AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,26 @@
|
||||
package apiresponse
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
type Envelope struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"msg"`
|
||||
Data any `json:"data,omitempty"`
|
||||
}
|
||||
|
||||
func OK(writer http.ResponseWriter, data any) {
|
||||
Write(writer, http.StatusOK, Envelope{Code: http.StatusOK, Message: "success", Data: data})
|
||||
}
|
||||
|
||||
func Error(writer http.ResponseWriter, status int, message string) {
|
||||
Write(writer, status, Envelope{Code: status, Message: message})
|
||||
}
|
||||
|
||||
func Write(writer http.ResponseWriter, status int, value any) {
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
writer.WriteHeader(status)
|
||||
_ = json.NewEncoder(writer).Encode(value)
|
||||
}
|
||||
Vendored
+23
@@ -0,0 +1,23 @@
|
||||
package cache
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
func Open(rawURL string) (*redis.Client, error) {
|
||||
if rawURL == "" {
|
||||
return nil, nil
|
||||
}
|
||||
options, err := redis.ParseURL(rawURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse redis URL: %w", err)
|
||||
}
|
||||
options.DialTimeout = 2 * time.Second
|
||||
options.ReadTimeout = time.Second
|
||||
options.WriteTimeout = time.Second
|
||||
options.PoolTimeout = 2 * time.Second
|
||||
return redis.NewClient(options), nil
|
||||
}
|
||||
@@ -0,0 +1,348 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Config struct {
|
||||
Environment string
|
||||
Server Server
|
||||
Database Database
|
||||
Redis Redis
|
||||
Security Security
|
||||
Auth Auth
|
||||
Credentials Credentials
|
||||
Upstream Upstream
|
||||
Audit Audit
|
||||
Outbox Outbox
|
||||
RuntimeData RuntimeData
|
||||
Shadow Shadow
|
||||
}
|
||||
|
||||
type Server struct {
|
||||
Address string
|
||||
ReadHeaderTimeout time.Duration
|
||||
IdleTimeout time.Duration
|
||||
ShutdownTimeout time.Duration
|
||||
MaxBodyBytes int64
|
||||
}
|
||||
|
||||
type Database struct {
|
||||
URL string
|
||||
MaxConns int32
|
||||
MinConns int32
|
||||
}
|
||||
|
||||
type Redis struct {
|
||||
CriticalURL string
|
||||
CacheURL string
|
||||
}
|
||||
|
||||
type Security struct {
|
||||
BootstrapAPIKey string
|
||||
BootstrapAPIKeyEnabled bool
|
||||
}
|
||||
|
||||
type Auth struct {
|
||||
SessionTTL time.Duration
|
||||
TOTPChallengeTTL time.Duration
|
||||
MaxFailures int
|
||||
LockDuration time.Duration
|
||||
LoginRateLimitMax int // 单 IP 滑动窗口内的最大登录尝试次数
|
||||
LoginRateLimitWindow time.Duration // 登录限流滑动窗口
|
||||
}
|
||||
|
||||
type Credentials struct {
|
||||
MasterKey string
|
||||
KEKVersion int
|
||||
KEKKeyring string
|
||||
AllowPrivateProviderURL bool
|
||||
AllowPrivateToolURL bool
|
||||
AllowPrivateWebhookURL bool
|
||||
ProviderRefreshInterval time.Duration
|
||||
}
|
||||
|
||||
type Upstream struct {
|
||||
BaseURL string
|
||||
APIKey string
|
||||
FallbackEnabled bool
|
||||
ResponseHeaderTimeout time.Duration
|
||||
MaxRetries int
|
||||
RetryBackoff time.Duration
|
||||
CircuitThreshold int
|
||||
CircuitOpenDuration time.Duration
|
||||
}
|
||||
|
||||
type Audit struct {
|
||||
QueueSize int
|
||||
BatchSize int
|
||||
FlushInterval time.Duration
|
||||
Retention time.Duration
|
||||
UsageRetention time.Duration
|
||||
PartitionMonthsAhead int
|
||||
MaintenanceInterval time.Duration
|
||||
}
|
||||
|
||||
type Outbox struct {
|
||||
Stream string
|
||||
BatchSize int
|
||||
PollInterval time.Duration
|
||||
Lease time.Duration
|
||||
MaxAttempts int
|
||||
MaxBackoff time.Duration
|
||||
StreamMaxLength int64
|
||||
MarkerTTL time.Duration
|
||||
}
|
||||
|
||||
type RuntimeData struct {
|
||||
ContentPolicyRefreshInterval time.Duration
|
||||
PricingRefreshInterval time.Duration
|
||||
}
|
||||
|
||||
type Shadow struct {
|
||||
BaseURL string
|
||||
APIKey string
|
||||
SampleRate float64
|
||||
Timeout time.Duration
|
||||
MaxBodyBytes int64
|
||||
MaxConcurrent int
|
||||
}
|
||||
|
||||
func Load() (Config, error) {
|
||||
cfg := Config{
|
||||
Environment: env("APP_ENV", "local"),
|
||||
Server: Server{
|
||||
Address: env("HTTP_ADDR", "127.0.0.1:8080"),
|
||||
ReadHeaderTimeout: duration("HTTP_READ_HEADER_TIMEOUT", 5*time.Second),
|
||||
IdleTimeout: duration("HTTP_IDLE_TIMEOUT", 120*time.Second),
|
||||
ShutdownTimeout: duration("HTTP_SHUTDOWN_TIMEOUT", 20*time.Second),
|
||||
MaxBodyBytes: int64Value("HTTP_MAX_BODY_BYTES", 32<<20),
|
||||
},
|
||||
Database: Database{
|
||||
URL: strings.TrimSpace(os.Getenv("DATABASE_URL")),
|
||||
MaxConns: int32(intValue("DATABASE_MAX_CONNS", 40)),
|
||||
MinConns: int32(intValue("DATABASE_MIN_CONNS", 4)),
|
||||
},
|
||||
Redis: Redis{
|
||||
CriticalURL: strings.TrimSpace(os.Getenv("REDIS_CRITICAL_URL")),
|
||||
CacheURL: strings.TrimSpace(os.Getenv("REDIS_CACHE_URL")),
|
||||
},
|
||||
Security: Security{
|
||||
BootstrapAPIKey: strings.TrimSpace(os.Getenv("GATEWAY_BOOTSTRAP_API_KEY")),
|
||||
BootstrapAPIKeyEnabled: boolValue("GATEWAY_BOOTSTRAP_API_KEY_ENABLED", false),
|
||||
},
|
||||
Auth: Auth{
|
||||
SessionTTL: duration("AUTH_SESSION_TTL", 12*time.Hour),
|
||||
TOTPChallengeTTL: duration("AUTH_TOTP_CHALLENGE_TTL", 5*time.Minute),
|
||||
MaxFailures: intValue("LOGIN_MAX_FAILURES", 5),
|
||||
LockDuration: duration("LOGIN_LOCK_DURATION", 15*time.Minute),
|
||||
LoginRateLimitMax: intValue("LOGIN_RATE_LIMIT_MAX", 30),
|
||||
LoginRateLimitWindow: duration("LOGIN_RATE_LIMIT_WINDOW", 5*time.Minute),
|
||||
},
|
||||
Credentials: Credentials{
|
||||
MasterKey: strings.TrimSpace(os.Getenv("CREDENTIAL_MASTER_KEY")),
|
||||
KEKVersion: intValue("CREDENTIAL_KEK_VERSION", 1),
|
||||
KEKKeyring: strings.TrimSpace(os.Getenv("CREDENTIAL_KEK_KEYRING")),
|
||||
AllowPrivateProviderURL: boolValue("ALLOW_PRIVATE_PROVIDER_URLS", false),
|
||||
AllowPrivateToolURL: boolValue("ALLOW_PRIVATE_TOOL_URLS", false),
|
||||
AllowPrivateWebhookURL: boolValue("ALLOW_PRIVATE_WEBHOOK_URLS", false),
|
||||
ProviderRefreshInterval: duration("PROVIDER_REFRESH_INTERVAL", 5*time.Second),
|
||||
},
|
||||
Upstream: Upstream{
|
||||
BaseURL: strings.TrimRight(env("UPSTREAM_BASE_URL", "https://api.openai.com"), "/"), APIKey: strings.TrimSpace(os.Getenv("UPSTREAM_API_KEY")),
|
||||
FallbackEnabled: boolValue("UPSTREAM_FALLBACK_ENABLED", true), ResponseHeaderTimeout: duration("UPSTREAM_RESPONSE_HEADER_TIMEOUT", 60*time.Second),
|
||||
MaxRetries: intValue("UPSTREAM_MAX_RETRIES", 2), RetryBackoff: duration("UPSTREAM_RETRY_BACKOFF", 50*time.Millisecond),
|
||||
CircuitThreshold: intValue("UPSTREAM_CIRCUIT_THRESHOLD", 5), CircuitOpenDuration: duration("UPSTREAM_CIRCUIT_OPEN_DURATION", 30*time.Second),
|
||||
},
|
||||
Audit: Audit{
|
||||
QueueSize: intValue("AUDIT_QUEUE_SIZE", 4096), BatchSize: intValue("AUDIT_BATCH_SIZE", 200),
|
||||
FlushInterval: duration("AUDIT_FLUSH_INTERVAL", time.Second), Retention: duration("AUDIT_RETENTION", 90*24*time.Hour),
|
||||
UsageRetention: duration("USAGE_RETENTION", 730*24*time.Hour), PartitionMonthsAhead: intValue("AUDIT_PARTITION_MONTHS_AHEAD", 3),
|
||||
MaintenanceInterval: duration("AUDIT_MAINTENANCE_INTERVAL", 6*time.Hour),
|
||||
},
|
||||
Outbox: Outbox{
|
||||
Stream: env("OUTBOX_STREAM", "gateway:{outbox}:events"), BatchSize: intValue("OUTBOX_BATCH_SIZE", 100),
|
||||
PollInterval: duration("OUTBOX_POLL_INTERVAL", 500*time.Millisecond), Lease: duration("OUTBOX_LEASE", 30*time.Second),
|
||||
MaxAttempts: intValue("OUTBOX_MAX_ATTEMPTS", 10), MaxBackoff: duration("OUTBOX_MAX_BACKOFF", 5*time.Minute),
|
||||
StreamMaxLength: int64Value("OUTBOX_STREAM_MAX_LENGTH", 100_000), MarkerTTL: duration("OUTBOX_MARKER_TTL", 30*24*time.Hour),
|
||||
},
|
||||
RuntimeData: RuntimeData{
|
||||
ContentPolicyRefreshInterval: duration("CONTENT_POLICY_REFRESH_INTERVAL", 30*time.Second),
|
||||
PricingRefreshInterval: duration("PRICING_REFRESH_INTERVAL", 30*time.Second),
|
||||
},
|
||||
Shadow: Shadow{
|
||||
BaseURL: strings.TrimRight(strings.TrimSpace(os.Getenv("SHADOW_BASE_URL")), "/"),
|
||||
APIKey: strings.TrimSpace(os.Getenv("SHADOW_API_KEY")), SampleRate: floatValue("SHADOW_SAMPLE_RATE", 0),
|
||||
Timeout: duration("SHADOW_TIMEOUT", 20*time.Second), MaxBodyBytes: int64Value("SHADOW_MAX_BODY_BYTES", 2<<20),
|
||||
MaxConcurrent: intValue("SHADOW_MAX_CONCURRENT", 16),
|
||||
},
|
||||
}
|
||||
|
||||
return cfg, cfg.Validate()
|
||||
}
|
||||
|
||||
func (c Config) Validate() error {
|
||||
var errs []error
|
||||
if c.Server.MaxBodyBytes <= 0 {
|
||||
errs = append(errs, errors.New("HTTP_MAX_BODY_BYTES must be positive"))
|
||||
}
|
||||
if c.Database.MinConns < 0 || c.Database.MaxConns < 1 || c.Database.MinConns > c.Database.MaxConns {
|
||||
errs = append(errs, errors.New("database pool sizes are invalid"))
|
||||
}
|
||||
if c.Auth.SessionTTL < 5*time.Minute || c.Auth.TOTPChallengeTTL < time.Minute || c.Auth.TOTPChallengeTTL > 15*time.Minute || c.Auth.MaxFailures < 1 || c.Auth.LockDuration < time.Minute {
|
||||
errs = append(errs, errors.New("authentication limits are invalid"))
|
||||
}
|
||||
if c.Auth.LoginRateLimitMax < 1 || c.Auth.LoginRateLimitWindow < time.Second {
|
||||
errs = append(errs, errors.New("login rate limit is invalid"))
|
||||
}
|
||||
if c.Credentials.KEKVersion < 1 {
|
||||
errs = append(errs, errors.New("CREDENTIAL_KEK_VERSION must be positive"))
|
||||
}
|
||||
if c.Credentials.ProviderRefreshInterval < time.Second || c.Credentials.ProviderRefreshInterval > time.Minute {
|
||||
errs = append(errs, errors.New("PROVIDER_REFRESH_INTERVAL must be between 1s and 1m"))
|
||||
}
|
||||
if err := validateHTTPURL(c.Upstream.BaseURL); err != nil {
|
||||
errs = append(errs, fmt.Errorf("UPSTREAM_BASE_URL: %w", err))
|
||||
}
|
||||
if c.Upstream.ResponseHeaderTimeout < time.Second || c.Upstream.ResponseHeaderTimeout > 10*time.Minute || c.Upstream.MaxRetries < 0 || c.Upstream.MaxRetries > 5 || c.Upstream.RetryBackoff < 0 || c.Upstream.RetryBackoff > 5*time.Second || c.Upstream.CircuitThreshold < 1 || c.Upstream.CircuitThreshold > 100 || c.Upstream.CircuitOpenDuration < time.Second || c.Upstream.CircuitOpenDuration > 10*time.Minute {
|
||||
errs = append(errs, errors.New("upstream resilience settings are invalid"))
|
||||
}
|
||||
if c.Audit.QueueSize < 100 || c.Audit.QueueSize > 1_000_000 || c.Audit.BatchSize < 1 || c.Audit.BatchSize > c.Audit.QueueSize || c.Audit.FlushInterval < 100*time.Millisecond || c.Audit.FlushInterval > time.Minute {
|
||||
errs = append(errs, errors.New("audit buffering settings are invalid"))
|
||||
}
|
||||
if c.Audit.Retention < 24*time.Hour || c.Audit.Retention > 10*365*24*time.Hour || c.Audit.UsageRetention < c.Audit.Retention || c.Audit.UsageRetention > 10*365*24*time.Hour || c.Audit.PartitionMonthsAhead < 1 || c.Audit.PartitionMonthsAhead > 24 || c.Audit.MaintenanceInterval < time.Hour || c.Audit.MaintenanceInterval > 7*24*time.Hour {
|
||||
errs = append(errs, errors.New("audit retention settings are invalid"))
|
||||
}
|
||||
if !strings.Contains(c.Outbox.Stream, "{outbox}") || c.Outbox.BatchSize < 1 || c.Outbox.BatchSize > 1000 || c.Outbox.PollInterval < 50*time.Millisecond || c.Outbox.PollInterval > time.Minute || c.Outbox.Lease < 5*time.Second || c.Outbox.Lease > 10*time.Minute || c.Outbox.MaxAttempts < 1 || c.Outbox.MaxAttempts > 100 || c.Outbox.MaxBackoff < time.Second || c.Outbox.MaxBackoff > time.Hour || c.Outbox.StreamMaxLength < 1000 || c.Outbox.StreamMaxLength > 100_000_000 || c.Outbox.MarkerTTL < 24*time.Hour || c.Outbox.MarkerTTL > 365*24*time.Hour {
|
||||
errs = append(errs, errors.New("outbox delivery settings are invalid"))
|
||||
}
|
||||
if c.RuntimeData.ContentPolicyRefreshInterval < time.Second || c.RuntimeData.ContentPolicyRefreshInterval > 10*time.Minute || c.RuntimeData.PricingRefreshInterval < time.Second || c.RuntimeData.PricingRefreshInterval > 10*time.Minute {
|
||||
errs = append(errs, errors.New("runtime data refresh settings must be between 1s and 10m"))
|
||||
}
|
||||
if c.Shadow.SampleRate < 0 || c.Shadow.SampleRate > 1 || c.Shadow.Timeout < time.Second || c.Shadow.Timeout > 2*time.Minute || c.Shadow.MaxBodyBytes < 1024 || c.Shadow.MaxBodyBytes > 32<<20 || c.Shadow.MaxConcurrent < 1 || c.Shadow.MaxConcurrent > 1000 {
|
||||
errs = append(errs, errors.New("shadow traffic settings are invalid"))
|
||||
}
|
||||
if c.Shadow.BaseURL != "" {
|
||||
if err := validateHTTPURL(c.Shadow.BaseURL); err != nil {
|
||||
errs = append(errs, fmt.Errorf("SHADOW_BASE_URL: %w", err))
|
||||
}
|
||||
if c.Shadow.APIKey == "" {
|
||||
errs = append(errs, errors.New("SHADOW_API_KEY is required when SHADOW_BASE_URL is set"))
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
func (c Config) ValidateRuntime() error {
|
||||
var errs []error
|
||||
if strings.EqualFold(c.Environment, "production") {
|
||||
if c.Database.URL == "" {
|
||||
errs = append(errs, errors.New("DATABASE_URL is required in production"))
|
||||
}
|
||||
if c.Redis.CriticalURL == "" {
|
||||
errs = append(errs, errors.New("REDIS_CRITICAL_URL is required in production"))
|
||||
}
|
||||
if c.Security.BootstrapAPIKeyEnabled {
|
||||
if len(c.Security.BootstrapAPIKey) < 32 {
|
||||
errs = append(errs, errors.New("GATEWAY_BOOTSTRAP_API_KEY must contain at least 32 characters in production"))
|
||||
}
|
||||
// Reject the documented placeholder and other obviously weak values
|
||||
// so an operator cannot accidentally deploy with the example secret.
|
||||
switch strings.ToLower(c.Security.BootstrapAPIKey) {
|
||||
case "change-me-in-production", "local-development-key", "changeme", "change-me", "password", "secret":
|
||||
errs = append(errs, errors.New("GATEWAY_BOOTSTRAP_API_KEY is set to a known weak default; generate a strong random value"))
|
||||
}
|
||||
}
|
||||
if c.Upstream.FallbackEnabled && c.Upstream.APIKey == "" {
|
||||
errs = append(errs, errors.New("UPSTREAM_API_KEY is required in production"))
|
||||
}
|
||||
if c.Credentials.MasterKey == "" {
|
||||
errs = append(errs, errors.New("CREDENTIAL_MASTER_KEY is required in production"))
|
||||
}
|
||||
}
|
||||
return errors.Join(errs...)
|
||||
}
|
||||
|
||||
func validateHTTPURL(raw string) error {
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" || u.User != nil {
|
||||
return errors.New("must be an absolute http(s) URL without user info")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func env(key, fallback string) string {
|
||||
if value := strings.TrimSpace(os.Getenv(key)); value != "" {
|
||||
return value
|
||||
}
|
||||
return fallback
|
||||
}
|
||||
|
||||
func duration(key string, fallback time.Duration) time.Duration {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := time.ParseDuration(value)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func intValue(key string, fallback int) int {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.Atoi(value)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func int64Value(key string, fallback int64) int64 {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.ParseInt(value, 10, 64)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func floatValue(key string, fallback float64) float64 {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.ParseFloat(value, 64)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
func boolValue(key string, fallback bool) bool {
|
||||
value := strings.TrimSpace(os.Getenv(key))
|
||||
if value == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.ParseBool(value)
|
||||
if err != nil {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
@@ -0,0 +1,45 @@
|
||||
package config
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestProductionRequiresCoreDependencies(t *testing.T) {
|
||||
t.Setenv("APP_ENV", "production")
|
||||
t.Setenv("DATABASE_URL", "")
|
||||
t.Setenv("REDIS_CRITICAL_URL", "")
|
||||
t.Setenv("GATEWAY_BOOTSTRAP_API_KEY", "short")
|
||||
t.Setenv("UPSTREAM_API_KEY", "")
|
||||
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected structural configuration error: %v", err)
|
||||
}
|
||||
if err := cfg.ValidateRuntime(); err == nil {
|
||||
t.Fatal("expected production validation error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalDefaultsAreValid(t *testing.T) {
|
||||
t.Setenv("APP_ENV", "local")
|
||||
t.Setenv("UPSTREAM_BASE_URL", "https://example.com")
|
||||
if _, err := Load(); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestProductionCanDisableBootstrapCompatibility(t *testing.T) {
|
||||
t.Setenv("APP_ENV", "production")
|
||||
t.Setenv("DATABASE_URL", "postgres://gateway:secret@db.example/gateway?sslmode=require")
|
||||
t.Setenv("REDIS_CRITICAL_URL", "rediss://redis.example/0")
|
||||
t.Setenv("GATEWAY_BOOTSTRAP_API_KEY", "")
|
||||
t.Setenv("GATEWAY_BOOTSTRAP_API_KEY_ENABLED", "false")
|
||||
t.Setenv("UPSTREAM_FALLBACK_ENABLED", "false")
|
||||
t.Setenv("CREDENTIAL_MASTER_KEY", "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=")
|
||||
t.Setenv("UPSTREAM_BASE_URL", "https://example.com")
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := cfg.ValidateRuntime(); err != nil {
|
||||
t.Fatalf("disabled bootstrap compatibility should not require a key: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package cryptox
|
||||
|
||||
import (
|
||||
"crypto/aes"
|
||||
"crypto/cipher"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
var ErrKeyUnavailable = errors.New("encryption key unavailable")
|
||||
|
||||
type AESGCM struct {
|
||||
aead cipher.AEAD
|
||||
version int
|
||||
purpose string
|
||||
}
|
||||
|
||||
func NewAESGCM(encodedKey string, version int, purpose string) (*AESGCM, error) {
|
||||
if encodedKey == "" {
|
||||
return nil, nil
|
||||
}
|
||||
key, err := base64.StdEncoding.DecodeString(encodedKey)
|
||||
if err != nil {
|
||||
key, err = base64.RawStdEncoding.DecodeString(encodedKey)
|
||||
}
|
||||
if err != nil || len(key) != 32 {
|
||||
return nil, errors.New("encryption key must be base64-encoded 32 bytes")
|
||||
}
|
||||
if version < 1 || purpose == "" {
|
||||
return nil, errors.New("encryption version and purpose are required")
|
||||
}
|
||||
block, err := aes.NewCipher(key)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
aead, err := cipher.NewGCM(block)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &AESGCM{aead: aead, version: version, purpose: purpose}, nil
|
||||
}
|
||||
|
||||
func (c *AESGCM) Encrypt(plaintext []byte) ([]byte, int, error) {
|
||||
if c == nil {
|
||||
return nil, 0, ErrKeyUnavailable
|
||||
}
|
||||
nonce := make([]byte, c.aead.NonceSize())
|
||||
if _, err := rand.Read(nonce); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
ciphertext := c.aead.Seal(nil, nonce, plaintext, c.additionalData())
|
||||
return append(nonce, ciphertext...), c.version, nil
|
||||
}
|
||||
|
||||
func (c *AESGCM) Decrypt(encrypted []byte, version int) ([]byte, error) {
|
||||
if c == nil {
|
||||
return nil, ErrKeyUnavailable
|
||||
}
|
||||
if version != c.version {
|
||||
return nil, fmt.Errorf("encryption key version %d is not loaded", version)
|
||||
}
|
||||
nonceSize := c.aead.NonceSize()
|
||||
if len(encrypted) <= nonceSize {
|
||||
return nil, errors.New("encrypted value is truncated")
|
||||
}
|
||||
plaintext, err := c.aead.Open(nil, encrypted[:nonceSize], encrypted[nonceSize:], c.additionalData())
|
||||
if err != nil {
|
||||
return nil, errors.New("encrypted value authentication failed")
|
||||
}
|
||||
return plaintext, nil
|
||||
}
|
||||
|
||||
func (c *AESGCM) additionalData() []byte {
|
||||
return []byte(fmt.Sprintf("ai-gateway/%s/v%d", c.purpose, c.version))
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
package cryptox
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type Cipher interface {
|
||||
Encrypt([]byte) ([]byte, int, error)
|
||||
Decrypt([]byte, int) ([]byte, error)
|
||||
}
|
||||
|
||||
type Keyring struct {
|
||||
activeVersion int
|
||||
ciphers map[int]*AESGCM
|
||||
}
|
||||
|
||||
func NewKeyring(activeKey string, activeVersion int, encodedKeyring, purpose string) (*Keyring, error) {
|
||||
if activeVersion < 1 || purpose == "" {
|
||||
return nil, errors.New("encryption version and purpose are required")
|
||||
}
|
||||
encoded := make(map[int]string)
|
||||
if strings.TrimSpace(encodedKeyring) != "" {
|
||||
var values map[string]string
|
||||
if err := json.Unmarshal([]byte(encodedKeyring), &values); err != nil {
|
||||
return nil, fmt.Errorf("encryption keyring must be a JSON object: %w", err)
|
||||
}
|
||||
for rawVersion, key := range values {
|
||||
version, err := strconv.Atoi(rawVersion)
|
||||
if err != nil || version < 1 || strings.TrimSpace(key) == "" {
|
||||
return nil, fmt.Errorf("invalid encryption keyring version %q", rawVersion)
|
||||
}
|
||||
encoded[version] = strings.TrimSpace(key)
|
||||
}
|
||||
}
|
||||
if strings.TrimSpace(activeKey) != "" {
|
||||
if existing, ok := encoded[activeVersion]; ok && existing != strings.TrimSpace(activeKey) {
|
||||
return nil, fmt.Errorf("active encryption key version %d is configured twice with different values", activeVersion)
|
||||
}
|
||||
encoded[activeVersion] = strings.TrimSpace(activeKey)
|
||||
}
|
||||
if len(encoded) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
if _, ok := encoded[activeVersion]; !ok {
|
||||
return nil, fmt.Errorf("active encryption key version %d is not loaded", activeVersion)
|
||||
}
|
||||
keyring := &Keyring{activeVersion: activeVersion, ciphers: make(map[int]*AESGCM, len(encoded))}
|
||||
for version, key := range encoded {
|
||||
cipher, err := NewAESGCM(key, version, purpose)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encryption key version %d: %w", version, err)
|
||||
}
|
||||
keyring.ciphers[version] = cipher
|
||||
}
|
||||
return keyring, nil
|
||||
}
|
||||
|
||||
func (k *Keyring) Encrypt(plaintext []byte) ([]byte, int, error) {
|
||||
if k == nil {
|
||||
return nil, 0, ErrKeyUnavailable
|
||||
}
|
||||
return k.ciphers[k.activeVersion].Encrypt(plaintext)
|
||||
}
|
||||
|
||||
func (k *Keyring) Decrypt(encrypted []byte, version int) ([]byte, error) {
|
||||
if k == nil {
|
||||
return nil, ErrKeyUnavailable
|
||||
}
|
||||
cipher, ok := k.ciphers[version]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("encryption key version %d is not loaded", version)
|
||||
}
|
||||
return cipher.Decrypt(encrypted, version)
|
||||
}
|
||||
|
||||
func (k *Keyring) ActiveVersion() int {
|
||||
if k == nil {
|
||||
return 0
|
||||
}
|
||||
return k.activeVersion
|
||||
}
|
||||
|
||||
func (k *Keyring) Versions() []int {
|
||||
if k == nil {
|
||||
return nil
|
||||
}
|
||||
versions := make([]int, 0, len(k.ciphers))
|
||||
for version := range k.ciphers {
|
||||
versions = append(versions, version)
|
||||
}
|
||||
sort.Ints(versions)
|
||||
return versions
|
||||
}
|
||||
|
||||
var _ Cipher = (*Keyring)(nil)
|
||||
@@ -0,0 +1,37 @@
|
||||
package cryptox
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestKeyringDecryptsOldVersionAndEncryptsActiveVersion(t *testing.T) {
|
||||
oldKey := base64.StdEncoding.EncodeToString(make([]byte, 32))
|
||||
newBytes := make([]byte, 32)
|
||||
newBytes[0] = 1
|
||||
newKey := base64.StdEncoding.EncodeToString(newBytes)
|
||||
oldCipher, err := NewAESGCM(oldKey, 1, "test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
oldEncrypted, _, err := oldCipher.Encrypt([]byte("secret"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keyring, err := NewKeyring(newKey, 2, fmt.Sprintf(`{"1":%q}`, oldKey), "test")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
plaintext, err := keyring.Decrypt(oldEncrypted, 1)
|
||||
if err != nil || string(plaintext) != "secret" {
|
||||
t.Fatalf("old key was not usable: %q, %v", plaintext, err)
|
||||
}
|
||||
newEncrypted, version, err := keyring.Encrypt([]byte("new-secret"))
|
||||
if err != nil || version != 2 {
|
||||
t.Fatalf("active version was not used: version=%d error=%v", version, err)
|
||||
}
|
||||
if _, err := oldCipher.Decrypt(newEncrypted, version); err == nil {
|
||||
t.Fatal("old cipher must not decrypt the new version")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package database
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
func Open(ctx context.Context, cfg config.Database) (*pgxpool.Pool, error) {
|
||||
if cfg.URL == "" {
|
||||
return nil, nil
|
||||
}
|
||||
poolConfig, err := pgxpool.ParseConfig(cfg.URL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("parse database configuration: %w", err)
|
||||
}
|
||||
poolConfig.MaxConns = cfg.MaxConns
|
||||
poolConfig.MinConns = cfg.MinConns
|
||||
poolConfig.MaxConnLifetime = 30 * time.Minute
|
||||
poolConfig.MaxConnIdleTime = 5 * time.Minute
|
||||
poolConfig.HealthCheckPeriod = 30 * time.Second
|
||||
|
||||
pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create database pool: %w", err)
|
||||
}
|
||||
return pool, nil
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
package health
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type Probe func(context.Context) error
|
||||
|
||||
type Dependency struct {
|
||||
Name string
|
||||
Required bool
|
||||
Probe Probe
|
||||
}
|
||||
|
||||
type Result struct {
|
||||
Status string `json:"status"`
|
||||
Required bool `json:"required"`
|
||||
LatencyMS float64 `json:"latency_ms"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
type Report struct {
|
||||
Ready bool `json:"ready"`
|
||||
Status string `json:"status"`
|
||||
Components map[string]Result `json:"components"`
|
||||
}
|
||||
|
||||
type Checker struct {
|
||||
Timeout time.Duration
|
||||
Dependencies []Dependency
|
||||
}
|
||||
|
||||
func (c Checker) Readiness(ctx context.Context) Report {
|
||||
report := Report{Ready: true, Status: "ok", Components: make(map[string]Result, len(c.Dependencies))}
|
||||
type namedResult struct {
|
||||
name string
|
||||
result Result
|
||||
}
|
||||
results := make(chan namedResult, len(c.Dependencies))
|
||||
var wg sync.WaitGroup
|
||||
|
||||
for _, dependency := range c.Dependencies {
|
||||
dependency := dependency
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
result := Result{Status: "ok", Required: dependency.Required}
|
||||
if dependency.Probe == nil {
|
||||
result.Status = "not_configured"
|
||||
results <- namedResult{dependency.Name, result}
|
||||
return
|
||||
}
|
||||
probeCtx, cancel := context.WithTimeout(ctx, c.Timeout)
|
||||
defer cancel()
|
||||
started := time.Now()
|
||||
err := dependency.Probe(probeCtx)
|
||||
result.LatencyMS = float64(time.Since(started).Microseconds()) / 1000
|
||||
if err != nil {
|
||||
result.Status = "unavailable"
|
||||
result.Error = err.Error()
|
||||
}
|
||||
results <- namedResult{dependency.Name, result}
|
||||
}()
|
||||
}
|
||||
|
||||
go func() {
|
||||
wg.Wait()
|
||||
close(results)
|
||||
}()
|
||||
for item := range results {
|
||||
report.Components[item.name] = item.result
|
||||
if item.result.Required && item.result.Status != "ok" {
|
||||
report.Ready = false
|
||||
}
|
||||
if !item.result.Required && item.result.Status == "unavailable" && report.Status == "ok" {
|
||||
report.Status = "degraded"
|
||||
}
|
||||
}
|
||||
if !report.Ready {
|
||||
report.Status = "unavailable"
|
||||
}
|
||||
return report
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package health
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestOptionalFailureDegradesWithoutBlockingReadiness(t *testing.T) {
|
||||
checker := Checker{
|
||||
Timeout: time.Second,
|
||||
Dependencies: []Dependency{
|
||||
{Name: "postgres", Required: true, Probe: func(context.Context) error { return nil }},
|
||||
{Name: "redis_cache", Required: false, Probe: func(context.Context) error { return errors.New("offline") }},
|
||||
},
|
||||
}
|
||||
report := checker.Readiness(context.Background())
|
||||
if !report.Ready || report.Status != "degraded" {
|
||||
t.Fatalf("unexpected report: %+v", report)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRequiredFailureBlocksReadiness(t *testing.T) {
|
||||
checker := Checker{
|
||||
Timeout: time.Second,
|
||||
Dependencies: []Dependency{
|
||||
{Name: "postgres", Required: true, Probe: func(context.Context) error { return errors.New("offline") }},
|
||||
},
|
||||
}
|
||||
report := checker.Readiness(context.Background())
|
||||
if report.Ready || report.Status != "unavailable" {
|
||||
t.Fatalf("unexpected report: %+v", report)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package httpserver
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"runtime/debug"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"aigateway.local/core/internal/gateway"
|
||||
"aigateway.local/core/internal/platform/config"
|
||||
"aigateway.local/core/internal/platform/health"
|
||||
)
|
||||
|
||||
type Dependencies struct {
|
||||
Config config.Config
|
||||
Logger *slog.Logger
|
||||
Checker health.Checker
|
||||
Gateway http.Handler
|
||||
Control http.Handler
|
||||
Version string
|
||||
StartedAt time.Time
|
||||
BootstrapUses func() uint64
|
||||
ExtraMetrics func() string
|
||||
}
|
||||
|
||||
type metrics struct {
|
||||
requests atomic.Uint64
|
||||
panics atomic.Uint64
|
||||
}
|
||||
|
||||
func New(dependencies Dependencies) *http.Server {
|
||||
mux := http.NewServeMux()
|
||||
stats := &metrics{}
|
||||
mux.HandleFunc("GET /healthz", func(writer http.ResponseWriter, _ *http.Request) {
|
||||
writeJSON(writer, http.StatusOK, map[string]any{"status": "ok", "version": dependencies.Version})
|
||||
})
|
||||
mux.HandleFunc("GET /readyz", func(writer http.ResponseWriter, request *http.Request) {
|
||||
report := dependencies.Checker.Readiness(request.Context())
|
||||
status := http.StatusOK
|
||||
if !report.Ready {
|
||||
status = http.StatusServiceUnavailable
|
||||
}
|
||||
writeJSON(writer, status, report)
|
||||
})
|
||||
mux.HandleFunc("GET /metrics", func(writer http.ResponseWriter, _ *http.Request) {
|
||||
writer.Header().Set("Content-Type", "text/plain; version=0.0.4")
|
||||
bootstrapUses := uint64(0)
|
||||
if dependencies.BootstrapUses != nil {
|
||||
bootstrapUses = dependencies.BootstrapUses()
|
||||
}
|
||||
_, _ = fmt.Fprintf(writer, "gateway_http_requests_total %d\ngateway_http_panics_total %d\ngateway_bootstrap_api_key_uses_total %d\ngateway_uptime_seconds %.0f\n", stats.requests.Load(), stats.panics.Load(), bootstrapUses, time.Since(dependencies.StartedAt).Seconds())
|
||||
if dependencies.ExtraMetrics != nil {
|
||||
_, _ = fmt.Fprint(writer, dependencies.ExtraMetrics())
|
||||
}
|
||||
})
|
||||
mux.Handle("/v1/", dependencies.Gateway)
|
||||
if dependencies.Control != nil {
|
||||
mux.Handle("/api/", dependencies.Control)
|
||||
}
|
||||
|
||||
handler := recoverMiddleware(dependencies.Logger, stats, requestIDMiddleware(stats, mux))
|
||||
return &http.Server{
|
||||
Addr: dependencies.Config.Server.Address,
|
||||
Handler: handler,
|
||||
ReadHeaderTimeout: dependencies.Config.Server.ReadHeaderTimeout,
|
||||
IdleTimeout: dependencies.Config.Server.IdleTimeout,
|
||||
// WriteTimeout intentionally remains zero: SSE responses may be long lived.
|
||||
}
|
||||
}
|
||||
|
||||
func requestIDMiddleware(stats *metrics, next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
stats.requests.Add(1)
|
||||
requestID := request.Header.Get("X-Request-ID")
|
||||
if requestID == "" || len(requestID) > 128 {
|
||||
requestID = newRequestID()
|
||||
}
|
||||
writer.Header().Set("X-Request-ID", requestID)
|
||||
next.ServeHTTP(writer, request.WithContext(gateway.WithRequestID(request.Context(), requestID)))
|
||||
})
|
||||
}
|
||||
|
||||
func recoverMiddleware(logger *slog.Logger, stats *metrics, next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
stats.panics.Add(1)
|
||||
logger.Error("http handler panic", "request_id", gateway.RequestID(request.Context()), "panic", recovered, "stack", string(debug.Stack()))
|
||||
writeJSON(writer, http.StatusInternalServerError, map[string]string{"error": "internal server error"})
|
||||
}
|
||||
}()
|
||||
next.ServeHTTP(writer, request)
|
||||
})
|
||||
}
|
||||
|
||||
func writeJSON(writer http.ResponseWriter, status int, value any) {
|
||||
writer.Header().Set("Content-Type", "application/json")
|
||||
writer.WriteHeader(status)
|
||||
_ = json.NewEncoder(writer).Encode(value)
|
||||
}
|
||||
|
||||
func newRequestID() string {
|
||||
buffer := make([]byte, 12)
|
||||
if _, err := rand.Read(buffer); err != nil {
|
||||
return fmt.Sprintf("req_%x", time.Now().UnixNano())
|
||||
}
|
||||
return "req_" + hex.EncodeToString(buffer)
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
package id
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"fmt"
|
||||
)
|
||||
|
||||
func NewUUID() (string, error) {
|
||||
value := make([]byte, 16)
|
||||
if _, err := rand.Read(value); err != nil {
|
||||
return "", err
|
||||
}
|
||||
value[6] = (value[6] & 0x0f) | 0x40
|
||||
value[8] = (value[8] & 0x3f) | 0x80
|
||||
return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x",
|
||||
value[0:4], value[4:6], value[6:8], value[8:10], value[10:16]), nil
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package legacyid
|
||||
|
||||
import (
|
||||
"crypto/sha1"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Namespace is stable for the lifetime of this migration lineage. Changing it
|
||||
// would generate different UUIDs for the same legacy records.
|
||||
const Namespace = "e2464f95-0c8d-5a9c-9c1d-d5bca69bb09f"
|
||||
|
||||
// UUID returns an RFC 4122 UUIDv5 for a legacy record. SHA-1 is required by the
|
||||
// UUIDv5 standard here and is not used for credentials or signatures.
|
||||
func UUID(sourceSystem, entityType, legacyID string) (string, error) {
|
||||
sourceSystem = strings.TrimSpace(sourceSystem)
|
||||
entityType = strings.TrimSpace(entityType)
|
||||
legacyID = strings.TrimSpace(legacyID)
|
||||
if sourceSystem == "" || entityType == "" || legacyID == "" {
|
||||
return "", errors.New("legacy ID components must not be empty")
|
||||
}
|
||||
if len(sourceSystem) > 64 || len(entityType) > 128 || len(legacyID) > 256 {
|
||||
return "", errors.New("legacy ID component is too long")
|
||||
}
|
||||
namespace, err := hex.DecodeString(strings.ReplaceAll(Namespace, "-", ""))
|
||||
if err != nil || len(namespace) != 16 {
|
||||
return "", errors.New("invalid UUID namespace")
|
||||
}
|
||||
name := sourceSystem + "/" + entityType + "/" + legacyID
|
||||
digest := sha1.Sum(append(namespace, []byte(name)...))
|
||||
digest[6] = digest[6]&0x0f | 0x50
|
||||
digest[8] = digest[8]&0x3f | 0x80
|
||||
return fmt.Sprintf("%x-%x-%x-%x-%x", digest[0:4], digest[4:6], digest[6:8], digest[8:10], digest[10:16]), nil
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
package legacyid
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestUUIDIsStableAndEntityScoped(t *testing.T) {
|
||||
department, err := UUID("python-v1", "departments", "42")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if department != "028ac0aa-75d4-5ce6-b20d-15a343df445e" {
|
||||
t.Fatalf("unexpected deterministic UUID %q", department)
|
||||
}
|
||||
user, err := UUID("python-v1", "users", "42")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if user != "de8453de-93b5-5ab9-8284-ed099894296b" || user == department {
|
||||
t.Fatalf("entity namespace was not applied: %q", user)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUUIDRejectsIncompleteIdentity(t *testing.T) {
|
||||
if _, err := UUID("python-v1", "users", " "); err == nil {
|
||||
t.Fatal("expected empty legacy ID to be rejected")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package migrate
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
const advisoryLockID int64 = 6720240805
|
||||
|
||||
type Migration struct {
|
||||
Version string
|
||||
Filename string
|
||||
SQL string
|
||||
Checksum string
|
||||
}
|
||||
|
||||
func Load(directory string) ([]Migration, error) {
|
||||
entries, err := os.ReadDir(directory)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read migration directory: %w", err)
|
||||
}
|
||||
var filenames []string
|
||||
for _, entry := range entries {
|
||||
if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".sql") {
|
||||
filenames = append(filenames, entry.Name())
|
||||
}
|
||||
}
|
||||
sort.Strings(filenames)
|
||||
migrations := make([]Migration, 0, len(filenames))
|
||||
for _, filename := range filenames {
|
||||
body, err := os.ReadFile(filepath.Join(directory, filename))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read migration %s: %w", filename, err)
|
||||
}
|
||||
version, _, ok := strings.Cut(filename, "_")
|
||||
if !ok || version == "" {
|
||||
return nil, fmt.Errorf("migration %s must start with a version and underscore", filename)
|
||||
}
|
||||
digest := sha256.Sum256(body)
|
||||
migrations = append(migrations, Migration{
|
||||
Version: version, Filename: filename, SQL: string(body), Checksum: hex.EncodeToString(digest[:]),
|
||||
})
|
||||
}
|
||||
return migrations, nil
|
||||
}
|
||||
|
||||
func Apply(ctx context.Context, pool *pgxpool.Pool, migrations []Migration) error {
|
||||
connection, err := pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("acquire migration connection: %w", err)
|
||||
}
|
||||
defer connection.Release()
|
||||
if _, err := connection.Exec(ctx, "SELECT pg_advisory_lock($1)", advisoryLockID); err != nil {
|
||||
return fmt.Errorf("lock migrations: %w", err)
|
||||
}
|
||||
defer func() { _, _ = connection.Exec(context.Background(), "SELECT pg_advisory_unlock($1)", advisoryLockID) }()
|
||||
|
||||
if _, err := connection.Exec(ctx, `
|
||||
CREATE SCHEMA IF NOT EXISTS gateway;
|
||||
CREATE TABLE IF NOT EXISTS gateway.schema_migrations (
|
||||
version text PRIMARY KEY,
|
||||
filename text NOT NULL,
|
||||
checksum text NOT NULL,
|
||||
applied_at timestamptz NOT NULL DEFAULT clock_timestamp()
|
||||
)`); err != nil {
|
||||
return fmt.Errorf("initialize migration table: %w", err)
|
||||
}
|
||||
|
||||
for _, migration := range migrations {
|
||||
var existingChecksum string
|
||||
err := connection.QueryRow(ctx,
|
||||
"SELECT checksum FROM gateway.schema_migrations WHERE version = $1", migration.Version,
|
||||
).Scan(&existingChecksum)
|
||||
if err == nil {
|
||||
if existingChecksum != migration.Checksum {
|
||||
return fmt.Errorf("migration %s checksum changed after application", migration.Filename)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if !errors.Is(err, pgx.ErrNoRows) {
|
||||
return fmt.Errorf("check migration %s: %w", migration.Filename, err)
|
||||
}
|
||||
transaction, err := connection.Begin(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("begin migration %s: %w", migration.Filename, err)
|
||||
}
|
||||
if _, err := transaction.Exec(ctx, migration.SQL); err != nil {
|
||||
_ = transaction.Rollback(ctx)
|
||||
return fmt.Errorf("apply migration %s: %w", migration.Filename, err)
|
||||
}
|
||||
if _, err := transaction.Exec(ctx,
|
||||
"INSERT INTO gateway.schema_migrations (version, filename, checksum) VALUES ($1, $2, $3)",
|
||||
migration.Version, migration.Filename, migration.Checksum,
|
||||
); err != nil {
|
||||
_ = transaction.Rollback(ctx)
|
||||
return fmt.Errorf("record migration %s: %w", migration.Filename, err)
|
||||
}
|
||||
if err := transaction.Commit(ctx); err != nil {
|
||||
return fmt.Errorf("commit migration %s: %w", migration.Filename, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user