Crussell/backend/auth/jwt.go

package auth

import (
	"context"
	"crypto/rand"
	"errors"
	"fmt"
	"log"
	"log/slog"
	"time"

	"crussell/clock"
	"crussell/db"

	"github.com/jackc/pgx/v5"

	"github.com/go-chi/jwtauth/v5"
)

var TokenAuth *jwtauth.JWTAuth

// RefreshTokenLifetime is how long an issued refresh token stays valid before
// it expires (90 days). Minted rows get expires_at = NOW() + RefreshTokenLifetime
// (see GenerateRefreshToken / GenerateRefreshTokenInFamily), and the cleanup
// jobs that expire refresh-token families and retain revoked tokens
// (handlers/scheduling CleanupExpiredRefreshTokens, jobs SweepSquareWebhookEvents)
// reference the SAME constant so the SQL window can never drift from the mint.
const RefreshTokenLifetime = 90 * 24 * time.Hour

// refreshTokenLifetimeDays is RefreshTokenLifetime expressed in whole days, fed
// to the SQL make_interval(days => ...) calls in the mint queries.
const refreshTokenLifetimeDays = int64(RefreshTokenLifetime / (24 * time.Hour))

// refreshTokenReuseGrace is how long after a rotation a used-token replay is
// treated as a BENIGN concurrent refresh (two tabs sharing one refresh token in
// localStorage both refreshing on load) instead of theft. A replay inside the
// grace window gets the generic error but does NOT kill the rotation family and
// does NOT raise the refresh_token_reuse alert — only a replay after the window
// has elapsed is treated as theft (see VerifyRefreshToken).
const refreshTokenReuseGrace = 30 * time.Second

// refreshTokenReuseGraceSecs is the grace window in whole seconds for the SQL
// make_interval(secs => ...) comparison in VerifyRefreshToken's reuse branch.
const refreshTokenReuseGraceSecs = int64(refreshTokenReuseGrace / time.Second)

// accessTokenFamilyClaim is the JWT claim that binds an access token to the
// refresh-token rotation family it was minted alongside. VerifyToken rejects an
// access token whose family_id no longer exists in refresh_tokens, so when
// reuse detection DELETEs a family every access token minted by that lineage
// dies immediately instead of remaining valid for its 1-hour TTL (HIGH 1).
const accessTokenFamilyClaim = "family_id"

// AuthResponse is the response structure for login/refresh endpoints.
// RefreshToken is a 90-day opaque, DB-hashed, single-use credential; the
// client stores it and presents it (Bearer) to POST /api/refresh-token in
// exchange for a fresh access token + a rotated refresh token. It is returned
// in JSON so the SPA can persist it — omitting it would make the refresh flow
// unusable — but it is never logged and never returned by any other endpoint.
type AuthResponse struct {
	Token        string `json:"token"`
	JTI          string `json:"jti"`
	RefreshToken string `json:"refreshToken,omitempty"`
}

// generateJTI generates a UUID v4 string using crypto/rand
func generateJTI() (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("failed to generate JTI: %w", err)
	}
	b[6] = (b[6] & 0x0f) | 0x40
	b[8] = (b[8] & 0x3f) | 0x80
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x",
		b[0:4], b[4:6], b[6:8], b[8:10], b[10:16]), nil
}

// RevokeJTI adds a JTI to the revoked set in PostgreSQL.
// Returns an error if the operation fails.
func RevokeJTI(ctx context.Context, jti string, expiresAt time.Time) error {
	if db.Conn == nil {
		return fmt.Errorf("revoke JTI: db.Conn is nil")
	}
	tx, err := db.Conn.Begin(ctx)
	if err != nil {
		return fmt.Errorf("revoke JTI: begin transaction: %w", err)
	}
	defer func() {
		if err := tx.Rollback(ctx); err != nil && !errors.Is(err, pgx.ErrTxClosed) {
			slog.Error("failed to rollback transaction", "err", err)
		}
	}()

	_, err = tx.Exec(ctx,
		`INSERT INTO revoked_jtis (jti, expires_at) VALUES ($1, $2)
		 ON CONFLICT (jti) DO NOTHING`,
		jti, expiresAt)
	if err != nil {
		return fmt.Errorf("revoke JTI %s: %w", jti, err)
	}

	if err := tx.Commit(ctx); err != nil {
		return fmt.Errorf("revoke JTI: commit transaction: %w", err)
	}

	return nil
}

// IsJTIRevoked checks if a JTI is in the revoked set via PostgreSQL.
// Returns false if the DB is not initialized (unit tests, startup) — treating
// the token as valid is the safer default for availability over security during
// startup, and revoked checks are quickly re-evaluated on each request.
func IsJTIRevoked(ctx context.Context, jti string) bool {
	if db.Conn == nil {
		return false
	}
	var exists bool
	err := db.Conn.QueryRow(ctx,
		`SELECT EXISTS(SELECT 1 FROM revoked_jtis WHERE jti = $1 AND expires_at > NOW())`,
		jti).Scan(&exists)
	if err != nil {
		return false
	}
	return exists
}

// CleanupRevokedJTIs removes expired entries from PostgreSQL and returns the count of deleted rows.
func CleanupRevokedJTIs(ctx context.Context) (int, error) {
	if db.Conn == nil {
		return 0, nil
	}
	tx, err := db.Conn.Begin(ctx)
	if err != nil {
		log.Printf("WARN: Failed to begin transaction for JTI cleanup: %v", err)
		return 0, err
	}
	defer func() {
		if err := tx.Rollback(ctx); err != nil && !errors.Is(err, pgx.ErrTxClosed) {
			slog.Error("failed to rollback transaction", "err", err)
		}
	}()

	tag, err := tx.Exec(ctx,
		`DELETE FROM revoked_jtis WHERE expires_at < NOW()`)
	if err != nil {
		log.Printf("WARN: Failed to cleanup revoked JTIs: %v", err)
		return 0, err
	}

	if err := tx.Commit(ctx); err != nil {
		log.Printf("WARN: Failed to commit transaction for JTI cleanup: %v", err)
		return 0, err
	}

	return int(tag.RowsAffected()), nil
}

func InitJWT(secret string) {
	TokenAuth = jwtauth.New("HS256", []byte(secret), nil)
}

// GenerateToken creates a JWT with user_id, role, and a unique jti claim.
// Returns the token string, the JTI, and any error. The token carries NO
// family_id claim, so it is exempt from the family-alive check in VerifyToken —
// kept for callers minting transient/test tokens. Production login and refresh
// paths use GenerateTokenForFamily so a killed rotation family cannot keep an
// access token alive (HIGH 1).
func GenerateToken(userID string, role string) (string, string, error) {
	return GenerateTokenForFamily(userID, role, "")
}

// GenerateTokenForFamily mints an access token carrying a family_id claim that
// binds it to a refresh-token rotation family. VerifyToken rejects a token whose
// family_id no longer exists in refresh_tokens, so when reuse detection DELETEs
// the whole family (VerifyRefreshToken), every access token minted by that
// lineage dies immediately instead of remaining valid for its 1-hour TTL.
// familyID == "" mints an unbound token (same as GenerateToken).
func GenerateTokenForFamily(userID string, role string, familyID string) (string, string, error) {
	jti, err := generateJTI()
	if err != nil {
		return "", "", err
	}

	claims := map[string]any{
		"user_id": userID,
		"role":    role,
		"jti":     jti,
		"exp":     clock.Now().Add(1 * time.Hour).Unix(), // 1 hour
	}
	if familyID != "" {
		claims[accessTokenFamilyClaim] = familyID
	}

	_, tokenString, err := TokenAuth.Encode(claims)
	return tokenString, jti, err
}

// VerifyToken validates JWT and returns user_id, role, and jti
func VerifyToken(tokenString string, ctx context.Context) (userID string, role string, jti string, err error) {
	token, err := jwtauth.VerifyToken(TokenAuth, tokenString)
	if err != nil {
		return "", "", "", err
	}

	var uidVal any
	if err := token.Get("user_id", &uidVal); err != nil {
		return "", "", "", fmt.Errorf("invalid user_id claim")
	}
	userID, ok := uidVal.(string)
	if !ok {
		return "", "", "", fmt.Errorf("invalid user_id claim")
	}

	var roleVal any
	if err := token.Get("role", &roleVal); err != nil {
		return "", "", "", fmt.Errorf("invalid role claim")
	}
	role, ok = roleVal.(string)
	if !ok {
		return "", "", "", fmt.Errorf("invalid role claim")
	}

	var jtiVal any
	if err := token.Get("jti", &jtiVal); err != nil {
		return "", "", "", fmt.Errorf("invalid jti claim")
	}
	jti, ok = jtiVal.(string)
	if !ok || jti == "" {
		return "", "", "", fmt.Errorf("invalid jti claim")
	}

	if IsJTIRevoked(ctx, jti) {
		return "", "", "", fmt.Errorf("token revoked")
	}

	// HIGH 1: reject an access token bound (family_id claim) to a rotation
	// family that reuse detection has killed. A token minted at rotation carries
	// family_id; when that family no longer exists in refresh_tokens the token is
	// dead even though its JTI was never revoked — closing the up-to-1-hour
	// window where a stolen refresh token's freshly-minted access token could
	// still hit money endpoints (gift-card buy, saved-card booking payment, tip).
	if err := verifyFamilyAlive(ctx, token, userID); err != nil {
		return "", "", "", err
	}

	return userID, role, jti, nil
}

// verifyFamilyAlive rejects access tokens bound to a rotation family that no
// longer exists in refresh_tokens. A token WITHOUT a family_id claim is unbound
// (minted via GenerateToken — tests/legacy callers) and passes. Fails closed on
// a live-DB query error: when the family cannot be confirmed alive the safe
// default for money endpoints is to refuse. Mirrors IsJTIRevoked's nil-db
// availability default (startup/unit tests).
func verifyFamilyAlive(ctx context.Context, token jwtClaimGetter, userID string) error {
	var familyVal any
	if err := token.Get(accessTokenFamilyClaim, &familyVal); err != nil {
		return nil
	}
	familyID, ok := familyVal.(string)
	if !ok || familyID == "" {
		return nil
	}
	if db.Conn == nil {
		return nil
	}
	var exists bool
	err := db.Conn.QueryRow(ctx,
		`SELECT EXISTS(
			SELECT 1 FROM refresh_tokens WHERE family_id = $1 AND user_id = $2
		)`, familyID, userID).Scan(&exists)
	if err != nil {
		return fmt.Errorf("token revoked")
	}
	if !exists {
		return fmt.Errorf("token revoked")
	}
	return nil
}

// jwtClaimGetter is the minimal subset of jwt.Token needed to read a claim
// (the token returned by jwtauth.VerifyToken). Kept as an interface so the
// lestrrat-go/jwx dependency stays out of this file's imports.
type jwtClaimGetter interface {
	Get(string, interface{}) error
}

// generateRefreshTokenString creates a cryptographically random opaque refresh token
func generateRefreshTokenString() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("failed to generate refresh token: %w", err)
	}
	return fmt.Sprintf("%x", b), nil
}

// GenerateRefreshToken creates a refresh token stored in the database. Returns
// the opaque token string to return to the client plus the id of the rotation
// family the new row was created in (so the caller can bind the matching access
// token to it via GenerateTokenForFamily — see HIGH 1).
func GenerateRefreshToken(ctx context.Context, userID string, role string) (string, string, error) {
	token, err := generateRefreshTokenString()
	if err != nil {
		return "", "", err
	}

	// Store hashed version in DB with the shared RefreshTokenLifetime expiry
	query := `
		INSERT INTO refresh_tokens (user_id, token_hash, role, family_id, expires_at)
		VALUES ($1, encode(sha256($2::bytea), 'hex'), $3, gen_random_uuid(), NOW() + make_interval(days => $4))
		RETURNING id, family_id`

	tx, err := db.Conn.Begin(ctx)
	if err != nil {
		return "", "", fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer func() {
		if err := tx.Rollback(ctx); err != nil && !errors.Is(err, pgx.ErrTxClosed) {
			slog.Error("failed to rollback transaction", "err", err)
		}
	}()

	var tokenID int64
	var familyID string
	err = tx.QueryRow(ctx, query, userID, token, role, refreshTokenLifetimeDays).Scan(&tokenID, &familyID)
	if err != nil {
		return "", "", fmt.Errorf("failed to store refresh token: %w", err)
	}

	if err := tx.Commit(ctx); err != nil {
		return "", "", fmt.Errorf("failed to commit transaction: %w", err)
	}

	return token, familyID, nil
}

// GenerateRefreshTokenInFamily creates a refresh token in the SAME rotation
// family as its parent (the family_id returned by VerifyRefreshToken). Rotation
// must mint the descendant in the parent's family so a replayed (already-used)
// ancestor can revoke the ENTIRE lineage — the descendant included — instead of
// leaving a fresh 90-day token alive after theft is detected.
func GenerateRefreshTokenInFamily(ctx context.Context, userID string, role string, familyID string) (string, error) {
	token, err := generateRefreshTokenString()
	if err != nil {
		return "", err
	}

	// Store hashed version in DB with the shared RefreshTokenLifetime expiry, in the given family
	query := `
		INSERT INTO refresh_tokens (user_id, token_hash, role, family_id, expires_at)
		VALUES ($1, encode(sha256($2::bytea), 'hex'), $3, $4, NOW() + make_interval(days => $5))
		RETURNING id`

	tx, err := db.Conn.Begin(ctx)
	if err != nil {
		return "", fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer func() {
		if err := tx.Rollback(ctx); err != nil && !errors.Is(err, pgx.ErrTxClosed) {
			slog.Error("failed to rollback transaction", "err", err)
		}
	}()

	var tokenID int64
	err = tx.QueryRow(ctx, query, userID, token, role, familyID, refreshTokenLifetimeDays).Scan(&tokenID)
	if err != nil {
		return "", fmt.Errorf("failed to store refresh token: %w", err)
	}

	if err := tx.Commit(ctx); err != nil {
		return "", fmt.Errorf("failed to commit transaction: %w", err)
	}

	return token, nil
}

// VerifyRefreshToken checks a refresh token and returns user details if valid.
// The token is consumed (marked used) upon successful verification — rotation —
// and its family_id is returned so the caller can mint the descendant in the
// SAME family. If an ALREADY-ROTATED token is presented again after the
// refreshTokenReuseGrace window (a replay: the attacker rotated it, then the
// victim replayed it), the entire rotation family is revoked (the descendant
// minted at rotation dies too, and so does every access token bound to the
// family — HIGH 1) and a critical admin notification (reason
// 'refresh_token_reuse') is raised. A replay WITHIN the grace window is a
// benign concurrent refresh (two tabs sharing one localStorage refresh token
// refreshing on load): it gets the generic error but kills NOTHING and raises
// NO alert, so the legitimately-rotated session survives. The caller always gets
// the generic "invalid or expired refresh token" error so reuse is never leaked.
func VerifyRefreshToken(ctx context.Context, tokenString string) (userID string, role string, familyID string, err error) {
	query := `
		UPDATE refresh_tokens SET used_at = NOW()
		WHERE token_hash = encode(sha256($1::bytea), 'hex')
		AND expires_at > NOW()
		AND NOT revoked
		AND used_at IS NULL
		RETURNING user_id, role, family_id`

	tx, err := db.Conn.Begin(ctx)
	if err != nil {
		return "", "", "", fmt.Errorf("failed to begin transaction: %w", err)
	}
	defer func() {
		if err := tx.Rollback(ctx); err != nil && !errors.Is(err, pgx.ErrTxClosed) {
			slog.Error("failed to rollback transaction", "err", err)
		}
	}()

	err = tx.QueryRow(ctx, query, tokenString).Scan(&userID, &role, &familyID)
	if err == nil {
		// Rotation: the token is marked used (kept in the row) so a later
		// replay can be detected, and its family_id is returned.
		if err := tx.Commit(ctx); err != nil {
			return "", "", "", fmt.Errorf("failed to commit transaction: %w", err)
		}
		return userID, role, familyID, nil
	}
	if !errors.Is(err, pgx.ErrNoRows) {
		return "", "", "", fmt.Errorf("failed to verify refresh token: %w", err)
	}

	// The rotation UPDATE matched nothing: the token is expired, revoked, or
	// never issued — OR it was already used (replayed). A used token replayed
	// AFTER the grace window is theft: the descendant minted at rotation would
	// otherwise stay valid for 90 days. A used token replayed WITHIN the grace
	// window is a benign concurrent refresh (two tabs, one shared refresh token)
	// and falls through to the generic error below — no family kill, no alert.
	var reusedUserID, reusedFamilyID string
	reuseErr := tx.QueryRow(ctx, `
		SELECT user_id, family_id FROM refresh_tokens
		WHERE token_hash = encode(sha256($1::bytea), 'hex')
		  AND used_at IS NOT NULL
		  AND used_at < NOW() - make_interval(secs => $2)
	`, tokenString, refreshTokenReuseGraceSecs).Scan(&reusedUserID, &reusedFamilyID)

	if reuseErr == nil {
		// (i) Revoke the ENTIRE family — the reused token and every descendant.
		if _, err := tx.Exec(ctx, `DELETE FROM refresh_tokens WHERE family_id = $1`, reusedFamilyID); err != nil {
			slog.Error("CRITICAL: refresh token reuse detected but family revocation failed", "userID", reusedUserID, "familyID", reusedFamilyID, "err", err)
		}
		// (ii) Surface the theft in the admin notification centre. The NOT
		// EXISTS guard keeps ONE alert per reused family until an admin
		// acknowledges it — mirroring insertCriticalPaymentNotification.
		if _, err := tx.Exec(ctx, `
			INSERT INTO admin_notifications (reason, user_id, created_at)
			SELECT 'refresh_token_reuse', $1, NOW()
			WHERE NOT EXISTS (
				SELECT 1 FROM admin_notifications an
				WHERE an.reason = 'refresh_token_reuse'
				  AND an.user_id = $1
				  AND an.acknowledged_at IS NULL
			)
		`, reusedUserID); err != nil {
			slog.Error("CRITICAL: refresh token reuse detected but admin alert insert failed", "userID", reusedUserID, "err", err)
		}
		// Commit the family revocation + alert — NOT the deferred rollback.
		if err := tx.Commit(ctx); err != nil {
			return "", "", "", fmt.Errorf("failed to commit transaction: %w", err)
		}
		// (iii) CRITICAL log; (iv) generic error — never leak that reuse was seen.
		slog.Error("CRITICAL: refresh token reuse detected — rotation family revoked", "userID", reusedUserID, "familyID", reusedFamilyID)
		return "", "", "", fmt.Errorf("invalid or expired refresh token")
	}
	if !errors.Is(reuseErr, pgx.ErrNoRows) {
		return "", "", "", fmt.Errorf("failed to verify refresh token: %w", reuseErr)
	}

	// Never-issued / expired / revoked token — indistinguishable from a replay
	// to the client, as before.
	return "", "", "", fmt.Errorf("invalid or expired refresh token")
}