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
type AuthResponse struct {
	Token        string `json:"token"`
	JTI          string `json:"jti"`
	RefreshToken string `json:"-"`
}

// 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 := TokenAuth.Decode(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
}

// VerifyRefreshToken checks a refresh token and returns user details if valid
// The token is consumed (deleted) upon successful verification, implementing rotation.
func VerifyRefreshToken(ctx context.Context, tokenString string) (userID string, role string, err error) {
	query := `
		DELETE FROM refresh_tokens
		WHERE token_hash = encode(sha256($1::bytea), 'hex')
		AND expires_at > NOW()
		AND NOT revoked
		RETURNING user_id, role`

	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)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return "", "", fmt.Errorf("invalid or expired refresh token")
		}
		return "", "", fmt.Errorf("failed to verify refresh token: %w", err)
	}

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

	// Token was consumed (DELETE returned it) — this is rotation
	// If a token is used twice, the second DELETE returns no rows = invalid
	return userID, role, nil
}