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

// 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
func GenerateToken(userID string, role string) (string, string, error) {
	jti, err := generateJTI()
	if err != nil {
		return "", "", err
	}

	_, tokenString, err := TokenAuth.Encode(map[string]any{
		"user_id": userID,
		"role":    role,
		"jti":     jti,
		"exp":     clock.Now().Add(1 * time.Hour).Unix(), // 1 hour
	})
	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")
	}

	return userID, role, jti, nil
}

// 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
func GenerateRefreshToken(ctx context.Context, userID string, role string) (string, error) {
	token, err := generateRefreshTokenString()
	if err != nil {
		return "", err
	}

	// Store hashed version in DB with 90-day expiry
	query := `
		INSERT INTO refresh_tokens (user_id, token_hash, role, expires_at)
		VALUES ($1, encode(sha256($2::bytea), 'hex'), $3, NOW() + INTERVAL '90 days')
		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).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
}

// 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 90-day 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() + INTERVAL '90 days')
		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).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 (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 a critical admin
// notification (reason 'refresh_token_reuse') is raised. 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 is theft:
	// the descendant minted at rotation would otherwise stay valid for 90 days.
	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
	`, tokenString).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")
}