silo-server/internal/auth/oauth_store.go

package auth

import (
	"context"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// OAuthSession is one in-flight OAuth authorization-code exchange.
//
// State is the primary key — an HMAC-signed string generated by the
// host on /init and verified on /callback. ProviderState is opaque JSON
// the plugin returned from InitAuthorize and that the host round-trips
// back on ExchangeCode (e.g., PKCE verifier, OIDC nonce). LinkingUserID
// is empty when the flow is a fresh login; non-empty when an existing
// user is attaching another identity from /me/account.
type OAuthSession struct {
	State         string
	InstallID     string
	RedirectURI   string
	LinkingUserID string
	ProviderState []byte
	NextURL       string
	CreatedAt     time.Time
	ExpiresAt     time.Time
}

// OAuthStore is the persistence interface used by the OAuth handlers.
// Defined as an interface so handler tests can supply InMemoryOAuthStore.
type OAuthStore interface {
	Insert(ctx context.Context, s OAuthSession) error
	GetAndDelete(ctx context.Context, state string) (OAuthSession, error)
	DeleteExpired(ctx context.Context, now time.Time) (int, error)
}

// ErrOAuthSessionNotFound is returned by GetAndDelete when no row matches.
var ErrOAuthSessionNotFound = errors.New("oauth_session not found")

type OAuthCompletion struct {
	Code         string
	AccessToken  string
	RefreshToken string
	ExpiresIn    int
	NextURL      string
	CreatedAt    time.Time
	ExpiresAt    time.Time
}

type OAuthCompletionStore interface {
	InsertCompletion(ctx context.Context, c OAuthCompletion) error
	GetAndDeleteCompletion(ctx context.Context, code string) (OAuthCompletion, error)
	DeleteExpiredCompletions(ctx context.Context, now time.Time) (int, error)
}

var ErrOAuthCompletionNotFound = errors.New("oauth_completion not found")

// PGOAuthStore is the Postgres-backed OAuthStore.
type PGOAuthStore struct {
	pool          *pgxpool.Pool
	completionKey [32]byte
}

func NewPGOAuthStore(pool *pgxpool.Pool, completionSecret ...[]byte) *PGOAuthStore {
	var material []byte
	if len(completionSecret) > 0 {
		material = completionSecret[0]
	}
	if len(material) == 0 {
		material = []byte("silo-oauth-completion-default")
	}
	key := sha256.Sum256(append([]byte("silo/oauth-completion/v1:"), material...))
	return &PGOAuthStore{pool: pool, completionKey: key}
}

func (s *PGOAuthStore) Insert(ctx context.Context, sess OAuthSession) error {
	if err := validateSession(&sess); err != nil {
		return err
	}
	_, err := s.pool.Exec(ctx, `
		INSERT INTO oauth_session (state, install_id, redirect_uri, linking_user_id, provider_state, next_url, expires_at)
		VALUES ($1, $2, $3, NULLIF($4, ''), $5, $6, $7)
	`, sess.State, sess.InstallID, sess.RedirectURI, sess.LinkingUserID, sess.ProviderState, sess.NextURL, sess.ExpiresAt)
	if err != nil {
		return fmt.Errorf("insert oauth_session: %w", err)
	}
	return nil
}

func (s *PGOAuthStore) GetAndDelete(ctx context.Context, state string) (OAuthSession, error) {
	row := s.pool.QueryRow(ctx, `
		DELETE FROM oauth_session WHERE state = $1
		RETURNING state, install_id, redirect_uri, COALESCE(linking_user_id, ''), provider_state, next_url, created_at, expires_at
	`, state)
	var out OAuthSession
	if err := row.Scan(&out.State, &out.InstallID, &out.RedirectURI, &out.LinkingUserID, &out.ProviderState, &out.NextURL, &out.CreatedAt, &out.ExpiresAt); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return OAuthSession{}, ErrOAuthSessionNotFound
		}
		return OAuthSession{}, fmt.Errorf("get_and_delete oauth_session: %w", err)
	}
	return out, nil
}

func (s *PGOAuthStore) DeleteExpired(ctx context.Context, now time.Time) (int, error) {
	tag, err := s.pool.Exec(ctx, `DELETE FROM oauth_session WHERE expires_at < $1`, now)
	if err != nil {
		return 0, fmt.Errorf("delete expired oauth_session: %w", err)
	}
	return int(tag.RowsAffected()), nil
}

func (s *PGOAuthStore) InsertCompletion(ctx context.Context, c OAuthCompletion) error {
	if err := validateCompletion(&c); err != nil {
		return err
	}
	codeHash := oauthCompletionCodeHash(c.Code)
	tokenCiphertext, err := s.encryptCompletionTokens(c, codeHash)
	if err != nil {
		return fmt.Errorf("encrypt oauth_completion tokens: %w", err)
	}
	_, err = s.pool.Exec(ctx, `
		INSERT INTO oauth_completion (code_hash, token_ciphertext, expires_in, next_url, expires_at)
		VALUES ($1, $2, $3, $4, $5)
	`, codeHash, tokenCiphertext, c.ExpiresIn, c.NextURL, c.ExpiresAt)
	if err != nil {
		return fmt.Errorf("insert oauth_completion: %w", err)
	}
	return nil
}

func (s *PGOAuthStore) GetAndDeleteCompletion(ctx context.Context, code string) (OAuthCompletion, error) {
	codeHash := oauthCompletionCodeHash(code)
	row := s.pool.QueryRow(ctx, `
		DELETE FROM oauth_completion WHERE code_hash = $1 AND expires_at >= now()
		RETURNING token_ciphertext, expires_in, next_url, created_at, expires_at
	`, codeHash)
	var out OAuthCompletion
	var tokenCiphertext string
	if err := row.Scan(&tokenCiphertext, &out.ExpiresIn, &out.NextURL, &out.CreatedAt, &out.ExpiresAt); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return OAuthCompletion{}, ErrOAuthCompletionNotFound
		}
		return OAuthCompletion{}, fmt.Errorf("get_and_delete oauth_completion: %w", err)
	}
	out.Code = code
	if err := s.decryptCompletionTokens(tokenCiphertext, codeHash, &out); err != nil {
		return OAuthCompletion{}, fmt.Errorf("decrypt oauth_completion tokens: %w", err)
	}
	return out, nil
}

func (s *PGOAuthStore) DeleteExpiredCompletions(ctx context.Context, now time.Time) (int, error) {
	tag, err := s.pool.Exec(ctx, `DELETE FROM oauth_completion WHERE expires_at < $1`, now)
	if err != nil {
		return 0, fmt.Errorf("delete expired oauth_completion: %w", err)
	}
	return int(tag.RowsAffected()), nil
}

type oauthCompletionTokenPayload struct {
	AccessToken  string `json:"access_token"`
	RefreshToken string `json:"refresh_token"`
}

func (s *PGOAuthStore) encryptCompletionTokens(c OAuthCompletion, codeHash string) (string, error) {
	block, err := aes.NewCipher(s.completionKey[:])
	if err != nil {
		return "", err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return "", err
	}
	plaintext, err := json.Marshal(oauthCompletionTokenPayload{
		AccessToken:  c.AccessToken,
		RefreshToken: c.RefreshToken,
	})
	if err != nil {
		return "", err
	}
	sealed := gcm.Seal(nonce, nonce, plaintext, []byte(codeHash))
	return base64.RawURLEncoding.EncodeToString(sealed), nil
}

func (s *PGOAuthStore) decryptCompletionTokens(ciphertext, codeHash string, out *OAuthCompletion) error {
	sealed, err := base64.RawURLEncoding.DecodeString(ciphertext)
	if err != nil {
		return err
	}
	block, err := aes.NewCipher(s.completionKey[:])
	if err != nil {
		return err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return err
	}
	if len(sealed) < gcm.NonceSize() {
		return fmt.Errorf("ciphertext too short")
	}
	nonce, body := sealed[:gcm.NonceSize()], sealed[gcm.NonceSize():]
	plaintext, err := gcm.Open(nil, nonce, body, []byte(codeHash))
	if err != nil {
		return err
	}
	var payload oauthCompletionTokenPayload
	if err := json.Unmarshal(plaintext, &payload); err != nil {
		return err
	}
	out.AccessToken = payload.AccessToken
	out.RefreshToken = payload.RefreshToken
	return nil
}

func oauthCompletionCodeHash(code string) string {
	sum := sha256.Sum256([]byte(code))
	return hex.EncodeToString(sum[:])
}

// InMemoryOAuthStore is a process-local OAuthStore for tests.
type InMemoryOAuthStore struct {
	mu          sync.Mutex
	rows        map[string]OAuthSession
	completions map[string]OAuthCompletion
}

func NewInMemoryOAuthStore() *InMemoryOAuthStore {
	return &InMemoryOAuthStore{
		rows:        make(map[string]OAuthSession),
		completions: make(map[string]OAuthCompletion),
	}
}

func (s *InMemoryOAuthStore) Insert(_ context.Context, sess OAuthSession) error {
	if err := validateSession(&sess); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, exists := s.rows[sess.State]; exists {
		return fmt.Errorf("insert oauth_session: state %q already exists", sess.State)
	}
	if sess.CreatedAt.IsZero() {
		sess.CreatedAt = time.Now().UTC()
	}
	s.rows[sess.State] = sess
	return nil
}

func (s *InMemoryOAuthStore) GetAndDelete(_ context.Context, state string) (OAuthSession, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	sess, ok := s.rows[state]
	if !ok {
		return OAuthSession{}, ErrOAuthSessionNotFound
	}
	delete(s.rows, state)
	return sess, nil
}

func (s *InMemoryOAuthStore) DeleteExpired(_ context.Context, now time.Time) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	n := 0
	for k, v := range s.rows {
		if v.ExpiresAt.Before(now) {
			delete(s.rows, k)
			n++
		}
	}
	return n, nil
}

func (s *InMemoryOAuthStore) InsertCompletion(_ context.Context, c OAuthCompletion) error {
	if err := validateCompletion(&c); err != nil {
		return err
	}
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, exists := s.completions[c.Code]; exists {
		return fmt.Errorf("insert oauth_completion: code %q already exists", c.Code)
	}
	if c.CreatedAt.IsZero() {
		c.CreatedAt = time.Now().UTC()
	}
	s.completions[c.Code] = c
	return nil
}

func (s *InMemoryOAuthStore) GetAndDeleteCompletion(_ context.Context, code string) (OAuthCompletion, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	c, ok := s.completions[code]
	if !ok || c.ExpiresAt.Before(time.Now().UTC()) {
		delete(s.completions, code)
		return OAuthCompletion{}, ErrOAuthCompletionNotFound
	}
	delete(s.completions, code)
	return c, nil
}

func (s *InMemoryOAuthStore) DeleteExpiredCompletions(_ context.Context, now time.Time) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	n := 0
	for k, v := range s.completions {
		if v.ExpiresAt.Before(now) {
			delete(s.completions, k)
			n++
		}
	}
	return n, nil
}

func validateSession(sess *OAuthSession) error {
	if sess.State == "" || sess.InstallID == "" {
		return fmt.Errorf("oauth_session: state and install_id required")
	}
	if sess.ExpiresAt.IsZero() {
		return fmt.Errorf("oauth_session: expires_at required")
	}
	if sess.RedirectURI == "" {
		return fmt.Errorf("oauth_session: redirect_uri required")
	}
	if sess.NextURL == "" {
		sess.NextURL = "/"
	}
	if len(sess.ProviderState) == 0 {
		sess.ProviderState = []byte("{}")
	}
	return nil
}

func validateCompletion(c *OAuthCompletion) error {
	if c.Code == "" {
		return fmt.Errorf("oauth_completion: code required")
	}
	if c.AccessToken == "" || c.RefreshToken == "" {
		return fmt.Errorf("oauth_completion: tokens required")
	}
	if c.ExpiresIn <= 0 {
		return fmt.Errorf("oauth_completion: expires_in required")
	}
	if c.ExpiresAt.IsZero() {
		return fmt.Errorf("oauth_completion: expires_at required")
	}
	if c.NextURL == "" {
		c.NextURL = "/"
	}
	return nil
}