llmgw/internal/store/claims.go

package store

import (
	"context"
	"crypto/rand"
	"database/sql"
	"errors"
	"time"
)

// ClaimTTL is how long an unclaimed link stays usable.
//
// Short enough that a link forwarded to the wrong place is unlikely to still
// work when it gets there; long enough to survive a weekend.
const ClaimTTL = 72 * time.Hour

// ErrClaimed says a claim link is gone: used, expired, or never real. The three
// are not distinguished, because saying which would confirm that a token was
// once valid.
var ErrClaimed = errors.New("this link is no longer valid")

// Claim is a pending one-time delivery of a key.
type Claim struct {
	KeyID     string
	ExpiresAt time.Time
}

// NewClaimToken mints the token that goes in the link.
//
// The same 256 bits as an API key, for the same reason: it is a bearer
// credential for as long as it lives.
func NewClaimToken() (string, error) {
	raw := make([]byte, keyBytes)
	if _, err := rand.Read(raw); err != nil {
		return "", err
	}
	return encodeBase62(raw), nil
}

// CreateClaim records a sealed key for later collection.
//
// The secret is sealed with the token, which is not stored, so what lands in
// the database cannot be opened by anything that can read the database.
func (s *DB) CreateClaim(ctx context.Context, keyID, token, secret string, ttl time.Duration) error {
	sealed, err := seal(token, []byte(secret))
	if err != nil {
		return err
	}
	if ttl <= 0 {
		ttl = ClaimTTL
	}
	now := time.Now()
	_, err = s.db.ExecContext(ctx,
		`INSERT INTO key_claims (token_hash, key_id, sealed, created_at, expires_at)
		 VALUES (?, ?, ?, ?, ?)`,
		HashKey(token), keyID, sealed, millis(now), millis(now.Add(ttl)))
	return err
}

// PeekClaim reports whether a link is still good, without consuming it.
//
// The page behind the link is shown before it is used, because link previewers
// follow URLs: a GET that consumed the claim would mean the first chat client
// to unfurl it burns the key.
func (s *DB) PeekClaim(ctx context.Context, token string) (Claim, error) {
	row := s.db.QueryRowContext(ctx,
		`SELECT key_id, expires_at, claimed_at FROM key_claims WHERE token_hash = ?`,
		HashKey(token))

	var c Claim
	var expires int64
	var claimed sql.NullInt64
	if err := row.Scan(&c.KeyID, &expires, &claimed); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return Claim{}, ErrClaimed
		}
		return Claim{}, err
	}
	c.ExpiresAt = fromMillis(expires)
	if claimed.Valid || time.Now().After(c.ExpiresAt) {
		return Claim{}, ErrClaimed
	}
	return c, nil
}

// ConsumeClaim opens the sealed key and marks the link used.
//
// The update is the lock: exactly one caller can move claimed_at from NULL, so
// two people opening the link at the same moment produce one winner rather than
// two copies of the key.
func (s *DB) ConsumeClaim(ctx context.Context, token string) (string, error) {
	hash := HashKey(token)
	now := time.Now()

	q, tx, err := s.begin(ctx)
	if err != nil {
		return "", err
	}
	defer func() { _ = tx.Rollback() }()

	var sealed []byte
	var expires int64
	err = q.QueryRowContext(ctx,
		`SELECT sealed, expires_at FROM key_claims WHERE token_hash = ? AND claimed_at IS NULL`,
		hash).Scan(&sealed, &expires)
	if errors.Is(err, sql.ErrNoRows) {
		return "", ErrClaimed
	}
	if err != nil {
		return "", err
	}
	if now.After(fromMillis(expires)) {
		return "", ErrClaimed
	}

	secret, err := open(token, sealed)
	if err != nil {
		return "", ErrClaimed
	}

	// Blank the sealed value as well as marking it claimed, so a used link
	// leaves nothing behind even before the row is purged.
	res, err := q.ExecContext(ctx,
		`UPDATE key_claims SET claimed_at = ?, sealed = x'' WHERE token_hash = ? AND claimed_at IS NULL`,
		millis(now), hash)
	if err != nil {
		return "", err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		return "", ErrClaimed
	}
	if err := tx.Commit(); err != nil {
		return "", err
	}
	return string(secret), nil
}

// PurgeClaims deletes rows that can no longer be used.
//
// Housekeeping rather than security — an expired row is already refused and its
// contents are already unopenable — but a table that only grows is its own
// problem.
func (s *DB) PurgeClaims(ctx context.Context, before time.Time) (int64, error) {
	res, err := s.db.ExecContext(ctx,
		`DELETE FROM key_claims WHERE expires_at < ? OR claimed_at IS NOT NULL`, millis(before))
	if err != nil {
		return 0, err
	}
	n, _ := res.RowsAffected()
	return n, nil
}