photoplaces/backend/internal/repository/refresh_tokens.go

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
}

func (r *RefreshTokenRepo) Delete(ctx context.Context, tokenHash string) error {
	_, err := r.pool.Exec(ctx,
		`DELETE FROM refresh_tokens WHERE token_hash = $1`,
		tokenHash,
	)
	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
}