feat: harden prepaid billing and commercial operations

1 files changed, 71 insertions, 11 deletions

diff --git a/internal/controlplane/store.go b/internal/controlplane/store.go
index 833f846..fb48d78 100644
--- a/internal/controlplane/store.go
+++ b/internal/controlplane/store.go
@@ -2,7 +2,9 @@ package controlplane
 
 import (
 	"context"
+	"crypto/sha256"
 	_ "embed"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -11,6 +13,7 @@ import (
 
 	"aigw/internal/security"
 
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/redis/go-redis/v9"
 )
@@ -20,12 +23,15 @@ var schemaSQL string
 
 var ErrRedisDisabled = errors.New("Redis propagation is disabled")
 
+const migrationVersion int64 = 2026080504
+
 type Options struct {
-	DatabaseURL     string
-	RedisURL        string
-	CredentialKey   string
-	RedisChannel    string
-	VersionCacheKey string
+	DatabaseURL            string
+	RedisURL               string
+	CredentialKey          string
+	PreviousCredentialKeys []string
+	RedisChannel           string
+	VersionCacheKey        string
 }
 
 type Store struct {
@@ -37,7 +43,7 @@ type Store struct {
 }
 
 func NewStore(ctx context.Context, options Options) (*Store, error) {
-	cipher, err := security.NewCredentialCipher(options.CredentialKey)
+	cipher, err := security.NewCredentialKeyring(options.CredentialKey, options.PreviousCredentialKeys)
 	if err != nil {
 		return nil, err
 	}
@@ -76,11 +82,10 @@ func (s *Store) RedisEnabled() bool {
 	return s.redis != nil
 }
 
+func (s *Store) Ping(ctx context.Context) error { return s.db.Ping(ctx) }
+
 func (s *Store) Migrate(ctx context.Context) error {
-	if _, err := s.db.Exec(ctx, schemaSQL); err != nil {
-		return fmt.Errorf("apply control-plane schema: %w", err)
-	}
-	return nil
+	return applySchema(ctx, s.db)
 }
 
 func MigrateDatabase(ctx context.Context, databaseURL string) error {
@@ -89,12 +94,67 @@ func MigrateDatabase(ctx context.Context, databaseURL string) error {
 		return fmt.Errorf("configure PostgreSQL: %w", err)
 	}
 	defer db.Close()
-	if _, err := db.Exec(ctx, schemaSQL); err != nil {
+	return applySchema(ctx, db)
+}
+
+func MigrationStatusDatabase(ctx context.Context, databaseURL string) (MigrationStatus, error) {
+	db, err := pgxpool.New(ctx, databaseURL)
+	if err != nil {
+		return MigrationStatus{}, fmt.Errorf("configure PostgreSQL: %w", err)
+	}
+	defer db.Close()
+	var result MigrationStatus
+	err = db.QueryRow(ctx, `SELECT version,name,checksum,applied_at FROM schema_migrations ORDER BY version DESC LIMIT 1`).Scan(&result.Version, &result.Name, &result.Checksum, &result.AppliedAt)
+	if err != nil {
+		return MigrationStatus{}, fmt.Errorf("read migration status: %w", err)
+	}
+	return result, nil
+}
+
+func applySchema(ctx context.Context, db *pgxpool.Pool) error {
+	tx, err := db.Begin(ctx)
+	if err != nil {
+		return fmt.Errorf("begin migration: %w", err)
+	}
+	defer tx.Rollback(ctx)
+	if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1)`, migrationVersion); err != nil {
+		return err
+	}
+	if _, err := tx.Exec(ctx, schemaSQL); err != nil {
 		return fmt.Errorf("apply control-plane schema: %w", err)
 	}
+	hash := sha256.Sum256([]byte(schemaSQL))
+	checksum := hex.EncodeToString(hash[:])
+	var existing string
+	err = tx.QueryRow(ctx, `SELECT checksum FROM schema_migrations WHERE version=$1`, migrationVersion).Scan(&existing)
+	if err == nil && existing != checksum {
+		return fmt.Errorf("migration %d checksum changed; deploy an explicit new migration version", migrationVersion)
+	}
+	if !errors.Is(err, pgx.ErrNoRows) && err != nil {
+		return err
+	}
+	if _, err := tx.Exec(ctx, `INSERT INTO schema_migrations(version,name,checksum) VALUES ($1,$2,$3) ON CONFLICT DO NOTHING`, migrationVersion, "commercial-control-plane", checksum); err != nil {
+		return err
+	}
+	if err := tx.Commit(ctx); err != nil {
+		return fmt.Errorf("commit migration: %w", err)
+	}
 	return nil
 }
 
+type MigrationStatus struct {
+	Version   int64     `json:"version"`
+	Name      string    `json:"name"`
+	Checksum  string    `json:"checksum"`
+	AppliedAt time.Time `json:"applied_at"`
+}
+
+func (s *Store) MigrationStatus(ctx context.Context) (MigrationStatus, error) {
+	var result MigrationStatus
+	err := s.db.QueryRow(ctx, `SELECT version,name,checksum,applied_at FROM schema_migrations ORDER BY version DESC LIMIT 1`).Scan(&result.Version, &result.Name, &result.Checksum, &result.AppliedAt)
+	return result, err
+}
+
 func (s *Store) DatabaseGeneration(ctx context.Context) (int64, error) {
 	var generation int64
 	err := s.db.QueryRow(ctx, `SELECT generation FROM control_state WHERE singleton = TRUE`).Scan(&generation)