pangolin/server/internal/auth/service.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/subtle"
	"errors"
	"fmt"
	"log/slog"
	"math/big"
	"strings"
	"time"

	"github.com/redis/go-redis/v9"
	"github.com/wangjia/pangolin/server/internal/apierr"
)

// maskEmail masks the local part of an email for safe logging (no full PII):
// "chenxin880812@gmail.com" -> "ch***@gmail.com". Used only in error logs.
func maskEmail(e string) string {
	at := strings.IndexByte(e, '@')
	if at <= 0 {
		return "***"
	}
	local := e[:at]
	if len(local) <= 2 {
		return "***" + e[at:]
	}
	return local[:2] + "***" + e[at:]
}

// Redis key helpers for verification codes (doc/03 §4: auth:code:{email}).
func codeKey(email string) string         { return "auth:code:" + email }
func codeAttemptsKey(email string) string { return "auth:code:attempts:" + email }

// Rate-limit scopes.
const (
	scopeCodeEmail = "code:email" // per-email send limit
	scopeCodeIP    = "code:ip"    // per-IP send limit
	scopeLogin     = "login"      // per-email login-failure counter
)

// ServiceConfig tunes the auth service. Zero values fall back to the documented
// defaults (doc/02 §4.1, doc/06 §3).
type ServiceConfig struct {
	CodeTTL         time.Duration // verification code lifetime (default 10m)
	CodeMaxAttempts int           // max verify attempts before code is burned (default 5)
	TrialDays       int           // auto PRO trial length (default 7)

	EmailPerMinute  int           // code sends per email per minute (default 1)
	IPPerHour       int           // code sends per IP per hour (default 10)
	LoginFailMax    int           // failed logins before lock (default 5)
	LoginLockWindow time.Duration // lock / failure-window length (default 15m)
}

func (c *ServiceConfig) withDefaults() {
	if c.CodeTTL <= 0 {
		c.CodeTTL = 10 * time.Minute
	}
	if c.CodeMaxAttempts <= 0 {
		c.CodeMaxAttempts = 5
	}
	if c.TrialDays <= 0 {
		c.TrialDays = 7
	}
	if c.EmailPerMinute <= 0 {
		c.EmailPerMinute = 1
	}
	if c.IPPerHour <= 0 {
		c.IPPerHour = 10
	}
	if c.LoginFailMax <= 0 {
		c.LoginFailMax = 5
	}
	if c.LoginLockWindow <= 0 {
		c.LoginLockWindow = 15 * time.Minute
	}
}

// DeviceMeta is the client-reported device identity carried on login/register so
// the device can be registered (devices table) and, later, bound to a session.
type DeviceMeta struct {
	DeviceID      string // client-generated stable UUID (secure storage)
	Name          string // host/model name
	Platform      string // ios|android|windows|macos|linux
	ClientVersion string // app version (stored from P2 onward)
}

// DeviceRegistrar registers the logging-in device and returns its internal id
// (for session binding). Defined consumer-side to avoid an import cycle;
// devices.Service is adapted to it in main wiring. Best-effort (see recordLogin).
type DeviceRegistrar interface {
	RegisterDevice(ctx context.Context, userID int64, meta DeviceMeta) (deviceID int64, err error)
	// CheckDeviceLimit reports over-limit status after registration; nil means the
	// account is within its plan cap. Surfaced on the login response (non-hard-reject).
	CheckDeviceLimit(ctx context.Context, userID int64) (*DeviceLimit, error)
}

// DeviceLimit is the over-limit signal on a successful login: the account's active
// device count exceeds its plan cap. Login still succeeds (tokens issued); the
// client must present "remove a device" UX before proceeding. nil = within cap.
type DeviceLimit struct {
	MaxDevices int           `json:"max_devices"`
	Devices    []DeviceBrief `json:"devices"`
}

// DeviceBrief is a minimal device view for the over-limit picker.
type DeviceBrief struct {
	UUID     string  `json:"uuid"`
	Name     string  `json:"name"`
	Platform string  `json:"platform"`
	LastSeen *string `json:"last_seen"` // RFC 3339 UTC; null when never seen
}

// SessionStore persists login sessions bound to a refresh JTI. Satisfied by
// sessions.Store; injected so auth stays decoupled.
type SessionStore interface {
	Create(ctx context.Context, userID, deviceID int64, jti, ip, clientVersion string) error
	Rotate(ctx context.Context, oldJTI, newJTI string) error
	Revoke(ctx context.Context, jti string) error
}

// Service is the auth business layer: code issuance, registration, login, and
// token refresh. It is safe for concurrent use.
type Service struct {
	store    UserStore
	rdb      *redis.Client
	rl       *RateLimiter
	tokens   *TokenManager
	mailer   Mailer
	cfg      ServiceConfig
	now      func() time.Time
	devReg   DeviceRegistrar // nil until wired; registration is best-effort
	sessions SessionStore    // nil until wired; session recording is best-effort
	referral ReferralHook    // nil until wired; referral attribution is best-effort
}

// ReferralHook is invoked after a successful registration to attribute the new
// account to the invite code it registered with (best-effort; implemented by
// reward.Service). Defined consumer-side to avoid an import cycle.
type ReferralHook interface {
	OnRegister(ctx context.Context, inviteeID int64, inviteCode, deviceUUID string)
}

// SetDeviceRegistrar / SetSessionStore / SetReferralHook wire collaborators
// after construction (main keeps auth, devices, sessions and reward decoupled).
// Call once during startup.
func (s *Service) SetDeviceRegistrar(r DeviceRegistrar) { s.devReg = r }
func (s *Service) SetSessionStore(st SessionStore)      { s.sessions = st }
func (s *Service) SetReferralHook(h ReferralHook)       { s.referral = h }

// recordLogin registers the device and binds a session to the freshly-issued
// refresh JTI. Best-effort: a registrar/session error (device cap, transient DB)
// is logged but never fails login — the user must always be able to get in
// (notably: free-plan reinstall churns the device UUID, so a hard cap here would
// lock users out).
func (s *Service) recordLogin(ctx context.Context, userID int64, refreshJTI, ip string, meta DeviceMeta) *DeviceLimit {
	if meta.DeviceID == "" {
		return nil
	}
	var deviceID int64
	var limit *DeviceLimit
	if s.devReg != nil {
		id, err := s.devReg.RegisterDevice(ctx, userID, meta)
		if err != nil {
			slog.Warn("auth: device register failed (login proceeds)", "uid", userID, "err", err)
		}
		deviceID = id
		// 超限闸(非硬拒登):设备已注册,若活跃设备数超套餐上限则返回信号让客户端处理;
		// 检查失败也绝不阻断登录。
		if dl, err := s.devReg.CheckDeviceLimit(ctx, userID); err != nil {
			slog.Warn("auth: device-limit check failed (login proceeds)", "uid", userID, "err", err)
		} else {
			limit = dl
		}
	}
	if s.sessions != nil && deviceID > 0 && refreshJTI != "" {
		if err := s.sessions.Create(ctx, userID, deviceID, refreshJTI, ip, meta.ClientVersion); err != nil {
			slog.Warn("auth: session create failed (login proceeds)", "uid", userID, "err", err)
		}
	}
	return limit
}

// NewService wires the auth service. now may be nil (defaults to time.Now).
func NewService(store UserStore, rdb *redis.Client, rl *RateLimiter, tokens *TokenManager, mailer Mailer, cfg ServiceConfig, now func() time.Time) *Service {
	cfg.withDefaults()
	if now == nil {
		now = time.Now
	}
	return &Service{
		store:  store,
		rdb:    rdb,
		rl:     rl,
		tokens: tokens,
		mailer: mailer,
		cfg:    cfg,
		now:    now,
	}
}

// SendCode applies rate-limiting and disposable-domain checks, generates a
// 6-digit code, stores it in Redis (TTL CodeTTL), and dispatches it
// asynchronously. retryAfter is non-zero only when a rate limit was hit.
func (s *Service) SendCode(ctx context.Context, rawEmail, ip string) (retryAfter time.Duration, apiErr *apierr.Error) {
	email := NormalizeEmail(rawEmail)
	if !ValidEmail(email) {
		return 0, ErrInvalidRequest
	}
	if IsDisposable(email) {
		return 0, ErrEmailDisposable
	}

	// Per-email limit: 1/min by default.
	ok, ra, err := s.rl.Allow(ctx, scopeCodeEmail, email, s.cfg.EmailPerMinute, time.Minute)
	if err != nil {
		return 0, ErrInternal
	}
	if !ok {
		return ra, ErrRateLimited
	}
	// Per-IP limit: e.g. 10/h. Skipped when IP is unknown.
	if ip != "" {
		ok, ra, err = s.rl.Allow(ctx, scopeCodeIP, ip, s.cfg.IPPerHour, time.Hour)
		if err != nil {
			return 0, ErrInternal
		}
		if !ok {
			return ra, ErrRateLimited
		}
	}

	// Anti-enumeration: a registered email gets a "you already have an account"
	// notice instead of a verification code; the API response (204) is identical
	// either way, so this endpoint can't be used to discover which emails are
	// registered. (Critical for a circumvention tool: a registered email implies
	// the person uses the service.)
	if _, err := s.store.GetUserByEmail(ctx, email); err == nil {
		go func(to string) {
			sendCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
			defer cancel()
			if err := s.mailer.SendAlreadyRegistered(sendCtx, to); err != nil {
				slog.Error("auth: send already-registered notice failed",
					"email", maskEmail(to), "err", err)
			}
		}(email)
		return 0, nil
	} else if !errors.Is(err, ErrNotFound) {
		return 0, ErrInternal
	}

	code, err := genNumericCode(6)
	if err != nil {
		return 0, ErrInternal
	}

	pipe := s.rdb.Pipeline()
	pipe.Set(ctx, codeKey(email), code, s.cfg.CodeTTL)
	pipe.Del(ctx, codeAttemptsKey(email)) // reset attempt counter for the new code
	if _, err := pipe.Exec(ctx); err != nil {
		return 0, ErrInternal
	}

	// Dispatch asynchronously; the request must not block on SMTP. A detached
	// context is used so request cancellation doesn't abort delivery.
	go func(to, c string) {
		sendCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
		defer cancel()
		if err := s.mailer.SendCode(sendCtx, to, c); err != nil {
			slog.Error("auth: send verification code failed",
				"email", maskEmail(to), "err", err)
		}
	}(email, code)

	return 0, nil
}

// Register verifies the code (one-time), creates the account plus a 7-day PRO
// trial in a single transaction, and returns a fresh token pair.
func (s *Service) Register(ctx context.Context, rawEmail, code, password, ip string, device DeviceMeta, inviteCode string) (*TokenPair, *apierr.Error) {
	email := NormalizeEmail(rawEmail)
	if !ValidEmail(email) || len(password) < 8 || len(code) != 6 {
		return nil, ErrInvalidRequest
	}

	// Verify the code with brute-force protection.
	if apiErr := s.verifyCode(ctx, email, code); apiErr != nil {
		return nil, apiErr
	}

	pwHash, err := HashPassword(password)
	if err != nil {
		return nil, ErrInternal
	}

	user, err := s.store.CreateUserWithTrial(ctx, email, pwHash, s.cfg.TrialDays)
	if err != nil {
		if errors.Is(err, ErrEmailTaken) {
			// Anti-enumeration: do NOT reveal the email is already registered.
			// A registered email never receives a code (SendCode mails a login
			// reminder instead), so this branch is normally unreachable; if hit
			// (e.g. a race), return the same generic error as a wrong/expired
			// code rather than "email exists".
			return nil, ErrCodeInvalid
		}
		return nil, ErrInternal
	}

	pair, jti, err := s.tokens.IssueWithJTI(ctx, user.ID, user.UUID)
	if err != nil {
		return nil, ErrInternal
	}
	s.recordLogin(ctx, user.ID, jti, ip, device)
	if s.referral != nil {
		s.referral.OnRegister(ctx, user.ID, inviteCode, device.DeviceID) // best-effort
	}
	return pair, nil
}

// verifyCode checks the supplied code against Redis. The code is consumed
// (deleted) on success, and after CodeMaxAttempts failed tries it is burned to
// stop brute forcing the 6-digit space. All failure modes return ErrCodeInvalid
// to avoid distinguishing wrong / expired / used.
func (s *Service) verifyCode(ctx context.Context, email, code string) *apierr.Error {
	stored, err := s.rdb.Get(ctx, codeKey(email)).Result()
	if err == redis.Nil {
		return ErrCodeInvalid
	}
	if err != nil {
		return ErrInternal
	}

	// Count this attempt (TTL bounded by the code lifetime).
	attempts, err := s.rdb.Incr(ctx, codeAttemptsKey(email)).Result()
	if err != nil {
		return ErrInternal
	}
	if attempts == 1 {
		_ = s.rdb.Expire(ctx, codeAttemptsKey(email), s.cfg.CodeTTL).Err()
	}
	if attempts > int64(s.cfg.CodeMaxAttempts) {
		// Burn the code and the counter.
		_ = s.rdb.Del(ctx, codeKey(email), codeAttemptsKey(email)).Err()
		return ErrCodeInvalid
	}

	if subtle.ConstantTimeCompare([]byte(stored), []byte(code)) != 1 {
		return ErrCodeInvalid
	}

	// Success: consume the code (one-time use).
	_ = s.rdb.Del(ctx, codeKey(email), codeAttemptsKey(email)).Err()
	return nil
}

// Login authenticates email+password with constant-time behaviour and a
// failure-count lock. retryAfter is non-zero only when the account is locked.
// LoginOutcome is the result of a password login: either issued tokens, or a
// short-lived PendingToken when the account has TOTP enabled and must complete
// the second step at /auth/login/totp.
type LoginOutcome struct {
	Tokens       *TokenPair
	PendingToken string
	DeviceLimit  *DeviceLimit // non-nil when active devices exceed the plan cap
}

// TOTP pending-login token (Redis): maps a one-time pending token → user id.
const (
	totpPendingPrefix = "totp:pending:"
	totpPendingTTL    = 5 * time.Minute
)

func (s *Service) Login(ctx context.Context, rawEmail, password, ip string, device DeviceMeta) (*LoginOutcome, time.Duration, *apierr.Error) {
	email := NormalizeEmail(rawEmail)
	if email == "" || password == "" {
		return nil, 0, ErrInvalidRequest
	}

	// Lock check.
	count, ttl, err := s.rl.FailureCount(ctx, scopeLogin, email)
	if err != nil {
		return nil, 0, ErrInternal
	}
	if count >= int64(s.cfg.LoginFailMax) && ttl > 0 {
		return nil, ttl, ErrAccountLocked
	}

	user, err := s.store.GetUserByEmail(ctx, email)
	if err != nil {
		if errors.Is(err, ErrNotFound) {
			// Spend equivalent CPU so timing doesn't reveal account existence.
			ConstantTimeReject(password)
			_, _ = s.rl.RecordFailure(ctx, scopeLogin, email, s.cfg.LoginLockWindow)
			return nil, 0, ErrInvalidCredentials
		}
		return nil, 0, ErrInternal
	}

	valid, verr := VerifyPassword(user.PwHash, password)
	if verr != nil || !valid {
		_, _ = s.rl.RecordFailure(ctx, scopeLogin, email, s.cfg.LoginLockWindow)
		return nil, 0, ErrInvalidCredentials
	}

	if user.Status == "banned" {
		return nil, 0, ErrAccountBanned
	}

	// Success: clear the failure counter.
	_ = s.rl.ClearFailures(ctx, scopeLogin, email)

	// Two-factor gate: when enabled, don't issue tokens yet — return a short-lived
	// pending token the client exchanges with a TOTP code at /auth/login/totp.
	if user.TOTPEnabled {
		pending, perr := newToken16()
		if perr != nil {
			return nil, 0, ErrInternal
		}
		if err := s.rdb.Set(ctx, totpPendingPrefix+pending, user.ID, totpPendingTTL).Err(); err != nil {
			return nil, 0, ErrInternal
		}
		return &LoginOutcome{PendingToken: pending}, 0, nil
	}

	pair, jti, err := s.tokens.IssueWithJTI(ctx, user.ID, user.UUID)
	if err != nil {
		return nil, 0, ErrInternal
	}
	dl := s.recordLogin(ctx, user.ID, jti, ip, device)
	return &LoginOutcome{Tokens: pair, DeviceLimit: dl}, 0, nil
}

// Logout revokes the given refresh token. Idempotent: an empty, unknown, or
// malformed token is a no-op (the client clears its own session regardless).
func (s *Service) Logout(ctx context.Context, refreshToken string) {
	if refreshToken == "" {
		return
	}
	// Mark the bound session revoked before dropping the Redis whitelist entry.
	if s.sessions != nil {
		if jti, err := s.tokens.ParseRefreshJTI(refreshToken); err == nil {
			_ = s.sessions.Revoke(ctx, jti)
		}
	}
	_ = s.tokens.RevokeRefresh(ctx, refreshToken)
}

// Refresh validates and rotates a refresh token, moving the bound session to the
// new JTI.
func (s *Service) Refresh(ctx context.Context, refreshToken string) (*TokenPair, *apierr.Error) {
	if refreshToken == "" {
		return nil, ErrInvalidRequest
	}
	pair, oldJTI, newJTI, err := s.tokens.RefreshWithJTI(ctx, refreshToken)
	if err != nil {
		if errors.Is(err, ErrInvalidTokenSentinel) {
			return nil, ErrInvalidToken
		}
		// Parse / signature / expiry failures all map to an opaque invalid-token.
		return nil, ErrInvalidToken
	}
	if s.sessions != nil {
		_ = s.sessions.Rotate(ctx, oldJTI, newJTI)
	}
	return pair, nil
}

// genNumericCode returns an n-digit numeric string drawn from crypto/rand.
func genNumericCode(n int) (string, error) {
	const digits = "0123456789"
	b := make([]byte, n)
	for i := range b {
		idx, err := rand.Int(rand.Reader, big.NewInt(int64(len(digits))))
		if err != nil {
			return "", fmt.Errorf("auth: gen code: %w", err)
		}
		b[i] = digits[idx.Int64()]
	}
	return string(b), nil
}