package repository import ( "context" "crypto/sha256" "encoding/hex" "fmt" "time" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" "gogs.fxtmmsk.ru/foxtime/photoplaces/backend/internal/models" ) type RefreshTokenRepo struct { pool *pgxpool.Pool } func NewRefreshTokenRepo(pool *pgxpool.Pool) *RefreshTokenRepo { return &RefreshTokenRepo{pool: pool} } func hashToken(token string) string { sum := sha256.Sum256([]byte(token)) return hex.EncodeToString(sum[:]) } func (r *RefreshTokenRepo) Create(ctx context.Context, userID, plainToken string, expiresAt time.Time) error { tokenHash := hashToken(plainToken) _, err := r.pool.Exec(ctx, `INSERT INTO refresh_tokens (user_id, token_hash, expires_at) VALUES ($1, $2, $3)`, userID, tokenHash, expiresAt, ) if err != nil { return fmt.Errorf("create refresh token: %w", err) } return nil } func (r *RefreshTokenRepo) GetValid(ctx context.Context, plainToken string) (*models.RefreshToken, error) { tokenHash := hashToken(plainToken) row := r.pool.QueryRow(ctx, `SELECT id, user_id, token_hash, expires_at, created_at, revoked_at, replaced_at FROM refresh_tokens WHERE token_hash = $1 AND expires_at > now() AND revoked_at IS NULL AND (replaced_at IS NULL OR replaced_at > now() - interval '30 seconds')`, tokenHash, ) var rt models.RefreshToken err := row.Scan(&rt.ID, &rt.UserID, &rt.TokenHash, &rt.ExpiresAt, &rt.CreatedAt, &rt.RevokedAt, &rt.ReplacedAt) if err != nil { if err == pgx.ErrNoRows { return nil, nil } return nil, fmt.Errorf("get refresh token: %w", err) } return &rt, nil } // GetValidAndLock atomically reads a valid token and locks the row FOR UPDATE. // Returns the token and a cleanup func that releases the lock on error. func (r *RefreshTokenRepo) GetValidAndLock(ctx context.Context, plainToken string) (*models.RefreshToken, func(), error) { tx, err := r.pool.Begin(ctx) if err != nil { return nil, nil, fmt.Errorf("begin tx: %w", err) } tokenHash := hashToken(plainToken) row := tx.QueryRow(ctx, `SELECT id, user_id, token_hash, expires_at, created_at, revoked_at, replaced_at FROM refresh_tokens WHERE token_hash = $1 AND expires_at > now() AND revoked_at IS NULL AND (replaced_at IS NULL OR replaced_at > now() - interval '30 seconds') FOR UPDATE`, tokenHash, ) var rt models.RefreshToken err = row.Scan(&rt.ID, &rt.UserID, &rt.TokenHash, &rt.ExpiresAt, &rt.CreatedAt, &rt.RevokedAt, &rt.ReplacedAt) if err != nil { tx.Rollback(ctx) if err == pgx.ErrNoRows { return nil, nil, nil } return nil, nil, fmt.Errorf("get and lock refresh token: %w", err) } cleanup := func() { tx.Rollback(ctx) } return &rt, cleanup, nil } // RotateToken atomically replaces the old refresh token and creates a new one // within a transaction. Uses soft-replace (sets replaced_at) to allow a grace // period where the old token is still accepted. func (r *RefreshTokenRepo) RotateToken(ctx context.Context, oldPlainToken, newPlainToken, userID string, expiresAt time.Time) error { tx, err := r.pool.Begin(ctx) if err != nil { return fmt.Errorf("begin tx: %w", err) } defer tx.Rollback(ctx) oldHash := hashToken(oldPlainToken) tag, err := tx.Exec(ctx, `UPDATE refresh_tokens SET replaced_at = now() WHERE token_hash = $1 AND replaced_at IS NULL`, oldHash, ) if err != nil { return fmt.Errorf("replace old token: %w", err) } if tag.RowsAffected() == 0 { return fmt.Errorf("old token not found or already replaced") } newHash := hashToken(newPlainToken) _, err = tx.Exec(ctx, `INSERT INTO refresh_tokens (user_id, token_hash, expires_at) VALUES ($1, $2, $3)`, userID, newHash, expiresAt, ) if err != nil { return fmt.Errorf("insert new token: %w", err) } return tx.Commit(ctx) } func (r *RefreshTokenRepo) GetRevoked(ctx context.Context, plainToken string) (*models.RefreshToken, error) { tokenHash := hashToken(plainToken) row := r.pool.QueryRow(ctx, `SELECT id, user_id, token_hash, expires_at, created_at, revoked_at, replaced_at FROM refresh_tokens WHERE token_hash = $1 AND revoked_at IS NOT NULL`, tokenHash, ) var rt models.RefreshToken err := row.Scan(&rt.ID, &rt.UserID, &rt.TokenHash, &rt.ExpiresAt, &rt.CreatedAt, &rt.RevokedAt, &rt.ReplacedAt) if err != nil { if err == pgx.ErrNoRows { return nil, nil } return nil, fmt.Errorf("get revoked refresh token: %w", err) } return &rt, nil } func (r *RefreshTokenRepo) Revoke(ctx context.Context, tokenHash string) error { _, err := r.pool.Exec(ctx, `UPDATE refresh_tokens SET revoked_at = now() WHERE token_hash = $1 AND revoked_at IS NULL`, tokenHash, ) return err } func (r *RefreshTokenRepo) RevokeAllForUser(ctx context.Context, userID string) error { _, err := r.pool.Exec(ctx, `UPDATE refresh_tokens SET revoked_at = now() WHERE user_id = $1 AND revoked_at IS NULL`, userID, ) return err } // ReplaceIfExists soft-replaces a token (sets replaced_at) and returns true if a row was updated. // Returns false if the token was already replaced by another request. func (r *RefreshTokenRepo) ReplaceIfExists(ctx context.Context, tokenHash string) (bool, error) { tag, err := r.pool.Exec(ctx, `UPDATE refresh_tokens SET replaced_at = now() WHERE token_hash = $1 AND replaced_at IS NULL`, tokenHash, ) if err != nil { return false, fmt.Errorf("replace if exists: %w", err) } return tag.RowsAffected() > 0, nil } func (r *RefreshTokenRepo) CleanupExpired(ctx context.Context) error { _, err := r.pool.Exec(ctx, `DELETE FROM refresh_tokens WHERE expires_at < now() - interval '1 day' OR (replaced_at IS NOT NULL AND replaced_at < now() - interval '1 hour')`, ) return err }