pangolin/server/internal/auth/token.go

package auth

import (
	"context"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"time"

	"github.com/golang-jwt/jwt/v4"
	"github.com/google/uuid"
	"github.com/redis/go-redis/v9"
)

// Token type discriminator carried in the `typ` claim so an access token can
// never be replayed as a refresh token and vice-versa.
const (
	typAccess  = "access"
	typRefresh = "refresh"
)

// Default lifetimes (doc/02 §4.1): access 15 min, refresh 30 days.
const (
	DefaultAccessTTL  = 15 * time.Minute
	DefaultRefreshTTL = 30 * 24 * time.Hour
)

// refreshKeyPrefix is the Redis whitelist prefix for refresh-token JTIs.
const refreshKeyPrefix = "jwt:refresh:"

// Claims is the JWT payload. Subject holds the user UUID; UID carries the
// numeric primary key so downstream middleware/stores avoid a DB round-trip.
type Claims struct {
	jwt.RegisteredClaims
	UID int64  `json:"uid"`
	Typ string `json:"typ"`
}

// TokenPair is the issued access/refresh pair plus the access lifetime seconds.
type TokenPair struct {
	AccessToken  string
	RefreshToken string
	ExpiresIn    int // access_token validity in seconds
}

// TokenManager signs tokens with one RSA private key (identified by kid) and
// verifies with a set of public keys keyed by kid — accepting both the current
// and previous keys to support zero-downtime key rotation. Refresh tokens are
// whitelisted in Redis so logout/ban takes effect immediately.
type TokenManager struct {
	signKey    *rsa.PrivateKey
	signKID    string
	verifyKeys map[string]*rsa.PublicKey // kid -> public key (current + old)
	accessTTL  time.Duration
	refreshTTL time.Duration
	rdb        *redis.Client
	now        func() time.Time
}

// TokenConfig configures a TokenManager.
type TokenConfig struct {
	// SignKey is the active RSA private key used to sign new tokens.
	SignKey *rsa.PrivateKey
	// SignKID is the key id written into the JWT header.
	SignKID string
	// VerifyKeys maps kid -> public key. Must contain SignKID; may contain
	// additional (older) keys still accepted during rotation. If nil, the
	// public part of SignKey under SignKID is used.
	VerifyKeys map[string]*rsa.PublicKey
	AccessTTL  time.Duration
	RefreshTTL time.Duration
	// Now is an optional clock override for tests.
	Now func() time.Time
}

// NewTokenManager constructs a TokenManager, validating the key material.
func NewTokenManager(rdb *redis.Client, cfg TokenConfig) (*TokenManager, error) {
	if cfg.SignKey == nil {
		return nil, errors.New("auth: token manager requires a signing key")
	}
	if cfg.SignKID == "" {
		return nil, errors.New("auth: token manager requires a signing key id (kid)")
	}
	verify := cfg.VerifyKeys
	if verify == nil {
		verify = map[string]*rsa.PublicKey{}
	}
	if _, ok := verify[cfg.SignKID]; !ok {
		verify[cfg.SignKID] = &cfg.SignKey.PublicKey
	}
	accessTTL := cfg.AccessTTL
	if accessTTL <= 0 {
		accessTTL = DefaultAccessTTL
	}
	refreshTTL := cfg.RefreshTTL
	if refreshTTL <= 0 {
		refreshTTL = DefaultRefreshTTL
	}
	now := cfg.Now
	if now == nil {
		now = time.Now
	}
	return &TokenManager{
		signKey:    cfg.SignKey,
		signKID:    cfg.SignKID,
		verifyKeys: verify,
		accessTTL:  accessTTL,
		refreshTTL: refreshTTL,
		rdb:        rdb,
		now:        now,
	}, nil
}

// IssueWithJTI mints a fresh access+refresh pair, whitelists the refresh JTI in
// Redis, and returns that JTI so the caller can bind a session to the refresh
// token. Issue wraps this when the JTI isn't needed.
func (tm *TokenManager) IssueWithJTI(ctx context.Context, userID int64, userUUID string) (*TokenPair, string, error) {
	now := tm.now()

	access, _, err := tm.sign(userID, userUUID, typAccess, tm.accessTTL, now)
	if err != nil {
		return nil, "", err
	}
	refresh, refreshJTI, err := tm.sign(userID, userUUID, typRefresh, tm.refreshTTL, now)
	if err != nil {
		return nil, "", err
	}
	if err := tm.whitelist(ctx, refreshJTI, userID); err != nil {
		return nil, "", err
	}
	return &TokenPair{
		AccessToken:  access,
		RefreshToken: refresh,
		ExpiresIn:    int(tm.accessTTL.Seconds()),
	}, refreshJTI, nil
}

// Issue mints a fresh access+refresh pair (refresh JTI discarded).
func (tm *TokenManager) Issue(ctx context.Context, userID int64, userUUID string) (*TokenPair, error) {
	pair, _, err := tm.IssueWithJTI(ctx, userID, userUUID)
	return pair, err
}

// sign builds and signs one token, returning the compact string and its JTI.
func (tm *TokenManager) sign(userID int64, userUUID, typ string, ttl time.Duration, now time.Time) (string, string, error) {
	jti := uuid.NewString()
	claims := Claims{
		RegisteredClaims: jwt.RegisteredClaims{
			Subject:   userUUID,
			ID:        jti,
			IssuedAt:  jwt.NewNumericDate(now),
			ExpiresAt: jwt.NewNumericDate(now.Add(ttl)),
		},
		UID: userID,
		Typ: typ,
	}
	tok := jwt.NewWithClaims(jwt.SigningMethodRS256, claims)
	tok.Header["kid"] = tm.signKID
	signed, err := tok.SignedString(tm.signKey)
	if err != nil {
		return "", "", fmt.Errorf("auth: sign token: %w", err)
	}
	return signed, jti, nil
}

// keyfunc resolves the verification key from the token's kid header and rejects
// any algorithm other than RS256.
func (tm *TokenManager) keyfunc(t *jwt.Token) (interface{}, error) {
	if _, ok := t.Method.(*jwt.SigningMethodRSA); !ok {
		return nil, fmt.Errorf("auth: unexpected signing method %q", t.Header["alg"])
	}
	kid, _ := t.Header["kid"].(string)
	if kid == "" {
		return nil, errors.New("auth: token missing kid")
	}
	pub, ok := tm.verifyKeys[kid]
	if !ok {
		return nil, fmt.Errorf("auth: unknown kid %q", kid)
	}
	return pub, nil
}

// parse validates signature, expiry, and the expected token type. Expiry is
// checked manually against tm.now so tests can inject a clock without mutating
// the package-global jwt.TimeFunc.
func (tm *TokenManager) parse(tokenStr, wantTyp string) (*Claims, error) {
	claims := &Claims{}
	parser := jwt.NewParser(
		jwt.WithValidMethods([]string{"RS256"}),
		jwt.WithoutClaimsValidation(), // we validate exp/iat ourselves below
	)
	if _, err := parser.ParseWithClaims(tokenStr, claims, tm.keyfunc); err != nil {
		return nil, err
	}

	now := tm.now()
	if claims.ExpiresAt == nil || now.After(claims.ExpiresAt.Time) {
		return nil, errors.New("auth: token expired")
	}
	if claims.IssuedAt != nil && now.Add(time.Minute).Before(claims.IssuedAt.Time) {
		return nil, errors.New("auth: token used before issued")
	}
	if claims.Typ != wantTyp {
		return nil, fmt.Errorf("auth: token type %q, want %q", claims.Typ, wantTyp)
	}
	return claims, nil
}

// ParseAccess validates an access token and returns its claims.
func (tm *TokenManager) ParseAccess(tokenStr string) (*Claims, error) {
	return tm.parse(tokenStr, typAccess)
}

// RefreshWithJTI validates+rotates a refresh token (single-use) and returns the
// old and new refresh JTIs so the caller can move the bound session. Refresh
// wraps this when the JTIs aren't needed.
func (tm *TokenManager) RefreshWithJTI(ctx context.Context, refreshToken string) (pair *TokenPair, oldJTI, newJTI string, err error) {
	claims, err := tm.parse(refreshToken, typRefresh)
	if err != nil {
		return nil, "", "", err
	}

	// Whitelist check + single-use rotation: DEL returns the number of keys
	// removed; 0 means the JTI was absent (already rotated, logged out, or
	// banned) → reject.
	removed, err := tm.rdb.Del(ctx, refreshKeyPrefix+claims.ID).Result()
	if err != nil {
		return nil, "", "", fmt.Errorf("auth: refresh whitelist del: %w", err)
	}
	if removed == 0 {
		return nil, "", "", ErrInvalidTokenSentinel
	}

	pair, newJTI, err = tm.IssueWithJTI(ctx, claims.UID, claims.Subject)
	if err != nil {
		return nil, "", "", err
	}
	return pair, claims.ID, newJTI, nil
}

// Refresh validates+rotates a refresh token (JTIs discarded).
func (tm *TokenManager) Refresh(ctx context.Context, refreshToken string) (*TokenPair, error) {
	pair, _, _, err := tm.RefreshWithJTI(ctx, refreshToken)
	return pair, err
}

// Revoke removes a single refresh JTI from the whitelist (logout).
func (tm *TokenManager) Revoke(ctx context.Context, jti string) error {
	return tm.rdb.Del(ctx, refreshKeyPrefix+jti).Err()
}

// ParseRefreshJTI returns a refresh token's JTI (for session lookup on logout).
func (tm *TokenManager) ParseRefreshJTI(refreshToken string) (string, error) {
	claims, err := tm.parse(refreshToken, typRefresh)
	if err != nil {
		return "", err
	}
	return claims.ID, nil
}

// RevokeRefresh parses a refresh token and revokes its JTI. Used by logout.
// Returns an error only when the token is structurally invalid; a missing or
// already-rotated JTI is a no-op (logout is idempotent).
func (tm *TokenManager) RevokeRefresh(ctx context.Context, refreshToken string) error {
	claims, err := tm.parse(refreshToken, typRefresh)
	if err != nil {
		return err
	}
	return tm.Revoke(ctx, claims.ID)
}

// whitelist stores the refresh JTI with the refresh TTL.
func (tm *TokenManager) whitelist(ctx context.Context, jti string, userID int64) error {
	if err := tm.rdb.Set(ctx, refreshKeyPrefix+jti, userID, tm.refreshTTL).Err(); err != nil {
		return fmt.Errorf("auth: refresh whitelist set: %w", err)
	}
	return nil
}

// ErrInvalidTokenSentinel is returned by Refresh when the token is structurally
// valid but no longer whitelisted. Callers map it to ErrInvalidToken.
var ErrInvalidTokenSentinel = errors.New("auth: refresh token not in whitelist")

// --------------------------------------------------------------------------
// PEM loading helpers (used by wiring/config to build a TokenManager)
// --------------------------------------------------------------------------

// LoadPrivateKeyPEM parses a PEM-encoded RSA private key (PKCS#1 or PKCS#8).
func LoadPrivateKeyPEM(pemBytes []byte) (*rsa.PrivateKey, error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, errors.New("auth: no PEM block in private key")
	}
	if key, err := x509.ParsePKCS1PrivateKey(block.Bytes); err == nil {
		return key, nil
	}
	parsed, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("auth: parse private key: %w", err)
	}
	rsaKey, ok := parsed.(*rsa.PrivateKey)
	if !ok {
		return nil, errors.New("auth: private key is not RSA")
	}
	return rsaKey, nil
}

// LoadPublicKeyPEM parses a PEM-encoded RSA public key (PKIX or PKCS#1).
func LoadPublicKeyPEM(pemBytes []byte) (*rsa.PublicKey, error) {
	block, _ := pem.Decode(pemBytes)
	if block == nil {
		return nil, errors.New("auth: no PEM block in public key")
	}
	if pub, err := x509.ParsePKIXPublicKey(block.Bytes); err == nil {
		if rsaPub, ok := pub.(*rsa.PublicKey); ok {
			return rsaPub, nil
		}
		return nil, errors.New("auth: public key is not RSA")
	}
	rsaPub, err := x509.ParsePKCS1PublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("auth: parse public key: %w", err)
	}
	return rsaPub, nil
}