Cyrene_v2/internal/auth/refresh.go

package auth

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"time"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/logger"
)

// Opaque refresh-token parameters and the canonical schema. The refresh secret
// is 32 random bytes (256 bits), so a plain unsalted SHA-256 digest is a sound
// one-way store: there is no low-entropy input to brute-force, and the digest is
// looked up through a unique index, never compared byte-by-byte in Go.
const (
	refreshTokenBytes = 32
	refreshIDBytes    = 16
	refreshTimeLayout = time.RFC3339Nano

	createRefreshTableSQL = `
CREATE TABLE IF NOT EXISTS refresh_tokens (
    id         TEXT PRIMARY KEY,
    user_id    TEXT NOT NULL,
    role       TEXT NOT NULL,
    token_hash TEXT NOT NULL UNIQUE,
    expires_at TEXT NOT NULL,
    revoked_at TEXT,
    created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
    updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
)`
	createRefreshUserIndexSQL = `CREATE INDEX IF NOT EXISTS idx_refresh_tokens_user_id
    ON refresh_tokens (user_id)`
	createRefreshExpiryIndexSQL = `CREATE INDEX IF NOT EXISTS idx_refresh_tokens_expires_at
    ON refresh_tokens (expires_at)`
	insertRefreshSQL = `INSERT INTO refresh_tokens (id, user_id, role, token_hash, expires_at)
    VALUES (?, ?, ?, ?, ?)`
	selectRefreshByHashSQL = `SELECT id, user_id, role, expires_at, revoked_at
    FROM refresh_tokens WHERE token_hash = ?`
	revokeRefreshByIDSQL = `UPDATE refresh_tokens SET revoked_at = ?, updated_at = ?
    WHERE id = ? AND revoked_at IS NULL`
	revokeRefreshByHashSQL = `UPDATE refresh_tokens SET revoked_at = ?, updated_at = ?
    WHERE token_hash = ? AND revoked_at IS NULL`
	deleteRefreshAllSQL    = `DELETE FROM refresh_tokens`
	deleteRefreshByUserSQL = `DELETE FROM refresh_tokens WHERE user_id = ?`
	pruneRefreshSQL        = `DELETE FROM refresh_tokens WHERE expires_at < ? OR (revoked_at IS NOT NULL AND revoked_at < ?)`
)

// RefreshToken is the stored metadata of a refresh grant. It never carries the
// secret: only the SHA-256 hash of the opaque token is persisted.
type RefreshToken struct {
	ID        string
	UserID    string
	Role      Role
	ExpiresAt time.Time
}

// RotateOptions configures a refresh call. Rotate enables single-use rotation:
// the presented token is revoked and replaced by a fresh one. When false the
// presented token stays valid and no replacement is minted (NewToken is empty),
// which is the opt-out for clients that cannot persist a rotated token.
type RotateOptions struct {
	TTL    time.Duration
	Rotate bool
}

// RefreshStore persists opaque refresh tokens in SQLite. Only the SHA-256 hash
// of a token is stored; the plaintext is returned once by Issue/Rotate and never
// written to the database or a log. It is safe for concurrent use by virtue of
// the underlying *sql.DB.
type RefreshStore struct {
	db  *sql.DB
	log *slog.Logger
}

// NewRefreshStore builds a SQLite-backed refresh-token store over db, ensuring
// the refresh_tokens table and indexes exist. The DDL mirrors migration 0008
// and is idempotent, so the store is self-contained for tests while the
// migration remains the authoritative path for an existing database. db must be
// live; nil returns ErrNilDB. log may be nil.
func NewRefreshStore(ctx context.Context, db *sql.DB, log *slog.Logger) (*RefreshStore, error) {
	if db == nil {
		return nil, ErrNilDB
	}
	for _, ddl := range []string{createRefreshTableSQL, createRefreshUserIndexSQL, createRefreshExpiryIndexSQL} {
		if _, err := db.ExecContext(ctx, ddl); err != nil {
			return nil, fmt.Errorf("auth: ensure refresh_tokens schema: %w", err)
		}
	}
	if log == nil {
		log = slog.Default()
	}
	return &RefreshStore{db: db, log: logger.WithSource(log, source)}, nil
}

// Issue mints a new opaque refresh token for userID/role, stores only its hash,
// and returns the plaintext (visible exactly once) plus its expiry. Expired or
// long-revoked rows are pruned best-effort so the table stays bounded.
func (s *RefreshStore) Issue(ctx context.Context, userID string, role Role, ttl time.Duration) (string, time.Time, error) {
	if strings.TrimSpace(userID) == "" {
		return "", time.Time{}, fmt.Errorf("%w: empty user id", ErrInvalidGrant)
	}
	if !role.Valid() {
		return "", time.Time{}, fmt.Errorf("%w: %q", ErrInvalidRole, role)
	}
	if ttl <= 0 {
		return "", time.Time{}, fmt.Errorf("%w: non-positive ttl", ErrInvalidGrant)
	}
	plaintext, hash, err := newRefreshToken()
	if err != nil {
		return "", time.Time{}, err
	}
	id, err := newRefreshID()
	if err != nil {
		return "", time.Time{}, err
	}
	expires := time.Now().UTC().Add(ttl)
	if _, err := s.db.ExecContext(ctx, insertRefreshSQL, id, userID, string(role), hash, formatRefreshTime(expires)); err != nil {
		return "", time.Time{}, fmt.Errorf("auth: insert refresh token: %w", err)
	}
	s.pruneExpired(ctx)
	return plaintext, expires, nil
}

// Rotate validates plaintext (known, unrevoked, unexpired) and, when opts.Rotate
// is true, atomically revokes it and issues a replacement. It returns the grant
// metadata and the new plaintext token (empty when rotation is disabled). An
// unknown, revoked, or expired token yields ErrInvalidGrant. The transaction
// makes rotation single-use under concurrent replay: the second caller's revoke
// affects no row and fails closed.
func (s *RefreshStore) Rotate(ctx context.Context, plaintext string, opts RotateOptions) (RefreshToken, string, error) {
	if strings.TrimSpace(plaintext) == "" {
		return RefreshToken{}, "", fmt.Errorf("%w: empty token", ErrInvalidGrant)
	}
	if opts.Rotate && opts.TTL <= 0 {
		return RefreshToken{}, "", fmt.Errorf("%w: non-positive ttl", ErrInvalidGrant)
	}
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return RefreshToken{}, "", fmt.Errorf("auth: begin refresh rotate: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	grant, revokedAt, err := selectRefreshGrant(ctx, tx, hashRefreshToken(plaintext))
	if err != nil {
		return RefreshToken{}, "", err
	}
	if revokedAt.Valid {
		return RefreshToken{}, "", fmt.Errorf("%w: revoked", ErrInvalidGrant)
	}
	now := time.Now().UTC()
	if !grant.ExpiresAt.After(now) {
		return RefreshToken{}, "", fmt.Errorf("%w: expired", ErrInvalidGrant)
	}
	if !opts.Rotate {
		if err := tx.Commit(); err != nil {
			return RefreshToken{}, "", fmt.Errorf("auth: commit refresh: %w", err)
		}
		return grant, "", nil
	}
	if err := revokeRefreshRow(ctx, tx, grant.ID, now); err != nil {
		return RefreshToken{}, "", err
	}
	newPlaintext, newHash, err := newRefreshToken()
	if err != nil {
		return RefreshToken{}, "", err
	}
	newID, err := newRefreshID()
	if err != nil {
		return RefreshToken{}, "", err
	}
	newExpires := now.Add(opts.TTL)
	if _, err := tx.ExecContext(ctx, insertRefreshSQL, newID, grant.UserID, string(grant.Role), newHash, formatRefreshTime(newExpires)); err != nil {
		return RefreshToken{}, "", fmt.Errorf("auth: insert rotated refresh token: %w", err)
	}
	if err := tx.Commit(); err != nil {
		return RefreshToken{}, "", fmt.Errorf("auth: commit refresh rotate: %w", err)
	}
	return RefreshToken{ID: newID, UserID: grant.UserID, Role: grant.Role, ExpiresAt: newExpires}, newPlaintext, nil
}

// Revoke marks the refresh token matching plaintext as revoked, returning true
// when a live token was found. It is idempotent: revoking an unknown or already
// revoked token is a no-op returning false.
func (s *RefreshStore) Revoke(ctx context.Context, plaintext string) (bool, error) {
	if strings.TrimSpace(plaintext) == "" {
		return false, nil
	}
	now := formatRefreshTime(time.Now().UTC())
	res, err := s.db.ExecContext(ctx, revokeRefreshByHashSQL, now, now, hashRefreshToken(plaintext))
	if err != nil {
		return false, fmt.Errorf("auth: revoke refresh token: %w", err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return false, fmt.Errorf("auth: revoke refresh token: %w", err)
	}
	return n > 0, nil
}

// RevokeAll deletes every refresh token and returns how many rows were removed.
// It is the global invalidation used when the JWT signing secret is rotated:
// credential rotation must invalidate every long-lived grant, not only the
// short-lived access tokens. Deleting (rather than marking revoked) is exact and
// leaves no valid row behind.
func (s *RefreshStore) RevokeAll(ctx context.Context) (int64, error) {
	res, err := s.db.ExecContext(ctx, deleteRefreshAllSQL)
	if err != nil {
		return 0, fmt.Errorf("auth: revoke all refresh tokens: %w", err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("auth: revoke all refresh tokens: %w", err)
	}
	return n, nil
}

// RevokeByUser deletes every refresh token belonging to userID and returns how
// many rows were removed. It backs logout by access token, where the individual
// refresh token is unknown: the account is signed out on every device.
func (s *RefreshStore) RevokeByUser(ctx context.Context, userID string) (int64, error) {
	if strings.TrimSpace(userID) == "" {
		return 0, fmt.Errorf("%w: empty user id", ErrInvalidGrant)
	}
	res, err := s.db.ExecContext(ctx, deleteRefreshByUserSQL, userID)
	if err != nil {
		return 0, fmt.Errorf("auth: revoke user refresh tokens: %w", err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return 0, fmt.Errorf("auth: revoke user refresh tokens: %w", err)
	}
	return n, nil
}

// pruneExpired deletes rows that are expired or were revoked long ago. It is
// best-effort: a failure is logged and never fails the Issue that triggered it.
func (s *RefreshStore) pruneExpired(ctx context.Context) {
	cutoff := formatRefreshTime(time.Now().UTC())
	if _, err := s.db.ExecContext(ctx, pruneRefreshSQL, cutoff, cutoff); err != nil {
		s.log.Warn("refresh token prune failed", "err", err)
	}
}

// selectRefreshGrant loads the grant for hash, returning ErrInvalidGrant when no
// row matches. The query runs on tx so the caller's rotate transaction observes
// a stable row.
func selectRefreshGrant(ctx context.Context, tx *sql.Tx, hash string) (RefreshToken, sql.NullString, error) {
	var (
		grant     RefreshToken
		role      string
		expiresAt string
		revokedAt sql.NullString
	)
	err := tx.QueryRowContext(ctx, selectRefreshByHashSQL, hash).
		Scan(&grant.ID, &grant.UserID, &role, &expiresAt, &revokedAt)
	if errors.Is(err, sql.ErrNoRows) {
		return RefreshToken{}, sql.NullString{}, fmt.Errorf("%w: unknown token", ErrInvalidGrant)
	}
	if err != nil {
		return RefreshToken{}, sql.NullString{}, fmt.Errorf("auth: lookup refresh token: %w", err)
	}
	parsedRole, err := ParseRole(role)
	if err != nil {
		return RefreshToken{}, sql.NullString{}, fmt.Errorf("auth: stored refresh role: %w", err)
	}
	expiry, err := time.Parse(refreshTimeLayout, expiresAt)
	if err != nil {
		return RefreshToken{}, sql.NullString{}, fmt.Errorf("auth: parse refresh expiry: %w", err)
	}
	grant.Role = parsedRole
	grant.ExpiresAt = expiry
	return grant, revokedAt, nil
}

// revokeRefreshRow marks the grant row revoked inside tx, failing closed with
// ErrInvalidGrant when no live row was updated (a concurrent rotate won the race).
func revokeRefreshRow(ctx context.Context, tx *sql.Tx, id string, now time.Time) error {
	stamp := formatRefreshTime(now)
	res, err := tx.ExecContext(ctx, revokeRefreshByIDSQL, stamp, stamp, id)
	if err != nil {
		return fmt.Errorf("auth: revoke rotated refresh token: %w", err)
	}
	n, err := res.RowsAffected()
	if err != nil {
		return fmt.Errorf("auth: revoke rotated refresh token: %w", err)
	}
	if n == 0 {
		return fmt.Errorf("%w: already used", ErrInvalidGrant)
	}
	return nil
}

// newRefreshToken returns a fresh opaque token and its stored hash.
func newRefreshToken() (string, string, error) {
	raw := make([]byte, refreshTokenBytes)
	if _, err := rand.Read(raw); err != nil {
		return "", "", fmt.Errorf("auth: generate refresh token: %w", err)
	}
	token := base64.RawURLEncoding.EncodeToString(raw)
	return token, hashRefreshToken(token), nil
}

// newRefreshID returns a random 128-bit hex row id.
func newRefreshID() (string, error) {
	raw := make([]byte, refreshIDBytes)
	if _, err := rand.Read(raw); err != nil {
		return "", fmt.Errorf("auth: generate refresh id: %w", err)
	}
	return hex.EncodeToString(raw), nil
}

// hashRefreshToken returns the hex SHA-256 of token, the only form persisted.
func hashRefreshToken(token string) string {
	sum := sha256.Sum256([]byte(token))
	return hex.EncodeToString(sum[:])
}

// formatRefreshTime renders t as UTC RFC3339Nano for storage.
func formatRefreshTime(t time.Time) string {
	return t.UTC().Format(refreshTimeLayout)
}