at-container-registry/pkg/appview/db/device_store.go

package db

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

	"github.com/google/uuid"
	"golang.org/x/crypto/bcrypt"
)

// Device represents an authorized device
type Device struct {
	ID         string    `json:"id"`
	DID        string    `json:"did"`
	Handle     string    `json:"handle"`
	Name       string    `json:"name"`
	SecretHash string    `json:"secret_hash"`
	IPAddress  string    `json:"ip_address"`
	Location   string    `json:"location"`
	UserAgent  string    `json:"user_agent"`
	CreatedAt  time.Time `json:"created_at"`
	LastUsed   time.Time `json:"last_used"`
}

// PendingAuthorization represents a device awaiting user approval
type PendingAuthorization struct {
	DeviceCode   string     `json:"device_code"`
	UserCode     string     `json:"user_code"`
	DeviceName   string     `json:"device_name"`
	IPAddress    string     `json:"ip_address"`
	UserAgent    string     `json:"user_agent"`
	ExpiresAt    time.Time  `json:"expires_at"`
	ApprovedDID  *string    `json:"approved_did"`
	ApprovedAt   *time.Time `json:"approved_at"`
	DeviceSecret *string    `json:"device_secret"`
}

// DeviceStore manages devices and pending authorizations with SQLite persistence
type DeviceStore struct {
	db *sql.DB

	// lastUsedWrite records when each device's last_used column was last
	// written, so the timestamp is refreshed at most once per
	// deviceLastUsedInterval instead of on every authentication.
	lastUsedWrite sync.Map
}

// deviceLastUsedInterval is how often a device's last_used timestamp is actually
// written. It is shown in the UI as "last used" and read by the MAU queries,
// neither of which needs better than minute resolution.
//
// Without it, every /auth/token call writes — so every docker push and pull,
// including each layer's re-auth. Cheap against a local file; a network round
// trip against a remote primary, on the authentication path.
const deviceLastUsedInterval = 5 * time.Minute

// NewDeviceStore creates a new SQLite-backed device store
func NewDeviceStore(db *sql.DB) *DeviceStore {
	return &DeviceStore{db: db}
}

// CreatePendingAuth creates a new pending device authorization
func (s *DeviceStore) CreatePendingAuth(deviceName, ip, userAgent string) (*PendingAuthorization, error) {
	// Generate device code (long, random)
	deviceCodeBytes := make([]byte, 32)
	if _, err := rand.Read(deviceCodeBytes); err != nil {
		return nil, fmt.Errorf("failed to generate device code: %w", err)
	}
	deviceCode := base64.RawURLEncoding.EncodeToString(deviceCodeBytes)

	// Generate user code (short, human-readable)
	userCode := generateUserCode()

	expiresAt := time.Now().Add(10 * time.Minute)

	_, err := s.db.Exec(`
		INSERT INTO pending_device_auth (device_code, user_code, device_name, ip_address, user_agent, expires_at, created_at)
		VALUES (?, ?, ?, ?, ?, ?, datetime('now'))
	`, deviceCode, userCode, deviceName, ip, userAgent, expiresAt)

	if err != nil {
		return nil, fmt.Errorf("failed to create pending auth: %w", err)
	}

	pending := &PendingAuthorization{
		DeviceCode: deviceCode,
		UserCode:   userCode,
		DeviceName: deviceName,
		IPAddress:  ip,
		UserAgent:  userAgent,
		ExpiresAt:  expiresAt,
	}

	return pending, nil
}

// GetPendingByUserCode retrieves a pending auth by user code
func (s *DeviceStore) GetPendingByUserCode(userCode string) (*PendingAuthorization, bool) {
	var pending PendingAuthorization

	err := s.db.QueryRow(`
		SELECT device_code, user_code, device_name, ip_address, user_agent, expires_at, approved_did, approved_at, device_secret
		FROM pending_device_auth
		WHERE user_code = ?
	`, userCode).Scan(
		&pending.DeviceCode,
		&pending.UserCode,
		&pending.DeviceName,
		&pending.IPAddress,
		&pending.UserAgent,
		&pending.ExpiresAt,
		&pending.ApprovedDID,
		&pending.ApprovedAt,
		&pending.DeviceSecret,
	)

	if err == sql.ErrNoRows {
		return nil, false
	}
	if err != nil {
		slog.Warn("Failed to query pending auth", "component", "device_store", "error", err)
		return nil, false
	}

	// Check if expired
	if time.Now().After(pending.ExpiresAt) {
		return nil, false
	}

	return &pending, true
}

// GetPendingByDeviceCode retrieves a pending auth by device code
func (s *DeviceStore) GetPendingByDeviceCode(deviceCode string) (*PendingAuthorization, bool) {
	var pending PendingAuthorization

	err := s.db.QueryRow(`
		SELECT device_code, user_code, device_name, ip_address, user_agent, expires_at, approved_did, approved_at, device_secret
		FROM pending_device_auth
		WHERE device_code = ?
	`, deviceCode).Scan(
		&pending.DeviceCode,
		&pending.UserCode,
		&pending.DeviceName,
		&pending.IPAddress,
		&pending.UserAgent,
		&pending.ExpiresAt,
		&pending.ApprovedDID,
		&pending.ApprovedAt,
		&pending.DeviceSecret,
	)

	if err == sql.ErrNoRows {
		return nil, false
	}
	if err != nil {
		slog.Warn("Failed to query pending auth", "component", "device_store", "error", err)
		return nil, false
	}

	// Check if expired
	if time.Now().After(pending.ExpiresAt) {
		return nil, false
	}

	return &pending, true
}

// ApprovePending approves a pending authorization and generates device secret
func (s *DeviceStore) ApprovePending(userCode, did, handle string) (deviceSecret string, err error) {
	// Start transaction
	tx, err := s.db.Begin()
	if err != nil {
		return "", fmt.Errorf("failed to start transaction: %w", err)
	}
	defer tx.Rollback()

	// Get pending auth
	var pending PendingAuthorization
	err = tx.QueryRow(`
		SELECT device_code, user_code, device_name, ip_address, user_agent, expires_at, approved_did
		FROM pending_device_auth
		WHERE user_code = ?
	`, userCode).Scan(
		&pending.DeviceCode,
		&pending.UserCode,
		&pending.DeviceName,
		&pending.IPAddress,
		&pending.UserAgent,
		&pending.ExpiresAt,
		&pending.ApprovedDID,
	)

	if err == sql.ErrNoRows {
		return "", fmt.Errorf("pending authorization not found")
	}
	if err != nil {
		return "", fmt.Errorf("failed to query pending auth: %w", err)
	}

	// Check expiration
	if time.Now().After(pending.ExpiresAt) {
		return "", fmt.Errorf("authorization expired")
	}

	// Check if already approved
	if pending.ApprovedDID != nil && *pending.ApprovedDID != "" {
		return "", fmt.Errorf("already approved")
	}

	// Generate device secret
	secretBytes := make([]byte, 32)
	if _, err := rand.Read(secretBytes); err != nil {
		return "", fmt.Errorf("failed to generate device secret: %w", err)
	}
	deviceSecret = "atcr_device_" + base64.RawURLEncoding.EncodeToString(secretBytes)

	// Hash for storage
	secretHashBytes, err := bcrypt.GenerateFromPassword([]byte(deviceSecret), bcrypt.DefaultCost)
	if err != nil {
		return "", fmt.Errorf("failed to hash device secret: %w", err)
	}
	secretHash := string(secretHashBytes)

	// Create device record
	deviceID := uuid.New().String()
	now := time.Now()

	_, err = tx.Exec(`
		INSERT INTO devices (id, did, handle, name, secret_hash, secret_lookup, ip_address, user_agent, created_at)
		VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
	`, deviceID, did, handle, pending.DeviceName, secretHash, deviceSecretLookup(deviceSecret),
		pending.IPAddress, pending.UserAgent, now)

	if err != nil {
		return "", fmt.Errorf("failed to create device: %w", err)
	}

	// Update pending auth to mark as approved
	_, err = tx.Exec(`
		UPDATE pending_device_auth
		SET approved_did = ?, approved_at = ?, device_secret = ?
		WHERE user_code = ?
	`, did, now, deviceSecret, userCode)

	if err != nil {
		return "", fmt.Errorf("failed to update pending auth: %w", err)
	}

	// Commit transaction
	if err := tx.Commit(); err != nil {
		return "", fmt.Errorf("failed to commit transaction: %w", err)
	}

	return deviceSecret, nil
}

// deviceSecretLookup derives the stored verifier for a device secret.
//
// Plain SHA-256 is deliberate, and it is the verifier rather than merely an
// index. Device secrets are 32 bytes from crypto/rand (see the secret
// construction in ApprovePending), so recovering one from its digest means a
// SHA-256 preimage or a 2^256 search. bcrypt's work factor exists to make
// guessing expensive when the input space is small enough to enumerate, which
// a 256-bit random token is not — it buys nothing here and costs ~65ms per
// comparison. This is the same pattern used for API tokens and session
// cookies generally.
//
// The entropy in the generated secret is therefore load-bearing for the whole
// auth model. If that generation is ever weakened, this must change with it.
func deviceSecretLookup(secret string) string {
	sum := sha256.Sum256([]byte(secret))
	return hex.EncodeToString(sum[:])
}

// ValidateDeviceSecret validates a device secret and returns the device.
//
// Fast path: fetch the single row whose secret_lookup matches. That match is
// the authentication — see deviceSecretLookup.
//
// Legacy path: rows created before migration 0028 have no lookup value, and
// their plaintext is not recoverable from the bcrypt hash, so they are scanned
// and bcrypt-verified once, then backfilled. bcrypt exists here solely to carry
// those rows across; once every device has authenticated once, this branch and
// the bcrypt dependency can be deleted.
//
// The scan used to be unconditional, which made every authentication O(number
// of devices) in bcrypt comparisons. At cost 10 and 244 devices that was ~15.8s
// of CPU per /auth/token, past Docker's client deadline, and it grew with every
// device registered.
func (s *DeviceStore) ValidateDeviceSecret(secret string) (*Device, error) {
	if device, err := s.validateBySecretLookup(secret); err != nil {
		return nil, err
	} else if device != nil {
		return device, nil
	}

	// Only rows that have not been backfilled still need scanning.
	rows, err := s.db.Query(`
		SELECT id, did, handle, name, secret_hash, ip_address, location, user_agent, created_at, last_used
		FROM devices
		WHERE secret_lookup IS NULL OR secret_lookup = ''
	`)
	if err != nil {
		return nil, fmt.Errorf("failed to query devices: %w", err)
	}
	defer rows.Close()

	// Record the match and finish iterating before writing. Issuing the
	// backfill UPDATE inside this loop deadlocks: the open cursor holds a
	// connection, and on a single-connection pool the write waits on a
	// connection that only the cursor can release.
	var matched *Device

	for rows.Next() {
		var device Device
		var lastUsed sql.NullTime
		var location sql.NullString

		err := rows.Scan(
			&device.ID,
			&device.DID,
			&device.Handle,
			&device.Name,
			&device.SecretHash,
			&device.IPAddress,
			&location,
			&device.UserAgent,
			&device.CreatedAt,
			&lastUsed,
		)
		if err != nil {
			continue
		}

		if lastUsed.Valid {
			device.LastUsed = lastUsed.Time
		}
		if location.Valid {
			device.Location = location.String
		}

		// Check if this device's hash matches the secret
		if err := bcrypt.CompareHashAndPassword([]byte(device.SecretHash), []byte(secret)); err == nil {
			matched = &device
			break
		}
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("failed to scan devices: %w", err)
	}
	// Release the connection before writing.
	rows.Close()

	if matched == nil {
		return nil, fmt.Errorf("invalid device secret")
	}

	// Backfill so this device never takes the scan path again.
	if _, err := s.db.Exec(
		`UPDATE devices SET secret_lookup = ? WHERE id = ?`,
		deviceSecretLookup(secret), matched.ID,
	); err != nil {
		slog.Warn("Failed to backfill device secret_lookup",
			"component", "db/devices", "deviceID", matched.ID, "error", err)
	}

	// Update last used asynchronously
	go s.UpdateLastUsed(matched.SecretHash)

	return matched, nil
}

// validateBySecretLookup resolves a secret via the indexed lookup column.
// Returns (nil, nil) when there is no indexed match, so the caller can fall
// back to scanning rows that predate migration 0028.
//
// The SHA-256 match IS the authentication. Presenting a value that hashes to a
// stored digest requires either a SHA-256 preimage or guessing the 256 bits of
// crypto/rand entropy in the secret, so no second factor is needed here; bcrypt
// is retained only on the legacy path below, where it is the sole verifier
// available for rows that have no lookup value yet.
//
// The comparison happens in SQL and so is not constant-time, which is
// immaterial: exploiting the timing would require iterating candidate secrets,
// and every candidate is a 256-bit guess.
func (s *DeviceStore) validateBySecretLookup(secret string) (*Device, error) {
	var device Device
	var lastUsed sql.NullTime
	var location sql.NullString

	err := s.db.QueryRow(`
		SELECT id, did, handle, name, secret_hash, ip_address, location, user_agent, created_at, last_used
		FROM devices
		WHERE secret_lookup = ?
	`, deviceSecretLookup(secret)).Scan(
		&device.ID, &device.DID, &device.Handle, &device.Name, &device.SecretHash,
		&device.IPAddress, &location, &device.UserAgent, &device.CreatedAt, &lastUsed,
	)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("failed to look up device by secret: %w", err)
	}

	if lastUsed.Valid {
		device.LastUsed = lastUsed.Time
	}
	if location.Valid {
		device.Location = location.String
	}

	go s.UpdateLastUsed(device.SecretHash)

	return &device, nil
}

// ListDevices returns all devices for a DID
func (s *DeviceStore) ListDevices(did string) []*Device {
	rows, err := s.db.Query(`
		SELECT id, did, handle, name, ip_address, location, user_agent, created_at, last_used
		FROM devices
		WHERE did = ?
		ORDER BY created_at DESC
	`, did)

	if err != nil {
		return []*Device{}
	}
	defer rows.Close()

	var devices []*Device
	for rows.Next() {
		var device Device
		var lastUsed sql.NullTime
		var location sql.NullString

		err := rows.Scan(
			&device.ID,
			&device.DID,
			&device.Handle,
			&device.Name,
			&device.IPAddress,
			&location,
			&device.UserAgent,
			&device.CreatedAt,
			&lastUsed,
		)
		if err != nil {
			continue
		}

		if lastUsed.Valid {
			device.LastUsed = lastUsed.Time
		}
		if location.Valid {
			device.Location = location.String
		}

		devices = append(devices, &device)
	}

	return devices
}

// RevokeDevice removes a device
func (s *DeviceStore) RevokeDevice(did, deviceID string) error {
	result, err := s.db.Exec(`
		DELETE FROM devices
		WHERE did = ? AND id = ?
	`, did, deviceID)

	if err != nil {
		return fmt.Errorf("failed to revoke device: %w", err)
	}

	rows, _ := result.RowsAffected()
	if rows == 0 {
		return fmt.Errorf("device not found")
	}

	return nil
}

// UpdateLastUsed updates the last used timestamp, at most once per
// deviceLastUsedInterval per device.
//
// Callers invoke this on every successful authentication. The throttle state is
// per-process and lost on restart, which costs one extra write per device per
// boot.
func (s *DeviceStore) UpdateLastUsed(secretHash string) {
	now := time.Now()
	if prev, ok := s.lastUsedWrite.Load(secretHash); ok {
		if last, ok := prev.(time.Time); ok && now.Sub(last) < deviceLastUsedInterval {
			return
		}
	}
	// Stamp before writing rather than after: a slow or failing write must not
	// let every concurrent layer upload through to pile on more of them.
	s.lastUsedWrite.Store(secretHash, now)

	_, err := s.db.Exec(`
		UPDATE devices
		SET last_used = ?
		WHERE secret_hash = ?
	`, now, secretHash)

	if err != nil {
		slog.Warn("Failed to update device last used timestamp", "component", "device_store", "error", err)
	}
}

// CleanupExpired removes expired pending authorizations
func (s *DeviceStore) CleanupExpired() {
	result, err := s.db.Exec(`
		DELETE FROM pending_device_auth
		WHERE expires_at < datetime('now')
	`)

	if err != nil {
		slog.Warn("Failed to cleanup expired pending auths", "component", "device_store", "error", err)
		return
	}

	deleted, _ := result.RowsAffected()
	if deleted > 0 {
		slog.Info("Cleaned up expired pending device auths", "count", deleted)
	}
}

// CleanupExpiredContext is a context-aware version for background workers
func (s *DeviceStore) CleanupExpiredContext(ctx context.Context) error {
	result, err := s.db.ExecContext(ctx, `
		DELETE FROM pending_device_auth
		WHERE expires_at < datetime('now')
	`)

	if err != nil {
		return fmt.Errorf("failed to cleanup expired pending auths: %w", err)
	}

	deleted, _ := result.RowsAffected()
	if deleted > 0 {
		slog.Info("Cleaned up expired pending device auths", "count", deleted)
	}

	return nil
}

// generateUserCode creates a short, human-readable code
// Format: XXXX-XXXX (e.g., "WDJB-MJHT")
// Character set: A-Z excluding ambiguous chars (0, O, I, 1, L)
func generateUserCode() string {
	chars := "ABCDEFGHJKMNPQRSTUVWXYZ23456789"
	code := make([]byte, 8)
	if _, err := rand.Read(code); err != nil {
		// Fallback to timestamp-based generation if crypto rand fails
		now := time.Now().UnixNano()
		for i := range code {
			code[i] = byte(now >> (i * 8))
		}
	}
	for i := range code {
		code[i] = chars[int(code[i])%len(chars)]
	}
	return string(code[:4]) + "-" + string(code[4:])
}