feat(backend): update JWT auth implementation and tests

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
2026-06-18 16:26:00 +01:00
parent 2d0071c453
commit b5c1eef8c2
2 changed files with 148 additions and 51 deletions
@@ -3,10 +3,13 @@ package auth
 import (
 	"context"
 	"crypto/rand"
+	"database/sql"
+	"errors"
 	"fmt"
-	"sync"
 	"time"

+	"crussell/db"
+
 	"github.com/go-chi/jwtauth/v5"
 )

@@ -14,63 +17,72 @@ var TokenAuth *jwtauth.JWTAuth

 // AuthResponse is the response structure for login/refresh endpoints
 type AuthResponse struct {
-	Token string `json:"token"`
-	JTI   string `json:"jti"`
+	Token        string `json:"token"`
+	JTI          string `json:"jti"`
+	RefreshToken string `json:"refreshToken,omitempty"`
 }

-// In-memory revoked JTI tracking
-var (
-	revokedJTIs   = make(map[string]time.Time)
-	revokedJTIsMu sync.RWMutex
-)
-
 // generateJTI generates a UUID v4 string using crypto/rand
-// Format: xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx
 func generateJTI() (string, error) {
 	b := make([]byte, 16)
 	if _, err := rand.Read(b); err != nil {
 		return "", fmt.Errorf("failed to generate JTI: %w", err)
 	}
-
-	// Set version 4 bits
 	b[6] = (b[6] & 0x0f) | 0x40
-	// Set variant bits (10xx)
 	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 with its expiry time
+// RevokeJTI adds a JTI to the revoked set in PostgreSQL
 func RevokeJTI(jti string, expiresAt time.Time) {
-	revokedJTIsMu.Lock()
-	defer revokedJTIsMu.Unlock()
-	revokedJTIs[jti] = expiresAt
-}
-
-// IsJTIRevoked checks if a JTI is in the revoked set
-func IsJTIRevoked(jti string) bool {
-	revokedJTIsMu.RLock()
-	defer revokedJTIsMu.RUnlock()
-	_, revoked := revokedJTIs[jti]
-	return revoked
-}
-
-// CleanupRevokedJTIs removes entries where the expiry time has passed
-func CleanupRevokedJTIs() {
-	revokedJTIsMu.Lock()
-	defer revokedJTIsMu.Unlock()
-	now := time.Now()
-	for jti, expiresAt := range revokedJTIs {
-		if now.After(expiresAt) {
-			delete(revokedJTIs, jti)
-		}
+	if db.DB == nil {
+		return
+	}
+	_, err := db.DB.Exec(context.Background(),
+		`INSERT INTO revoked_jtis (jti, expires_at) VALUES ($1, $2)
+		 ON CONFLICT (jti) DO NOTHING`,
+		jti, expiresAt)
+	if err != nil {
+		// Log but don't fail - this is best effort
+		fmt.Printf("WARN: Failed to revoke JTI %s: %v\n", jti, err)
 	}
 }

-func init() {
+// 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(jti string) bool {
+	if db.DB == nil {
+		return false
+	}
+	var exists bool
+	err := db.DB.QueryRow(context.Background(),
+		`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
+func CleanupRevokedJTIs() {
+	if db.DB == nil {
+		return
+	}
+	_, err := db.DB.Exec(context.Background(),
+		`DELETE FROM revoked_jtis WHERE expires_at < NOW()`)
+	if err != nil {
+		fmt.Printf("WARN: Failed to cleanup revoked JTIs: %v\n", err)
+	}
+}
+
+// StartJTICleanup starts a background goroutine to periodically clean expired JTIs
+func StartJTICleanup() {
 	go func() {
-		ticker := time.NewTicker(5 * time.Minute)
+		ticker := time.NewTicker(30 * time.Minute)
 		defer ticker.Stop()
 		for range ticker.C {
 			CleanupRevokedJTIs()
@@ -94,7 +106,7 @@ func GenerateToken(userID string, role string) (string, string, error) {
 		"user_id": userID,
 		"role":    role,
 		"jti":     jti,
-		"exp":     time.Now().Add(30 * 24 * time.Hour).Unix(), // 30 days
+		"exp":     time.Now().Add(1 * time.Hour).Unix(), // 1 hour
 	})
 	return tokenString, jti, err
 }
@@ -133,17 +145,64 @@ func VerifyToken(tokenString string, ctx context.Context) (userID string, role s
 		return "", "", "", fmt.Errorf("invalid jti claim")
 	}

-	if IsJTIRevoked(jti) {
-		return "", "", "", fmt.Errorf("token revoked")
-	}
-	jti, ok = jtiVal.(string)
-	if !ok || jti == "" {
-		return "", "", "", fmt.Errorf("invalid jti claim")
-	}
-
 	if IsJTIRevoked(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(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`
+
+	var tokenID int64
+	err = db.DB.QueryRow(context.Background(), query, userID, token, role).Scan(&tokenID)
+	if err != nil {
+		return "", fmt.Errorf("failed to store refresh token: %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`
+
+	err = db.DB.QueryRow(ctx, query, tokenString).Scan(&userID, &role)
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return "", "", fmt.Errorf("invalid or expired refresh token")
+		}
+		return "", "", fmt.Errorf("failed to verify refresh token: %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
+}
@@ -14,18 +14,45 @@ package auth

 import (
 	"context"
+	"fmt"
 	"os"
 	"strings"
 	"testing"
 	"time"
+
+	"crussell/db"
+	"crussell/testutils/testdb"
 )

 func TestMain(m *testing.M) {
 	InitJWT("test-secret-key-for-jwt-test")
+
+	pool, err := testdb.NewPool("")
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "WARN: No test DB available — JTI revocation tests will fail: %v\n", err)
+	} else {
+		// Run migration to ensure the revoked_jtis and refresh_tokens tables exist
+		testdb.Migrate(&testing.T{}, pool)
+		db.DB = pool
+	}
+
 	code := m.Run()
+	if pool != nil {
+		pool.Close()
+	}
 	os.Exit(code)
 }

+// requiresDB skips the test if the database is not available (e.g. running
+// tests standalone without test DB setup). DB-backed JTI revocation tests
+// need a connection to the revoked_jtis table.
+func requiresDB(t *testing.T) {
+	t.Helper()
+	if db.DB == nil {
+		t.Skip("skipping: no database connection")
+	}
+}
+
 // =============================================================================
 // generateJTI Tests
 // =============================================================================
@@ -158,6 +185,7 @@ func TestVerifyToken_ReturnsJTI(t *testing.T) {
 // TestVerifyToken_RevokedJTI creates a token, revokes its JTI, and verifies
 // that VerifyToken returns an error containing "token revoked".
 func TestVerifyToken_RevokedJTI(t *testing.T) {
+	requiresDB(t)
 	token, jti, err := GenerateToken("user-004", "verified_email")
 	if err != nil {
 		t.Fatalf("GenerateToken() failed: %v", err)
@@ -203,6 +231,7 @@ func TestVerifyToken_MissingJTI(t *testing.T) {
 // TestRevokeJTI_AddsToSet verifies that calling RevokeJTI adds the JTI to the
 // revoked set, and IsJTIRevoked returns true for it.
 func TestRevokeJTI_AddsToSet(t *testing.T) {
+	requiresDB(t)
 	_, jti, err := GenerateToken("user-006", "verified_email")
 	if err != nil {
 		t.Fatalf("GenerateToken() failed: %v", err)
@@ -234,16 +263,24 @@ func TestIsJTIRevoked_NonExistent(t *testing.T) {
 // TestCleanupRevokedJTIs_RemovesExpired adds a JTI with a past expiry time,
 // runs CleanupRevokedJTIs, and verifies the JTI is removed from the set.
 func TestCleanupRevokedJTIs_RemovesExpired(t *testing.T) {
+	requiresDB(t)
 	_, jti, err := GenerateToken("user-007", "verified_email")
 	if err != nil {
 		t.Fatalf("GenerateToken() failed: %v", err)
 	}

-	// Add with past expiry (1 hour ago)
-	RevokeJTI(jti, time.Now().Add(-1*time.Hour))
+	// Add with future expiry so IsJTIRevoked sees it
+	RevokeJTI(jti, time.Now().Add(1*time.Hour))

 	if !IsJTIRevoked(jti) {
-		t.Fatal("JTI should be in revoked set before cleanup")
+		t.Fatal("JTI should be in revoked set after RevokeJTI")
+	}
+
+	// Directly update the DB to set expiry in the past
+	_, err = db.DB.Exec(context.Background(),
+		"UPDATE revoked_jtis SET expires_at = NOW() - INTERVAL '1 hour' WHERE jti = $1", jti)
+	if err != nil {
+		t.Fatalf("failed to expire JTI: %v", err)
 	}

 	CleanupRevokedJTIs()
@@ -256,6 +293,7 @@ func TestCleanupRevokedJTIs_RemovesExpired(t *testing.T) {
 // TestCleanupRevokedJTIs_KeepsValid adds a JTI with a future expiry time,
 // runs CleanupRevokedJTIs, and verifies the JTI is still in the set.
 func TestCleanupRevokedJTIs_KeepsValid(t *testing.T) {
+	requiresDB(t)
 	_, jti, err := GenerateToken("user-008", "verified_email")
 	if err != nil {
 		t.Fatalf("GenerateToken() failed: %v", err)