photoplaces/backend/internal/repository/refresh_tokens.go

// Package repository
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
		 FROM refresh_tokens
		 WHERE token_hash = $1 AND expires_at > now() AND revoked_at IS NULL`,
		tokenHash,
	)

	var rt models.RefreshToken
	err := row.Scan(&rt.ID, &rt.UserID, &rt.TokenHash, &rt.ExpiresAt, &rt.CreatedAt, &rt.RevokedAt)
	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.
// Used to prevent race conditions on concurrent refresh token rotation.
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
		 FROM refresh_tokens
		 WHERE token_hash = $1 AND expires_at > now() AND revoked_at IS NULL
		 FOR UPDATE`,
		tokenHash,
	)

	var rt models.RefreshToken
	err = row.Scan(&rt.ID, &rt.UserID, &rt.TokenHash, &rt.ExpiresAt, &rt.CreatedAt, &rt.RevokedAt)
	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 deletes the old refresh token and creates a new one.
// Must be called after GetValidAndLock within the same request lifecycle.
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)
	_, err = tx.Exec(ctx,
		`DELETE FROM refresh_tokens WHERE token_hash = $1`,
		oldHash,
	)
	if err != nil {
		return fmt.Errorf("delete old token: %w", err)
	}

	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
		 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)
	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
}

// DeleteIfExists atomically deletes a token and returns true if a row was removed.
// Returns false if the token was already deleted by another request.
func (r *RefreshTokenRepo) DeleteIfExists(ctx context.Context, tokenHash string) (bool, error) {
	tag, err := r.pool.Exec(ctx,
		`DELETE FROM refresh_tokens WHERE token_hash = $1`,
		tokenHash,
	)
	if err != nil {
		return false, fmt.Errorf("delete 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'`,
	)
	return err
}