refactor(auth): wrap JWT operations in transactions and migrate to clock.Now()
Wrap JTI revocation, cleanup, refresh token generation, and verification in explicit DB transactions with Begin/defer Rollback/Commit. Replace time.Now() with clock.Now() for testability. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
This commit is contained in:
+59
-6
@@ -5,9 +5,11 @@ import (
|
|||||||
"crypto/rand"
|
"crypto/rand"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"log"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"crussell/db"
|
"crussell/db"
|
||||||
|
"crussell/clock"
|
||||||
"github.com/jackc/pgx/v5"
|
"github.com/jackc/pgx/v5"
|
||||||
|
|
||||||
"github.com/go-chi/jwtauth/v5"
|
"github.com/go-chi/jwtauth/v5"
|
||||||
@@ -39,13 +41,25 @@ func RevokeJTI(ctx context.Context, jti string, expiresAt time.Time) {
|
|||||||
if db.Conn == nil {
|
if db.Conn == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
_, err := db.Conn.Exec(ctx,
|
tx, err := db.Conn.Begin(ctx)
|
||||||
|
if err != nil {
|
||||||
|
fmt.Printf("WARN: Failed to begin transaction for JTI revocation: %v\n", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer tx.Rollback(ctx)
|
||||||
|
|
||||||
|
_, err = tx.Exec(ctx,
|
||||||
`INSERT INTO revoked_jtis (jti, expires_at) VALUES ($1, $2)
|
`INSERT INTO revoked_jtis (jti, expires_at) VALUES ($1, $2)
|
||||||
ON CONFLICT (jti) DO NOTHING`,
|
ON CONFLICT (jti) DO NOTHING`,
|
||||||
jti, expiresAt)
|
jti, expiresAt)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
// Log but don't fail - this is best effort
|
// Log but don't fail - this is best effort
|
||||||
fmt.Printf("WARN: Failed to revoke JTI %s: %v\n", jti, err)
|
fmt.Printf("WARN: Failed to revoke JTI %s: %v\n", jti, err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := tx.Commit(ctx); err != nil {
|
||||||
|
fmt.Printf("WARN: Failed to commit transaction for JTI revocation: %v\n", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -72,10 +86,22 @@ func CleanupRevokedJTIs(ctx context.Context) {
|
|||||||
if db.Conn == nil {
|
if db.Conn == nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
_, err := db.Conn.Exec(ctx,
|
tx, err := db.Conn.Begin(ctx)
|
||||||
|
if err != nil {
|
||||||
|
fmt.Printf("WARN: Failed to begin transaction for JTI cleanup: %v\n", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer tx.Rollback(ctx)
|
||||||
|
|
||||||
|
_, err = tx.Exec(ctx,
|
||||||
`DELETE FROM revoked_jtis WHERE expires_at < NOW()`)
|
`DELETE FROM revoked_jtis WHERE expires_at < NOW()`)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
fmt.Printf("WARN: Failed to cleanup revoked JTIs: %v\n", err)
|
fmt.Printf("WARN: Failed to cleanup revoked JTIs: %v\n", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := tx.Commit(ctx); err != nil {
|
||||||
|
fmt.Printf("WARN: Failed to commit transaction for JTI cleanup: %v\n", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -85,7 +111,14 @@ func StartJTICleanup() {
|
|||||||
ticker := time.NewTicker(30 * time.Minute)
|
ticker := time.NewTicker(30 * time.Minute)
|
||||||
defer ticker.Stop()
|
defer ticker.Stop()
|
||||||
for range ticker.C {
|
for range ticker.C {
|
||||||
CleanupRevokedJTIs(context.Background())
|
func() {
|
||||||
|
defer func() {
|
||||||
|
if r := recover(); r != nil {
|
||||||
|
log.Printf("Panic recovered in JWT cleanup ticker: %v", r)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
CleanupRevokedJTIs(context.Background())
|
||||||
|
}()
|
||||||
}
|
}
|
||||||
}()
|
}()
|
||||||
}
|
}
|
||||||
@@ -106,7 +139,7 @@ func GenerateToken(userID string, role string) (string, string, error) {
|
|||||||
"user_id": userID,
|
"user_id": userID,
|
||||||
"role": role,
|
"role": role,
|
||||||
"jti": jti,
|
"jti": jti,
|
||||||
"exp": time.Now().Add(1 * time.Hour).Unix(), // 1 hour
|
"exp": clock.Now().Add(1 * time.Hour).Unix(), // 1 hour
|
||||||
})
|
})
|
||||||
return tokenString, jti, err
|
return tokenString, jti, err
|
||||||
}
|
}
|
||||||
@@ -175,12 +208,22 @@ func GenerateRefreshToken(ctx context.Context, userID string, role string) (stri
|
|||||||
VALUES ($1, encode(sha256($2::bytea), 'hex'), $3, NOW() + INTERVAL '90 days')
|
VALUES ($1, encode(sha256($2::bytea), 'hex'), $3, NOW() + INTERVAL '90 days')
|
||||||
RETURNING id`
|
RETURNING id`
|
||||||
|
|
||||||
|
tx, err := db.Conn.Begin(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("failed to begin transaction: %w", err)
|
||||||
|
}
|
||||||
|
defer tx.Rollback(ctx)
|
||||||
|
|
||||||
var tokenID int64
|
var tokenID int64
|
||||||
err = db.Conn.QueryRow(ctx, query, userID, token, role).Scan(&tokenID)
|
err = tx.QueryRow(ctx, query, userID, token, role).Scan(&tokenID)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return "", fmt.Errorf("failed to store refresh token: %w", err)
|
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
|
return token, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -194,7 +237,13 @@ func VerifyRefreshToken(ctx context.Context, tokenString string) (userID string,
|
|||||||
AND NOT revoked
|
AND NOT revoked
|
||||||
RETURNING user_id, role`
|
RETURNING user_id, role`
|
||||||
|
|
||||||
err = db.Conn.QueryRow(ctx, query, tokenString).Scan(&userID, &role)
|
tx, err := db.Conn.Begin(ctx)
|
||||||
|
if err != nil {
|
||||||
|
return "", "", fmt.Errorf("failed to begin transaction: %w", err)
|
||||||
|
}
|
||||||
|
defer tx.Rollback(ctx)
|
||||||
|
|
||||||
|
err = tx.QueryRow(ctx, query, tokenString).Scan(&userID, &role)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if errors.Is(err, pgx.ErrNoRows) {
|
if errors.Is(err, pgx.ErrNoRows) {
|
||||||
return "", "", fmt.Errorf("invalid or expired refresh token")
|
return "", "", fmt.Errorf("invalid or expired refresh token")
|
||||||
@@ -202,6 +251,10 @@ func VerifyRefreshToken(ctx context.Context, tokenString string) (userID string,
|
|||||||
return "", "", fmt.Errorf("failed to verify refresh token: %w", err)
|
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
|
// Token was consumed (DELETE returned it) — this is rotation
|
||||||
// If a token is used twice, the second DELETE returns no rows = invalid
|
// If a token is used twice, the second DELETE returns no rows = invalid
|
||||||
return userID, role, nil
|
return userID, role, nil
|
||||||
|
|||||||
@@ -18,6 +18,7 @@ import (
|
|||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
|
"crussell/clock"
|
||||||
"crussell/testutils/testtx"
|
"crussell/testutils/testtx"
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -160,7 +161,7 @@ func TestVerifyToken_RevokedJTI(t *testing.T) {
|
|||||||
t.Fatalf("GenerateToken() failed: %v", err)
|
t.Fatalf("GenerateToken() failed: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
RevokeJTI(ctx, jti, time.Now().Add(30*24*time.Hour))
|
RevokeJTI(ctx, jti, clock.Now().Add(30*24*time.Hour))
|
||||||
|
|
||||||
_, _, _, err = VerifyToken(token, ctx)
|
_, _, _, err = VerifyToken(token, ctx)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
@@ -178,7 +179,7 @@ func TestVerifyToken_MissingJTI(t *testing.T) {
|
|||||||
_, tokenString, err := TokenAuth.Encode(map[string]interface{}{
|
_, tokenString, err := TokenAuth.Encode(map[string]interface{}{
|
||||||
"user_id": "user-005",
|
"user_id": "user-005",
|
||||||
"role": "verified_email",
|
"role": "verified_email",
|
||||||
"exp": time.Now().Add(30 * 24 * time.Hour).Unix(),
|
"exp": clock.Now().Add(30 * 24 * time.Hour).Unix(),
|
||||||
})
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("failed to create token without JTI: %v", err)
|
t.Fatalf("failed to create token without JTI: %v", err)
|
||||||
@@ -210,7 +211,7 @@ func TestRevokeJTI_AddsToSet(t *testing.T) {
|
|||||||
t.Fatal("JTI should not be revoked before calling RevokeJTI")
|
t.Fatal("JTI should not be revoked before calling RevokeJTI")
|
||||||
}
|
}
|
||||||
|
|
||||||
RevokeJTI(ctx, jti, time.Now().Add(30*24*time.Hour))
|
RevokeJTI(ctx, jti, clock.Now().Add(30*24*time.Hour))
|
||||||
|
|
||||||
if !IsJTIRevoked(ctx, jti) {
|
if !IsJTIRevoked(ctx, jti) {
|
||||||
t.Error("expected IsJTIRevoked to return true after RevokeJTI")
|
t.Error("expected IsJTIRevoked to return true after RevokeJTI")
|
||||||
@@ -240,7 +241,7 @@ func TestCleanupRevokedJTIs_RemovesExpired(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Add with future expiry so IsJTIRevoked sees it
|
// Add with future expiry so IsJTIRevoked sees it
|
||||||
RevokeJTI(ctx, jti, time.Now().Add(1*time.Hour))
|
RevokeJTI(ctx, jti, clock.Now().Add(1*time.Hour))
|
||||||
|
|
||||||
if !IsJTIRevoked(ctx, jti) {
|
if !IsJTIRevoked(ctx, jti) {
|
||||||
t.Fatal("JTI should be in revoked set after RevokeJTI")
|
t.Fatal("JTI should be in revoked set after RevokeJTI")
|
||||||
@@ -270,7 +271,7 @@ func TestCleanupRevokedJTIs_KeepsValid(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
// Add with future expiry
|
// Add with future expiry
|
||||||
RevokeJTI(ctx, jti, time.Now().Add(30*24*time.Hour))
|
RevokeJTI(ctx, jti, clock.Now().Add(30*24*time.Hour))
|
||||||
|
|
||||||
if !IsJTIRevoked(ctx, jti) {
|
if !IsJTIRevoked(ctx, jti) {
|
||||||
t.Fatal("JTI should be in revoked set before cleanup")
|
t.Fatal("JTI should be in revoked set before cleanup")
|
||||||
|
|||||||
Reference in New Issue
Block a user