dpsapi/token_store.go

// SPDX-License-Identifier: GPL-2.0-or-later
// Copyright (C) 2026 Bryan Joshua Pedini

package main

import (
	"cmp"
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"slices"
	"time"

	sqlite3 "modernc.org/sqlite/lib"
)

// tokenBytes is the amount of random data in a token (256 bits)
const tokenBytes = 32

// tokenLength is the encoded token length: base64 RawURLEncoding of 32 bytes
const tokenLength = 43

// maxTokenAttempts bounds the retries on a primary-key collision
const maxTokenAttempts = 5

// Token store errors
var (
	ErrTokenNotFound  = errors.New("token not found")
	ErrTokenInactive  = errors.New("token already inactive")
	ErrTokenCollision = errors.New("could not generate a unique token")
)

// Token is one row of the tokens table
type Token struct {
	Token       string
	Description string
	Active      bool
	CreatedAt   time.Time
	DeletedAt   *time.Time
}

// GenerateToken returns a new random token from crypto/rand
func GenerateToken() (string, error) {
	return generateTokenFrom(rand.Reader)
}

// generateTokenFrom reads tokenBytes from r and encodes them with base64 RawURLEncoding
func generateTokenFrom(r io.Reader) (string, error) {
	buf := make([]byte, tokenBytes)
	if _, err := io.ReadFull(r, buf); err != nil {
		return "", fmt.Errorf("read random bytes: %w", err)
	}
	return base64.RawURLEncoding.EncodeToString(buf), nil
}

// TokenStore persists tokens in SQLite
type TokenStore struct {
	db          *sql.DB
	log         *slog.Logger
	now         func() time.Time
	generate    func() (string, error)
	maxAttempts int
}

// NewTokenStore creates a token store using the given clock and token generator
func NewTokenStore(db *sql.DB, log *slog.Logger, now func() time.Time, generate func() (string, error)) *TokenStore {
	return &TokenStore{
		db:          db,
		log:         log,
		now:         now,
		generate:    generate,
		maxAttempts: maxTokenAttempts,
	}
}

// CreateToken generates and stores a new active token with the given description
// On primary-key collision a new token is generated, up to maxAttempts times
func (s *TokenStore) CreateToken(ctx context.Context, description string) (Token, error) {
	defer traceFunc(ctx, s.log, "TokenStore.CreateToken")()

	for attempt := 1; attempt <= s.maxAttempts; attempt++ {
		// Generate a candidate token
		value, err := s.generate()
		if err != nil {
			return Token{}, fmt.Errorf("generate token: %w", err)
		}
		tok := Token{
			Token:       value,
			Description: description,
			Active:      true,
			CreatedAt:   s.now().UTC().Round(0),
		}
		createdAt := formatTime(tok.CreatedAt)

		// Insert it
		_, err = s.db.ExecContext(ctx,
			`INSERT INTO tokens (token, description, active, created_at, deleted_at) VALUES (?, ?, 1, ?, NULL)`,
			value, description, createdAt)
		if err == nil {
			trace(ctx, s.log, "row inserted", "table", "tokens", "token_id", tokenID(value),
				"description", description, "active", true, "created_at", createdAt, "deleted_at", nil)
			return tok, nil
		}
		if isSQLiteCode(err, sqlite3.SQLITE_CONSTRAINT_PRIMARYKEY) {
			s.log.DebugContext(ctx, "generated token collided with an existing token, retrying", "attempt", attempt)
			continue
		}
		return Token{}, fmt.Errorf("insert token: %w", err)
	}
	return Token{}, ErrTokenCollision
}

// rowScanner is implemented by *sql.Row and *sql.Rows
type rowScanner interface {
	Scan(dest ...any) error
}

// tokenColumns is the column list matching scanToken
const tokenColumns = `token, description, active, created_at, deleted_at`

// scanToken reads one tokens row
func scanToken(sc rowScanner) (Token, error) {
	var (
		tok       Token
		active    int64
		createdAt string
		deletedAt sql.NullString
	)
	if err := sc.Scan(&tok.Token, &tok.Description, &active, &createdAt, &deletedAt); err != nil {
		return Token{}, err
	}
	tok.Active = active == 1
	created, err := parseTime(createdAt)
	if err != nil {
		return Token{}, err
	}
	tok.CreatedAt = created
	if deletedAt.Valid {
		deleted, err := parseTime(deletedAt.String)
		if err != nil {
			return Token{}, err
		}
		tok.DeletedAt = &deleted
	}
	return tok, nil
}

// traceTokenRow emits the row-level TRACE record for a token row; the secret is replaced by token_id
func (s *TokenStore) traceTokenRow(ctx context.Context, msg string, tok Token) {
	var deletedAt any
	if tok.DeletedAt != nil {
		deletedAt = formatTime(*tok.DeletedAt)
	}
	trace(ctx, s.log, msg, "table", "tokens", "token_id", tokenID(tok.Token), "description", tok.Description,
		"active", tok.Active, "created_at", formatTime(tok.CreatedAt), "deleted_at", deletedAt)
}

// ListTokens returns the active tokens, or every token (active and inactive) when activeOnly is false,
// ordered by creation time, then by token
func (s *TokenStore) ListTokens(ctx context.Context, activeOnly bool) ([]Token, error) {
	defer traceFunc(ctx, s.log, "TokenStore.ListTokens")()

	query := `SELECT ` + tokenColumns + ` FROM tokens`
	if activeOnly {
		query += ` WHERE active = 1`
	}
	rows, err := s.db.QueryContext(ctx, query)
	if err != nil {
		return nil, fmt.Errorf("query tokens: %w", err)
	}
	defer rows.Close()

	tokens := []Token{}
	for rows.Next() {
		tok, err := scanToken(rows)
		if err != nil {
			return nil, fmt.Errorf("read token row: %w", err)
		}
		s.traceTokenRow(ctx, "row read", tok)
		tokens = append(tokens, tok)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate token rows: %w", err)
	}

	// Sort in Go: RFC3339Nano strings with variable fraction length do not sort lexically
	slices.SortFunc(tokens, func(a, b Token) int {
		if c := a.CreatedAt.Compare(b.CreatedAt); c != 0 {
			return c
		}
		return cmp.Compare(a.Token, b.Token)
	})
	return tokens, nil
}

// GetToken returns the token row matching value exactly, or ErrTokenNotFound
func (s *TokenStore) GetToken(ctx context.Context, value string) (Token, error) {
	defer traceFunc(ctx, s.log, "TokenStore.GetToken")()

	row := s.db.QueryRowContext(ctx, `SELECT `+tokenColumns+` FROM tokens WHERE token = ?`, value)
	tok, err := scanToken(row)
	if errors.Is(err, sql.ErrNoRows) {
		trace(ctx, s.log, "row not found", "table", "tokens", "token_id", tokenID(value))
		return Token{}, ErrTokenNotFound
	}
	if err != nil {
		return Token{}, fmt.Errorf("query token: %w", err)
	}
	s.traceTokenRow(ctx, "row read", tok)
	return tok, nil
}

// DeactivateToken sets active=0 and deleted_at=now in a single transaction
// It returns ErrTokenNotFound or ErrTokenInactive without modifying anything
func (s *TokenStore) DeactivateToken(ctx context.Context, value string) error {
	defer traceFunc(ctx, s.log, "TokenStore.DeactivateToken")()

	// Begin an immediate (write-locked) transaction
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin transaction: %w", err)
	}
	defer tx.Rollback()

	// Read the current state
	var active int64
	err = tx.QueryRowContext(ctx, `SELECT active FROM tokens WHERE token = ?`, value).Scan(&active)
	if errors.Is(err, sql.ErrNoRows) {
		trace(ctx, s.log, "row not found", "table", "tokens", "token_id", tokenID(value))
		return ErrTokenNotFound
	}
	if err != nil {
		return fmt.Errorf("query token: %w", err)
	}
	if active != 1 {
		trace(ctx, s.log, "row already inactive", "table", "tokens", "token_id", tokenID(value), "active", false)
		return ErrTokenInactive
	}

	// Flip the state
	deletedAt := formatTime(s.now())
	if _, err := tx.ExecContext(ctx,
		`UPDATE tokens SET active = 0, deleted_at = ? WHERE token = ? AND active = 1`,
		deletedAt, value); err != nil {
		return fmt.Errorf("update token: %w", err)
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit transaction: %w", err)
	}
	trace(ctx, s.log, "row updated", "table", "tokens", "token_id", tokenID(value),
		"active_from", true, "active_to", false, "deleted_at", deletedAt)
	return nil
}