fix(control): 加固幂等恢复与配置并发
This commit is contained in:
55 files changed
+5477
-233
No files matched your search
+18
-1
@@ -32,6 +32,7 @@ type App struct {
|
||||
runtimeState RuntimeStateLoader
|
||||
controlRuntime ControlRuntime
|
||||
startupBootstrap StartupBootstrap
|
||||
startupRecovery StartupRecovery
|
||||
requestLogs RequestLogRuntime
|
||||
listen func(network, address string) (net.Listener, error)
|
||||
|
||||
@@ -53,6 +54,11 @@ type StartupBootstrap interface {
|
||||
EnsureInitialState(context.Context) error
|
||||
}
|
||||
|
||||
// StartupRecovery restores durable post-commit control-plane side effects.
|
||||
type StartupRecovery interface {
|
||||
DrainCommittedOperations(context.Context) error
|
||||
}
|
||||
|
||||
// ControlRuntime runs background control-plane maintenance.
|
||||
type ControlRuntime interface {
|
||||
Run(context.Context)
|
||||
@@ -73,6 +79,7 @@ type AppParams struct {
|
||||
Encryption encryption.Service
|
||||
DB *gorm.DB
|
||||
StartupBootstrap StartupBootstrap
|
||||
StartupRecovery StartupRecovery `optional:"true"`
|
||||
RuntimeState RuntimeStateLoader
|
||||
ControlRuntime ControlRuntime
|
||||
RequestLogs RequestLogRuntime
|
||||
@@ -106,6 +113,7 @@ func NewApp(params AppParams) *App {
|
||||
runtimeState: params.RuntimeState,
|
||||
controlRuntime: params.ControlRuntime,
|
||||
startupBootstrap: params.StartupBootstrap,
|
||||
startupRecovery: params.StartupRecovery,
|
||||
requestLogs: params.RequestLogs,
|
||||
listen: net.Listen,
|
||||
serveErrors: make(chan error, 1),
|
||||
@@ -123,7 +131,11 @@ func (a *App) Start() error {
|
||||
if err := i18n.Init(); err != nil {
|
||||
return fmt.Errorf("initialize i18n: %w", err)
|
||||
}
|
||||
if err := storage.AutoMigrate(a.db); err != nil {
|
||||
if err := storage.AutoMigrateWithEncryption(
|
||||
a.db,
|
||||
a.encryption,
|
||||
a.config.DataDir,
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := a.startupBootstrap.EnsureInitialState(context.Background()); err != nil {
|
||||
@@ -132,6 +144,11 @@ func (a *App) Start() error {
|
||||
if err := a.runtimeState.Load(context.Background()); err != nil {
|
||||
return fmt.Errorf("load runtime state: %w", err)
|
||||
}
|
||||
if a.startupRecovery != nil {
|
||||
if err := a.startupRecovery.DrainCommittedOperations(context.Background()); err != nil {
|
||||
return fmt.Errorf("recover committed control operations: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
address := net.JoinHostPort(a.config.Server.Host, strconv.Itoa(a.config.Server.Port))
|
||||
listener, err := a.listen("tcp", address)
|
||||
|
||||
+127
-1
@@ -45,6 +45,12 @@ func noopStartupBootstrap(context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
type startupRecoveryFunc func(context.Context) error
|
||||
|
||||
func (f startupRecoveryFunc) DrainCommittedOperations(ctx context.Context) error {
|
||||
return f(ctx)
|
||||
}
|
||||
|
||||
type requestLogRuntimeFake struct {
|
||||
startFunc func() error
|
||||
stopFunc func(context.Context) error
|
||||
@@ -302,7 +308,7 @@ func TestAppStartMigratesDatabaseAndServesHTTP(t *testing.T) {
|
||||
}
|
||||
for _, table := range []string{
|
||||
"groups", "upstream_keys", "access_keys", "request_logs", "usage_stats",
|
||||
"model_prices", "system_settings", "jobs", "schema_info",
|
||||
"model_prices", "system_settings", "jobs", "control_operations", "schema_info",
|
||||
} {
|
||||
if !db.Migrator().HasTable(table) {
|
||||
t.Errorf("Start() did not migrate table %q", table)
|
||||
@@ -319,6 +325,78 @@ func TestAppStartMigratesDatabaseAndServesHTTP(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppStartUsesEncryptionForSchemaV1UpgradeBeforeRuntimeLoad(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("storage.Open() error = %v", err)
|
||||
}
|
||||
keyService, err := encryption.NewService("schema-v1-app-test-master-key")
|
||||
if err != nil {
|
||||
t.Fatalf("encryption.NewService() error = %v", err)
|
||||
}
|
||||
for _, statement := range []string{
|
||||
`CREATE TABLE schema_info (version integer PRIMARY KEY)`,
|
||||
`INSERT INTO schema_info(version) VALUES (1)`,
|
||||
`CREATE TABLE access_keys (
|
||||
id integer PRIMARY KEY AUTOINCREMENT,
|
||||
name varchar(255) NOT NULL,
|
||||
key_value text NOT NULL,
|
||||
key_hash varchar(128) NOT NULL UNIQUE,
|
||||
status varchar(32) NOT NULL DEFAULT 'active'
|
||||
CHECK (status IN ('active','disabled')),
|
||||
filters json,
|
||||
rpm_limit integer NOT NULL DEFAULT 0,
|
||||
daily_cost_limit real NOT NULL DEFAULT 0,
|
||||
monthly_cost_limit real NOT NULL DEFAULT 0,
|
||||
created_at datetime,
|
||||
updated_at datetime
|
||||
)`,
|
||||
} {
|
||||
if err := db.Exec(statement).Error; err != nil {
|
||||
t.Fatalf("create schema v1 fixture: %v", err)
|
||||
}
|
||||
}
|
||||
const plaintext = "sk-gl-00112233445566778899aabbccdd7f2a"
|
||||
ciphertext, err := keyService.Encrypt(plaintext)
|
||||
if err != nil {
|
||||
t.Fatalf("Encrypt() error = %v", err)
|
||||
}
|
||||
if err := db.Exec(`INSERT INTO access_keys
|
||||
(name,key_value,key_hash,status,filters,rpm_limit,daily_cost_limit,monthly_cost_limit)
|
||||
VALUES (?,?,?,?,?,?,?,?)`,
|
||||
"legacy", ciphertext, keyService.Hash(plaintext), "active",
|
||||
`{"groups":[],"protocols":[],"models":[]}`, 0, 0, 0,
|
||||
).Error; err != nil {
|
||||
t.Fatalf("insert legacy AccessKey: %v", err)
|
||||
}
|
||||
|
||||
loadErr := errors.New("stop after schema upgrade")
|
||||
application := NewApp(AppParams{
|
||||
Engine: mustNewEngine(t),
|
||||
Config: testConfig(t),
|
||||
Encryption: keyService,
|
||||
DB: db,
|
||||
StartupBootstrap: startupBootstrapFunc(noopStartupBootstrap),
|
||||
RuntimeState: runtimeStateLoaderFunc(func(context.Context) error {
|
||||
return loadErr
|
||||
}),
|
||||
ControlRuntime: newControlRuntimeFake(nil, false),
|
||||
RequestLogs: newRequestLogRuntimeFake(nil, nil),
|
||||
})
|
||||
cleanupApp(t, application)
|
||||
|
||||
if err := application.Start(); !errors.Is(err, loadErr) {
|
||||
t.Fatalf("Start() error = %v, want runtime load sentinel", err)
|
||||
}
|
||||
var suffix string
|
||||
if err := db.Raw("SELECT key_suffix FROM access_keys WHERE name = 'legacy'").Scan(&suffix).Error; err != nil {
|
||||
t.Fatalf("read migrated suffix: %v", err)
|
||||
}
|
||||
if suffix != "7f2a" {
|
||||
t.Fatalf("migrated suffix = %q, want 7f2a", suffix)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppStartBootstrapsAfterMigrationBeforeRuntimeLoad(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
@@ -362,6 +440,54 @@ func TestAppStartBootstrapsAfterMigrationBeforeRuntimeLoad(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppStartDrainsCommittedOperationsAfterRuntimeLoadBeforeListen(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("storage.Open() error = %v", err)
|
||||
}
|
||||
keyService, err := encryption.NewService("test-master-key")
|
||||
if err != nil {
|
||||
t.Fatalf("encryption.NewService() error = %v", err)
|
||||
}
|
||||
|
||||
var order []string
|
||||
recoveryErr := errors.New("durable recovery failed")
|
||||
application := NewApp(AppParams{
|
||||
Engine: mustNewEngine(t),
|
||||
Config: testConfig(t),
|
||||
Encryption: keyService,
|
||||
DB: db,
|
||||
StartupBootstrap: startupBootstrapFunc(noopStartupBootstrap),
|
||||
RuntimeState: runtimeStateLoaderFunc(func(context.Context) error {
|
||||
order = append(order, "load")
|
||||
return nil
|
||||
}),
|
||||
StartupRecovery: startupRecoveryFunc(func(context.Context) error {
|
||||
order = append(order, "recover")
|
||||
return recoveryErr
|
||||
}),
|
||||
ControlRuntime: newControlRuntimeFake(nil, false),
|
||||
RequestLogs: newRequestLogRuntimeFake(nil, nil),
|
||||
})
|
||||
listenCalled := false
|
||||
application.listen = func(string, string) (net.Listener, error) {
|
||||
listenCalled = true
|
||||
return nil, errors.New("unexpected listen")
|
||||
}
|
||||
cleanupApp(t, application)
|
||||
|
||||
err = application.Start()
|
||||
if !errors.Is(err, recoveryErr) {
|
||||
t.Fatalf("Start() error = %v, want recovery failure", err)
|
||||
}
|
||||
if want := []string{"load", "recover"}; !slices.Equal(order, want) {
|
||||
t.Fatalf("startup order = %#v, want %#v", order, want)
|
||||
}
|
||||
if listenCalled {
|
||||
t.Fatal("Start() listened before committed operations drained")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAppStartRejectsBootstrapFailureBeforeRuntimeLoadAndListen(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
|
||||
@@ -111,6 +111,7 @@ func BuildContainer() (*dig.Container, error) {
|
||||
gateway.NewHandler,
|
||||
control.NewService,
|
||||
func(service *control.Service) app.StartupBootstrap { return service },
|
||||
func(service *control.Service) app.StartupRecovery { return service },
|
||||
control.NewServer,
|
||||
func(db *gorm.DB, manager *state.Manager, registry *state.KeyRegistry) app.RuntimeStateLoader {
|
||||
return stateloader.New(db, manager, registry)
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"gpt-load/internal/platform/canonicaljson"
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/protocol"
|
||||
"gpt-load/internal/state"
|
||||
stateloader "gpt-load/internal/state/loader"
|
||||
)
|
||||
|
||||
type accessKeyCreateDigestBody struct {
|
||||
Name string `json:"name"`
|
||||
Filters AccessKeyFilters `json:"filters"`
|
||||
RPMLimit int64 `json:"rpm_limit"`
|
||||
}
|
||||
|
||||
func (s *Service) CreateAccessKeyIdempotent(
|
||||
ctx context.Context,
|
||||
idempotencyKey string,
|
||||
request AccessKeyCreateRequest,
|
||||
) (AccessKeyCreateResult, error) {
|
||||
name, err := normalizeAccessKeyName(request.Name)
|
||||
if err != nil {
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
filters, err := normalizeAccessKeyFilters(request.Filters)
|
||||
if err != nil {
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
if err := validateAccessKeyProtocolsForCreate(filters.Protocols); err != nil {
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
rpmLimit, err := normalizeRPMLimit(request.RPMLimit, 0)
|
||||
if err != nil {
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
digestFilters := canonicalAccessKeyFilterSet(filters)
|
||||
canonicalBody, err := canonicalIdempotencyBody(accessKeyCreateDigestBody{
|
||||
Name: name, Filters: digestFilters, RPMLimit: rpmLimit,
|
||||
})
|
||||
if err != nil {
|
||||
return AccessKeyCreateResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
digest, err := buildIdempotencyDigest(idempotencyDigestInput{
|
||||
Version: 1,
|
||||
Method: "POST",
|
||||
OperationKind: operationKindAccessKeyCreate,
|
||||
PathTemplate: "/api/access-keys",
|
||||
ResourceLocator: "new",
|
||||
AuthScopeID: idempotencyAuthScopeID,
|
||||
CanonicalBody: canonicalBody,
|
||||
})
|
||||
if err != nil {
|
||||
return AccessKeyCreateResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
|
||||
operationResult, err := s.executeIdempotentOperation(ctx, idempotentOperationInput{
|
||||
IdempotencyKey: idempotencyKey,
|
||||
DigestVersion: 1,
|
||||
RequestDigest: digest.Digest,
|
||||
Kind: operationKindAccessKeyCreate,
|
||||
Mutate: func(tx *gorm.DB) (idempotentMutationResult, error) {
|
||||
if err := validateFilterGroupReferences(tx, filters.Groups); err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
row, plaintext, err := s.newAccessKeyRow(name, filters, rpmLimit)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
if err := tx.Create(&row).Error; err != nil {
|
||||
return idempotentMutationResult{}, app_errors.ParseDBError(err)
|
||||
}
|
||||
metadata, err := mapAccessKeyMetadataRow(accessKeyMetadataRow{
|
||||
ID: row.ID, Name: row.Name, KeySuffix: row.KeySuffix,
|
||||
Status: row.Status, Filters: row.Filters, RPMLimit: row.RPMLimit,
|
||||
CreatedAt: row.CreatedAt, UpdatedAt: row.UpdatedAt,
|
||||
})
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
input, err := stateloader.BuildCompileInput(ctx, tx)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
if _, err := state.Compile(input); err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
canonicalResult, err := canonicaljson.Marshal(metadata)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, fmt.Errorf(
|
||||
"encode AccessKey operation result: %w",
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
return idempotentMutationResult{
|
||||
ResourceIdentity: fmt.Sprintf("access-key:%d", row.ID),
|
||||
CanonicalResult: canonicalResult,
|
||||
Ephemeral: plaintext,
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
var metadata AccessKeyMetadata
|
||||
if err := json.Unmarshal(operationResult.CanonicalResult, &metadata); err != nil {
|
||||
return AccessKeyCreateResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
result := AccessKeyCreateResult{
|
||||
AccessKeyMetadata: metadata,
|
||||
Replayed: operationResult.Replayed,
|
||||
}
|
||||
if plaintext, ok := operationResult.Ephemeral.(string); ok &&
|
||||
!operationResult.Replayed {
|
||||
result.Key = plaintext
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func canonicalAccessKeyFilterSet(filters AccessKeyFilters) AccessKeyFilters {
|
||||
result := AccessKeyFilters{
|
||||
Groups: append([]uint(nil), filters.Groups...),
|
||||
Protocols: append([]protocol.Protocol(nil), filters.Protocols...),
|
||||
Models: append([]string(nil), filters.Models...),
|
||||
}
|
||||
sort.Slice(result.Groups, func(left, right int) bool {
|
||||
return result.Groups[left] < result.Groups[right]
|
||||
})
|
||||
sort.Slice(result.Protocols, func(left, right int) bool {
|
||||
return string(result.Protocols[left]) < string(result.Protocols[right])
|
||||
})
|
||||
sort.Strings(result.Models)
|
||||
return result
|
||||
}
|
||||
+269
-61
@@ -9,6 +9,7 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
@@ -55,13 +56,47 @@ type AccessKeyUpdateRequest struct {
|
||||
RPMLimit OptionalRPMLimit `json:"rpm_limit"`
|
||||
}
|
||||
|
||||
type AccessKeyResponse struct {
|
||||
ID uint `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Key string `json:"key"`
|
||||
Status state.AccessKeyStatus `json:"status"`
|
||||
Filters AccessKeyFilters `json:"filters"`
|
||||
RPMLimit int64 `json:"rpm_limit"`
|
||||
type AccessKeyMetadata struct {
|
||||
ID uint `json:"id"`
|
||||
Name string `json:"name"`
|
||||
MaskedKey string `json:"masked_key"`
|
||||
Status state.AccessKeyStatus `json:"status"`
|
||||
Filters AccessKeyFilters `json:"filters"`
|
||||
RPMLimit int64 `json:"rpm_limit"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
type AccessKeyCreateResult struct {
|
||||
AccessKeyMetadata
|
||||
Key string `json:"key,omitempty"`
|
||||
Replayed bool `json:"replayed"`
|
||||
}
|
||||
|
||||
// AccessKeyResponse is retained as the create-response compatibility name.
|
||||
type AccessKeyResponse = AccessKeyCreateResult
|
||||
|
||||
type AccessKeyOption struct {
|
||||
ID uint `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Status state.AccessKeyStatus `json:"status"`
|
||||
}
|
||||
|
||||
type AccessKeyRevealResult struct {
|
||||
ID uint `json:"id"`
|
||||
Key string `json:"key"`
|
||||
RevealedAt time.Time `json:"revealed_at"`
|
||||
}
|
||||
|
||||
type accessKeyMetadataRow struct {
|
||||
ID uint
|
||||
Name string
|
||||
KeySuffix string
|
||||
Status string
|
||||
Filters models.JSON
|
||||
RPMLimit int64
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
|
||||
const accessKeyPrefix = "sk-gl-"
|
||||
@@ -85,36 +120,37 @@ func (s *Service) newAccessKeyRow(
|
||||
return models.AccessKey{}, "", fmt.Errorf("encrypt access key: %w", err)
|
||||
}
|
||||
return models.AccessKey{
|
||||
Name: name,
|
||||
KeyValue: ciphertext,
|
||||
KeyHash: s.encryption.Hash(plaintext),
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(encodedFilters),
|
||||
RPMLimit: rpmLimit,
|
||||
Name: name,
|
||||
KeyValue: ciphertext,
|
||||
KeyHash: s.encryption.Hash(plaintext),
|
||||
KeySuffix: plaintext[len(plaintext)-4:],
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(encodedFilters),
|
||||
RPMLimit: rpmLimit,
|
||||
}, plaintext, nil
|
||||
}
|
||||
|
||||
func (s *Service) CreateAccessKey(
|
||||
ctx context.Context,
|
||||
request AccessKeyCreateRequest,
|
||||
) (AccessKeyResponse, error) {
|
||||
) (AccessKeyCreateResult, error) {
|
||||
name, err := normalizeAccessKeyName(request.Name)
|
||||
if err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
filters, err := normalizeAccessKeyFilters(request.Filters)
|
||||
if err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
if err := validateAccessKeyProtocolsForCreate(filters.Protocols); err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
rpmLimit, err := normalizeRPMLimit(request.RPMLimit, 0)
|
||||
if err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
|
||||
var result AccessKeyResponse
|
||||
var result AccessKeyCreateResult
|
||||
_, err = s.writeConfig(ctx, func(tx *gorm.DB) error {
|
||||
if err := validateFilterGroupReferences(tx, filters.Groups); err != nil {
|
||||
return err
|
||||
@@ -126,41 +162,46 @@ func (s *Service) CreateAccessKey(
|
||||
if err := tx.Create(&row).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
result = AccessKeyResponse{
|
||||
ID: row.ID, Name: row.Name, Key: plaintext,
|
||||
Status: state.AccessKeyStatusActive, Filters: filters, RPMLimit: row.RPMLimit,
|
||||
metadata, err := mapAccessKeyMetadataRow(accessKeyMetadataRow{
|
||||
ID: row.ID, Name: row.Name, KeySuffix: row.KeySuffix,
|
||||
Status: row.Status, Filters: row.Filters, RPMLimit: row.RPMLimit,
|
||||
CreatedAt: row.CreatedAt, UpdatedAt: row.UpdatedAt,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
result = AccessKeyCreateResult{
|
||||
AccessKeyMetadata: metadata,
|
||||
Key: plaintext,
|
||||
}
|
||||
return nil
|
||||
}, nil)
|
||||
if err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyCreateResult{}, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *Service) ListAccessKeys(ctx context.Context) ([]AccessKeyResponse, error) {
|
||||
var rows []models.AccessKey
|
||||
if err := s.db.WithContext(ctx).Order("id ASC").Find(&rows).Error; err != nil {
|
||||
func (s *Service) ListAccessKeys(ctx context.Context) ([]AccessKeyMetadata, error) {
|
||||
var rows []accessKeyMetadataRow
|
||||
if err := s.db.WithContext(ctx).
|
||||
Model(&models.AccessKey{}).
|
||||
Select(
|
||||
"id", "name", "key_suffix", "status", "filters", "rpm_limit",
|
||||
"created_at", "updated_at",
|
||||
).
|
||||
Order("id ASC").
|
||||
Scan(&rows).Error; err != nil {
|
||||
return nil, app_errors.ParseDBError(err)
|
||||
}
|
||||
|
||||
result := make([]AccessKeyResponse, 0, len(rows))
|
||||
result := make([]AccessKeyMetadata, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
status := state.AccessKeyStatus(row.Status)
|
||||
if status != state.AccessKeyStatusActive && status != state.AccessKeyStatusDisabled {
|
||||
return nil, fmt.Errorf("access key %d has invalid status", row.ID)
|
||||
}
|
||||
filters, err := decodeStoredAccessKeyFilters(row.Filters)
|
||||
metadata, err := mapAccessKeyMetadataRow(row)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("decode access key %d filters: %w", row.ID, err)
|
||||
return nil, err
|
||||
}
|
||||
plaintext, err := s.encryption.Decrypt(row.KeyValue)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("decrypt access key %d: %w", row.ID, err)
|
||||
}
|
||||
result = append(result, AccessKeyResponse{
|
||||
ID: row.ID, Name: row.Name, Key: plaintext, Status: status, Filters: filters, RPMLimit: row.RPMLimit,
|
||||
})
|
||||
result = append(result, metadata)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -169,51 +210,60 @@ func (s *Service) UpdateAccessKey(
|
||||
ctx context.Context,
|
||||
id uint,
|
||||
request AccessKeyUpdateRequest,
|
||||
) (AccessKeyResponse, error) {
|
||||
) (AccessKeyMetadata, error) {
|
||||
if id == 0 || (request.Name == nil && request.Status == nil && request.Filters == nil && !request.RPMLimit.Set) {
|
||||
return AccessKeyResponse{}, app_errors.ErrBadRequest
|
||||
return AccessKeyMetadata{}, app_errors.ErrBadRequest
|
||||
}
|
||||
if _, err := normalizeRPMLimit(request.RPMLimit, 0); err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyMetadata{}, err
|
||||
}
|
||||
|
||||
var name *string
|
||||
if request.Name != nil {
|
||||
normalized, err := normalizeAccessKeyName(*request.Name)
|
||||
if err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyMetadata{}, err
|
||||
}
|
||||
name = &normalized
|
||||
}
|
||||
if request.Status != nil &&
|
||||
*request.Status != state.AccessKeyStatusActive &&
|
||||
*request.Status != state.AccessKeyStatusDisabled {
|
||||
return AccessKeyResponse{}, app_errors.ErrValidation
|
||||
return AccessKeyMetadata{}, app_errors.ErrValidation
|
||||
}
|
||||
var filters *AccessKeyFilters
|
||||
var encodedFilters []byte
|
||||
if request.Filters != nil {
|
||||
normalized, err := normalizeAccessKeyFilters(request.Filters)
|
||||
if err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyMetadata{}, err
|
||||
}
|
||||
encoded, err := json.Marshal(normalized)
|
||||
if err != nil {
|
||||
return AccessKeyResponse{}, fmt.Errorf("encode access key filters: %w", err)
|
||||
return AccessKeyMetadata{}, fmt.Errorf("encode access key filters: %w", err)
|
||||
}
|
||||
filters = &normalized
|
||||
encodedFilters = encoded
|
||||
}
|
||||
|
||||
var result AccessKeyResponse
|
||||
var result AccessKeyMetadata
|
||||
_, err := s.writeConfig(ctx, func(tx *gorm.DB) error {
|
||||
var row models.AccessKey
|
||||
if err := tx.First(&row, id).Error; err != nil {
|
||||
var row accessKeyMetadataRow
|
||||
if err := tx.Model(&models.AccessKey{}).
|
||||
Select(
|
||||
"id", "name", "key_suffix", "status", "filters", "rpm_limit",
|
||||
"created_at", "updated_at",
|
||||
).
|
||||
Where("id = ?", id).
|
||||
Take(&row).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
plaintext, err := s.encryption.Decrypt(row.KeyValue)
|
||||
if err != nil {
|
||||
return fmt.Errorf("decrypt access key %d: %w", row.ID, err)
|
||||
if !validAccessKeySuffix(row.KeySuffix) {
|
||||
return fmt.Errorf(
|
||||
"access key %d has invalid persisted suffix: %w",
|
||||
row.ID,
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
currentFilters, err := decodeStoredAccessKeyFilters(row.Filters)
|
||||
if err != nil {
|
||||
@@ -223,6 +273,13 @@ func (s *Service) UpdateAccessKey(
|
||||
if err := validateAccessKeyProtocolUpdate(currentFilters.Protocols, filters.Protocols); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := validateAccessKeyGroupUpdate(
|
||||
tx,
|
||||
currentFilters.Groups,
|
||||
filters.Groups,
|
||||
); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
status := state.AccessKeyStatus(row.Status)
|
||||
if status != state.AccessKeyStatusActive && status != state.AccessKeyStatusDisabled {
|
||||
@@ -239,9 +296,6 @@ func (s *Service) UpdateAccessKey(
|
||||
updates["status"] = string(status)
|
||||
}
|
||||
if filters != nil {
|
||||
if err := validateFilterGroupReferences(tx, filters.Groups); err != nil {
|
||||
return err
|
||||
}
|
||||
currentFilters = *filters
|
||||
updates["filters"] = models.JSON(encodedFilters)
|
||||
}
|
||||
@@ -249,20 +303,141 @@ func (s *Service) UpdateAccessKey(
|
||||
row.RPMLimit = request.RPMLimit.Value
|
||||
updates["rpm_limit"] = row.RPMLimit
|
||||
}
|
||||
if err := tx.Model(&row).Updates(updates).Error; err != nil {
|
||||
if err := tx.Model(&models.AccessKey{}).
|
||||
Where("id = ?", row.ID).
|
||||
Updates(updates).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
result = AccessKeyResponse{
|
||||
ID: row.ID, Name: row.Name, Key: plaintext, Status: status, Filters: currentFilters, RPMLimit: row.RPMLimit,
|
||||
if err := tx.Model(&models.AccessKey{}).
|
||||
Select(
|
||||
"id", "name", "key_suffix", "status", "filters", "rpm_limit",
|
||||
"created_at", "updated_at",
|
||||
).
|
||||
Where("id = ?", row.ID).
|
||||
Take(&row).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
return nil
|
||||
result, err = mapAccessKeyMetadataRow(row)
|
||||
return err
|
||||
}, nil)
|
||||
if err != nil {
|
||||
return AccessKeyResponse{}, err
|
||||
return AccessKeyMetadata{}, err
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *Service) ListAccessKeyOptions(ctx context.Context) ([]AccessKeyOption, error) {
|
||||
var rows []AccessKeyOption
|
||||
if err := s.db.WithContext(ctx).
|
||||
Model(&models.AccessKey{}).
|
||||
Select("id", "name", "status").
|
||||
Order("id ASC").
|
||||
Scan(&rows).Error; err != nil {
|
||||
return nil, app_errors.ParseDBError(err)
|
||||
}
|
||||
for _, row := range rows {
|
||||
if row.Status != state.AccessKeyStatusActive &&
|
||||
row.Status != state.AccessKeyStatusDisabled {
|
||||
return nil, fmt.Errorf(
|
||||
"access key %d has invalid status: %w",
|
||||
row.ID,
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
}
|
||||
return rows, nil
|
||||
}
|
||||
|
||||
func (s *Service) RevealAccessKey(
|
||||
ctx context.Context,
|
||||
id uint,
|
||||
) (AccessKeyRevealResult, error) {
|
||||
if id == 0 {
|
||||
return AccessKeyRevealResult{}, app_errors.ErrBadRequest
|
||||
}
|
||||
var row struct {
|
||||
ID uint
|
||||
KeyValue string
|
||||
}
|
||||
if err := s.db.WithContext(ctx).
|
||||
Model(&models.AccessKey{}).
|
||||
Select("id", "key_value").
|
||||
Where("id = ?", id).
|
||||
Take(&row).Error; err != nil {
|
||||
return AccessKeyRevealResult{}, app_errors.ParseDBError(err)
|
||||
}
|
||||
plaintext, err := s.encryption.Decrypt(row.KeyValue)
|
||||
if err != nil || !validAccessKeyPlaintext(plaintext) {
|
||||
return AccessKeyRevealResult{}, fmt.Errorf(
|
||||
"reveal access key %d: %w",
|
||||
row.ID,
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
return AccessKeyRevealResult{
|
||||
ID: row.ID,
|
||||
Key: plaintext,
|
||||
RevealedAt: s.now().UTC(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func mapAccessKeyMetadataRow(row accessKeyMetadataRow) (AccessKeyMetadata, error) {
|
||||
status := state.AccessKeyStatus(row.Status)
|
||||
if status != state.AccessKeyStatusActive && status != state.AccessKeyStatusDisabled {
|
||||
return AccessKeyMetadata{}, fmt.Errorf(
|
||||
"access key %d has invalid status: %w",
|
||||
row.ID,
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
if !validAccessKeySuffix(row.KeySuffix) {
|
||||
return AccessKeyMetadata{}, fmt.Errorf(
|
||||
"access key %d has invalid persisted suffix: %w",
|
||||
row.ID,
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
filters, err := decodeStoredAccessKeyFilters(row.Filters)
|
||||
if err != nil {
|
||||
return AccessKeyMetadata{}, fmt.Errorf(
|
||||
"decode access key %d filters: %w",
|
||||
row.ID,
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
return AccessKeyMetadata{
|
||||
ID: row.ID, Name: row.Name,
|
||||
MaskedKey: accessKeyPrefix + "••••••••" + row.KeySuffix,
|
||||
Status: status, Filters: filters, RPMLimit: row.RPMLimit,
|
||||
CreatedAt: row.CreatedAt.UTC(), UpdatedAt: row.UpdatedAt.UTC(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func validAccessKeySuffix(value string) bool {
|
||||
if len(value) != 4 {
|
||||
return false
|
||||
}
|
||||
for _, character := range []byte(value) {
|
||||
if character < '0' || character > '9' && character < 'a' || character > 'f' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func validAccessKeyPlaintext(value string) bool {
|
||||
if len(value) != len(accessKeyPrefix)+32 ||
|
||||
!strings.HasPrefix(value, accessKeyPrefix) {
|
||||
return false
|
||||
}
|
||||
for _, character := range []byte(strings.TrimPrefix(value, accessKeyPrefix)) {
|
||||
if character < '0' || character > '9' && character < 'a' || character > 'f' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *Service) DeleteAccessKey(ctx context.Context, id uint) error {
|
||||
if id == 0 {
|
||||
return app_errors.ErrBadRequest
|
||||
@@ -385,6 +560,39 @@ func validateFilterGroupReferences(tx *gorm.DB, groupIDs []uint) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateAccessKeyGroupUpdate(
|
||||
tx *gorm.DB,
|
||||
current []uint,
|
||||
requested []uint,
|
||||
) error {
|
||||
if len(requested) == 0 {
|
||||
return nil
|
||||
}
|
||||
currentSet := make(map[uint]struct{}, len(current))
|
||||
for _, groupID := range current {
|
||||
currentSet[groupID] = struct{}{}
|
||||
}
|
||||
var existing []uint
|
||||
if err := tx.Model(&models.Group{}).
|
||||
Where("id IN ?", requested).
|
||||
Pluck("id", &existing).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
existingSet := make(map[uint]struct{}, len(existing))
|
||||
for _, groupID := range existing {
|
||||
existingSet[groupID] = struct{}{}
|
||||
}
|
||||
for _, groupID := range requested {
|
||||
if _, exists := existingSet[groupID]; exists {
|
||||
continue
|
||||
}
|
||||
if _, historical := currentSet[groupID]; !historical {
|
||||
return app_errors.ErrValidation
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func decodeStoredAccessKeyFilters(raw models.JSON) (AccessKeyFilters, error) {
|
||||
if len(raw) == 0 {
|
||||
return normalizeAccessKeyFilters(nil)
|
||||
|
||||
@@ -0,0 +1,268 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gpt-load/internal/platform/encryption"
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/protocol"
|
||||
"gpt-load/internal/state"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
type decryptCountingEncryption struct {
|
||||
encryption.Service
|
||||
decryptCalls int
|
||||
}
|
||||
|
||||
func (spy *decryptCountingEncryption) Decrypt(ciphertext string) (string, error) {
|
||||
spy.decryptCalls++
|
||||
return spy.Service.Decrypt(ciphertext)
|
||||
}
|
||||
|
||||
func TestAccessKeyMetadataListAndUpdateNeverDecryptCiphertext(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.random = bytes.NewReader(make([]byte, 16))
|
||||
created, err := fixture.service.CreateAccessKey(
|
||||
t.Context(),
|
||||
AccessKeyCreateRequest{Name: "metadata-only"},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateAccessKey() error = %v", err)
|
||||
}
|
||||
spy := &decryptCountingEncryption{Service: fixture.encryption}
|
||||
fixture.service.encryption = spy
|
||||
if err := fixture.db.Model(&models.AccessKey{}).
|
||||
Where("id = ?", created.ID).
|
||||
UpdateColumn("key_value", "intentionally-corrupt").Error; err != nil {
|
||||
t.Fatalf("corrupt ciphertext: %v", err)
|
||||
}
|
||||
|
||||
list, err := fixture.service.ListAccessKeys(t.Context())
|
||||
if err != nil {
|
||||
t.Fatalf("ListAccessKeys() error = %v", err)
|
||||
}
|
||||
if len(list) != 1 || list[0].ID != created.ID ||
|
||||
list[0].MaskedKey != "sk-gl-••••••••0000" {
|
||||
t.Fatalf("ListAccessKeys() = %#v", list)
|
||||
}
|
||||
encoded, err := json.Marshal(list)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal metadata: %v", err)
|
||||
}
|
||||
if strings.Contains(string(encoded), `"key"`) ||
|
||||
strings.Contains(string(encoded), "intentionally-corrupt") {
|
||||
t.Fatalf("metadata response exposes secret field: %s", encoded)
|
||||
}
|
||||
|
||||
updated, err := fixture.service.UpdateAccessKey(
|
||||
t.Context(),
|
||||
created.ID,
|
||||
AccessKeyUpdateRequest{Name: stringPointer("renamed")},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("UpdateAccessKey() error = %v", err)
|
||||
}
|
||||
if updated.Name != "renamed" || updated.MaskedKey != list[0].MaskedKey {
|
||||
t.Fatalf("UpdateAccessKey() = %#v", updated)
|
||||
}
|
||||
if spy.decryptCalls != 0 {
|
||||
t.Fatalf("metadata list/update decrypt calls = %d, want 0", spy.decryptCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreateAccessKeyIdempotentReplaysMetadataWithoutPersistingSecret(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.random = bytes.NewReader(make([]byte, 16))
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x61}, 16))
|
||||
request := AccessKeyCreateRequest{Name: "idempotent-client"}
|
||||
const key = "118f47a2-9c35-4d6e-8b1a-1234567890ab"
|
||||
|
||||
first, err := fixture.service.CreateAccessKeyIdempotent(t.Context(), key, request)
|
||||
if err != nil {
|
||||
t.Fatalf("first CreateAccessKeyIdempotent() error = %v", err)
|
||||
}
|
||||
if first.Key == "" || first.Replayed {
|
||||
t.Fatalf("first result = %#v", first)
|
||||
}
|
||||
replayed, err := fixture.service.CreateAccessKeyIdempotent(t.Context(), key, request)
|
||||
if err != nil {
|
||||
t.Fatalf("replay CreateAccessKeyIdempotent() error = %v", err)
|
||||
}
|
||||
if replayed.Key != "" || !replayed.Replayed ||
|
||||
!reflect.DeepEqual(replayed.AccessKeyMetadata, first.AccessKeyMetadata) {
|
||||
t.Fatalf("replay result = %#v, first = %#v", replayed, first)
|
||||
}
|
||||
var operation models.ControlOperation
|
||||
if err := fixture.db.First(&operation).Error; err != nil {
|
||||
t.Fatalf("read operation: %v", err)
|
||||
}
|
||||
if bytes.Contains(operation.CanonicalResult, []byte(first.Key)) {
|
||||
t.Fatal("operation ledger contains one-time AccessKey secret")
|
||||
}
|
||||
|
||||
_, err = fixture.service.CreateAccessKeyIdempotent(
|
||||
t.Context(),
|
||||
key,
|
||||
AccessKeyCreateRequest{Name: "different-client"},
|
||||
)
|
||||
var apiErr *app_errors.APIError
|
||||
if !errors.As(err, &apiErr) ||
|
||||
apiErr.Code != app_errors.ErrIdempotencyKeyReused.Code {
|
||||
t.Fatalf("different request error = %v", err)
|
||||
}
|
||||
var count int64
|
||||
if err := fixture.db.Model(&models.AccessKey{}).Count(&count).Error; err != nil {
|
||||
t.Fatalf("count AccessKeys: %v", err)
|
||||
}
|
||||
if count != 1 {
|
||||
t.Fatalf("AccessKey count = %d, want 1", count)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccessKeyMetadataFailsClosedForInvalidPersistedSuffix(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
created, err := fixture.service.CreateAccessKey(
|
||||
t.Context(),
|
||||
AccessKeyCreateRequest{Name: "invalid-suffix"},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateAccessKey() error = %v", err)
|
||||
}
|
||||
if err := fixture.db.Exec("PRAGMA ignore_check_constraints = ON").Error; err != nil {
|
||||
t.Fatalf("disable check constraints: %v", err)
|
||||
}
|
||||
if err := fixture.db.Model(&models.AccessKey{}).
|
||||
Where("id = ?", created.ID).
|
||||
UpdateColumn("key_suffix", "ZZZZ").Error; err != nil {
|
||||
t.Fatalf("set invalid suffix: %v", err)
|
||||
}
|
||||
if err := fixture.db.Exec("PRAGMA ignore_check_constraints = OFF").Error; err != nil {
|
||||
t.Fatalf("restore check constraints: %v", err)
|
||||
}
|
||||
|
||||
if list, err := fixture.service.ListAccessKeys(t.Context()); err == nil || list != nil {
|
||||
t.Fatalf("ListAccessKeys() = %#v, %v, want fail closed", list, err)
|
||||
}
|
||||
if _, err := fixture.service.UpdateAccessKey(
|
||||
t.Context(),
|
||||
created.ID,
|
||||
AccessKeyUpdateRequest{Name: stringPointer("must-not-change")},
|
||||
); err == nil {
|
||||
t.Fatal("UpdateAccessKey() error = nil for invalid suffix")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRevealAccessKeyIsTheExplicitMetadataDecryptPath(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
created, err := fixture.service.CreateAccessKey(
|
||||
t.Context(),
|
||||
AccessKeyCreateRequest{Name: "reveal"},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("CreateAccessKey() error = %v", err)
|
||||
}
|
||||
spy := &decryptCountingEncryption{Service: fixture.encryption}
|
||||
fixture.service.encryption = spy
|
||||
|
||||
revealed, err := fixture.service.RevealAccessKey(t.Context(), created.ID)
|
||||
if err != nil {
|
||||
t.Fatalf("RevealAccessKey() error = %v", err)
|
||||
}
|
||||
if revealed.ID != created.ID || revealed.Key != created.Key ||
|
||||
revealed.RevealedAt.IsZero() {
|
||||
t.Fatalf("RevealAccessKey() = %#v", revealed)
|
||||
}
|
||||
if spy.decryptCalls != 1 {
|
||||
t.Fatalf("RevealAccessKey() decrypt calls = %d, want 1", spy.decryptCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAccessKeyOptionsContainsOnlySelectorMetadata(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
first, err := fixture.service.CreateAccessKey(
|
||||
t.Context(),
|
||||
AccessKeyCreateRequest{Name: "first-option"},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("create first option: %v", err)
|
||||
}
|
||||
second, err := fixture.service.CreateAccessKey(
|
||||
t.Context(),
|
||||
AccessKeyCreateRequest{Name: "second-option"},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("create second option: %v", err)
|
||||
}
|
||||
if _, err := fixture.service.UpdateAccessKey(
|
||||
t.Context(),
|
||||
second.ID,
|
||||
AccessKeyUpdateRequest{Status: accessKeyStatusPointer(state.AccessKeyStatusDisabled)},
|
||||
); err != nil {
|
||||
t.Fatalf("disable second option: %v", err)
|
||||
}
|
||||
|
||||
options, err := fixture.service.ListAccessKeyOptions(t.Context())
|
||||
if err != nil {
|
||||
t.Fatalf("ListAccessKeyOptions() error = %v", err)
|
||||
}
|
||||
if len(options) != 2 ||
|
||||
options[0] != (AccessKeyOption{
|
||||
ID: first.ID, Name: "first-option", Status: state.AccessKeyStatusActive,
|
||||
}) ||
|
||||
options[1] != (AccessKeyOption{
|
||||
ID: second.ID, Name: "second-option", Status: state.AccessKeyStatusDisabled,
|
||||
}) {
|
||||
t.Fatalf("options = %#v", options)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccessKeyFilterUpdateCanPreserveOrRemoveButNotAddDanglingScope(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
row := models.AccessKey{
|
||||
Name: "historical-scope",
|
||||
KeyValue: "ciphertext",
|
||||
KeyHash: "historical-scope-hash",
|
||||
KeySuffix: "cafe",
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(`{"groups":[999],"protocols":["openai-response"],"models":[]}`),
|
||||
}
|
||||
if err := fixture.db.Create(&row).Error; err != nil {
|
||||
t.Fatalf("create historical scope: %v", err)
|
||||
}
|
||||
|
||||
preserved := AccessKeyFilters{
|
||||
Groups: []uint{999}, Protocols: []protocol.Protocol{protocol.OpenAIResponse},
|
||||
}
|
||||
if _, err := fixture.service.UpdateAccessKey(
|
||||
t.Context(),
|
||||
row.ID,
|
||||
AccessKeyUpdateRequest{Filters: &preserved},
|
||||
); err != nil {
|
||||
t.Fatalf("preserve historical scope: %v", err)
|
||||
}
|
||||
removed := AccessKeyFilters{}
|
||||
if _, err := fixture.service.UpdateAccessKey(
|
||||
t.Context(),
|
||||
row.ID,
|
||||
AccessKeyUpdateRequest{Filters: &removed},
|
||||
); err != nil {
|
||||
t.Fatalf("remove historical scope: %v", err)
|
||||
}
|
||||
if _, err := fixture.service.UpdateAccessKey(
|
||||
t.Context(),
|
||||
row.ID,
|
||||
AccessKeyUpdateRequest{Filters: &preserved},
|
||||
); err == nil {
|
||||
t.Fatal("re-add removed dangling/reserved scope error = nil")
|
||||
}
|
||||
}
|
||||
|
||||
func accessKeyStatusPointer(value state.AccessKeyStatus) *state.AccessKeyStatus {
|
||||
return &value
|
||||
}
|
||||
@@ -170,7 +170,7 @@ func TestAccessKeyCreateRejectsReservedProtocolWithoutPublishing(t *testing.T) {
|
||||
|
||||
func TestAccessKeyHistoricalReservedProtocolPreservedAcrossReadAndNameUpdate(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
row, plaintext := seedHistoricalAccessKeyWithReservedProtocol(
|
||||
row, _ := seedHistoricalAccessKeyWithReservedProtocol(
|
||||
t,
|
||||
fixture,
|
||||
[]protocol.Protocol{protocol.OpenAIResponse, protocol.OpenAI},
|
||||
@@ -183,7 +183,7 @@ func TestAccessKeyHistoricalReservedProtocolPreservedAcrossReadAndNameUpdate(t *
|
||||
wantProtocols := []protocol.Protocol{protocol.OpenAIResponse, protocol.OpenAI}
|
||||
if len(listed) != 1 ||
|
||||
listed[0].ID != row.ID ||
|
||||
listed[0].Key != plaintext ||
|
||||
listed[0].MaskedKey != "sk-gl-••••••••abab" ||
|
||||
!reflect.DeepEqual(listed[0].Filters.Protocols, wantProtocols) {
|
||||
t.Fatalf("ListAccessKeys() = %#v, want historical protocols preserved", listed)
|
||||
}
|
||||
@@ -205,7 +205,7 @@ func TestAccessKeyHistoricalReservedProtocolPreservedAcrossReadAndNameUpdate(t *
|
||||
t.Fatalf("name-only UpdateAccessKey() error = %v", err)
|
||||
}
|
||||
if updated.Name != "renamed legacy" ||
|
||||
updated.Key != plaintext ||
|
||||
updated.MaskedKey != "sk-gl-••••••••abab" ||
|
||||
!reflect.DeepEqual(updated.Filters.Protocols, wantProtocols) {
|
||||
t.Fatalf("name-only UpdateAccessKey() = %#v", updated)
|
||||
}
|
||||
@@ -357,7 +357,7 @@ func TestAccessKeyFiltersRejectInvalidCurrentInputWithoutPublishing(t *testing.T
|
||||
}
|
||||
}
|
||||
|
||||
func TestListAccessKeysReturnsPlaintextInIDOrderAndFailsClosed(t *testing.T) {
|
||||
func TestListAccessKeysReturnsMaskedMetadataInIDOrderWithoutDecrypting(t *testing.T) {
|
||||
initControlI18n(t)
|
||||
fixture := newServiceFixture(t)
|
||||
randomBytes := make([]byte, 32)
|
||||
@@ -378,7 +378,11 @@ func TestListAccessKeysReturnsPlaintextInIDOrderAndFailsClosed(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("ListAccessKeys() error = %v", err)
|
||||
}
|
||||
if len(listed) != 2 || listed[0].ID != first.ID || listed[0].Key != first.Key || listed[1].ID != second.ID || listed[1].Key != second.Key {
|
||||
if len(listed) != 2 ||
|
||||
listed[0].ID != first.ID ||
|
||||
listed[0].MaskedKey != "sk-gl-••••••••0e0f" ||
|
||||
listed[1].ID != second.ID ||
|
||||
listed[1].MaskedKey != "sk-gl-••••••••1e1f" {
|
||||
t.Fatalf("ListAccessKeys() = %#v", listed)
|
||||
}
|
||||
|
||||
@@ -386,8 +390,13 @@ func TestListAccessKeysReturnsPlaintextInIDOrderAndFailsClosed(t *testing.T) {
|
||||
if err := fixture.db.Model(&models.AccessKey{}).Where("id = ?", second.ID).UpdateColumn("key_value", corruptCiphertext).Error; err != nil {
|
||||
t.Fatalf("corrupt second ciphertext: %v", err)
|
||||
}
|
||||
if partial, err := fixture.service.ListAccessKeys(context.Background()); err == nil || partial != nil {
|
||||
t.Fatalf("ListAccessKeys() = %#v, %v, want nil/error", partial, err)
|
||||
if afterCorruption, err := fixture.service.ListAccessKeys(context.Background()); err != nil ||
|
||||
!reflect.DeepEqual(afterCorruption, listed) {
|
||||
t.Fatalf(
|
||||
"ListAccessKeys() after ciphertext corruption = %#v, %v, want unchanged metadata",
|
||||
afterCorruption,
|
||||
err,
|
||||
)
|
||||
}
|
||||
|
||||
engine := gin.New()
|
||||
@@ -396,7 +405,7 @@ func TestListAccessKeysReturnsPlaintextInIDOrderAndFailsClosed(t *testing.T) {
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/access-keys", nil)
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
engine.ServeHTTP(recorder, request)
|
||||
if recorder.Code != http.StatusInternalServerError || !strings.Contains(recorder.Body.String(), "INTERNAL_SERVER_ERROR") {
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("GET access keys = %d %s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
for _, forbidden := range []string{first.Key, corruptCiphertext} {
|
||||
@@ -429,7 +438,9 @@ func TestUpdateAccessKeyPreservesCredentialAcrossPointerPatches(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("name-only UpdateAccessKey() error = %v", err)
|
||||
}
|
||||
if updated.Name != "after" || updated.Status != state.AccessKeyStatusActive || updated.Key != created.Key {
|
||||
if updated.Name != "after" ||
|
||||
updated.Status != state.AccessKeyStatusActive ||
|
||||
updated.MaskedKey != created.MaskedKey {
|
||||
t.Fatalf("name-only update = %#v", updated)
|
||||
}
|
||||
assertAccessKeyCredentialUnchanged(t, fixture.db, created.ID, original)
|
||||
@@ -557,7 +568,7 @@ func TestUpdateAccessKeyStatusAndDeletePublishWithoutMutatingRegistry(t *testing
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpdateAccessKeyRollsBackWhenCredentialCannotDecrypt(t *testing.T) {
|
||||
func TestUpdateAccessKeyDoesNotReadCredentialCiphertext(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.random = bytes.NewReader(make([]byte, 16))
|
||||
created, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "before"})
|
||||
@@ -568,12 +579,16 @@ func TestUpdateAccessKeyRollsBackWhenCredentialCannotDecrypt(t *testing.T) {
|
||||
t.Fatalf("corrupt AccessKey: %v", err)
|
||||
}
|
||||
before := fixture.manager.Current().Revision
|
||||
if _, err := fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{Name: stringPointer("after")}); err == nil {
|
||||
t.Fatal("UpdateAccessKey() error = nil")
|
||||
if _, err := fixture.service.UpdateAccessKey(
|
||||
context.Background(),
|
||||
created.ID,
|
||||
AccessKeyUpdateRequest{Name: stringPointer("after")},
|
||||
); err != nil {
|
||||
t.Fatalf("UpdateAccessKey() error = %v", err)
|
||||
}
|
||||
row := loadAccessKeyRow(t, fixture.db, created.ID)
|
||||
if row.Name != "before" || fixture.manager.Current().Revision != before {
|
||||
t.Fatalf("failed update changed row/revision: name=%q revision=%d", row.Name, fixture.manager.Current().Revision)
|
||||
if row.Name != "after" || fixture.manager.Current().Revision != before+1 {
|
||||
t.Fatalf("update row/revision: name=%q revision=%d", row.Name, fixture.manager.Current().Revision)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -586,7 +601,8 @@ func TestAccessKeyDanglingFiltersDoNotBlockUnrelatedUpdate(t *testing.T) {
|
||||
}
|
||||
if err := fixture.db.Create(&models.AccessKey{
|
||||
Name: "legacy", KeyValue: danglingCiphertext, KeyHash: fixture.encryption.Hash(danglingPlaintext),
|
||||
Status: string(state.AccessKeyStatusActive), Filters: models.JSON(`{"groups":[999],"protocols":[],"models":[]}`),
|
||||
KeySuffix: "0000", Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(`{"groups":[999],"protocols":[],"models":[]}`),
|
||||
}).Error; err != nil {
|
||||
t.Fatalf("create dangling AccessKey: %v", err)
|
||||
}
|
||||
@@ -738,6 +754,9 @@ func TestAccessKeyEndpointsDistinguishRPMLimit(t *testing.T) {
|
||||
request := httptest.NewRequest(method, path, strings.NewReader(payload))
|
||||
request.Header.Set("Authorization", "Bearer "+authKey)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if method == http.MethodPost && path == "/api/access-keys" {
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
}
|
||||
engine.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
@@ -909,11 +928,12 @@ func seedHistoricalAccessKeyWithReservedProtocol(
|
||||
t.Fatalf("encode historical AccessKey filters: %v", err)
|
||||
}
|
||||
row := models.AccessKey{
|
||||
Name: "legacy",
|
||||
KeyValue: ciphertext,
|
||||
KeyHash: fixture.encryption.Hash(plaintext),
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(encodedFilters),
|
||||
Name: "legacy",
|
||||
KeyValue: ciphertext,
|
||||
KeyHash: fixture.encryption.Hash(plaintext),
|
||||
KeySuffix: "abab",
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(encodedFilters),
|
||||
}
|
||||
if _, err := fixture.service.writeConfig(
|
||||
t.Context(),
|
||||
|
||||
@@ -238,11 +238,12 @@ func TestEnsureInitialStateWithExistingAccessKeyOnlyWritesMarker(t *testing.T) {
|
||||
t.Fatalf("Encrypt(existing key) error = %v", err)
|
||||
}
|
||||
existing := models.AccessKey{
|
||||
Name: "Existing",
|
||||
KeyValue: ciphertext,
|
||||
KeyHash: fixture.encryption.Hash(plaintext),
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(`{"groups":[],"protocols":[],"models":[]}`),
|
||||
Name: "Existing",
|
||||
KeyValue: ciphertext,
|
||||
KeyHash: fixture.encryption.Hash(plaintext),
|
||||
KeySuffix: "0000",
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(`{"groups":[],"protocols":[],"models":[]}`),
|
||||
}
|
||||
if err := fixture.db.Create(&existing).Error; err != nil {
|
||||
t.Fatalf("create existing AccessKey: %v", err)
|
||||
|
||||
@@ -30,18 +30,21 @@ func TestDeleteGroupRejectsActiveAndDisabledExplicitAccessKeyReferences(t *testi
|
||||
for _, row := range []models.AccessKey{
|
||||
{
|
||||
Name: "Active", KeyValue: "cipher-client-one", KeyHash: "client-hash-one",
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d],"protocols":[],"models":[]}`, groupID)),
|
||||
KeySuffix: "0001",
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d],"protocols":[],"models":[]}`, groupID)),
|
||||
},
|
||||
{
|
||||
Name: "Disabled", KeyValue: "cipher-client-two", KeyHash: "client-hash-two",
|
||||
Status: string(state.AccessKeyStatusDisabled),
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d,%d],"protocols":[],"models":[]}`, other.ID, groupID)),
|
||||
KeySuffix: "0002",
|
||||
Status: string(state.AccessKeyStatusDisabled),
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d,%d],"protocols":[],"models":[]}`, other.ID, groupID)),
|
||||
},
|
||||
{
|
||||
Name: "Unrestricted", KeyValue: "cipher-client-three", KeyHash: "client-hash-three",
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(`{"groups":[],"protocols":["openai"],"models":["gpt-4o"]}`),
|
||||
KeySuffix: "0003",
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(`{"groups":[],"protocols":["openai"],"models":["gpt-4o"]}`),
|
||||
},
|
||||
} {
|
||||
row := row
|
||||
@@ -194,7 +197,7 @@ func TestDeleteGroupCorruptAccessKeyFiltersPreservesDatabaseRegistryAndSnapshot(
|
||||
groupID := createGroupForKeyImport(t, fixture, "sk-delete-corrupt-filter")
|
||||
corrupt := models.AccessKey{
|
||||
Name: "Corrupt filters", KeyValue: "cipher-corrupt-filter", KeyHash: "hash-corrupt-filter",
|
||||
Status: string(state.AccessKeyStatusDisabled), Filters: models.JSON(`{}`),
|
||||
KeySuffix: "0004", Status: string(state.AccessKeyStatusDisabled), Filters: models.JSON(`{}`),
|
||||
}
|
||||
if err := fixture.db.Create(&corrupt).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
@@ -307,11 +310,12 @@ func TestDeleteGroupEndpointAuthenticationValidationNotFoundConflictAndSuccess(t
|
||||
{language: "ja-JP", message: "グループはアクセスキーから参照されています"},
|
||||
} {
|
||||
row := models.AccessKey{
|
||||
Name: "HTTP reference " + test.language,
|
||||
KeyValue: "cipher-http-" + test.language,
|
||||
KeyHash: "hash-http-" + test.language,
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d]}`, groupID)),
|
||||
Name: "HTTP reference " + test.language,
|
||||
KeyValue: "cipher-http-" + test.language,
|
||||
KeyHash: "hash-http-" + test.language,
|
||||
KeySuffix: "0005",
|
||||
Status: string(state.AccessKeyStatusActive),
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d]}`, groupID)),
|
||||
}
|
||||
if err := fixture.db.Create(&row).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
|
||||
@@ -0,0 +1,264 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strconv"
|
||||
|
||||
"github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"gpt-load/internal/platform/canonicaljson"
|
||||
"gpt-load/internal/platform/config"
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/protocol"
|
||||
"gpt-load/internal/state"
|
||||
stateloader "gpt-load/internal/state/loader"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
type groupCreateDigestBody struct {
|
||||
Name *string `json:"name"`
|
||||
UpstreamURL string `json:"upstream_url"`
|
||||
Protocols []protocol.Protocol `json:"protocols"`
|
||||
Models []GroupModel `json:"models"`
|
||||
Config config.Settings `json:"config"`
|
||||
Keys []string `json:"keys"`
|
||||
ConfirmSameUpstreamURL bool `json:"confirm_same_upstream_url"`
|
||||
}
|
||||
|
||||
type groupKeyImportDigestBody struct {
|
||||
Keys []string `json:"keys"`
|
||||
}
|
||||
|
||||
func (s *Service) CreateGroupIdempotent(
|
||||
ctx context.Context,
|
||||
idempotencyKey string,
|
||||
request GroupCreateRequest,
|
||||
) (GroupCreateResult, error) {
|
||||
normalized, err := s.normalizeGroupCreate(request)
|
||||
if err != nil {
|
||||
return GroupCreateResult{}, err
|
||||
}
|
||||
keyLines, err := normalizeIdempotencyKeyLines(request.Keys)
|
||||
if err != nil {
|
||||
return GroupCreateResult{}, err
|
||||
}
|
||||
protocols := append([]protocol.Protocol(nil), normalized.protocols...)
|
||||
sort.Slice(protocols, func(left, right int) bool {
|
||||
return string(protocols[left]) < string(protocols[right])
|
||||
})
|
||||
settings := make(config.Settings)
|
||||
if len(normalized.encodedConfig) > 0 {
|
||||
decoder := json.NewDecoder(bytes.NewReader(normalized.encodedConfig))
|
||||
decoder.UseNumber()
|
||||
if err := decoder.Decode(&settings); err != nil {
|
||||
return GroupCreateResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
}
|
||||
canonicalBody, err := canonicalIdempotencyBody(groupCreateDigestBody{
|
||||
Name: normalized.explicitName,
|
||||
UpstreamURL: normalized.upstreamURL,
|
||||
Protocols: protocols,
|
||||
Models: append([]GroupModel(nil), normalized.models...),
|
||||
Config: settings,
|
||||
Keys: keyLines,
|
||||
ConfirmSameUpstreamURL: normalized.confirmSameUpstreamURL,
|
||||
})
|
||||
if err != nil {
|
||||
return GroupCreateResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
digest, err := buildIdempotencyDigest(idempotencyDigestInput{
|
||||
Version: 1,
|
||||
Method: "POST",
|
||||
OperationKind: operationKindGroupCreate,
|
||||
PathTemplate: "/api/groups",
|
||||
ResourceLocator: "new",
|
||||
AuthScopeID: idempotencyAuthScopeID,
|
||||
CanonicalBody: canonicalBody,
|
||||
})
|
||||
if err != nil {
|
||||
return GroupCreateResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
if isLiteralPrivateHost(normalized.hostname) {
|
||||
logrus.WithField("host", normalized.hostname).
|
||||
Warn("Creating upstream group with a private or local host")
|
||||
}
|
||||
|
||||
operationResult, err := s.executeIdempotentOperation(ctx, idempotentOperationInput{
|
||||
IdempotencyKey: idempotencyKey,
|
||||
DigestVersion: 1,
|
||||
RequestDigest: digest.Digest,
|
||||
Kind: operationKindGroupCreate,
|
||||
Mutate: func(tx *gorm.DB) (idempotentMutationResult, error) {
|
||||
if !normalized.confirmSameUpstreamURL {
|
||||
conflicts, err := findGroupsByUpstreamURL(tx, normalized.upstreamURL)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
if len(conflicts) > 0 {
|
||||
return idempotentMutationResult{}, app_errors.NewAPIErrorWithData(
|
||||
app_errors.ErrUpstreamURLConflict,
|
||||
UpstreamURLConflictData{Groups: conflicts},
|
||||
)
|
||||
}
|
||||
}
|
||||
name, err := resolveGroupCreateName(
|
||||
tx,
|
||||
normalized.explicitName,
|
||||
normalized.hostname,
|
||||
)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
encodedProtocols, err := json.Marshal(normalized.protocols)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
encodedModels, err := json.Marshal(normalized.models)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
group := models.Group{
|
||||
Name: name,
|
||||
UpstreamURL: normalized.upstreamURL,
|
||||
Protocols: models.JSON(encodedProtocols),
|
||||
Models: models.JSON(encodedModels),
|
||||
Config: normalized.encodedConfig,
|
||||
Enabled: true,
|
||||
}
|
||||
if err := tx.Create(&group).Error; err != nil {
|
||||
return idempotentMutationResult{}, app_errors.ParseDBError(err)
|
||||
}
|
||||
entries, added, duplicated, err := s.persistUpstreamKeys(
|
||||
tx,
|
||||
group.ID,
|
||||
normalized.keys,
|
||||
)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
if err := state.ValidateKeyEntries(entries); err != nil {
|
||||
return idempotentMutationResult{}, fmt.Errorf(
|
||||
"validate created group keys: %w",
|
||||
err,
|
||||
)
|
||||
}
|
||||
input, err := stateloader.BuildCompileInput(ctx, tx)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
if _, err := state.Compile(input); err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
result := GroupCreateResult{
|
||||
GroupID: group.ID, GroupName: group.Name,
|
||||
KeysAdded: added, KeysDuplicated: duplicated,
|
||||
Models: append(make([]GroupModel, 0, len(normalized.models)), normalized.models...),
|
||||
}
|
||||
canonicalResult, err := canonicaljson.Marshal(result)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
return idempotentMutationResult{
|
||||
ResourceIdentity: "group:" + strconv.FormatUint(uint64(group.ID), 10),
|
||||
CanonicalResult: canonicalResult,
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return GroupCreateResult{}, err
|
||||
}
|
||||
var result GroupCreateResult
|
||||
if err := json.Unmarshal(operationResult.CanonicalResult, &result); err != nil {
|
||||
return GroupCreateResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *Service) ImportGroupKeysIdempotent(
|
||||
ctx context.Context,
|
||||
idempotencyKey string,
|
||||
groupID uint,
|
||||
request GroupKeyImportRequest,
|
||||
) (GroupKeyImportResult, error) {
|
||||
if groupID == 0 {
|
||||
return GroupKeyImportResult{}, app_errors.ErrValidation
|
||||
}
|
||||
keys, err := s.normalizeUpstreamKeys(request.Keys)
|
||||
if err != nil {
|
||||
return GroupKeyImportResult{}, err
|
||||
}
|
||||
keyLines, err := normalizeIdempotencyKeyLines(request.Keys)
|
||||
if err != nil {
|
||||
return GroupKeyImportResult{}, err
|
||||
}
|
||||
canonicalBody, err := canonicalIdempotencyBody(
|
||||
groupKeyImportDigestBody{Keys: keyLines},
|
||||
)
|
||||
if err != nil {
|
||||
return GroupKeyImportResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
resourceIdentity := "group:" + strconv.FormatUint(uint64(groupID), 10)
|
||||
digest, err := buildIdempotencyDigest(idempotencyDigestInput{
|
||||
Version: 1,
|
||||
Method: "POST",
|
||||
OperationKind: operationKindGroupKeyImport,
|
||||
PathTemplate: "/api/groups/:group_id/keys/import",
|
||||
ResourceLocator: resourceIdentity,
|
||||
AuthScopeID: idempotencyAuthScopeID,
|
||||
CanonicalBody: canonicalBody,
|
||||
})
|
||||
if err != nil {
|
||||
return GroupKeyImportResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
|
||||
operationResult, err := s.executeIdempotentOperation(ctx, idempotentOperationInput{
|
||||
IdempotencyKey: idempotencyKey,
|
||||
DigestVersion: 1,
|
||||
RequestDigest: digest.Digest,
|
||||
Kind: operationKindGroupKeyImport,
|
||||
Mutate: func(tx *gorm.DB) (idempotentMutationResult, error) {
|
||||
var group struct{ ID uint }
|
||||
if err := tx.Model(&models.Group{}).
|
||||
Select("id").
|
||||
Where("id = ?", groupID).
|
||||
Take(&group).Error; err != nil {
|
||||
return idempotentMutationResult{}, app_errors.ParseDBError(err)
|
||||
}
|
||||
entries, added, duplicated, err := s.persistUpstreamKeys(
|
||||
tx,
|
||||
groupID,
|
||||
keys,
|
||||
)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
if err := state.ValidateKeyEntries(entries); err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
result := GroupKeyImportResult{
|
||||
GroupID: groupID, KeysAdded: added, KeysDuplicated: duplicated,
|
||||
}
|
||||
canonicalResult, err := canonicaljson.Marshal(result)
|
||||
if err != nil {
|
||||
return idempotentMutationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
return idempotentMutationResult{
|
||||
ResourceIdentity: resourceIdentity,
|
||||
CanonicalResult: canonicalResult,
|
||||
}, nil
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return GroupKeyImportResult{}, err
|
||||
}
|
||||
var result GroupKeyImportResult
|
||||
if err := json.Unmarshal(operationResult.CanonicalResult, &result); err != nil {
|
||||
return GroupKeyImportResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"reflect"
|
||||
"testing"
|
||||
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/protocol"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
func TestCreateGroupIdempotentReplaysOriginalCountsAndPreservesKeyMultiplicity(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x71}, 16))
|
||||
request := GroupCreateRequest{
|
||||
Name: stringPointer("idempotent-group"),
|
||||
UpstreamURL: "https://idempotent.example.com/v1/",
|
||||
Protocols: []protocol.Protocol{protocol.OpenAI},
|
||||
Keys: " K \r\nK\n",
|
||||
}
|
||||
const key = "218f47a2-9c35-4d6e-8b1a-1234567890ab"
|
||||
|
||||
first, err := fixture.service.CreateGroupIdempotent(t.Context(), key, request)
|
||||
if err != nil {
|
||||
t.Fatalf("first CreateGroupIdempotent() error = %v", err)
|
||||
}
|
||||
if first.KeysAdded != 1 || first.KeysDuplicated != 1 {
|
||||
t.Fatalf("first result = %#v", first)
|
||||
}
|
||||
replayed, err := fixture.service.CreateGroupIdempotent(t.Context(), key, request)
|
||||
if err != nil {
|
||||
t.Fatalf("replay CreateGroupIdempotent() error = %v", err)
|
||||
}
|
||||
if !reflect.DeepEqual(replayed, first) {
|
||||
t.Fatalf("replay = %#v, first = %#v", replayed, first)
|
||||
}
|
||||
var groupCount, keyCount int64
|
||||
if err := fixture.db.Model(&models.Group{}).Count(&groupCount).Error; err != nil {
|
||||
t.Fatalf("count groups: %v", err)
|
||||
}
|
||||
if err := fixture.db.Model(&models.UpstreamKey{}).Count(&keyCount).Error; err != nil {
|
||||
t.Fatalf("count keys: %v", err)
|
||||
}
|
||||
if groupCount != 1 || keyCount != 1 {
|
||||
t.Fatalf("resource counts = group:%d key:%d, want 1/1", groupCount, keyCount)
|
||||
}
|
||||
|
||||
different := request
|
||||
different.Keys = "K"
|
||||
_, err = fixture.service.CreateGroupIdempotent(t.Context(), key, different)
|
||||
assertAPIErrorCode(t, err, app_errors.ErrIdempotencyKeyReused.Code)
|
||||
}
|
||||
|
||||
func TestImportGroupKeysIdempotentReplaysOriginalCountsWithoutSecondInsert(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
groupID := createGroupForKeyImport(t, fixture, "seed-existing")
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x72}, 16))
|
||||
request := GroupKeyImportRequest{Keys: " imported \n imported \n"}
|
||||
const key = "228f47a2-9c35-4d6e-8b1a-1234567890ab"
|
||||
|
||||
first, err := fixture.service.ImportGroupKeysIdempotent(
|
||||
t.Context(),
|
||||
key,
|
||||
groupID,
|
||||
request,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("first ImportGroupKeysIdempotent() error = %v", err)
|
||||
}
|
||||
if first.KeysAdded != 1 || first.KeysDuplicated != 1 {
|
||||
t.Fatalf("first result = %#v", first)
|
||||
}
|
||||
replayed, err := fixture.service.ImportGroupKeysIdempotent(
|
||||
t.Context(),
|
||||
key,
|
||||
groupID,
|
||||
request,
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("replay ImportGroupKeysIdempotent() error = %v", err)
|
||||
}
|
||||
if !reflect.DeepEqual(replayed, first) {
|
||||
t.Fatalf("replay = %#v, first = %#v", replayed, first)
|
||||
}
|
||||
var count int64
|
||||
if err := fixture.db.Model(&models.UpstreamKey{}).
|
||||
Where("group_id = ?", groupID).
|
||||
Count(&count).Error; err != nil {
|
||||
t.Fatalf("count group keys: %v", err)
|
||||
}
|
||||
if count != 2 {
|
||||
t.Fatalf("group key count = %d, want seed + imported = 2", count)
|
||||
}
|
||||
|
||||
_, err = fixture.service.ImportGroupKeysIdempotent(
|
||||
t.Context(),
|
||||
key,
|
||||
groupID,
|
||||
GroupKeyImportRequest{Keys: "imported"},
|
||||
)
|
||||
var apiErr *app_errors.APIError
|
||||
if !errors.As(err, &apiErr) ||
|
||||
apiErr.Code != app_errors.ErrIdempotencyKeyReused.Code {
|
||||
t.Fatalf("different multiplicity error = %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"unicode/utf8"
|
||||
|
||||
"gpt-load/internal/platform/canonicaljson"
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
)
|
||||
|
||||
const (
|
||||
idempotencyDigestDomain = "gpt-load/control-idempotency"
|
||||
idempotencyAuthScopeID = "control-admin-v1"
|
||||
)
|
||||
|
||||
type operationKind string
|
||||
|
||||
const (
|
||||
operationKindAccessKeyCreate operationKind = "access_key_create"
|
||||
operationKindGroupCreate operationKind = "group_create"
|
||||
operationKindGroupKeyImport operationKind = "group_key_import"
|
||||
)
|
||||
|
||||
type idempotencyDigestInput struct {
|
||||
Version uint
|
||||
Method string
|
||||
OperationKind operationKind
|
||||
PathTemplate string
|
||||
ResourceLocator string
|
||||
AuthScopeID string
|
||||
CanonicalBody []byte
|
||||
}
|
||||
|
||||
type idempotencyDigestResult struct {
|
||||
Digest [sha256.Size]byte
|
||||
FramedInput []byte
|
||||
}
|
||||
|
||||
func buildIdempotencyDigest(
|
||||
input idempotencyDigestInput,
|
||||
) (idempotencyDigestResult, error) {
|
||||
if input.Version != 1 {
|
||||
return idempotencyDigestResult{}, fmt.Errorf(
|
||||
"unsupported idempotency digest version %d: %w",
|
||||
input.Version,
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
if input.Method == "" || input.Method != strings.ToUpper(input.Method) {
|
||||
return idempotencyDigestResult{}, fmt.Errorf("idempotency method must be uppercase")
|
||||
}
|
||||
if !input.OperationKind.valid() ||
|
||||
input.PathTemplate == "" ||
|
||||
input.ResourceLocator == "" ||
|
||||
input.AuthScopeID == "" ||
|
||||
len(input.CanonicalBody) == 0 {
|
||||
return idempotencyDigestResult{}, fmt.Errorf("idempotency digest fields are incomplete")
|
||||
}
|
||||
if !utf8.Valid(input.CanonicalBody) {
|
||||
return idempotencyDigestResult{}, fmt.Errorf("canonical body is not valid UTF-8")
|
||||
}
|
||||
canonicalBody, err := canonicaljson.Canonicalize(input.CanonicalBody)
|
||||
if err != nil {
|
||||
return idempotencyDigestResult{}, fmt.Errorf("validate canonical body: %w", err)
|
||||
}
|
||||
if !bytes.Equal(canonicalBody, input.CanonicalBody) {
|
||||
return idempotencyDigestResult{}, fmt.Errorf("idempotency body is not canonical JSON")
|
||||
}
|
||||
|
||||
fields := [][]byte{
|
||||
[]byte(strconv.FormatUint(uint64(input.Version), 10)),
|
||||
[]byte(input.Method),
|
||||
[]byte(input.OperationKind),
|
||||
[]byte(input.PathTemplate),
|
||||
[]byte(input.ResourceLocator),
|
||||
[]byte(input.AuthScopeID),
|
||||
input.CanonicalBody,
|
||||
}
|
||||
var framed bytes.Buffer
|
||||
framed.WriteString(idempotencyDigestDomain)
|
||||
framed.WriteByte(0)
|
||||
for _, field := range fields {
|
||||
if len(field) > math.MaxUint32 {
|
||||
return idempotencyDigestResult{}, fmt.Errorf("idempotency digest field is too large")
|
||||
}
|
||||
var length [4]byte
|
||||
binary.BigEndian.PutUint32(length[:], uint32(len(field)))
|
||||
framed.Write(length[:])
|
||||
framed.Write(field)
|
||||
}
|
||||
inputBytes := framed.Bytes()
|
||||
return idempotencyDigestResult{
|
||||
Digest: sha256.Sum256(inputBytes),
|
||||
FramedInput: append([]byte(nil), inputBytes...),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (kind operationKind) valid() bool {
|
||||
switch kind {
|
||||
case operationKindAccessKeyCreate, operationKindGroupCreate, operationKindGroupKeyImport:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func canonicalIdempotencyBody(value any) ([]byte, error) {
|
||||
body, err := canonicaljson.Marshal(value)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode idempotency body: %w", err)
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func normalizeIdempotencyKeyLines(raw string) ([]string, error) {
|
||||
normalizedNewlines := strings.ReplaceAll(raw, "\r\n", "\n")
|
||||
normalizedNewlines = strings.ReplaceAll(normalizedNewlines, "\r", "\n")
|
||||
lines := make([]string, 0)
|
||||
for _, line := range strings.Split(normalizedNewlines, "\n") {
|
||||
normalized := strings.TrimSpace(line)
|
||||
if normalized == "" {
|
||||
continue
|
||||
}
|
||||
lines = append(lines, normalized)
|
||||
if len(lines) > maxUpstreamKeyLines {
|
||||
return nil, app_errors.ErrValidation
|
||||
}
|
||||
}
|
||||
if len(lines) == 0 {
|
||||
return nil, app_errors.ErrValidation
|
||||
}
|
||||
sort.Slice(lines, func(left, right int) bool {
|
||||
return bytes.Compare([]byte(lines[left]), []byte(lines[right])) < 0
|
||||
})
|
||||
return lines, nil
|
||||
}
|
||||
@@ -0,0 +1,188 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestIdempotencyDigestFixtures(t *testing.T) {
|
||||
type fixture struct {
|
||||
Name string `json:"name"`
|
||||
OperationKind operationKind `json:"operation_kind"`
|
||||
PathTemplate string `json:"path_template"`
|
||||
ResourceLocator string `json:"resource_locator"`
|
||||
CanonicalBody string `json:"canonical_body"`
|
||||
DomainHex string `json:"domain_hex"`
|
||||
SeparatorHex string `json:"separator_hex"`
|
||||
FramedInputHex string `json:"framed_input_hex"`
|
||||
ExpectedSHA256 string `json:"expected_sha256"`
|
||||
}
|
||||
paths, err := filepath.Glob("testdata/idempotency/*.json")
|
||||
if err != nil {
|
||||
t.Fatalf("glob fixtures: %v", err)
|
||||
}
|
||||
if len(paths) != 5 {
|
||||
t.Fatalf("fixture count = %d, want 5", len(paths))
|
||||
}
|
||||
digests := make(map[string]string, len(paths))
|
||||
for _, path := range paths {
|
||||
path := path
|
||||
t.Run(filepath.Base(path), func(t *testing.T) {
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatalf("read fixture: %v", err)
|
||||
}
|
||||
var value fixture
|
||||
if err := json.Unmarshal(raw, &value); err != nil {
|
||||
t.Fatalf("decode fixture: %v", err)
|
||||
}
|
||||
result, err := buildIdempotencyDigest(idempotencyDigestInput{
|
||||
Version: 1,
|
||||
Method: "POST",
|
||||
OperationKind: value.OperationKind,
|
||||
PathTemplate: value.PathTemplate,
|
||||
ResourceLocator: value.ResourceLocator,
|
||||
AuthScopeID: idempotencyAuthScopeID,
|
||||
CanonicalBody: []byte(value.CanonicalBody),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("buildIdempotencyDigest() error = %v", err)
|
||||
}
|
||||
if got := hex.EncodeToString([]byte(idempotencyDigestDomain)); got != value.DomainHex {
|
||||
t.Fatalf("domain hex = %s, want %s", got, value.DomainHex)
|
||||
}
|
||||
if value.SeparatorHex != "00" ||
|
||||
result.FramedInput[len(idempotencyDigestDomain)] != 0 {
|
||||
t.Fatalf("separator fixture/runtime = %q/%02x", value.SeparatorHex, result.FramedInput[len(idempotencyDigestDomain)])
|
||||
}
|
||||
if got := hex.EncodeToString(result.FramedInput); got != value.FramedInputHex {
|
||||
t.Fatalf("framed input hex = %s, want %s", got, value.FramedInputHex)
|
||||
}
|
||||
gotDigest := hex.EncodeToString(result.Digest[:])
|
||||
if gotDigest != value.ExpectedSHA256 {
|
||||
t.Fatalf("digest = %s, want %s", gotDigest, value.ExpectedSHA256)
|
||||
}
|
||||
digests[value.Name] = gotDigest
|
||||
})
|
||||
}
|
||||
for _, pair := range [][2]string{
|
||||
{"group-create-k", "group-create-k-twice"},
|
||||
{"group-key-import-k", "group-key-import-k-twice"},
|
||||
} {
|
||||
if digests[pair[0]] == digests[pair[1]] {
|
||||
t.Fatalf("%s and %s digests are equal", pair[0], pair[1])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdempotencyDigestV1MatchesAccessKeyCreateFixture(t *testing.T) {
|
||||
body := []byte(
|
||||
`{"filters":{"groups":[],"models":[],"protocols":[]},"name":"CI client","rpm_limit":0}`,
|
||||
)
|
||||
result, err := buildIdempotencyDigest(idempotencyDigestInput{
|
||||
Version: 1,
|
||||
Method: "POST",
|
||||
OperationKind: operationKindAccessKeyCreate,
|
||||
PathTemplate: "/api/access-keys",
|
||||
ResourceLocator: "new",
|
||||
AuthScopeID: "control-admin-v1",
|
||||
CanonicalBody: body,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("buildIdempotencyDigest() error = %v", err)
|
||||
}
|
||||
|
||||
const wantDomainHex = "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e6379"
|
||||
const wantInputHex = "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e637900000000013100000004504f5354000000116163636573735f6b65795f637265617465000000102f6170692f6163636573732d6b657973000000036e657700000010636f6e74726f6c2d61646d696e2d7631000000557b2266696c74657273223a7b2267726f757073223a5b5d2c226d6f64656c73223a5b5d2c2270726f746f636f6c73223a5b5d7d2c226e616d65223a22434920636c69656e74222c2272706d5f6c696d6974223a307d"
|
||||
const wantDigestHex = "7b22877fc5ce68bde59170f4906463bbbf873abff4cda536dd36c30da3a62b81"
|
||||
|
||||
if got := hex.EncodeToString([]byte(idempotencyDigestDomain)); got != wantDomainHex {
|
||||
t.Fatalf("domain hex = %s, want %s", got, wantDomainHex)
|
||||
}
|
||||
if got := hex.EncodeToString(result.FramedInput); got != wantInputHex {
|
||||
t.Fatalf("framed input hex = %s, want %s", got, wantInputHex)
|
||||
}
|
||||
if got := hex.EncodeToString(result.Digest[:]); got != wantDigestHex {
|
||||
t.Fatalf("digest hex = %s, want %s", got, wantDigestHex)
|
||||
}
|
||||
if result.FramedInput[len(idempotencyDigestDomain)] != 0 {
|
||||
t.Fatalf(
|
||||
"separator = %02x, want exactly one NUL after domain",
|
||||
result.FramedInput[len(idempotencyDigestDomain)],
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeIdempotencyKeyLinesPreservesSortedMultiplicity(t *testing.T) {
|
||||
one, err := normalizeIdempotencyKeyLines(" K \r\n")
|
||||
if err != nil {
|
||||
t.Fatalf("normalizeIdempotencyKeyLines(K) error = %v", err)
|
||||
}
|
||||
two, err := normalizeIdempotencyKeyLines("K\r\n K\n")
|
||||
if err != nil {
|
||||
t.Fatalf("normalizeIdempotencyKeyLines(K twice) error = %v", err)
|
||||
}
|
||||
if len(one) != 1 || one[0] != "K" {
|
||||
t.Fatalf("one = %#v, want [K]", one)
|
||||
}
|
||||
if len(two) != 2 || two[0] != "K" || two[1] != "K" {
|
||||
t.Fatalf("two = %#v, want [K K]", two)
|
||||
}
|
||||
|
||||
oneBody, err := canonicalIdempotencyBody(map[string]any{"keys": one})
|
||||
if err != nil {
|
||||
t.Fatalf("canonicalIdempotencyBody(one) error = %v", err)
|
||||
}
|
||||
twoBody, err := canonicalIdempotencyBody(map[string]any{"keys": two})
|
||||
if err != nil {
|
||||
t.Fatalf("canonicalIdempotencyBody(two) error = %v", err)
|
||||
}
|
||||
oneDigest, err := buildIdempotencyDigest(idempotencyDigestInput{
|
||||
Version: 1, Method: "POST", OperationKind: operationKindGroupKeyImport,
|
||||
PathTemplate: "/api/groups/:group_id/keys/import", ResourceLocator: "group:7",
|
||||
AuthScopeID: "control-admin-v1", CanonicalBody: oneBody,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("buildIdempotencyDigest(one) error = %v", err)
|
||||
}
|
||||
twoDigest, err := buildIdempotencyDigest(idempotencyDigestInput{
|
||||
Version: 1, Method: "POST", OperationKind: operationKindGroupKeyImport,
|
||||
PathTemplate: "/api/groups/:group_id/keys/import", ResourceLocator: "group:7",
|
||||
AuthScopeID: "control-admin-v1", CanonicalBody: twoBody,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("buildIdempotencyDigest(two) error = %v", err)
|
||||
}
|
||||
if oneDigest.Digest == twoDigest.Digest {
|
||||
t.Fatal("K and K\\nK produced the same digest")
|
||||
}
|
||||
}
|
||||
|
||||
func TestIdempotencyDigestRejectsUnsupportedVersionAndInvalidFields(t *testing.T) {
|
||||
valid := idempotencyDigestInput{
|
||||
Version: 1, Method: "POST", OperationKind: operationKindAccessKeyCreate,
|
||||
PathTemplate: "/api/access-keys", ResourceLocator: "new",
|
||||
AuthScopeID: "control-admin-v1", CanonicalBody: []byte(`{}`),
|
||||
}
|
||||
|
||||
unsupported := valid
|
||||
unsupported.Version = 2
|
||||
if _, err := buildIdempotencyDigest(unsupported); err == nil {
|
||||
t.Fatal("buildIdempotencyDigest(version 2) error = nil, want rejection")
|
||||
}
|
||||
|
||||
lowerMethod := valid
|
||||
lowerMethod.Method = "post"
|
||||
if _, err := buildIdempotencyDigest(lowerMethod); err == nil {
|
||||
t.Fatal("buildIdempotencyDigest(lowercase method) error = nil, want rejection")
|
||||
}
|
||||
|
||||
implicitLocator := valid
|
||||
implicitLocator.ResourceLocator = ""
|
||||
if _, err := buildIdempotencyDigest(implicitLocator); err == nil {
|
||||
t.Fatal("buildIdempotencyDigest(empty locator) error = nil, want rejection")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,484 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"gpt-load/internal/platform/canonicaljson"
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
stateloader "gpt-load/internal/state/loader"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
type operationStage string
|
||||
|
||||
const (
|
||||
operationStageDBCommitted operationStage = "db_committed"
|
||||
operationStageRegistryApplied operationStage = "registry_applied"
|
||||
operationStageSnapshotPublished operationStage = "snapshot_published"
|
||||
operationStageCompleted operationStage = "completed"
|
||||
)
|
||||
|
||||
type operationTransaction = *gorm.DB
|
||||
|
||||
type idempotentMutationResult struct {
|
||||
ResourceIdentity string
|
||||
CanonicalResult []byte
|
||||
Ephemeral any
|
||||
}
|
||||
|
||||
type idempotentOperationInput struct {
|
||||
IdempotencyKey string
|
||||
DigestVersion uint
|
||||
RequestDigest [32]byte
|
||||
Kind operationKind
|
||||
Mutate func(operationTransaction) (idempotentMutationResult, error)
|
||||
}
|
||||
|
||||
type idempotentOperationResult struct {
|
||||
OperationID string
|
||||
ResourceIdentity string
|
||||
CanonicalResult []byte
|
||||
Replayed bool
|
||||
Ephemeral any
|
||||
}
|
||||
|
||||
type operationErrorData struct {
|
||||
OperationID string `json:"operation_id"`
|
||||
OperationKind operationKind `json:"operation_kind"`
|
||||
}
|
||||
|
||||
type operationExpiredData struct {
|
||||
OperationID string `json:"operation_id"`
|
||||
OperationKind operationKind `json:"operation_kind"`
|
||||
ResourceIdentity string `json:"resource_identity"`
|
||||
CompletedAt time.Time `json:"completed_at"`
|
||||
}
|
||||
|
||||
func operationRequiredStages(kind operationKind) ([]operationStage, error) {
|
||||
var stages []operationStage
|
||||
switch kind {
|
||||
case operationKindAccessKeyCreate:
|
||||
stages = []operationStage{
|
||||
operationStageDBCommitted,
|
||||
operationStageSnapshotPublished,
|
||||
operationStageCompleted,
|
||||
}
|
||||
case operationKindGroupCreate:
|
||||
stages = []operationStage{
|
||||
operationStageDBCommitted,
|
||||
operationStageRegistryApplied,
|
||||
operationStageSnapshotPublished,
|
||||
operationStageCompleted,
|
||||
}
|
||||
case operationKindGroupKeyImport:
|
||||
stages = []operationStage{
|
||||
operationStageDBCommitted,
|
||||
operationStageRegistryApplied,
|
||||
operationStageCompleted,
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported control operation kind %q", kind)
|
||||
}
|
||||
return append([]operationStage(nil), stages...), nil
|
||||
}
|
||||
|
||||
func validateIdempotencyKey(value string) error {
|
||||
if len(value) != 36 ||
|
||||
value[8] != '-' ||
|
||||
value[13] != '-' ||
|
||||
value[18] != '-' ||
|
||||
value[23] != '-' ||
|
||||
value[14] != '4' {
|
||||
return fmt.Errorf("idempotency key must be a canonical lowercase UUID v4")
|
||||
}
|
||||
switch value[19] {
|
||||
case '8', '9', 'a', 'b':
|
||||
default:
|
||||
return fmt.Errorf("idempotency key must be a canonical lowercase UUID v4")
|
||||
}
|
||||
for index, character := range []byte(value) {
|
||||
if index == 8 || index == 13 || index == 18 || index == 23 {
|
||||
continue
|
||||
}
|
||||
if character < '0' || character > '9' && character < 'a' || character > 'f' {
|
||||
return fmt.Errorf("idempotency key must be a canonical lowercase UUID v4")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func newOperationID(random io.Reader) (string, error) {
|
||||
value := make([]byte, 16)
|
||||
if _, err := io.ReadFull(random, value); err != nil {
|
||||
return "", fmt.Errorf("generate operation id: %w", err)
|
||||
}
|
||||
value[6] = value[6]&0x0f | 0x40
|
||||
value[8] = value[8]&0x3f | 0x80
|
||||
encoded := hex.EncodeToString(value)
|
||||
return encoded[0:8] + "-" +
|
||||
encoded[8:12] + "-" +
|
||||
encoded[12:16] + "-" +
|
||||
encoded[16:20] + "-" +
|
||||
encoded[20:32], nil
|
||||
}
|
||||
|
||||
func (s *Service) executeIdempotentOperation(
|
||||
ctx context.Context,
|
||||
input idempotentOperationInput,
|
||||
) (idempotentOperationResult, error) {
|
||||
if err := validateIdempotencyKey(input.IdempotencyKey); err != nil {
|
||||
return idempotentOperationResult{}, app_errors.ErrInvalidIdempotencyKey
|
||||
}
|
||||
if input.DigestVersion != 1 || !input.Kind.valid() || input.Mutate == nil {
|
||||
return idempotentOperationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
|
||||
var existing models.ControlOperation
|
||||
query := s.db.WithContext(ctx).
|
||||
Where("idempotency_key = ?", input.IdempotencyKey).
|
||||
Take(&existing)
|
||||
if query.Error == nil {
|
||||
if err := validateOperationComparator(&existing, input); err != nil {
|
||||
return idempotentOperationResult{}, err
|
||||
}
|
||||
if blocked, recoveryErr := s.recoverPendingOperationsLocked(
|
||||
ctx,
|
||||
existing.CommitSequence,
|
||||
); recoveryErr != nil {
|
||||
return idempotentOperationResult{}, s.recoveryPendingError(*blocked)
|
||||
}
|
||||
return s.replayIdempotentOperationLocked(ctx, &existing, input)
|
||||
}
|
||||
if !errors.Is(query.Error, gorm.ErrRecordNotFound) {
|
||||
return idempotentOperationResult{}, app_errors.ParseDBError(query.Error)
|
||||
}
|
||||
if blocked, recoveryErr := s.recoverPendingOperationsLocked(ctx, 0); recoveryErr != nil {
|
||||
return idempotentOperationResult{}, s.recoveryPendingError(*blocked)
|
||||
}
|
||||
|
||||
requiredStages, err := operationRequiredStages(input.Kind)
|
||||
if err != nil {
|
||||
return idempotentOperationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
encodedStages, err := json.Marshal(requiredStages)
|
||||
if err != nil {
|
||||
return idempotentOperationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
operationID, err := newOperationID(s.operationRandom)
|
||||
if err != nil {
|
||||
return idempotentOperationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
|
||||
var mutationResult idempotentMutationResult
|
||||
var operation models.ControlOperation
|
||||
err = s.withControlTransaction(ctx, func(tx *gorm.DB) error {
|
||||
var mutationErr error
|
||||
mutationResult, mutationErr = input.Mutate(tx)
|
||||
if mutationErr != nil {
|
||||
return mutationErr
|
||||
}
|
||||
if mutationResult.ResourceIdentity == "" || len(mutationResult.CanonicalResult) == 0 {
|
||||
return fmt.Errorf(
|
||||
"idempotent mutation result is incomplete: %w",
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
canonicalResult, canonicalErr := canonicaljson.Canonicalize(
|
||||
mutationResult.CanonicalResult,
|
||||
)
|
||||
if canonicalErr != nil || !bytes.Equal(canonicalResult, mutationResult.CanonicalResult) {
|
||||
return fmt.Errorf(
|
||||
"idempotent mutation result is not canonical: %w",
|
||||
app_errors.ErrInternalServer,
|
||||
)
|
||||
}
|
||||
now := s.now().UTC()
|
||||
operation = models.ControlOperation{
|
||||
OperationID: operationID,
|
||||
IdempotencyKey: input.IdempotencyKey,
|
||||
DigestVersion: input.DigestVersion,
|
||||
RequestDigest: append([]byte(nil), input.RequestDigest[:]...),
|
||||
OperationKind: string(input.Kind),
|
||||
ResourceIdentity: mutationResult.ResourceIdentity,
|
||||
CanonicalResult: append([]byte(nil), mutationResult.CanonicalResult...),
|
||||
RequiredStages: models.JSON(encodedStages),
|
||||
LastCompletedStage: string(operationStageDBCommitted),
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if err := tx.Create(&operation).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return idempotentOperationResult{}, err
|
||||
}
|
||||
|
||||
if err := s.recoverOperationLocked(ctx, &operation); err != nil {
|
||||
s.wakeOperationRecovery()
|
||||
return idempotentOperationResult{}, s.operationIncompleteError(operation)
|
||||
}
|
||||
return idempotentOperationResult{
|
||||
OperationID: operation.OperationID,
|
||||
ResourceIdentity: operation.ResourceIdentity,
|
||||
CanonicalResult: append([]byte(nil), operation.CanonicalResult...),
|
||||
Ephemeral: mutationResult.Ephemeral,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (s *Service) replayIdempotentOperationLocked(
|
||||
ctx context.Context,
|
||||
operation *models.ControlOperation,
|
||||
input idempotentOperationInput,
|
||||
) (idempotentOperationResult, error) {
|
||||
if err := validateOperationComparator(operation, input); err != nil {
|
||||
return idempotentOperationResult{}, err
|
||||
}
|
||||
if operation.CompactedAt != nil {
|
||||
if operation.CompletedAt == nil {
|
||||
return idempotentOperationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
return idempotentOperationResult{}, app_errors.NewAPIErrorWithData(
|
||||
app_errors.ErrIdempotencyResultExpired,
|
||||
operationExpiredData{
|
||||
OperationID: operation.OperationID,
|
||||
OperationKind: operationKind(operation.OperationKind),
|
||||
ResourceIdentity: operation.ResourceIdentity,
|
||||
CompletedAt: operation.CompletedAt.UTC(),
|
||||
},
|
||||
)
|
||||
}
|
||||
if operation.LastCompletedStage != string(operationStageCompleted) {
|
||||
if err := s.recoverOperationLocked(ctx, operation); err != nil {
|
||||
s.wakeOperationRecovery()
|
||||
return idempotentOperationResult{}, s.operationIncompleteError(*operation)
|
||||
}
|
||||
}
|
||||
if operation.CompletedAt == nil || len(operation.CanonicalResult) == 0 {
|
||||
return idempotentOperationResult{}, app_errors.ErrInternalServer
|
||||
}
|
||||
return idempotentOperationResult{
|
||||
OperationID: operation.OperationID,
|
||||
ResourceIdentity: operation.ResourceIdentity,
|
||||
CanonicalResult: append([]byte(nil), operation.CanonicalResult...),
|
||||
Replayed: true,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func validateOperationComparator(
|
||||
operation *models.ControlOperation,
|
||||
input idempotentOperationInput,
|
||||
) error {
|
||||
if operation.DigestVersion != 1 ||
|
||||
input.DigestVersion != 1 ||
|
||||
len(operation.RequestDigest) != len(input.RequestDigest) ||
|
||||
!operationKind(operation.OperationKind).valid() {
|
||||
return app_errors.ErrInternalServer
|
||||
}
|
||||
if subtle.ConstantTimeCompare(operation.RequestDigest, input.RequestDigest[:]) != 1 {
|
||||
return app_errors.NewAPIErrorWithData(
|
||||
app_errors.ErrIdempotencyKeyReused,
|
||||
operationErrorData{
|
||||
OperationID: operation.OperationID,
|
||||
OperationKind: operationKind(operation.OperationKind),
|
||||
},
|
||||
)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) recoverOperationLocked(
|
||||
ctx context.Context,
|
||||
operation *models.ControlOperation,
|
||||
) error {
|
||||
if err := validateRecoverableOperation(operation); err != nil {
|
||||
return err
|
||||
}
|
||||
kind := operationKind(operation.OperationKind)
|
||||
wantStages, err := operationRequiredStages(kind)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var storedStages []operationStage
|
||||
if err := json.Unmarshal(operation.RequiredStages, &storedStages); err != nil ||
|
||||
!sameOperationStages(storedStages, wantStages) {
|
||||
return fmt.Errorf("invalid durable operation stage plan")
|
||||
}
|
||||
currentIndex := -1
|
||||
for index, stage := range storedStages {
|
||||
if string(stage) == operation.LastCompletedStage {
|
||||
currentIndex = index
|
||||
break
|
||||
}
|
||||
}
|
||||
if currentIndex < 0 {
|
||||
return fmt.Errorf("invalid durable operation current stage")
|
||||
}
|
||||
|
||||
for index := currentIndex + 1; index < len(storedStages); index++ {
|
||||
stage := storedStages[index]
|
||||
var stageErr error
|
||||
switch stage {
|
||||
case operationStageRegistryApplied:
|
||||
stageErr = s.recoverRegistryOperation(ctx, operation)
|
||||
case operationStageSnapshotPublished:
|
||||
compileInput, inputErr := stateloader.BuildCompileInput(ctx, s.db)
|
||||
if inputErr != nil {
|
||||
stageErr = inputErr
|
||||
} else {
|
||||
var matches bool
|
||||
matches, stageErr = s.manager.Matches(compileInput)
|
||||
if stageErr == nil && !matches {
|
||||
_, stageErr = s.publishSnapshot(compileInput)
|
||||
}
|
||||
}
|
||||
case operationStageCompleted:
|
||||
default:
|
||||
stageErr = fmt.Errorf("unsupported operation stage %q", stage)
|
||||
}
|
||||
if stageErr != nil {
|
||||
_ = s.recordOperationFailureLocked(ctx, operation, stage)
|
||||
return fmt.Errorf("recover operation stage %s", stage)
|
||||
}
|
||||
if err := s.advanceOperationStageLocked(ctx, operation, stage); err != nil {
|
||||
_ = s.recordOperationFailureLocked(ctx, operation, stage)
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) recoverRegistryOperation(
|
||||
ctx context.Context,
|
||||
operation *models.ControlOperation,
|
||||
) error {
|
||||
groupID, err := operationGroupID(operation)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
entries, err := stateloader.BuildGroupKeyEntries(ctx, s.db, groupID)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if s.registry.MatchesGroup(groupID, entries) {
|
||||
return nil
|
||||
}
|
||||
_, err = s.reconcileRegistryGroup(groupID, entries)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Service) advanceOperationStageLocked(
|
||||
ctx context.Context,
|
||||
operation *models.ControlOperation,
|
||||
stage operationStage,
|
||||
) error {
|
||||
if s.beforeAdvanceOperationStage != nil {
|
||||
if err := s.beforeAdvanceOperationStage(ctx, operation, stage); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
previous := operation.LastCompletedStage
|
||||
now := s.now().UTC()
|
||||
updates := map[string]any{
|
||||
"last_completed_stage": string(stage),
|
||||
"failed_stage": "",
|
||||
"updated_at": now,
|
||||
}
|
||||
if stage == operationStageCompleted {
|
||||
updates["completed_at"] = now
|
||||
}
|
||||
err := s.withControlTransaction(ctx, func(tx *gorm.DB) error {
|
||||
result := tx.Model(&models.ControlOperation{}).
|
||||
Where(
|
||||
"commit_sequence = ? AND last_completed_stage = ?",
|
||||
operation.CommitSequence,
|
||||
previous,
|
||||
).
|
||||
Updates(updates)
|
||||
if result.Error != nil {
|
||||
return app_errors.ParseDBError(result.Error)
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return fmt.Errorf("advance operation stage: durable stage changed")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
operation.LastCompletedStage = string(stage)
|
||||
operation.FailedStage = ""
|
||||
operation.UpdatedAt = now
|
||||
if stage == operationStageCompleted {
|
||||
operation.CompletedAt = &now
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) recordOperationFailureLocked(
|
||||
ctx context.Context,
|
||||
operation *models.ControlOperation,
|
||||
stage operationStage,
|
||||
) error {
|
||||
now := s.now().UTC()
|
||||
err := s.withControlTransaction(ctx, func(tx *gorm.DB) error {
|
||||
return tx.Model(&models.ControlOperation{}).
|
||||
Where("commit_sequence = ?", operation.CommitSequence).
|
||||
Updates(map[string]any{
|
||||
"failed_stage": string(stage),
|
||||
"updated_at": now,
|
||||
}).Error
|
||||
})
|
||||
if err == nil {
|
||||
operation.FailedStage = string(stage)
|
||||
operation.UpdatedAt = now
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Service) operationIncompleteError(
|
||||
operation models.ControlOperation,
|
||||
) error {
|
||||
return app_errors.NewAPIErrorWithData(
|
||||
app_errors.ErrControlOperationIncomplete,
|
||||
struct {
|
||||
OperationID string `json:"operation_id"`
|
||||
OperationKind operationKind `json:"operation_kind"`
|
||||
LastCompletedStage string `json:"last_completed_stage"`
|
||||
FailedStage string `json:"failed_stage"`
|
||||
CanReconcile bool `json:"can_reconcile"`
|
||||
}{
|
||||
OperationID: operation.OperationID,
|
||||
OperationKind: operationKind(operation.OperationKind),
|
||||
LastCompletedStage: operation.LastCompletedStage,
|
||||
FailedStage: operation.FailedStage,
|
||||
CanReconcile: true,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
func sameOperationStages(left, right []operationStage) bool {
|
||||
if len(left) != len(right) {
|
||||
return false
|
||||
}
|
||||
for index := range left {
|
||||
if left[index] != right[index] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
func TestOperationRequiredStagesAreKindSpecificAndDetached(t *testing.T) {
|
||||
tests := []struct {
|
||||
kind operationKind
|
||||
want []operationStage
|
||||
}{
|
||||
{
|
||||
kind: operationKindAccessKeyCreate,
|
||||
want: []operationStage{
|
||||
operationStageDBCommitted,
|
||||
operationStageSnapshotPublished,
|
||||
operationStageCompleted,
|
||||
},
|
||||
},
|
||||
{
|
||||
kind: operationKindGroupCreate,
|
||||
want: []operationStage{
|
||||
operationStageDBCommitted,
|
||||
operationStageRegistryApplied,
|
||||
operationStageSnapshotPublished,
|
||||
operationStageCompleted,
|
||||
},
|
||||
},
|
||||
{
|
||||
kind: operationKindGroupKeyImport,
|
||||
want: []operationStage{
|
||||
operationStageDBCommitted,
|
||||
operationStageRegistryApplied,
|
||||
operationStageCompleted,
|
||||
},
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(string(test.kind), func(t *testing.T) {
|
||||
got, err := operationRequiredStages(test.kind)
|
||||
if err != nil {
|
||||
t.Fatalf("operationRequiredStages() error = %v", err)
|
||||
}
|
||||
if !equalOperationStages(got, test.want) {
|
||||
t.Fatalf("stages = %#v, want %#v", got, test.want)
|
||||
}
|
||||
got[0] = operationStageCompleted
|
||||
again, err := operationRequiredStages(test.kind)
|
||||
if err != nil {
|
||||
t.Fatalf("second operationRequiredStages() error = %v", err)
|
||||
}
|
||||
if !equalOperationStages(again, test.want) {
|
||||
t.Fatalf("stage plan aliases caller: %#v", again)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
if _, err := operationRequiredStages(operationKind("unknown")); err == nil {
|
||||
t.Fatal("operationRequiredStages(unknown) error = nil, want rejection")
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateIdempotencyKeyRequiresCanonicalLowercaseUUIDV4(t *testing.T) {
|
||||
const valid = "018f47a2-9c35-4d6e-8b1a-1234567890ab"
|
||||
if err := validateIdempotencyKey(valid); err != nil {
|
||||
t.Fatalf("validateIdempotencyKey(valid) error = %v", err)
|
||||
}
|
||||
for _, invalid := range []string{
|
||||
"",
|
||||
" 018f47a2-9c35-4d6e-8b1a-1234567890ab",
|
||||
"018F47A2-9C35-4D6E-8B1A-1234567890AB",
|
||||
"018f47a2-9c35-3d6e-8b1a-1234567890ab",
|
||||
"018f47a2-9c35-4d6e-7b1a-1234567890ab",
|
||||
"018f47a29c354d6e8b1a1234567890ab",
|
||||
} {
|
||||
t.Run(invalid, func(t *testing.T) {
|
||||
if err := validateIdempotencyKey(invalid); err == nil {
|
||||
t.Fatalf("validateIdempotencyKey(%q) error = nil, want rejection", invalid)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewOperationIDUsesCanonicalUUIDV4WithoutSharingCredentialRandom(t *testing.T) {
|
||||
random := bytes.NewReader([]byte{
|
||||
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
|
||||
0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
|
||||
})
|
||||
got, err := newOperationID(random)
|
||||
if err != nil {
|
||||
t.Fatalf("newOperationID() error = %v", err)
|
||||
}
|
||||
if got != "00010203-0405-4607-8809-0a0b0c0d0e0f" {
|
||||
t.Fatalf("newOperationID() = %q", got)
|
||||
}
|
||||
if err := validateIdempotencyKey(got); err != nil {
|
||||
t.Fatalf("generated operation ID is not canonical UUID v4: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteIdempotentOperationCommitsResultAndReplaysWithoutMutation(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x42}, 16))
|
||||
const key = "018f47a2-9c35-4d6e-8b1a-1234567890ab"
|
||||
digest := sha256.Sum256([]byte("same request"))
|
||||
mutations := 0
|
||||
input := idempotentOperationInput{
|
||||
IdempotencyKey: key,
|
||||
DigestVersion: 1,
|
||||
RequestDigest: digest,
|
||||
Kind: operationKindAccessKeyCreate,
|
||||
Mutate: func(tx operationTransaction) (idempotentMutationResult, error) {
|
||||
mutations++
|
||||
return idempotentMutationResult{
|
||||
ResourceIdentity: "access-key:7",
|
||||
CanonicalResult: []byte(`{"id":7,"name":"client"}`),
|
||||
Ephemeral: "one-time-secret",
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
|
||||
first, err := fixture.service.executeIdempotentOperation(context.Background(), input)
|
||||
if err != nil {
|
||||
t.Fatalf("first executeIdempotentOperation() error = %v", err)
|
||||
}
|
||||
if first.Replayed || first.Ephemeral != "one-time-secret" ||
|
||||
string(first.CanonicalResult) != `{"id":7,"name":"client"}` {
|
||||
t.Fatalf("first result = %#v", first)
|
||||
}
|
||||
|
||||
replayInput := input
|
||||
replayInput.Mutate = func(operationTransaction) (idempotentMutationResult, error) {
|
||||
mutations++
|
||||
return idempotentMutationResult{}, errors.New("mutation reran")
|
||||
}
|
||||
replayed, err := fixture.service.executeIdempotentOperation(context.Background(), replayInput)
|
||||
if err != nil {
|
||||
t.Fatalf("replay executeIdempotentOperation() error = %v", err)
|
||||
}
|
||||
if !replayed.Replayed || replayed.Ephemeral != nil ||
|
||||
string(replayed.CanonicalResult) != string(first.CanonicalResult) {
|
||||
t.Fatalf("replay result = %#v", replayed)
|
||||
}
|
||||
if mutations != 1 {
|
||||
t.Fatalf("mutation calls = %d, want 1", mutations)
|
||||
}
|
||||
|
||||
var row models.ControlOperation
|
||||
if err := fixture.db.Where("idempotency_key = ?", key).Take(&row).Error; err != nil {
|
||||
t.Fatalf("read ControlOperation: %v", err)
|
||||
}
|
||||
if row.CommitSequence == 0 || row.ResourceIdentity != "access-key:7" ||
|
||||
row.LastCompletedStage != string(operationStageCompleted) ||
|
||||
row.CompletedAt == nil || row.CompactedAt != nil {
|
||||
t.Fatalf("operation row = %#v", row)
|
||||
}
|
||||
var stages []operationStage
|
||||
if err := json.Unmarshal(row.RequiredStages, &stages); err != nil {
|
||||
t.Fatalf("decode required stages: %v", err)
|
||||
}
|
||||
wantStages, _ := operationRequiredStages(operationKindAccessKeyCreate)
|
||||
if !equalOperationStages(stages, wantStages) {
|
||||
t.Fatalf("stored stages = %#v, want %#v", stages, wantStages)
|
||||
}
|
||||
if bytes.Contains(row.CanonicalResult, []byte("one-time-secret")) {
|
||||
t.Fatal("operation row persisted ephemeral secret")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteIdempotentOperationRejectsKeyReuseWithDifferentDigest(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x24}, 16))
|
||||
const key = "018f47a2-9c35-4d6e-8b1a-1234567890ab"
|
||||
firstDigest := sha256.Sum256([]byte("first"))
|
||||
input := idempotentOperationInput{
|
||||
IdempotencyKey: key, DigestVersion: 1, RequestDigest: firstDigest,
|
||||
Kind: operationKindAccessKeyCreate,
|
||||
Mutate: func(operationTransaction) (idempotentMutationResult, error) {
|
||||
return idempotentMutationResult{
|
||||
ResourceIdentity: "access-key:7",
|
||||
CanonicalResult: []byte(`{"id":7,"name":"client"}`),
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
if _, err := fixture.service.executeIdempotentOperation(t.Context(), input); err != nil {
|
||||
t.Fatalf("seed executeIdempotentOperation() error = %v", err)
|
||||
}
|
||||
|
||||
input.RequestDigest = sha256.Sum256([]byte("different"))
|
||||
_, err := fixture.service.executeIdempotentOperation(t.Context(), input)
|
||||
var apiErr *app_errors.APIError
|
||||
if !errors.As(err, &apiErr) || apiErr.Code != app_errors.ErrIdempotencyKeyReused.Code {
|
||||
t.Fatalf("different digest error = %v, want IDEMPOTENCY_KEY_REUSED", err)
|
||||
}
|
||||
}
|
||||
|
||||
func equalOperationStages(left, right []operationStage) bool {
|
||||
if len(left) != len(right) {
|
||||
return false
|
||||
}
|
||||
for index := range left {
|
||||
if left[index] != right[index] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
@@ -105,6 +105,9 @@ func (s *Service) writePriceTable(
|
||||
) error {
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
if blocked, err := s.recoverPendingOperationsLocked(ctx, 0); err != nil {
|
||||
return s.recoveryPendingError(*blocked)
|
||||
}
|
||||
|
||||
var table *pricing.Table
|
||||
err := s.withControlTransaction(ctx, func(tx *gorm.DB) error {
|
||||
|
||||
@@ -0,0 +1,259 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/sirupsen/logrus"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"gpt-load/internal/platform/canonicaljson"
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
const (
|
||||
operationResultRetention = 7 * 24 * time.Hour
|
||||
operationRecoveryInitialBackoff = 250 * time.Millisecond
|
||||
operationRecoveryMaximumBackoff = 30 * time.Second
|
||||
operationCompactionInterval = time.Hour
|
||||
)
|
||||
|
||||
type recoveryPendingData struct {
|
||||
OperationID string `json:"operation_id"`
|
||||
OperationKind operationKind `json:"operation_kind"`
|
||||
FailedStage string `json:"failed_stage"`
|
||||
RetryAfterMS int64 `json:"retry_after_ms"`
|
||||
}
|
||||
|
||||
// DrainCommittedOperations restores all known post-commit side effects in
|
||||
// commit order. It is called during startup before the process listens.
|
||||
func (s *Service) DrainCommittedOperations(ctx context.Context) error {
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
_, err := s.recoverPendingOperationsLocked(ctx, 0)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Service) recoverPendingOperationsLocked(
|
||||
ctx context.Context,
|
||||
beforeCommitSequence uint64,
|
||||
) (*models.ControlOperation, error) {
|
||||
query := s.db.WithContext(ctx).
|
||||
Where("completed_at IS NULL").
|
||||
Order("commit_sequence ASC")
|
||||
if beforeCommitSequence != 0 {
|
||||
query = query.Where("commit_sequence < ?", beforeCommitSequence)
|
||||
}
|
||||
var operations []models.ControlOperation
|
||||
if err := query.Find(&operations).Error; err != nil {
|
||||
return nil, app_errors.ParseDBError(err)
|
||||
}
|
||||
for index := range operations {
|
||||
operation := &operations[index]
|
||||
if err := s.recoverOperationLocked(ctx, operation); err != nil {
|
||||
return operation, err
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (s *Service) recoveryPendingError(
|
||||
operation models.ControlOperation,
|
||||
) error {
|
||||
return app_errors.NewAPIErrorWithData(
|
||||
app_errors.ErrControlRecoveryPending,
|
||||
recoveryPendingData{
|
||||
OperationID: operation.OperationID,
|
||||
OperationKind: operationKind(operation.OperationKind),
|
||||
FailedStage: operation.FailedStage,
|
||||
RetryAfterMS: operationRecoveryInitialBackoff.Milliseconds(),
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
// CompactCompletedOperations removes replay payloads after the guaranteed
|
||||
// retention period while preserving the permanent request comparator.
|
||||
func (s *Service) CompactCompletedOperations(
|
||||
ctx context.Context,
|
||||
now time.Time,
|
||||
) (int64, error) {
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
return s.compactCompletedOperationsLocked(ctx, now)
|
||||
}
|
||||
|
||||
func (s *Service) compactCompletedOperationsLocked(
|
||||
ctx context.Context,
|
||||
now time.Time,
|
||||
) (int64, error) {
|
||||
now = now.UTC()
|
||||
cutoff := now.Add(-operationResultRetention)
|
||||
var rowsAffected int64
|
||||
err := s.withControlTransaction(ctx, func(tx *gorm.DB) error {
|
||||
result := tx.Model(&models.ControlOperation{}).
|
||||
Where("completed_at IS NOT NULL").
|
||||
Where("completed_at <= ?", cutoff).
|
||||
Where("compacted_at IS NULL").
|
||||
Updates(map[string]any{
|
||||
"canonical_result": gorm.Expr("NULL"),
|
||||
"required_stages": gorm.Expr("NULL"),
|
||||
"last_completed_stage": "",
|
||||
"failed_stage": "",
|
||||
"compacted_at": now,
|
||||
"updated_at": now,
|
||||
})
|
||||
if result.Error != nil {
|
||||
return app_errors.ParseDBError(result.Error)
|
||||
}
|
||||
rowsAffected = result.RowsAffected
|
||||
return nil
|
||||
})
|
||||
return rowsAffected, err
|
||||
}
|
||||
|
||||
// RunOperationRecovery drains failures signaled by mutation requests and
|
||||
// periodically compacts terminal results. Retries are bounded and stop with
|
||||
// the application context.
|
||||
func (s *Service) RunOperationRecovery(ctx context.Context) {
|
||||
compactionTicker := time.NewTicker(operationCompactionInterval)
|
||||
defer compactionTicker.Stop()
|
||||
|
||||
var retryTimer *time.Timer
|
||||
var retry <-chan time.Time
|
||||
backoff := operationRecoveryInitialBackoff
|
||||
stopRetry := func() {
|
||||
if retryTimer != nil {
|
||||
if !retryTimer.Stop() {
|
||||
select {
|
||||
case <-retryTimer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
retryTimer = nil
|
||||
retry = nil
|
||||
}
|
||||
defer stopRetry()
|
||||
|
||||
drain := func() {
|
||||
if err := s.DrainCommittedOperations(ctx); err != nil {
|
||||
logrus.WithField("retry_after_ms", backoff.Milliseconds()).
|
||||
Warn("Control operation recovery remains pending")
|
||||
stopRetry()
|
||||
retryTimer = time.NewTimer(backoff)
|
||||
retry = retryTimer.C
|
||||
backoff = nextOperationRecoveryBackoff(backoff)
|
||||
return
|
||||
}
|
||||
stopRetry()
|
||||
backoff = operationRecoveryInitialBackoff
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-s.operationRecoveryWake:
|
||||
drain()
|
||||
case <-retry:
|
||||
retryTimer = nil
|
||||
retry = nil
|
||||
drain()
|
||||
case now := <-compactionTicker.C:
|
||||
if _, err := s.CompactCompletedOperations(ctx, now); err != nil &&
|
||||
ctx.Err() == nil {
|
||||
logrus.Warn("Control operation result compaction failed")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func nextOperationRecoveryBackoff(current time.Duration) time.Duration {
|
||||
if current >= operationRecoveryMaximumBackoff/2 {
|
||||
return operationRecoveryMaximumBackoff
|
||||
}
|
||||
return current * 2
|
||||
}
|
||||
|
||||
func (s *Service) wakeOperationRecovery() {
|
||||
if s.operationRecoveryWake == nil {
|
||||
return
|
||||
}
|
||||
select {
|
||||
case s.operationRecoveryWake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func validateRecoverableOperation(operation *models.ControlOperation) error {
|
||||
if operation == nil ||
|
||||
operation.DigestVersion != 1 ||
|
||||
len(operation.RequestDigest) != 32 ||
|
||||
!operationKind(operation.OperationKind).valid() ||
|
||||
operation.CompactedAt != nil {
|
||||
return fmt.Errorf("invalid durable operation comparator: %w", app_errors.ErrInternalServer)
|
||||
}
|
||||
if err := validateIdempotencyKey(operation.OperationID); err != nil {
|
||||
return fmt.Errorf("invalid durable operation id: %w", app_errors.ErrInternalServer)
|
||||
}
|
||||
if err := validateIdempotencyKey(operation.IdempotencyKey); err != nil {
|
||||
return fmt.Errorf("invalid durable idempotency key: %w", app_errors.ErrInternalServer)
|
||||
}
|
||||
if err := validateOperationResourceIdentity(
|
||||
operationKind(operation.OperationKind),
|
||||
operation.ResourceIdentity,
|
||||
); err != nil {
|
||||
return fmt.Errorf("invalid durable resource identity: %w", app_errors.ErrInternalServer)
|
||||
}
|
||||
if len(operation.CanonicalResult) == 0 {
|
||||
return fmt.Errorf("missing durable operation result: %w", app_errors.ErrInternalServer)
|
||||
}
|
||||
canonical, err := canonicaljson.Canonicalize(operation.CanonicalResult)
|
||||
if err != nil || !bytes.Equal(canonical, operation.CanonicalResult) {
|
||||
return fmt.Errorf("invalid durable operation result: %w", app_errors.ErrInternalServer)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func operationGroupID(operation *models.ControlOperation) (uint, error) {
|
||||
kind := operationKind(operation.OperationKind)
|
||||
if kind != operationKindGroupCreate && kind != operationKindGroupKeyImport {
|
||||
return 0, fmt.Errorf("operation %q has no group registry stage", kind)
|
||||
}
|
||||
return parseResourceIdentity(operation.ResourceIdentity, "group")
|
||||
}
|
||||
|
||||
func validateOperationResourceIdentity(kind operationKind, identity string) error {
|
||||
switch kind {
|
||||
case operationKindAccessKeyCreate:
|
||||
_, err := parseResourceIdentity(identity, "access-key")
|
||||
return err
|
||||
case operationKindGroupCreate, operationKindGroupKeyImport:
|
||||
_, err := parseResourceIdentity(identity, "group")
|
||||
return err
|
||||
default:
|
||||
return fmt.Errorf("unsupported operation kind")
|
||||
}
|
||||
}
|
||||
|
||||
func parseResourceIdentity(identity, prefix string) (uint, error) {
|
||||
wantPrefix := prefix + ":"
|
||||
if !strings.HasPrefix(identity, wantPrefix) {
|
||||
return 0, fmt.Errorf("resource identity must start with %q", wantPrefix)
|
||||
}
|
||||
raw := strings.TrimPrefix(identity, wantPrefix)
|
||||
if raw == "" || raw[0] == '0' {
|
||||
return 0, fmt.Errorf("resource identity id must be canonical base10")
|
||||
}
|
||||
value, err := strconv.ParseUint(raw, 10, strconv.IntSize)
|
||||
if err != nil || value == 0 ||
|
||||
strconv.FormatUint(value, 10) != raw {
|
||||
return 0, fmt.Errorf("resource identity id must be canonical base10")
|
||||
}
|
||||
return uint(value), nil
|
||||
}
|
||||
@@ -0,0 +1,384 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/protocol"
|
||||
"gpt-load/internal/state"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
func TestOperationReplayRepairsPostCommitRegistryFailureWithoutRerunningMutation(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x31}, 16))
|
||||
reconcileCalls := 0
|
||||
fixture.service.reconcileRegistryGroup = func(uint, []state.KeyEntry) (bool, error) {
|
||||
reconcileCalls++
|
||||
return false, errors.New("injected registry failure")
|
||||
}
|
||||
mutations := 0
|
||||
input := newDurableGroupOperationInput(
|
||||
t,
|
||||
fixture,
|
||||
"018f47a2-9c35-4d6e-8b1a-1234567890ab",
|
||||
&mutations,
|
||||
)
|
||||
|
||||
_, err := fixture.service.executeIdempotentOperation(t.Context(), input)
|
||||
assertAPIErrorCode(t, err, app_errors.ErrControlOperationIncomplete.Code)
|
||||
if mutations != 1 || reconcileCalls != 1 {
|
||||
t.Fatalf("mutation/reconcile calls = %d/%d, want 1/1", mutations, reconcileCalls)
|
||||
}
|
||||
|
||||
var operation models.ControlOperation
|
||||
if err := fixture.db.First(&operation).Error; err != nil {
|
||||
t.Fatalf("read operation: %v", err)
|
||||
}
|
||||
if operation.LastCompletedStage != string(operationStageDBCommitted) ||
|
||||
operation.FailedStage != string(operationStageRegistryApplied) {
|
||||
t.Fatalf("failed operation = %#v", operation)
|
||||
}
|
||||
var groupCount int64
|
||||
if err := fixture.db.Model(&models.Group{}).Count(&groupCount).Error; err != nil {
|
||||
t.Fatalf("count groups: %v", err)
|
||||
}
|
||||
if groupCount != 1 {
|
||||
t.Fatalf("committed group count = %d, want 1", groupCount)
|
||||
}
|
||||
|
||||
fixture.service.reconcileRegistryGroup = fixture.registry.ReconcileGroup
|
||||
replayed, err := fixture.service.executeIdempotentOperation(t.Context(), input)
|
||||
if err != nil {
|
||||
t.Fatalf("replay executeIdempotentOperation() error = %v", err)
|
||||
}
|
||||
if !replayed.Replayed || mutations != 1 {
|
||||
t.Fatalf("replay/mutation = %t/%d, want true/1", replayed.Replayed, mutations)
|
||||
}
|
||||
groupID := mustResourceGroupID(t, replayed.ResourceIdentity)
|
||||
if value, ok := fixture.registry.EncryptedValue(1); !ok || value != "cipher-one" {
|
||||
t.Fatalf("recovered registry key = %q, %t", value, ok)
|
||||
}
|
||||
if current := fixture.manager.Current(); current == nil || current.Groups[groupID].ID != groupID {
|
||||
t.Fatalf("recovered snapshot = %#v", current)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOperationReplayOnlyRecordsStageWhenRegistrySideEffectAlreadySatisfied(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x32}, 16))
|
||||
reconcileCalls := 0
|
||||
fixture.service.reconcileRegistryGroup = func(
|
||||
groupID uint,
|
||||
entries []state.KeyEntry,
|
||||
) (bool, error) {
|
||||
reconcileCalls++
|
||||
return fixture.registry.ReconcileGroup(groupID, entries)
|
||||
}
|
||||
failedStageWrite := false
|
||||
fixture.service.beforeAdvanceOperationStage = func(
|
||||
_ context.Context,
|
||||
_ *models.ControlOperation,
|
||||
stage operationStage,
|
||||
) error {
|
||||
if stage == operationStageRegistryApplied && !failedStageWrite {
|
||||
failedStageWrite = true
|
||||
return errors.New("injected stage write failure")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
mutations := 0
|
||||
input := newDurableGroupOperationInput(
|
||||
t,
|
||||
fixture,
|
||||
"028f47a2-9c35-4d6e-8b1a-1234567890ab",
|
||||
&mutations,
|
||||
)
|
||||
|
||||
_, err := fixture.service.executeIdempotentOperation(t.Context(), input)
|
||||
assertAPIErrorCode(t, err, app_errors.ErrControlOperationIncomplete.Code)
|
||||
if reconcileCalls != 1 {
|
||||
t.Fatalf("initial reconcile calls = %d, want 1", reconcileCalls)
|
||||
}
|
||||
if value, ok := fixture.registry.EncryptedValue(1); !ok || value != "cipher-one" {
|
||||
t.Fatalf("side effect did not complete before stage failure: %q, %t", value, ok)
|
||||
}
|
||||
|
||||
fixture.service.beforeAdvanceOperationStage = nil
|
||||
replayed, err := fixture.service.executeIdempotentOperation(t.Context(), input)
|
||||
if err != nil {
|
||||
t.Fatalf("replay executeIdempotentOperation() error = %v", err)
|
||||
}
|
||||
if !replayed.Replayed || mutations != 1 {
|
||||
t.Fatalf("replay/mutation = %t/%d, want true/1", replayed.Replayed, mutations)
|
||||
}
|
||||
if reconcileCalls != 1 {
|
||||
t.Fatalf("replay reapplied satisfied registry side effect, calls = %d", reconcileCalls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOperationBarrierBlocksNewMutationBehindOlderRecovery(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x33}, 32))
|
||||
fixture.service.reconcileRegistryGroup = func(uint, []state.KeyEntry) (bool, error) {
|
||||
return false, errors.New("registry remains unavailable")
|
||||
}
|
||||
firstMutations := 0
|
||||
first := newDurableGroupOperationInput(
|
||||
t,
|
||||
fixture,
|
||||
"038f47a2-9c35-4d6e-8b1a-1234567890ab",
|
||||
&firstMutations,
|
||||
)
|
||||
_, err := fixture.service.executeIdempotentOperation(t.Context(), first)
|
||||
assertAPIErrorCode(t, err, app_errors.ErrControlOperationIncomplete.Code)
|
||||
|
||||
secondMutations := 0
|
||||
second := idempotentOperationInput{
|
||||
IdempotencyKey: "048f47a2-9c35-4d6e-8b1a-1234567890ab",
|
||||
DigestVersion: 1,
|
||||
RequestDigest: [32]byte{4},
|
||||
Kind: operationKindAccessKeyCreate,
|
||||
Mutate: func(operationTransaction) (idempotentMutationResult, error) {
|
||||
secondMutations++
|
||||
return idempotentMutationResult{
|
||||
ResourceIdentity: "access-key:9",
|
||||
CanonicalResult: []byte(`{"id":9}`),
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
_, err = fixture.service.executeIdempotentOperation(t.Context(), second)
|
||||
assertAPIErrorCode(t, err, app_errors.ErrControlRecoveryPending.Code)
|
||||
if secondMutations != 0 {
|
||||
t.Fatalf("new mutation ran behind failed barrier %d times", secondMutations)
|
||||
}
|
||||
var operationCount int64
|
||||
if err := fixture.db.Model(&models.ControlOperation{}).Count(&operationCount).Error; err != nil {
|
||||
t.Fatalf("count operations: %v", err)
|
||||
}
|
||||
if operationCount != 1 {
|
||||
t.Fatalf("operation count = %d, want only older operation", operationCount)
|
||||
}
|
||||
assertRecoveryPendingDataShape(t, err)
|
||||
|
||||
ordinaryMutations := 0
|
||||
_, err = fixture.service.writeConfig(t.Context(), func(operationTransaction) error {
|
||||
ordinaryMutations++
|
||||
return nil
|
||||
}, nil)
|
||||
assertAPIErrorCode(t, err, app_errors.ErrControlRecoveryPending.Code)
|
||||
if ordinaryMutations != 0 {
|
||||
t.Fatalf("ordinary control mutation ran behind failed barrier %d times", ordinaryMutations)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDrainCommittedOperationsFailsClosedOnUnsupportedDigestVersion(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
stages, err := json.Marshal([]operationStage{
|
||||
operationStageDBCommitted,
|
||||
operationStageSnapshotPublished,
|
||||
operationStageCompleted,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("marshal stages: %v", err)
|
||||
}
|
||||
now := time.Date(2026, time.July, 28, 12, 0, 0, 0, time.UTC)
|
||||
row := models.ControlOperation{
|
||||
OperationID: "058f47a2-9c35-4d6e-8b1a-1234567890ab",
|
||||
IdempotencyKey: "068f47a2-9c35-4d6e-8b1a-1234567890ab",
|
||||
DigestVersion: 99,
|
||||
RequestDigest: bytes.Repeat([]byte{0x6}, 32),
|
||||
OperationKind: string(operationKindAccessKeyCreate),
|
||||
ResourceIdentity: "access-key:1",
|
||||
CanonicalResult: []byte(`{"id":1}`),
|
||||
RequiredStages: models.JSON(stages),
|
||||
LastCompletedStage: string(operationStageDBCommitted),
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
if err := fixture.db.Create(&row).Error; err != nil {
|
||||
t.Fatalf("seed unsupported operation: %v", err)
|
||||
}
|
||||
|
||||
err = fixture.service.DrainCommittedOperations(t.Context())
|
||||
if err == nil || !errors.Is(err, app_errors.ErrInternalServer) {
|
||||
t.Fatalf("DrainCommittedOperations() error = %v, want fail closed", err)
|
||||
}
|
||||
var stored models.ControlOperation
|
||||
if err := fixture.db.First(&stored).Error; err != nil {
|
||||
t.Fatalf("read operation: %v", err)
|
||||
}
|
||||
if stored.LastCompletedStage != string(operationStageDBCommitted) {
|
||||
t.Fatalf("unsupported operation advanced to %q", stored.LastCompletedStage)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOperationRecoveryBackoffIsExponentialAndBounded(t *testing.T) {
|
||||
backoff := operationRecoveryInitialBackoff
|
||||
want := []time.Duration{
|
||||
500 * time.Millisecond,
|
||||
time.Second,
|
||||
2 * time.Second,
|
||||
4 * time.Second,
|
||||
8 * time.Second,
|
||||
16 * time.Second,
|
||||
30 * time.Second,
|
||||
30 * time.Second,
|
||||
}
|
||||
for index, expected := range want {
|
||||
backoff = nextOperationRecoveryBackoff(backoff)
|
||||
if backoff != expected {
|
||||
t.Fatalf("backoff step %d = %s, want %s", index, backoff, expected)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompactCompletedOperationsKeepsPermanentComparatorTombstone(t *testing.T) {
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x34}, 16))
|
||||
fixture.service.now = func() time.Time {
|
||||
return time.Date(2026, time.July, 20, 12, 0, 0, 0, time.UTC)
|
||||
}
|
||||
mutations := 0
|
||||
input := idempotentOperationInput{
|
||||
IdempotencyKey: "078f47a2-9c35-4d6e-8b1a-1234567890ab",
|
||||
DigestVersion: 1,
|
||||
RequestDigest: [32]byte{7},
|
||||
Kind: operationKindAccessKeyCreate,
|
||||
Mutate: func(operationTransaction) (idempotentMutationResult, error) {
|
||||
mutations++
|
||||
return idempotentMutationResult{
|
||||
ResourceIdentity: "access-key:7",
|
||||
CanonicalResult: []byte(`{"id":7}`),
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
if _, err := fixture.service.executeIdempotentOperation(t.Context(), input); err != nil {
|
||||
t.Fatalf("execute operation: %v", err)
|
||||
}
|
||||
|
||||
compacted, err := fixture.service.CompactCompletedOperations(
|
||||
t.Context(),
|
||||
time.Date(2026, time.July, 28, 12, 0, 0, 0, time.UTC),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("CompactCompletedOperations() error = %v", err)
|
||||
}
|
||||
if compacted != 1 {
|
||||
t.Fatalf("compacted rows = %d, want 1", compacted)
|
||||
}
|
||||
var row models.ControlOperation
|
||||
if err := fixture.db.First(&row).Error; err != nil {
|
||||
t.Fatalf("read compacted operation: %v", err)
|
||||
}
|
||||
if row.CompactedAt == nil || len(row.CanonicalResult) != 0 ||
|
||||
len(row.RequiredStages) != 0 || row.LastCompletedStage != "" ||
|
||||
row.OperationID == "" || row.ResourceIdentity != "access-key:7" ||
|
||||
len(row.RequestDigest) != 32 {
|
||||
t.Fatalf("compacted row = %#v", row)
|
||||
}
|
||||
|
||||
_, err = fixture.service.executeIdempotentOperation(t.Context(), input)
|
||||
assertAPIErrorCode(t, err, app_errors.ErrIdempotencyResultExpired.Code)
|
||||
if mutations != 1 {
|
||||
t.Fatalf("expired replay reran mutation %d times", mutations)
|
||||
}
|
||||
input.RequestDigest = [32]byte{8}
|
||||
_, err = fixture.service.executeIdempotentOperation(t.Context(), input)
|
||||
assertAPIErrorCode(t, err, app_errors.ErrIdempotencyKeyReused.Code)
|
||||
}
|
||||
|
||||
func newDurableGroupOperationInput(
|
||||
t *testing.T,
|
||||
fixture serviceFixture,
|
||||
idempotencyKey string,
|
||||
mutations *int,
|
||||
) idempotentOperationInput {
|
||||
t.Helper()
|
||||
return idempotentOperationInput{
|
||||
IdempotencyKey: idempotencyKey,
|
||||
DigestVersion: 1,
|
||||
RequestDigest: [32]byte{byte(idempotencyKey[0])},
|
||||
Kind: operationKindGroupCreate,
|
||||
Mutate: func(tx operationTransaction) (idempotentMutationResult, error) {
|
||||
*mutations++
|
||||
group := models.Group{
|
||||
Name: "group-" + idempotencyKey[:4],
|
||||
UpstreamURL: "https://upstream.example.com",
|
||||
Protocols: models.JSON(`["` + string(protocol.OpenAI) + `"]`),
|
||||
Models: models.JSON(`[]`),
|
||||
Config: models.JSON(`{}`),
|
||||
Enabled: true,
|
||||
}
|
||||
if err := tx.Create(&group).Error; err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
key := models.UpstreamKey{
|
||||
GroupID: group.ID, KeyValue: "cipher-one", KeyHash: "hash-one",
|
||||
Status: models.UpstreamKeyStatusActive,
|
||||
}
|
||||
if err := tx.Create(&key).Error; err != nil {
|
||||
return idempotentMutationResult{}, err
|
||||
}
|
||||
result := []byte(fmt.Sprintf(`{"group_id":%d}`, group.ID))
|
||||
return idempotentMutationResult{
|
||||
ResourceIdentity: "group:" + strconv.FormatUint(uint64(group.ID), 10),
|
||||
CanonicalResult: result,
|
||||
}, nil
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func mustResourceGroupID(t *testing.T, identity string) uint {
|
||||
t.Helper()
|
||||
const prefix = "group:"
|
||||
if len(identity) <= len(prefix) || identity[:len(prefix)] != prefix {
|
||||
t.Fatalf("resource identity = %q", identity)
|
||||
}
|
||||
value, err := strconv.ParseUint(identity[len(prefix):], 10, 64)
|
||||
if err != nil {
|
||||
t.Fatalf("parse resource identity: %v", err)
|
||||
}
|
||||
return uint(value)
|
||||
}
|
||||
|
||||
func assertAPIErrorCode(t *testing.T, err error, want string) {
|
||||
t.Helper()
|
||||
var apiErr *app_errors.APIError
|
||||
if !errors.As(err, &apiErr) || apiErr.Code != want {
|
||||
t.Fatalf("error = %v, want API code %s", err, want)
|
||||
}
|
||||
}
|
||||
|
||||
func assertRecoveryPendingDataShape(t *testing.T, err error) {
|
||||
t.Helper()
|
||||
var apiErr *app_errors.APIError
|
||||
if !errors.As(err, &apiErr) {
|
||||
t.Fatalf("error = %v, want APIError", err)
|
||||
}
|
||||
raw, marshalErr := json.Marshal(apiErr.Data)
|
||||
if marshalErr != nil {
|
||||
t.Fatalf("marshal recovery data: %v", marshalErr)
|
||||
}
|
||||
var fields map[string]any
|
||||
if unmarshalErr := json.Unmarshal(raw, &fields); unmarshalErr != nil {
|
||||
t.Fatalf("unmarshal recovery data: %v", unmarshalErr)
|
||||
}
|
||||
for _, name := range []string{
|
||||
"operation_id", "operation_kind", "failed_stage", "retry_after_ms",
|
||||
} {
|
||||
if _, ok := fields[name]; !ok {
|
||||
t.Errorf("recovery data missing %q: %s", name, raw)
|
||||
}
|
||||
}
|
||||
if len(fields) != 4 {
|
||||
t.Fatalf("recovery data fields = %v, want exact frozen shape", fields)
|
||||
}
|
||||
}
|
||||
@@ -28,6 +28,10 @@ type keyMutationCoordinator interface {
|
||||
Do(uint, func())
|
||||
}
|
||||
|
||||
type operationRecoveryRuntime interface {
|
||||
RunOperationRecovery(context.Context)
|
||||
}
|
||||
|
||||
// RequestLogCleaner is the control-owned scheduling view of request log
|
||||
// retention. The requestlog package owns all cleanup semantics.
|
||||
type RequestLogCleaner interface {
|
||||
@@ -57,6 +61,7 @@ type Runtime struct {
|
||||
mutations keyMutationCoordinator
|
||||
validator validationSweep
|
||||
requestLogCleaner RequestLogCleaner
|
||||
operationRecovery operationRecoveryRuntime
|
||||
autoWeightInterval time.Duration
|
||||
validationInterval time.Duration
|
||||
validationJitter func() time.Duration
|
||||
@@ -72,12 +77,14 @@ func NewRuntime(
|
||||
encryptionService encryption.Service,
|
||||
dialects dialect.Set,
|
||||
requestLogCleaner RequestLogCleaner,
|
||||
operationRecovery *Service,
|
||||
) *Runtime {
|
||||
runtime := &Runtime{
|
||||
registry: registry,
|
||||
stats: stats,
|
||||
mutations: mutations,
|
||||
requestLogCleaner: requestLogCleaner,
|
||||
operationRecovery: operationRecovery,
|
||||
autoWeightInterval: autoWeightInterval,
|
||||
validationInterval: validationInterval,
|
||||
validationJitter: func() time.Duration {
|
||||
@@ -121,6 +128,13 @@ func (runtime *Runtime) Run(ctx context.Context) {
|
||||
runtime.runRetention(ctx, retentionTicker)
|
||||
}()
|
||||
}
|
||||
if runtime.operationRecovery != nil {
|
||||
wait.Add(1)
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
runtime.operationRecovery.RunOperationRecovery(ctx)
|
||||
}()
|
||||
}
|
||||
wait.Wait()
|
||||
}
|
||||
|
||||
|
||||
@@ -172,6 +172,38 @@ func (cleaner *controlledRequestLogCleaner) Sweep(ctx context.Context, now time.
|
||||
cleaner.returned <- struct{}{}
|
||||
}
|
||||
|
||||
type controlledOperationRecovery struct {
|
||||
started chan struct{}
|
||||
returned chan struct{}
|
||||
}
|
||||
|
||||
func (recovery *controlledOperationRecovery) RunOperationRecovery(ctx context.Context) {
|
||||
close(recovery.started)
|
||||
<-ctx.Done()
|
||||
close(recovery.returned)
|
||||
}
|
||||
|
||||
func TestRuntimeRunsOperationRecoveryUntilCancellation(t *testing.T) {
|
||||
recovery := &controlledOperationRecovery{
|
||||
started: make(chan struct{}),
|
||||
returned: make(chan struct{}),
|
||||
}
|
||||
runtime, _, _, created := newRuntimeHarness(
|
||||
newFakeAutoWeightRegistry(1),
|
||||
health.NewStatsStore(),
|
||||
newFakeValidationSweep(false),
|
||||
time.Now,
|
||||
)
|
||||
runtime.operationRecovery = recovery
|
||||
|
||||
cancel, done := startRuntime(t, runtime)
|
||||
awaitTickers(t, created)
|
||||
awaitSignal(t, recovery.started)
|
||||
cancel()
|
||||
awaitSignal(t, recovery.returned)
|
||||
awaitSignal(t, done)
|
||||
}
|
||||
|
||||
func TestRuntimeSweepsRequestLogsImmediatelyAndHourlyWithoutOverlap(t *testing.T) {
|
||||
base := time.Date(2026, time.July, 24, 12, 0, 0, 0, time.UTC)
|
||||
clock := &fakeRuntimeClock{now: base}
|
||||
|
||||
+140
-6
@@ -66,6 +66,8 @@ func (s *Server) RegisterRoutes(engine *gin.Engine) {
|
||||
api.POST("/groups/:group_id/models/discover", s.handleDiscoverGroupModels)
|
||||
api.POST("/models/discover", s.handleDiscoverModels)
|
||||
api.POST("/access-keys", s.handleCreateAccessKey)
|
||||
api.GET("/access-keys/options", s.handleListAccessKeyOptions)
|
||||
api.POST("/access-keys/:id/reveal", s.handleRevealAccessKey)
|
||||
api.GET("/access-keys", s.handleListAccessKeys)
|
||||
api.PUT("/access-keys/:id", s.handleUpdateAccessKey)
|
||||
api.DELETE("/access-keys/:id", s.handleDeleteAccessKey)
|
||||
@@ -77,21 +79,38 @@ func (s *Server) handleGetSettings(c *gin.Context) {
|
||||
writeServiceError(c, "get_settings", err)
|
||||
return
|
||||
}
|
||||
response.SuccessI18n(c, "common.success", result)
|
||||
representation, err := newSettingsWireRepresentation(
|
||||
i18n.Message(c, "common.success"),
|
||||
result.DTO(),
|
||||
)
|
||||
if err != nil {
|
||||
writeServiceError(c, "get_settings", err)
|
||||
return
|
||||
}
|
||||
writeSettingsRepresentation(c, representation)
|
||||
}
|
||||
|
||||
func (s *Server) handleUpdateSettings(c *gin.Context) {
|
||||
expectedETag, ok := requiredSettingsIfMatch(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var request SettingsUpdateRequest
|
||||
if err := bindStrictJSON(c, &request); err != nil {
|
||||
writeServiceError(c, "update_settings", mapControlJSONError(err))
|
||||
return
|
||||
}
|
||||
result, err := s.service.UpdateSettings(c.Request.Context(), request)
|
||||
result, err := s.service.UpdateSettingsIfMatch(
|
||||
c.Request.Context(),
|
||||
request,
|
||||
expectedETag,
|
||||
i18n.Message(c, "common.success"),
|
||||
)
|
||||
if err != nil {
|
||||
writeServiceError(c, "update_settings", err)
|
||||
return
|
||||
}
|
||||
response.SuccessI18n(c, "common.success", result)
|
||||
writeSettingsRepresentation(c, result)
|
||||
}
|
||||
|
||||
func (s *Server) handleListGroups(c *gin.Context) {
|
||||
@@ -117,12 +136,20 @@ func (s *Server) handleGetGroup(c *gin.Context) {
|
||||
}
|
||||
|
||||
func (s *Server) handleCreateGroup(c *gin.Context) {
|
||||
idempotencyKey, ok := requiredIdempotencyKey(c, "create_group")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var request GroupCreateRequest
|
||||
if err := bindStrictJSON(c, &request); err != nil {
|
||||
writeServiceError(c, "create_group", mapControlJSONError(err))
|
||||
return
|
||||
}
|
||||
result, err := s.service.CreateGroup(c.Request.Context(), request)
|
||||
result, err := s.service.CreateGroupIdempotent(
|
||||
c.Request.Context(),
|
||||
idempotencyKey,
|
||||
request,
|
||||
)
|
||||
if err != nil {
|
||||
writeServiceError(c, "create_group", err)
|
||||
return
|
||||
@@ -243,12 +270,21 @@ func (s *Server) handleImportGroupKeys(c *gin.Context) {
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
idempotencyKey, ok := requiredIdempotencyKey(c, "import_group_keys")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var request GroupKeyImportRequest
|
||||
if err := bindStrictJSON(c, &request); err != nil {
|
||||
writeServiceError(c, "import_group_keys", mapControlJSONError(err))
|
||||
return
|
||||
}
|
||||
result, err := s.service.ImportGroupKeys(c.Request.Context(), id, request)
|
||||
result, err := s.service.ImportGroupKeysIdempotent(
|
||||
c.Request.Context(),
|
||||
idempotencyKey,
|
||||
id,
|
||||
request,
|
||||
)
|
||||
if err != nil {
|
||||
writeServiceError(c, "import_group_keys", err)
|
||||
return
|
||||
@@ -288,16 +324,48 @@ func (s *Server) handleDiscoverModels(c *gin.Context) {
|
||||
}
|
||||
|
||||
func (s *Server) handleCreateAccessKey(c *gin.Context) {
|
||||
idempotencyKey, ok := requiredIdempotencyKey(c, "create_access_key")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
var request AccessKeyCreateRequest
|
||||
if err := bindStrictJSON(c, &request); err != nil {
|
||||
writeServiceError(c, "create_access_key", mapControlJSONError(err))
|
||||
return
|
||||
}
|
||||
result, err := s.service.CreateAccessKey(c.Request.Context(), request)
|
||||
result, err := s.service.CreateAccessKeyIdempotent(
|
||||
c.Request.Context(),
|
||||
idempotencyKey,
|
||||
request,
|
||||
)
|
||||
if err != nil {
|
||||
writeServiceError(c, "create_access_key", err)
|
||||
return
|
||||
}
|
||||
setSecretResponseHeaders(c)
|
||||
response.SuccessI18n(c, "common.success", result)
|
||||
}
|
||||
|
||||
func (s *Server) handleListAccessKeyOptions(c *gin.Context) {
|
||||
result, err := s.service.ListAccessKeyOptions(c.Request.Context())
|
||||
if err != nil {
|
||||
writeServiceError(c, "list_access_key_options", err)
|
||||
return
|
||||
}
|
||||
response.SuccessI18n(c, "common.success", result)
|
||||
}
|
||||
|
||||
func (s *Server) handleRevealAccessKey(c *gin.Context) {
|
||||
id, ok := accessKeyID(c)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
result, err := s.service.RevealAccessKey(c.Request.Context(), id)
|
||||
if err != nil {
|
||||
writeServiceError(c, "reveal_access_key", err)
|
||||
return
|
||||
}
|
||||
setSecretResponseHeaders(c)
|
||||
response.SuccessI18n(c, "common.success", result)
|
||||
}
|
||||
|
||||
@@ -396,6 +464,56 @@ func mapControlJSONError(err error) *app_errors.APIError {
|
||||
return app_errors.ErrInvalidJSON
|
||||
}
|
||||
|
||||
func requiredIdempotencyKey(
|
||||
c *gin.Context,
|
||||
operation string,
|
||||
) (string, bool) {
|
||||
values := c.Request.Header.Values("Idempotency-Key")
|
||||
if len(values) == 0 || len(values) == 1 && values[0] == "" {
|
||||
writeServiceError(c, operation, app_errors.ErrIdempotencyKeyRequired)
|
||||
return "", false
|
||||
}
|
||||
if len(values) != 1 || validateIdempotencyKey(values[0]) != nil {
|
||||
writeServiceError(c, operation, app_errors.ErrInvalidIdempotencyKey)
|
||||
return "", false
|
||||
}
|
||||
return values[0], true
|
||||
}
|
||||
|
||||
func requiredSettingsIfMatch(c *gin.Context) (string, bool) {
|
||||
values := c.Request.Header.Values("If-Match")
|
||||
if len(values) == 0 || len(values) == 1 && values[0] == "" {
|
||||
writeServiceError(c, "update_settings", app_errors.ErrSettingsPreconditionRequired)
|
||||
return "", false
|
||||
}
|
||||
if len(values) != 1 {
|
||||
writeServiceError(c, "update_settings", app_errors.ErrBadRequest)
|
||||
return "", false
|
||||
}
|
||||
token, valid := parseSettingsHeaderETag(values[0])
|
||||
if !valid {
|
||||
writeServiceError(c, "update_settings", app_errors.ErrBadRequest)
|
||||
return "", false
|
||||
}
|
||||
return token, true
|
||||
}
|
||||
|
||||
func writeSettingsRepresentation(
|
||||
c *gin.Context,
|
||||
representation settingsWireRepresentation,
|
||||
) {
|
||||
c.Header("ETag", representation.HeaderETag)
|
||||
c.Header("Content-Language", i18n.GetLanguageFromContext(c))
|
||||
c.Header("Vary", "Accept-Language")
|
||||
c.Header("Content-Encoding", "identity")
|
||||
c.Data(http.StatusOK, "application/json; charset=utf-8", representation.Body)
|
||||
}
|
||||
|
||||
func setSecretResponseHeaders(c *gin.Context) {
|
||||
c.Header("Cache-Control", "no-store")
|
||||
c.Header("Pragma", "no-cache")
|
||||
}
|
||||
|
||||
func accessKeyID(c *gin.Context) (uint, bool) {
|
||||
parsed, err := strconv.ParseUint(c.Param("id"), 10, strconv.IntSize)
|
||||
if err != nil || parsed == 0 {
|
||||
@@ -459,6 +577,22 @@ func serviceErrorMessageID(
|
||||
apiErr *app_errors.APIError,
|
||||
) string {
|
||||
switch apiErr.Code {
|
||||
case app_errors.ErrIdempotencyKeyRequired.Code:
|
||||
return "idempotency.required"
|
||||
case app_errors.ErrInvalidIdempotencyKey.Code:
|
||||
return "idempotency.invalid"
|
||||
case app_errors.ErrIdempotencyKeyReused.Code:
|
||||
return "idempotency.reused"
|
||||
case app_errors.ErrIdempotencyResultExpired.Code:
|
||||
return "idempotency.expired"
|
||||
case app_errors.ErrControlOperationIncomplete.Code:
|
||||
return "control.operation_incomplete"
|
||||
case app_errors.ErrControlRecoveryPending.Code:
|
||||
return "control.recovery_pending"
|
||||
case app_errors.ErrSettingsPreconditionRequired.Code:
|
||||
return "settings.precondition_required"
|
||||
case app_errors.ErrSettingsVersionConflict.Code:
|
||||
return "settings.version_conflict"
|
||||
case app_errors.ErrRequestTooLarge.Code:
|
||||
return "request_too_large"
|
||||
case app_errors.ErrBadRequest.Code, app_errors.ErrInvalidJSON.Code, app_errors.ErrValidation.Code:
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"gpt-load/internal/platform/config"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
func TestCreateAndImportEndpointsRequireCanonicalIdempotencyKeyBeforeMutation(t *testing.T) {
|
||||
initControlI18n(t)
|
||||
fixture := newServiceFixture(t)
|
||||
groupID := createGroupForKeyImport(t, fixture, "seed-idempotency-header")
|
||||
engine := gin.New()
|
||||
NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine)
|
||||
tests := []struct {
|
||||
name string
|
||||
path string
|
||||
body string
|
||||
}{
|
||||
{
|
||||
name: "AccessKey create",
|
||||
path: "/api/access-keys",
|
||||
body: `{"name":"client"}`,
|
||||
},
|
||||
{
|
||||
name: "Group create",
|
||||
path: "/api/groups",
|
||||
body: `{"upstream_url":"https://header.example.com","protocols":["openai"],"keys":"K"}`,
|
||||
},
|
||||
{
|
||||
name: "Group key import",
|
||||
path: "/api/groups/" + strconv.FormatUint(uint64(groupID), 10) + "/keys/import",
|
||||
body: `{"keys":"K"}`,
|
||||
},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
before := countCreateImportRows(t, fixture)
|
||||
for _, header := range []struct {
|
||||
value string
|
||||
wantStatus int
|
||||
wantCode string
|
||||
}{
|
||||
{wantStatus: http.StatusPreconditionRequired, wantCode: "IDEMPOTENCY_KEY_REQUIRED"},
|
||||
{
|
||||
value: "not-a-canonical-uuid",
|
||||
wantStatus: http.StatusBadRequest,
|
||||
wantCode: "INVALID_IDEMPOTENCY_KEY",
|
||||
},
|
||||
} {
|
||||
request := httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
test.path,
|
||||
strings.NewReader(test.body),
|
||||
)
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
if header.value != "" {
|
||||
request.Header.Set("Idempotency-Key", header.value)
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
engine.ServeHTTP(recorder, request)
|
||||
if recorder.Code != header.wantStatus ||
|
||||
!strings.Contains(recorder.Body.String(), header.wantCode) {
|
||||
t.Fatalf(
|
||||
"POST %s = %d %s, want %d/%s",
|
||||
test.path,
|
||||
recorder.Code,
|
||||
recorder.Body.String(),
|
||||
header.wantStatus,
|
||||
header.wantCode,
|
||||
)
|
||||
}
|
||||
if after := countCreateImportRows(t, fixture); after != before {
|
||||
t.Fatalf("rows changed from %#v to %#v", before, after)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccessKeyCreateReplayOptionsAndRevealWireContracts(t *testing.T) {
|
||||
initControlI18n(t)
|
||||
fixture := newServiceFixture(t)
|
||||
fixture.service.random = bytes.NewReader(make([]byte, 16))
|
||||
fixture.service.operationRandom = bytes.NewReader(bytes.Repeat([]byte{0x73}, 16))
|
||||
engine := gin.New()
|
||||
NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine)
|
||||
const idempotencyKey = "318f47a2-9c35-4d6e-8b1a-1234567890ab"
|
||||
|
||||
create := func() *httptest.ResponseRecorder {
|
||||
request := httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/api/access-keys",
|
||||
strings.NewReader(`{"name":"wire-client"}`),
|
||||
)
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Idempotency-Key", idempotencyKey)
|
||||
recorder := httptest.NewRecorder()
|
||||
engine.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
first := create()
|
||||
if first.Code != http.StatusOK ||
|
||||
first.Header().Get("Cache-Control") != "no-store" ||
|
||||
first.Header().Get("Pragma") != "no-cache" {
|
||||
t.Fatalf("first create = %d headers=%v body=%s", first.Code, first.Header(), first.Body)
|
||||
}
|
||||
var firstEnvelope struct {
|
||||
Data AccessKeyCreateResult `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(first.Body.Bytes(), &firstEnvelope); err != nil {
|
||||
t.Fatalf("decode first create: %v", err)
|
||||
}
|
||||
if firstEnvelope.Data.Key == "" || firstEnvelope.Data.Replayed {
|
||||
t.Fatalf("first create data = %#v", firstEnvelope.Data)
|
||||
}
|
||||
|
||||
replay := create()
|
||||
if replay.Code != http.StatusOK ||
|
||||
strings.Contains(replay.Body.String(), firstEnvelope.Data.Key) {
|
||||
t.Fatalf("replay create = %d %s", replay.Code, replay.Body)
|
||||
}
|
||||
var replayEnvelope struct {
|
||||
Data AccessKeyCreateResult `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(replay.Body.Bytes(), &replayEnvelope); err != nil {
|
||||
t.Fatalf("decode replay create: %v", err)
|
||||
}
|
||||
if replayEnvelope.Data.Key != "" || !replayEnvelope.Data.Replayed ||
|
||||
replayEnvelope.Data.ID != firstEnvelope.Data.ID {
|
||||
t.Fatalf("replay data = %#v", replayEnvelope.Data)
|
||||
}
|
||||
|
||||
optionsRequest := httptest.NewRequest(http.MethodGet, "/api/access-keys/options", nil)
|
||||
optionsRequest.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
options := httptest.NewRecorder()
|
||||
engine.ServeHTTP(options, optionsRequest)
|
||||
if options.Code != http.StatusOK ||
|
||||
strings.Contains(options.Body.String(), "masked_key") ||
|
||||
strings.Contains(options.Body.String(), `"key"`) {
|
||||
t.Fatalf("options = %d %s", options.Code, options.Body)
|
||||
}
|
||||
|
||||
revealRequest := httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
fmt.Sprintf("/api/access-keys/%d/reveal", firstEnvelope.Data.ID),
|
||||
nil,
|
||||
)
|
||||
revealRequest.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
reveal := httptest.NewRecorder()
|
||||
engine.ServeHTTP(reveal, revealRequest)
|
||||
if reveal.Code != http.StatusOK ||
|
||||
reveal.Header().Get("Cache-Control") != "no-store" ||
|
||||
reveal.Header().Get("Pragma") != "no-cache" ||
|
||||
!strings.Contains(reveal.Body.String(), firstEnvelope.Data.Key) {
|
||||
t.Fatalf("reveal = %d headers=%v body=%s", reveal.Code, reveal.Header(), reveal.Body)
|
||||
}
|
||||
}
|
||||
|
||||
type createImportRowCounts struct {
|
||||
groups int64
|
||||
upstream int64
|
||||
access int64
|
||||
operations int64
|
||||
}
|
||||
|
||||
func countCreateImportRows(
|
||||
t *testing.T,
|
||||
fixture serviceFixture,
|
||||
) createImportRowCounts {
|
||||
t.Helper()
|
||||
var result createImportRowCounts
|
||||
for model, target := range map[any]*int64{
|
||||
&models.Group{}: &result.groups,
|
||||
&models.UpstreamKey{}: &result.upstream,
|
||||
&models.AccessKey{}: &result.access,
|
||||
&models.ControlOperation{}: &result.operations,
|
||||
} {
|
||||
if err := fixture.db.Model(model).Count(target).Error; err != nil {
|
||||
t.Fatalf("count %T: %v", model, err)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -233,6 +233,15 @@ func TestControlMutationRejectsDuplicateJSONWithoutSideEffects(t *testing.T) {
|
||||
request := httptest.NewRequest(endpoint.method, endpoint.path, strings.NewReader(endpoint.body))
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if endpoint.path == "/api/groups" || endpoint.path == "/api/access-keys" {
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
}
|
||||
if endpoint.path == "/api/settings" {
|
||||
request.Header.Set(
|
||||
"If-Match",
|
||||
currentSettingsETag(t, engine, "test-auth-key", "zh-CN"),
|
||||
)
|
||||
}
|
||||
engine.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusBadRequest {
|
||||
@@ -502,6 +511,13 @@ func TestControlJSONBodyLimitAppliesToEveryJSONEndpoint(t *testing.T) {
|
||||
request.ContentLength = -1
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
if path == "/api/settings" {
|
||||
request.Header.Set(
|
||||
"If-Match",
|
||||
currentSettingsETag(t, engine, "test-auth-key", "zh-CN"),
|
||||
)
|
||||
}
|
||||
recorder := httptest.NewRecorder()
|
||||
|
||||
engine.ServeHTTP(recorder, request)
|
||||
@@ -554,6 +570,7 @@ func TestControlJSONBodyLimitContentLengthFastPathPreservesAuthenticationPriorit
|
||||
if auth.header != "" {
|
||||
request.Header.Set("Authorization", auth.header)
|
||||
}
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
recorder := httptest.NewRecorder()
|
||||
|
||||
engine.ServeHTTP(recorder, request)
|
||||
@@ -595,6 +612,7 @@ func TestControlJSONBodyLimitLocalizes413(t *testing.T) {
|
||||
request.ContentLength = -1
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Accept-Language", test.language)
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
recorder := httptest.NewRecorder()
|
||||
|
||||
engine.ServeHTTP(recorder, request)
|
||||
@@ -717,6 +735,7 @@ func TestCreateGroupEndpointReturnsSuccessAndConflictEnvelopes(t *testing.T) {
|
||||
))
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
recorder := httptest.NewRecorder()
|
||||
engine.ServeHTTP(recorder, request)
|
||||
if recorder.Code != http.StatusOK {
|
||||
@@ -739,6 +758,7 @@ func TestCreateGroupEndpointReturnsSuccessAndConflictEnvelopes(t *testing.T) {
|
||||
))
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
recorder = httptest.NewRecorder()
|
||||
engine.ServeHTTP(recorder, request)
|
||||
if recorder.Code != http.StatusConflict {
|
||||
@@ -775,6 +795,7 @@ func TestImportGroupKeysEndpointReturnsSuccessEnvelope(t *testing.T) {
|
||||
))
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
engine.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusOK {
|
||||
@@ -836,6 +857,7 @@ func TestImportGroupKeysEndpointRejectsUnknownFieldsAndInvalidGroupID(t *testing
|
||||
request := httptest.NewRequest(http.MethodPost, "/api/groups/"+test.groupID+"/keys/import", strings.NewReader(test.body))
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
engine.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusBadRequest {
|
||||
@@ -860,6 +882,7 @@ func TestImportGroupKeysEndpointReturnsGroupNotFound(t *testing.T) {
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("Accept-Language", "zh-CN")
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
engine.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusNotFound {
|
||||
@@ -904,6 +927,7 @@ func TestManagementWritesRejectUnknownFieldsAndMultipleJSONValues(t *testing.T)
|
||||
))
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
engine.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusBadRequest {
|
||||
@@ -927,6 +951,7 @@ func TestManagementWritesRejectUnknownFieldsAndMultipleJSONValues(t *testing.T)
|
||||
))
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
setRequiredTestIdempotencyHeader(request)
|
||||
engine.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code != http.StatusBadRequest {
|
||||
@@ -1404,7 +1429,9 @@ func TestUpdateAccessKeyRoutesParseIDsAndPreservePointerSemantics(t *testing.T)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
engine.ServeHTTP(recorder, request)
|
||||
if recorder.Code != http.StatusOK || !strings.Contains(recorder.Body.String(), created.Key) {
|
||||
if recorder.Code != http.StatusOK ||
|
||||
strings.Contains(recorder.Body.String(), created.Key) ||
|
||||
!strings.Contains(recorder.Body.String(), created.MaskedKey) {
|
||||
t.Fatalf("PUT access key = %d %s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
row := loadAccessKeyRow(t, fixture.db, created.ID)
|
||||
@@ -1943,6 +1970,12 @@ func TestSettingsHTTPContract(t *testing.T) {
|
||||
}
|
||||
request.Header.Set("Accept-Language", "en-US")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if test.method == http.MethodPut && test.auth != "" {
|
||||
request.Header.Set(
|
||||
"If-Match",
|
||||
currentSettingsETag(t, engine, strings.TrimPrefix(test.auth, "Bearer "), "en-US"),
|
||||
)
|
||||
}
|
||||
engine.ServeHTTP(recorder, request)
|
||||
if recorder.Code != test.wantStatus {
|
||||
t.Fatalf("response = %d %s, want %d", recorder.Code, recorder.Body.String(), test.wantStatus)
|
||||
@@ -1976,19 +2009,19 @@ func TestSettingsHTTPStableSuccessEnvelopeUpdateAndReset(t *testing.T) {
|
||||
NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine)
|
||||
|
||||
get := serveSettingsRequest(t, engine, http.MethodGet, "test-auth-key", "")
|
||||
const wantDefault = `{"code":0,"message":"Success","data":{"revision":1,"values":{"connect_timeout":15,"first_byte_timeout":120,"request_timeout":600,"stream_idle_timeout":300,"header_rules":{"set":{},"remove":[]},"inject_usage_options":true,"request_log_retention_days":7},"overrides":[]}}`
|
||||
const wantDefault = `{"code":0,"data":{"overrides":[],"values":{"connect_timeout":15,"first_byte_timeout":120,"header_rules":{"remove":[],"set":{}},"inject_usage_options":true,"request_log_retention_days":7,"request_timeout":600,"stream_idle_timeout":300}},"message":"Success"}`
|
||||
if get.Code != http.StatusOK || strings.TrimSpace(get.Body.String()) != wantDefault {
|
||||
t.Fatalf("default response = %d %s, want %s", get.Code, get.Body.String(), wantDefault)
|
||||
}
|
||||
|
||||
update := serveSettingsRequest(t, engine, http.MethodPut, "test-auth-key", `{"settings":{"request_timeout":900,"header_rules":{"set":{},"remove":[]}}}`)
|
||||
const wantUpdate = `{"code":0,"message":"Success","data":{"revision":2,"values":{"connect_timeout":15,"first_byte_timeout":120,"request_timeout":900,"stream_idle_timeout":300,"header_rules":{"set":{},"remove":[]},"inject_usage_options":true,"request_log_retention_days":7},"overrides":["header_rules","request_timeout"]}}`
|
||||
const wantUpdate = `{"code":0,"data":{"overrides":["header_rules","request_timeout"],"values":{"connect_timeout":15,"first_byte_timeout":120,"header_rules":{"remove":[],"set":{}},"inject_usage_options":true,"request_log_retention_days":7,"request_timeout":900,"stream_idle_timeout":300}},"message":"Success"}`
|
||||
if update.Code != http.StatusOK || strings.TrimSpace(update.Body.String()) != wantUpdate {
|
||||
t.Fatalf("update response = %d %s, want %s", update.Code, update.Body.String(), wantUpdate)
|
||||
}
|
||||
|
||||
reset := serveSettingsRequest(t, engine, http.MethodPut, "test-auth-key", `{"settings":{"request_timeout":null,"header_rules":null}}`)
|
||||
const wantReset = `{"code":0,"message":"Success","data":{"revision":3,"values":{"connect_timeout":15,"first_byte_timeout":120,"request_timeout":600,"stream_idle_timeout":300,"header_rules":{"set":{},"remove":[]},"inject_usage_options":true,"request_log_retention_days":7},"overrides":[]}}`
|
||||
const wantReset = wantDefault
|
||||
if reset.Code != http.StatusOK || strings.TrimSpace(reset.Body.String()) != wantReset {
|
||||
t.Fatalf("reset response = %d %s, want %s", reset.Code, reset.Body.String(), wantReset)
|
||||
}
|
||||
@@ -2011,6 +2044,7 @@ func TestSettingsHTTPBodyLimitRejectsBeforeMutation(t *testing.T) {
|
||||
request.Header.Set("Authorization", "Bearer test-auth-key")
|
||||
request.Header.Set("Accept-Language", "en-US")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
request.Header.Set("If-Match", currentSettingsETag(t, engine, "test-auth-key", "en-US"))
|
||||
engine.ServeHTTP(recorder, request)
|
||||
if recorder.Code != http.StatusRequestEntityTooLarge ||
|
||||
!strings.Contains(recorder.Body.String(), `"code":"REQUEST_TOO_LARGE"`) {
|
||||
@@ -2136,6 +2170,31 @@ func serveSettingsRequest(
|
||||
}
|
||||
request.Header.Set("Accept-Language", "en-US")
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
if method == http.MethodPut && authKey != "" {
|
||||
request.Header.Set("If-Match", currentSettingsETag(t, engine, authKey, "en-US"))
|
||||
}
|
||||
engine.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func currentSettingsETag(
|
||||
t *testing.T,
|
||||
engine *gin.Engine,
|
||||
authKey string,
|
||||
language string,
|
||||
) string {
|
||||
t.Helper()
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(http.MethodGet, "/api/settings", nil)
|
||||
request.Header.Set("Authorization", "Bearer "+authKey)
|
||||
request.Header.Set("Accept-Language", language)
|
||||
engine.ServeHTTP(recorder, request)
|
||||
if recorder.Code != http.StatusOK {
|
||||
t.Fatalf("get current Settings ETag = %d %s", recorder.Code, recorder.Body.String())
|
||||
}
|
||||
tag := recorder.Header().Get("ETag")
|
||||
if tag == "" {
|
||||
t.Fatal("get current Settings ETag returned no ETag")
|
||||
}
|
||||
return tag
|
||||
}
|
||||
+37
-15
@@ -19,6 +19,7 @@ import (
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/state"
|
||||
stateloader "gpt-load/internal/state/loader"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
const (
|
||||
@@ -27,19 +28,28 @@ const (
|
||||
)
|
||||
|
||||
type Service struct {
|
||||
db *gorm.DB
|
||||
manager *state.Manager
|
||||
registry *state.KeyRegistry
|
||||
priceRuntime *PriceRuntime
|
||||
encryption encryption.Service
|
||||
dialects dialect.Set
|
||||
requestLogs RequestLogReader
|
||||
usageStats UsageStatReader
|
||||
stats *health.StatsStore
|
||||
requestLogStats RequestLogStatsReader
|
||||
modelDiscoveryTimeout time.Duration
|
||||
random io.Reader
|
||||
now func() time.Time
|
||||
db *gorm.DB
|
||||
manager *state.Manager
|
||||
registry *state.KeyRegistry
|
||||
priceRuntime *PriceRuntime
|
||||
encryption encryption.Service
|
||||
dialects dialect.Set
|
||||
requestLogs RequestLogReader
|
||||
usageStats UsageStatReader
|
||||
stats *health.StatsStore
|
||||
requestLogStats RequestLogStatsReader
|
||||
modelDiscoveryTimeout time.Duration
|
||||
random io.Reader
|
||||
operationRandom io.Reader
|
||||
now func() time.Time
|
||||
publishSnapshot func(state.CompileInput) (*state.ConfigSnapshot, error)
|
||||
reconcileRegistryGroup func(uint, []state.KeyEntry) (bool, error)
|
||||
beforeAdvanceOperationStage func(
|
||||
context.Context,
|
||||
*models.ControlOperation,
|
||||
operationStage,
|
||||
) error
|
||||
operationRecoveryWake chan struct{}
|
||||
writeMu sync.RWMutex
|
||||
}
|
||||
|
||||
@@ -55,15 +65,21 @@ func NewService(
|
||||
stats *health.StatsStore,
|
||||
requestLogStats RequestLogStatsReader,
|
||||
) *Service {
|
||||
return &Service{
|
||||
service := &Service{
|
||||
db: db, manager: manager, registry: registry,
|
||||
priceRuntime: priceRuntime,
|
||||
encryption: encryptionService, dialects: dialects, requestLogs: requestLogs,
|
||||
usageStats: usageStats,
|
||||
stats: stats, requestLogStats: requestLogStats,
|
||||
modelDiscoveryTimeout: defaultModelDiscoveryTimeout,
|
||||
random: rand.Reader, now: time.Now,
|
||||
random: rand.Reader,
|
||||
operationRandom: rand.Reader,
|
||||
now: time.Now,
|
||||
operationRecoveryWake: make(chan struct{}, 1),
|
||||
}
|
||||
service.publishSnapshot = manager.Publish
|
||||
service.reconcileRegistryGroup = registry.ReconcileGroup
|
||||
return service
|
||||
}
|
||||
|
||||
func (s *Service) writeConfig(
|
||||
@@ -73,6 +89,9 @@ func (s *Service) writeConfig(
|
||||
) (*state.ConfigSnapshot, error) {
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
if blocked, err := s.recoverPendingOperationsLocked(ctx, 0); err != nil {
|
||||
return nil, s.recoveryPendingError(*blocked)
|
||||
}
|
||||
|
||||
var input state.CompileInput
|
||||
err := s.withControlTransaction(ctx, func(tx *gorm.DB) error {
|
||||
@@ -114,6 +133,9 @@ func (s *Service) writeKeyConfig(
|
||||
) error {
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
if blocked, err := s.recoverPendingOperationsLocked(ctx, 0); err != nil {
|
||||
return s.recoveryPendingError(*blocked)
|
||||
}
|
||||
|
||||
if err := s.withControlTransaction(ctx, mutate); err != nil {
|
||||
return err
|
||||
|
||||
+124
-21
@@ -15,6 +15,7 @@ import (
|
||||
|
||||
app_errors "gpt-load/internal/platform/errors"
|
||||
"gpt-load/internal/state"
|
||||
stateloader "gpt-load/internal/state/loader"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
@@ -39,6 +40,13 @@ type SettingsResponse struct {
|
||||
Overrides []string `json:"overrides"`
|
||||
}
|
||||
|
||||
func (response SettingsResponse) DTO() SettingsDTO {
|
||||
return canonicalizeSettingsDTO(SettingsDTO{
|
||||
Values: response.Values,
|
||||
Overrides: response.Overrides,
|
||||
})
|
||||
}
|
||||
|
||||
type SettingsUpdateRequest struct {
|
||||
Settings map[string]json.RawMessage `json:"settings"`
|
||||
}
|
||||
@@ -87,8 +95,16 @@ func (s *Service) GetSettings(ctx context.Context) (SettingsResponse, error) {
|
||||
if snapshot == nil {
|
||||
return SettingsResponse{}, app_errors.ErrInternalServer
|
||||
}
|
||||
return s.getSettingsWithSnapshot(ctx, s.db, snapshot)
|
||||
}
|
||||
|
||||
func (s *Service) getSettingsWithSnapshot(
|
||||
ctx context.Context,
|
||||
db *gorm.DB,
|
||||
snapshot *state.ConfigSnapshot,
|
||||
) (SettingsResponse, error) {
|
||||
var rows []models.SystemSetting
|
||||
if err := s.db.WithContext(ctx).Order("key ASC").Find(&rows).Error; err != nil {
|
||||
if err := db.WithContext(ctx).Order("key ASC").Find(&rows).Error; err != nil {
|
||||
return SettingsResponse{}, app_errors.ParseDBError(err)
|
||||
}
|
||||
return mapSettingsResponse(snapshot, rows), nil
|
||||
@@ -104,26 +120,7 @@ func (s *Service) UpdateSettings(
|
||||
}
|
||||
|
||||
_, err = s.writeConfig(ctx, func(tx *gorm.DB) error {
|
||||
for _, update := range updates {
|
||||
if update.value == nil {
|
||||
if err := tx.Where("key = ?", update.key).
|
||||
Delete(&models.SystemSetting{}).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err := tx.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "key"}},
|
||||
DoUpdates: clause.AssignmentColumns([]string{"value", "updated_at"}),
|
||||
}).Create(&models.SystemSetting{
|
||||
Key: update.key,
|
||||
Value: *update.value,
|
||||
UpdatedAt: s.now().UTC(),
|
||||
}).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
return s.applySettingUpdates(tx, updates)
|
||||
}, nil)
|
||||
if err != nil {
|
||||
return SettingsResponse{}, err
|
||||
@@ -131,6 +128,112 @@ func (s *Service) UpdateSettings(
|
||||
return s.GetSettings(ctx)
|
||||
}
|
||||
|
||||
func (s *Service) UpdateSettingsIfMatch(
|
||||
ctx context.Context,
|
||||
request SettingsUpdateRequest,
|
||||
expectedETag string,
|
||||
message string,
|
||||
) (settingsWireRepresentation, error) {
|
||||
updates, err := normalizeSettingUpdates(request)
|
||||
if err != nil {
|
||||
return settingsWireRepresentation{}, err
|
||||
}
|
||||
|
||||
s.writeMu.Lock()
|
||||
defer s.writeMu.Unlock()
|
||||
if blocked, err := s.recoverPendingOperationsLocked(ctx, 0); err != nil {
|
||||
return settingsWireRepresentation{}, s.recoveryPendingError(*blocked)
|
||||
}
|
||||
|
||||
var (
|
||||
input state.CompileInput
|
||||
result settingsWireRepresentation
|
||||
)
|
||||
err = s.withControlTransaction(ctx, func(tx *gorm.DB) error {
|
||||
currentInput, err := stateloader.BuildCompileInput(ctx, tx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
currentSnapshot, err := state.Compile(currentInput)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
currentSettings, err := s.getSettingsWithSnapshot(ctx, tx, currentSnapshot)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
currentRepresentation, err := newSettingsWireRepresentation(
|
||||
message,
|
||||
currentSettings.DTO(),
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if currentRepresentation.ETag != expectedETag {
|
||||
return app_errors.NewAPIErrorWithData(
|
||||
app_errors.ErrSettingsVersionConflict,
|
||||
SettingsConflictData{
|
||||
Settings: currentRepresentation.Settings,
|
||||
SettingsETag: currentRepresentation.ETag,
|
||||
},
|
||||
)
|
||||
}
|
||||
|
||||
if err := s.applySettingUpdates(tx, updates); err != nil {
|
||||
return err
|
||||
}
|
||||
input, err = stateloader.BuildCompileInput(ctx, tx)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
compiled, err := state.Compile(input)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
updatedSettings, err := s.getSettingsWithSnapshot(ctx, tx, compiled)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
result, err = newSettingsWireRepresentation(message, updatedSettings.DTO())
|
||||
return err
|
||||
})
|
||||
if err != nil {
|
||||
return settingsWireRepresentation{}, err
|
||||
}
|
||||
if _, err := s.manager.Publish(input); err != nil {
|
||||
return settingsWireRepresentation{}, newControlOperationError(
|
||||
stagePublishCommittedSnapshot,
|
||||
)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s *Service) applySettingUpdates(
|
||||
tx *gorm.DB,
|
||||
updates []persistedSettingUpdate,
|
||||
) error {
|
||||
for _, update := range updates {
|
||||
if update.value == nil {
|
||||
if err := tx.Where("key = ?", update.key).
|
||||
Delete(&models.SystemSetting{}).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
continue
|
||||
}
|
||||
if err := tx.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{{Name: "key"}},
|
||||
DoUpdates: clause.AssignmentColumns([]string{"value", "updated_at"}),
|
||||
}).Create(&models.SystemSetting{
|
||||
Key: update.key,
|
||||
Value: *update.value,
|
||||
UpdatedAt: s.now().UTC(),
|
||||
}).Error; err != nil {
|
||||
return app_errors.ParseDBError(err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func normalizeSettingUpdates(request SettingsUpdateRequest) ([]persistedSettingUpdate, error) {
|
||||
if len(request.Settings) == 0 {
|
||||
return nil, app_errors.ErrBadRequest
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"net/textproto"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"gpt-load/internal/platform/canonicaljson"
|
||||
)
|
||||
|
||||
const settingsWireETagDomain = "gpt-load/settings-wire-v1"
|
||||
|
||||
type SettingsDTO struct {
|
||||
Values SettingsValuesResponse `json:"values"`
|
||||
Overrides []string `json:"overrides"`
|
||||
}
|
||||
|
||||
type SettingsConflictData struct {
|
||||
Settings SettingsDTO `json:"settings"`
|
||||
SettingsETag string `json:"settings_etag"`
|
||||
}
|
||||
|
||||
type settingsWireRepresentation struct {
|
||||
Settings SettingsDTO
|
||||
Body []byte
|
||||
ETag string
|
||||
HeaderETag string
|
||||
}
|
||||
|
||||
func parseSettingsHeaderETag(value string) (string, bool) {
|
||||
if len(value) != 73 || value[0] != '"' || value[len(value)-1] != '"' {
|
||||
return "", false
|
||||
}
|
||||
token := value[1 : len(value)-1]
|
||||
if !strings.HasPrefix(token, "sha256-") {
|
||||
return "", false
|
||||
}
|
||||
for _, character := range token[len("sha256-"):] {
|
||||
if character < '0' || character > '9' {
|
||||
if character < 'a' || character > 'f' {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
}
|
||||
return token, true
|
||||
}
|
||||
|
||||
type settingsSuccessEnvelope struct {
|
||||
Code int `json:"code"`
|
||||
Message string `json:"message"`
|
||||
Data SettingsDTO `json:"data"`
|
||||
}
|
||||
|
||||
func newSettingsWireRepresentation(
|
||||
message string,
|
||||
settings SettingsDTO,
|
||||
) (settingsWireRepresentation, error) {
|
||||
settings = canonicalizeSettingsDTO(settings)
|
||||
body, err := canonicaljson.Marshal(settingsSuccessEnvelope{
|
||||
Code: 0,
|
||||
Message: message,
|
||||
Data: settings,
|
||||
})
|
||||
if err != nil {
|
||||
return settingsWireRepresentation{}, err
|
||||
}
|
||||
hash := sha256.New()
|
||||
_, _ = hash.Write([]byte(settingsWireETagDomain))
|
||||
_, _ = hash.Write([]byte{0})
|
||||
_, _ = hash.Write(body)
|
||||
tag := "sha256-" + hex.EncodeToString(hash.Sum(nil))
|
||||
return settingsWireRepresentation{
|
||||
Settings: settings,
|
||||
Body: body,
|
||||
ETag: tag,
|
||||
HeaderETag: `"` + tag + `"`,
|
||||
}, nil
|
||||
}
|
||||
|
||||
func canonicalizeSettingsDTO(settings SettingsDTO) SettingsDTO {
|
||||
setNames := make([]string, 0, len(settings.Values.HeaderRules.Set))
|
||||
for name := range settings.Values.HeaderRules.Set {
|
||||
setNames = append(setNames, name)
|
||||
}
|
||||
sort.Slice(setNames, func(left, right int) bool {
|
||||
leftFolded := strings.ToLower(setNames[left])
|
||||
rightFolded := strings.ToLower(setNames[right])
|
||||
if leftFolded == rightFolded {
|
||||
return setNames[left] < setNames[right]
|
||||
}
|
||||
return leftFolded < rightFolded
|
||||
})
|
||||
set := make(map[string]string, len(setNames))
|
||||
for _, name := range setNames {
|
||||
set[textproto.CanonicalMIMEHeaderKey(name)] = settings.Values.HeaderRules.Set[name]
|
||||
}
|
||||
|
||||
removeByIdentity := make(map[string]string, len(settings.Values.HeaderRules.Remove))
|
||||
for _, name := range settings.Values.HeaderRules.Remove {
|
||||
canonicalName := textproto.CanonicalMIMEHeaderKey(name)
|
||||
removeByIdentity[strings.ToLower(canonicalName)] = canonicalName
|
||||
}
|
||||
remove := make([]string, 0, len(removeByIdentity))
|
||||
for _, name := range removeByIdentity {
|
||||
remove = append(remove, name)
|
||||
}
|
||||
sort.Strings(remove)
|
||||
|
||||
overrideSet := make(map[string]struct{}, len(settings.Overrides))
|
||||
for _, override := range settings.Overrides {
|
||||
overrideSet[override] = struct{}{}
|
||||
}
|
||||
overrides := make([]string, 0, len(overrideSet))
|
||||
for override := range overrideSet {
|
||||
overrides = append(overrides, override)
|
||||
}
|
||||
sort.Strings(overrides)
|
||||
|
||||
settings.Values.HeaderRules.Set = set
|
||||
settings.Values.HeaderRules.Remove = remove
|
||||
settings.Overrides = overrides
|
||||
return settings
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
type settingsETagFixture struct {
|
||||
DomainHex string `json:"domain_hex"`
|
||||
CanonicalResponseHex string `json:"canonical_response_hex"`
|
||||
FramedInputHex string `json:"framed_input_hex"`
|
||||
ExpectedDigest string `json:"expected_digest"`
|
||||
}
|
||||
|
||||
func TestSettingsWireETagMatchesCheckedInFixture(t *testing.T) {
|
||||
raw, err := os.ReadFile("testdata/settings-etag-v1.json")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var fixture settingsETagFixture
|
||||
if err := json.Unmarshal(raw, &fixture); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
representation, err := newSettingsWireRepresentation("Success", SettingsDTO{
|
||||
Values: SettingsValuesResponse{
|
||||
ConnectTimeout: 15,
|
||||
FirstByteTimeout: 120,
|
||||
RequestTimeout: 600,
|
||||
StreamIdleTimeout: 300,
|
||||
HeaderRules: HeaderRulesResponse{
|
||||
Set: map[string]string{},
|
||||
Remove: []string{},
|
||||
},
|
||||
InjectUsageOptions: true,
|
||||
RequestLogRetentionDays: 7,
|
||||
},
|
||||
Overrides: []string{},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
domain, err := hex.DecodeString(fixture.DomainHex)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wantBody, err := hex.DecodeString(fixture.CanonicalResponseHex)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wantFramed, err := hex.DecodeString(fixture.FramedInputHex)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
framed := append(append(append([]byte{}, domain...), 0), representation.Body...)
|
||||
if !bytes.Equal(representation.Body, wantBody) {
|
||||
t.Fatalf("canonical body = %s, want %s", representation.Body, wantBody)
|
||||
}
|
||||
if !bytes.Equal(framed, wantFramed) {
|
||||
t.Fatalf("framed input = %x, want %x", framed, wantFramed)
|
||||
}
|
||||
if len(wantFramed) <= len(domain) || wantFramed[len(domain)] != 0 ||
|
||||
bytes.Contains(wantFramed[len(domain)+1:], []byte{0}) {
|
||||
t.Fatalf("fixture separator is not exactly one NUL: %x", wantFramed)
|
||||
}
|
||||
digest := sha256.Sum256(framed)
|
||||
if got := hex.EncodeToString(digest[:]); got != fixture.ExpectedDigest {
|
||||
t.Fatalf("fixture digest = %s, want %s", got, fixture.ExpectedDigest)
|
||||
}
|
||||
if representation.ETag != "sha256-"+fixture.ExpectedDigest {
|
||||
t.Fatalf("ETag = %q", representation.ETag)
|
||||
}
|
||||
if representation.HeaderETag != `"`+representation.ETag+`"` {
|
||||
t.Fatalf("header ETag = %q", representation.HeaderETag)
|
||||
}
|
||||
if bytes.Contains(representation.Body, []byte(`"revision"`)) ||
|
||||
bytes.Contains(representation.Body, []byte(`"settings_etag"`)) {
|
||||
t.Fatalf("wire body exposes process identity: %s", representation.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsDTOCanonicalizesHeaderRulesAndOverrides(t *testing.T) {
|
||||
dto := canonicalizeSettingsDTO(SettingsDTO{
|
||||
Values: SettingsValuesResponse{
|
||||
HeaderRules: HeaderRulesResponse{
|
||||
Set: map[string]string{
|
||||
"x-zed": "last",
|
||||
"A-test": "first",
|
||||
},
|
||||
Remove: []string{"x-old", "X-OLD", "a-remove"},
|
||||
},
|
||||
},
|
||||
Overrides: []string{"request_timeout", "connect_timeout", "request_timeout"},
|
||||
})
|
||||
if got := strings.Join(dto.Values.HeaderRules.Remove, ","); got != "A-Remove,X-Old" {
|
||||
t.Fatalf("remove = %q", got)
|
||||
}
|
||||
if _, exists := dto.Values.HeaderRules.Set["x-zed"]; exists ||
|
||||
dto.Values.HeaderRules.Set["X-Zed"] != "last" ||
|
||||
dto.Values.HeaderRules.Set["A-Test"] != "first" {
|
||||
t.Fatalf("set = %#v", dto.Values.HeaderRules.Set)
|
||||
}
|
||||
if got := strings.Join(dto.Overrides, ","); got != "connect_timeout,request_timeout" {
|
||||
t.Fatalf("overrides = %q", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,285 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
|
||||
"gpt-load/internal/platform/config"
|
||||
"gpt-load/internal/protocol"
|
||||
"gpt-load/internal/state"
|
||||
stateloader "gpt-load/internal/state/loader"
|
||||
)
|
||||
|
||||
func TestSettingsHTTPUsesCanonicalLocalizedStrongRepresentation(t *testing.T) {
|
||||
initControlI18n(t)
|
||||
fixture := newServiceFixture(t)
|
||||
engine := gin.New()
|
||||
NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine)
|
||||
|
||||
english := serveSettingsOCCRequest(
|
||||
t, engine, http.MethodGet, "test-auth-key", "en-US", "", "",
|
||||
)
|
||||
if english.Code != http.StatusOK {
|
||||
t.Fatalf("English GET = %d %s", english.Code, english.Body.String())
|
||||
}
|
||||
assertSettingsRepresentationHeaders(t, english, "en-US")
|
||||
const wantEnglish = `{"code":0,"data":{"overrides":[],"values":{"connect_timeout":15,"first_byte_timeout":120,"header_rules":{"remove":[],"set":{}},"inject_usage_options":true,"request_log_retention_days":7,"request_timeout":600,"stream_idle_timeout":300}},"message":"Success"}`
|
||||
if english.Body.String() != wantEnglish {
|
||||
t.Fatalf("English body = %s, want %s", english.Body.String(), wantEnglish)
|
||||
}
|
||||
|
||||
chinese := serveSettingsOCCRequest(
|
||||
t, engine, http.MethodGet, "test-auth-key", "zh-CN", "", "",
|
||||
)
|
||||
if chinese.Code != http.StatusOK {
|
||||
t.Fatalf("Chinese GET = %d %s", chinese.Code, chinese.Body.String())
|
||||
}
|
||||
assertSettingsRepresentationHeaders(t, chinese, "zh-CN")
|
||||
if chinese.Header().Get("ETag") == english.Header().Get("ETag") {
|
||||
t.Fatalf("localized representations share ETag %q", english.Header().Get("ETag"))
|
||||
}
|
||||
if !strings.Contains(chinese.Body.String(), `"message":"操作成功"`) {
|
||||
t.Fatalf("Chinese body = %s", chinese.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsHTTPRequiresStrongIfMatchAndReturnsLatestConflict(t *testing.T) {
|
||||
initControlI18n(t)
|
||||
fixture := newServiceFixture(t)
|
||||
engine := gin.New()
|
||||
NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine)
|
||||
|
||||
before := fixture.manager.Current()
|
||||
missing := serveSettingsOCCRequest(
|
||||
t,
|
||||
engine,
|
||||
http.MethodPut,
|
||||
"test-auth-key",
|
||||
"en-US",
|
||||
"",
|
||||
`{"settings":{"request_timeout":`,
|
||||
)
|
||||
if missing.Code != http.StatusPreconditionRequired ||
|
||||
!strings.Contains(missing.Body.String(), `"code":"SETTINGS_PRECONDITION_REQUIRED"`) {
|
||||
t.Fatalf("missing If-Match = %d %s", missing.Code, missing.Body.String())
|
||||
}
|
||||
if fixture.manager.Current() != before {
|
||||
t.Fatal("missing If-Match published a snapshot")
|
||||
}
|
||||
|
||||
malformed := serveSettingsOCCRequest(
|
||||
t,
|
||||
engine,
|
||||
http.MethodPut,
|
||||
"test-auth-key",
|
||||
"en-US",
|
||||
`W/"sha256-`+strings.Repeat("0", 64)+`"`,
|
||||
`{"settings":{"request_timeout":900}}`,
|
||||
)
|
||||
if malformed.Code != http.StatusBadRequest ||
|
||||
!strings.Contains(malformed.Body.String(), `"code":"BAD_REQUEST"`) {
|
||||
t.Fatalf("weak If-Match = %d %s", malformed.Code, malformed.Body.String())
|
||||
}
|
||||
|
||||
latest := serveSettingsOCCRequest(
|
||||
t, engine, http.MethodGet, "test-auth-key", "en-US", "", "",
|
||||
)
|
||||
conflict := serveSettingsOCCRequest(
|
||||
t,
|
||||
engine,
|
||||
http.MethodPut,
|
||||
"test-auth-key",
|
||||
"en-US",
|
||||
`"sha256-`+strings.Repeat("0", 64)+`"`,
|
||||
`{"settings":{"request_timeout":900}}`,
|
||||
)
|
||||
if conflict.Code != http.StatusPreconditionFailed {
|
||||
t.Fatalf("conflict = %d %s", conflict.Code, conflict.Body.String())
|
||||
}
|
||||
var envelope struct {
|
||||
Code string `json:"code"`
|
||||
Data struct {
|
||||
Settings SettingsDTO `json:"settings"`
|
||||
SettingsETag string `json:"settings_etag"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(conflict.Body.Bytes(), &envelope); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if envelope.Code != "SETTINGS_VERSION_CONFLICT" ||
|
||||
envelope.Data.SettingsETag != strings.Trim(latest.Header().Get("ETag"), `"`) ||
|
||||
envelope.Data.Settings.Values.RequestTimeout != 600 ||
|
||||
len(envelope.Data.Settings.Overrides) != 0 {
|
||||
t.Fatalf("conflict body = %s", conflict.Body.String())
|
||||
}
|
||||
if fixture.manager.Current() != before {
|
||||
t.Fatal("version conflict published a snapshot")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsHTTPMatchingWriteReturnsCurrentGETAndAllowsABA(t *testing.T) {
|
||||
initControlI18n(t)
|
||||
fixture := newServiceFixture(t)
|
||||
engine := gin.New()
|
||||
NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine)
|
||||
|
||||
initial := serveSettingsOCCRequest(
|
||||
t, engine, http.MethodGet, "test-auth-key", "en-US", "", "",
|
||||
)
|
||||
updated := serveSettingsOCCRequest(
|
||||
t,
|
||||
engine,
|
||||
http.MethodPut,
|
||||
"test-auth-key",
|
||||
"en-US",
|
||||
initial.Header().Get("ETag"),
|
||||
`{"settings":{"request_timeout":900}}`,
|
||||
)
|
||||
if updated.Code != http.StatusOK {
|
||||
t.Fatalf("update = %d %s", updated.Code, updated.Body.String())
|
||||
}
|
||||
assertSettingsRepresentationHeaders(t, updated, "en-US")
|
||||
if updated.Header().Get("ETag") == initial.Header().Get("ETag") {
|
||||
t.Fatalf("changed settings retained ETag %q", updated.Header().Get("ETag"))
|
||||
}
|
||||
current := serveSettingsOCCRequest(
|
||||
t, engine, http.MethodGet, "test-auth-key", "en-US", "", "",
|
||||
)
|
||||
if current.Header().Get("ETag") != updated.Header().Get("ETag") ||
|
||||
!bytes.Equal(current.Body.Bytes(), updated.Body.Bytes()) {
|
||||
t.Fatalf("GET/PUT representation differs:\nGET %q %s\nPUT %q %s",
|
||||
current.Header().Get("ETag"), current.Body.String(),
|
||||
updated.Header().Get("ETag"), updated.Body.String())
|
||||
}
|
||||
|
||||
reset := serveSettingsOCCRequest(
|
||||
t,
|
||||
engine,
|
||||
http.MethodPut,
|
||||
"test-auth-key",
|
||||
"en-US",
|
||||
updated.Header().Get("ETag"),
|
||||
`{"settings":{"request_timeout":null}}`,
|
||||
)
|
||||
if reset.Code != http.StatusOK ||
|
||||
reset.Header().Get("ETag") != initial.Header().Get("ETag") ||
|
||||
!bytes.Equal(reset.Body.Bytes(), initial.Body.Bytes()) {
|
||||
t.Fatalf("A-B-A response = %d %q %s, want %q %s",
|
||||
reset.Code, reset.Header().Get("ETag"), reset.Body.String(),
|
||||
initial.Header().Get("ETag"), initial.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSettingsETagIgnoresOtherResourcesAndSurvivesRuntimeReload(t *testing.T) {
|
||||
initControlI18n(t)
|
||||
fixture := newServiceFixture(t)
|
||||
engine := gin.New()
|
||||
NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine)
|
||||
initial := serveSettingsOCCRequest(
|
||||
t, engine, http.MethodGet, "test-auth-key", "en-US", "", "",
|
||||
)
|
||||
|
||||
groupName := "ETag unrelated group"
|
||||
if _, err := fixture.service.CreateGroup(t.Context(), GroupCreateRequest{
|
||||
Name: &groupName,
|
||||
UpstreamURL: "https://example.com/v1",
|
||||
Protocols: []protocol.Protocol{protocol.OpenAI},
|
||||
Keys: "sk-unrelated-etag-key",
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
afterGroup := serveSettingsOCCRequest(
|
||||
t, engine, http.MethodGet, "test-auth-key", "en-US", "", "",
|
||||
)
|
||||
if afterGroup.Header().Get("ETag") != initial.Header().Get("ETag") ||
|
||||
!bytes.Equal(afterGroup.Body.Bytes(), initial.Body.Bytes()) {
|
||||
t.Fatalf("unrelated mutation changed Settings representation")
|
||||
}
|
||||
|
||||
reloadedManager := state.NewManager()
|
||||
reloadedRegistry := state.NewKeyRegistry()
|
||||
if err := stateloader.New(fixture.db, reloadedManager, reloadedRegistry).Load(t.Context()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
reloadedService := NewService(
|
||||
fixture.db,
|
||||
reloadedManager,
|
||||
reloadedRegistry,
|
||||
fixture.priceRuntime,
|
||||
fixture.encryption,
|
||||
fixture.service.dialects,
|
||||
fixture.service.requestLogs,
|
||||
fixture.service.usageStats,
|
||||
fixture.stats,
|
||||
fixture.requestLogStats,
|
||||
)
|
||||
reloadedEngine := gin.New()
|
||||
NewServer(&config.Config{AuthKey: "test-auth-key"}, reloadedService).RegisterRoutes(reloadedEngine)
|
||||
afterReload := serveSettingsOCCRequest(
|
||||
t, reloadedEngine, http.MethodGet, "test-auth-key", "en-US", "", "",
|
||||
)
|
||||
if afterReload.Header().Get("ETag") != initial.Header().Get("ETag") ||
|
||||
!bytes.Equal(afterReload.Body.Bytes(), initial.Body.Bytes()) {
|
||||
t.Fatalf("reload changed Settings representation:\ninitial %q %s\nreload %q %s",
|
||||
initial.Header().Get("ETag"), initial.Body.String(),
|
||||
afterReload.Header().Get("ETag"), afterReload.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func assertSettingsRepresentationHeaders(
|
||||
t *testing.T,
|
||||
recorder *httptest.ResponseRecorder,
|
||||
language string,
|
||||
) {
|
||||
t.Helper()
|
||||
tag := recorder.Header().Get("ETag")
|
||||
if len(tag) != len(`"sha256-`)+64+len(`"`) ||
|
||||
!strings.HasPrefix(tag, `"sha256-`) ||
|
||||
!strings.HasSuffix(tag, `"`) {
|
||||
t.Fatalf("ETag = %q", tag)
|
||||
}
|
||||
if got := recorder.Header().Get("Content-Type"); got != "application/json; charset=utf-8" {
|
||||
t.Fatalf("Content-Type = %q", got)
|
||||
}
|
||||
if got := recorder.Header().Get("Content-Encoding"); got != "identity" {
|
||||
t.Fatalf("Content-Encoding = %q", got)
|
||||
}
|
||||
if got := recorder.Header().Get("Content-Language"); got != language {
|
||||
t.Fatalf("Content-Language = %q, want %q", got, language)
|
||||
}
|
||||
if got := recorder.Header().Get("Vary"); got != "Accept-Language" {
|
||||
t.Fatalf("Vary = %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func serveSettingsOCCRequest(
|
||||
t *testing.T,
|
||||
engine *gin.Engine,
|
||||
method string,
|
||||
authKey string,
|
||||
language string,
|
||||
ifMatch string,
|
||||
body string,
|
||||
) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
recorder := httptest.NewRecorder()
|
||||
request := httptest.NewRequest(method, "/api/settings", strings.NewReader(body))
|
||||
if authKey != "" {
|
||||
request.Header.Set("Authorization", "Bearer "+authKey)
|
||||
}
|
||||
if language != "" {
|
||||
request.Header.Set("Accept-Language", language)
|
||||
}
|
||||
if ifMatch != "" {
|
||||
request.Header.Set("If-Match", ifMatch)
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
engine.ServeHTTP(recorder, request)
|
||||
return recorder
|
||||
}
|
||||
@@ -1,9 +1,12 @@
|
||||
package control
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -21,10 +24,19 @@ import (
|
||||
)
|
||||
|
||||
var (
|
||||
controlI18nOnce sync.Once
|
||||
controlI18nErr error
|
||||
controlI18nOnce sync.Once
|
||||
controlI18nErr error
|
||||
testIdempotencySequence atomic.Uint64
|
||||
)
|
||||
|
||||
func setRequiredTestIdempotencyHeader(request *http.Request) {
|
||||
sequence := testIdempotencySequence.Add(1)
|
||||
request.Header.Set(
|
||||
"Idempotency-Key",
|
||||
fmt.Sprintf("00000000-0000-4000-8000-%012x", sequence),
|
||||
)
|
||||
}
|
||||
|
||||
type serviceFixture struct {
|
||||
db *gorm.DB
|
||||
manager *state.Manager
|
||||
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"name": "access-key-create",
|
||||
"operation_kind": "access_key_create",
|
||||
"path_template": "/api/access-keys",
|
||||
"resource_locator": "new",
|
||||
"canonical_body": "{\"filters\":{\"groups\":[],\"models\":[],\"protocols\":[]},\"name\":\"CI client\",\"rpm_limit\":0}",
|
||||
"domain_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e6379",
|
||||
"separator_hex": "00",
|
||||
"framed_input_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e637900000000013100000004504f5354000000116163636573735f6b65795f637265617465000000102f6170692f6163636573732d6b657973000000036e657700000010636f6e74726f6c2d61646d696e2d7631000000557b2266696c74657273223a7b2267726f757073223a5b5d2c226d6f64656c73223a5b5d2c2270726f746f636f6c73223a5b5d7d2c226e616d65223a22434920636c69656e74222c2272706d5f6c696d6974223a307d",
|
||||
"expected_sha256": "7b22877fc5ce68bde59170f4906463bbbf873abff4cda536dd36c30da3a62b81"
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"name": "group-create-k-twice",
|
||||
"operation_kind": "group_create",
|
||||
"path_template": "/api/groups",
|
||||
"resource_locator": "new",
|
||||
"canonical_body": "{\"config\":{},\"confirm_same_upstream_url\":false,\"keys\":[\"K\",\"K\"],\"models\":[],\"name\":\"fixture\",\"protocols\":[\"openai\"],\"upstream_url\":\"https://fixture.example.com/v1\"}",
|
||||
"domain_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e6379",
|
||||
"separator_hex": "00",
|
||||
"framed_input_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e637900000000013100000004504f53540000000c67726f75705f6372656174650000000b2f6170692f67726f757073000000036e657700000010636f6e74726f6c2d61646d696e2d7631000000a47b22636f6e666967223a7b7d2c22636f6e6669726d5f73616d655f757073747265616d5f75726c223a66616c73652c226b657973223a5b224b222c224b225d2c226d6f64656c73223a5b5d2c226e616d65223a2266697874757265222c2270726f746f636f6c73223a5b226f70656e6169225d2c22757073747265616d5f75726c223a2268747470733a2f2f666978747572652e6578616d706c652e636f6d2f7631227d",
|
||||
"expected_sha256": "701aa5aba8c91c53e4df345026b64cf72d2210fb3400a0a4c3fe3ca1e9660d30"
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"name": "group-create-k",
|
||||
"operation_kind": "group_create",
|
||||
"path_template": "/api/groups",
|
||||
"resource_locator": "new",
|
||||
"canonical_body": "{\"config\":{},\"confirm_same_upstream_url\":false,\"keys\":[\"K\"],\"models\":[],\"name\":\"fixture\",\"protocols\":[\"openai\"],\"upstream_url\":\"https://fixture.example.com/v1\"}",
|
||||
"domain_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e6379",
|
||||
"separator_hex": "00",
|
||||
"framed_input_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e637900000000013100000004504f53540000000c67726f75705f6372656174650000000b2f6170692f67726f757073000000036e657700000010636f6e74726f6c2d61646d696e2d7631000000a07b22636f6e666967223a7b7d2c22636f6e6669726d5f73616d655f757073747265616d5f75726c223a66616c73652c226b657973223a5b224b225d2c226d6f64656c73223a5b5d2c226e616d65223a2266697874757265222c2270726f746f636f6c73223a5b226f70656e6169225d2c22757073747265616d5f75726c223a2268747470733a2f2f666978747572652e6578616d706c652e636f6d2f7631227d",
|
||||
"expected_sha256": "20a145b24354734cf86e4b4ba0b7a5b850ae2d8877e8fedf91aef0cb98880862"
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"name": "group-key-import-k-twice",
|
||||
"operation_kind": "group_key_import",
|
||||
"path_template": "/api/groups/:group_id/keys/import",
|
||||
"resource_locator": "group:7",
|
||||
"canonical_body": "{\"keys\":[\"K\",\"K\"]}",
|
||||
"domain_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e6379",
|
||||
"separator_hex": "00",
|
||||
"framed_input_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e637900000000013100000004504f53540000001067726f75705f6b65795f696d706f7274000000212f6170692f67726f7570732f3a67726f75705f69642f6b6579732f696d706f72740000000767726f75703a3700000010636f6e74726f6c2d61646d696e2d7631000000127b226b657973223a5b224b222c224b225d7d",
|
||||
"expected_sha256": "8e98a779323ab4b2652c7814126307c14aa106e435b5abdf5a3f98c8be1f264b"
|
||||
}
|
||||
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"name": "group-key-import-k",
|
||||
"operation_kind": "group_key_import",
|
||||
"path_template": "/api/groups/:group_id/keys/import",
|
||||
"resource_locator": "group:7",
|
||||
"canonical_body": "{\"keys\":[\"K\"]}",
|
||||
"domain_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e6379",
|
||||
"separator_hex": "00",
|
||||
"framed_input_hex": "6770742d6c6f61642f636f6e74726f6c2d6964656d706f74656e637900000000013100000004504f53540000001067726f75705f6b65795f696d706f7274000000212f6170692f67726f7570732f3a67726f75705f69642f6b6579732f696d706f72740000000767726f75703a3700000010636f6e74726f6c2d61646d696e2d76310000000e7b226b657973223a5b224b225d7d",
|
||||
"expected_sha256": "33ed46d786c43e583a1aab96104d834e26f8d46d402cf8fcacb83ec7e87373d4"
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"domain_hex": "6770742d6c6f61642f73657474696e67732d776972652d7631",
|
||||
"canonical_response_hex": "7b22636f6465223a302c2264617461223a7b226f7665727269646573223a5b5d2c2276616c756573223a7b22636f6e6e6563745f74696d656f7574223a31352c2266697273745f627974655f74696d656f7574223a3132302c226865616465725f72756c6573223a7b2272656d6f7665223a5b5d2c22736574223a7b7d7d2c22696e6a6563745f75736167655f6f7074696f6e73223a747275652c22726571756573745f6c6f675f726574656e74696f6e5f64617973223a372c22726571756573745f74696d656f7574223a3630302c2273747265616d5f69646c655f74696d656f7574223a3330307d7d2c226d657373616765223a2253756363657373227d",
|
||||
"framed_input_hex": "6770742d6c6f61642f73657474696e67732d776972652d7631007b22636f6465223a302c2264617461223a7b226f7665727269646573223a5b5d2c2276616c756573223a7b22636f6e6e6563745f74696d656f7574223a31352c2266697273745f627974655f74696d656f7574223a3132302c226865616465725f72756c6573223a7b2272656d6f7665223a5b5d2c22736574223a7b7d7d2c22696e6a6563745f75736167655f6f7074696f6e73223a747275652c22726571756573745f6c6f675f726574656e74696f6e5f64617973223a372c22726571756573745f74696d656f7574223a3630302c2273747265616d5f69646c655f74696d656f7574223a3330307d7d2c226d657373616765223a2253756363657373227d",
|
||||
"expected_digest": "227bd1218eee98e141755f98a14e372540a8856964664484c63ce35eecad8306"
|
||||
}
|
||||
@@ -0,0 +1,278 @@
|
||||
// Package canonicaljson implements the restricted JSON Canonicalization Scheme
|
||||
// subset used by GPT-Load control-plane contracts. Contract numbers are signed
|
||||
// 64-bit integers; floating-point values are rejected.
|
||||
package canonicaljson
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"reflect"
|
||||
"sort"
|
||||
"strconv"
|
||||
"unicode/utf16"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
// Marshal encodes a typed value as canonical JSON.
|
||||
func Marshal(value any) ([]byte, error) {
|
||||
if err := rejectUnsupportedTypedValue(reflect.ValueOf(value)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
raw, err := json.Marshal(value)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal canonical JSON input: %w", err)
|
||||
}
|
||||
return Canonicalize(raw)
|
||||
}
|
||||
|
||||
// Canonicalize validates and canonicalizes one JSON value.
|
||||
func Canonicalize(raw []byte) ([]byte, error) {
|
||||
decoder := json.NewDecoder(bytes.NewReader(raw))
|
||||
decoder.UseNumber()
|
||||
value, err := decodeUniqueValue(decoder)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("decode canonical JSON input: %w", err)
|
||||
}
|
||||
if _, err := decoder.Token(); err == nil {
|
||||
return nil, fmt.Errorf("decode canonical JSON input: multiple values")
|
||||
} else if err != io.EOF {
|
||||
return nil, fmt.Errorf("decode canonical JSON trailing input: %w", err)
|
||||
}
|
||||
|
||||
var output bytes.Buffer
|
||||
if err := appendCanonical(&output, value); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return output.Bytes(), nil
|
||||
}
|
||||
|
||||
func decodeUniqueValue(decoder *json.Decoder) (any, error) {
|
||||
token, err := decoder.Token()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
delimiter, isDelimiter := token.(json.Delim)
|
||||
if !isDelimiter {
|
||||
return token, nil
|
||||
}
|
||||
|
||||
switch delimiter {
|
||||
case '{':
|
||||
object := make(map[string]any)
|
||||
for decoder.More() {
|
||||
keyToken, err := decoder.Token()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key, ok := keyToken.(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("object key is not a string")
|
||||
}
|
||||
if _, duplicate := object[key]; duplicate {
|
||||
return nil, fmt.Errorf("duplicate object key %q", key)
|
||||
}
|
||||
value, err := decodeUniqueValue(decoder)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
object[key] = value
|
||||
}
|
||||
end, err := decoder.Token()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if end != json.Delim('}') {
|
||||
return nil, fmt.Errorf("object is not closed")
|
||||
}
|
||||
return object, nil
|
||||
case '[':
|
||||
array := make([]any, 0)
|
||||
for decoder.More() {
|
||||
value, err := decodeUniqueValue(decoder)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
array = append(array, value)
|
||||
}
|
||||
end, err := decoder.Token()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if end != json.Delim(']') {
|
||||
return nil, fmt.Errorf("array is not closed")
|
||||
}
|
||||
return array, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unexpected JSON delimiter %q", delimiter)
|
||||
}
|
||||
}
|
||||
|
||||
func rejectUnsupportedTypedValue(value reflect.Value) error {
|
||||
if !value.IsValid() {
|
||||
return nil
|
||||
}
|
||||
switch value.Kind() {
|
||||
case reflect.Interface, reflect.Pointer:
|
||||
if value.IsNil() {
|
||||
return nil
|
||||
}
|
||||
return rejectUnsupportedTypedValue(value.Elem())
|
||||
case reflect.Float32, reflect.Float64:
|
||||
return fmt.Errorf("canonical JSON floats are not supported")
|
||||
case reflect.Map:
|
||||
if value.Type().Key().Kind() != reflect.String {
|
||||
return fmt.Errorf("canonical JSON object keys must be strings")
|
||||
}
|
||||
iter := value.MapRange()
|
||||
for iter.Next() {
|
||||
if err := rejectUnsupportedTypedValue(iter.Value()); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case reflect.Array, reflect.Slice:
|
||||
if value.Type().Elem().Kind() == reflect.Uint8 {
|
||||
return nil
|
||||
}
|
||||
for index := 0; index < value.Len(); index++ {
|
||||
if err := rejectUnsupportedTypedValue(value.Index(index)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
case reflect.Struct:
|
||||
if _, ok := value.Interface().(json.Number); ok {
|
||||
return nil
|
||||
}
|
||||
for index := 0; index < value.NumField(); index++ {
|
||||
if value.Type().Field(index).IsExported() {
|
||||
if err := rejectUnsupportedTypedValue(value.Field(index)); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
case reflect.Func, reflect.Chan, reflect.Complex64, reflect.Complex128, reflect.UnsafePointer:
|
||||
return fmt.Errorf("canonical JSON does not support %s", value.Kind())
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func appendCanonical(output *bytes.Buffer, value any) error {
|
||||
switch typed := value.(type) {
|
||||
case nil:
|
||||
output.WriteString("null")
|
||||
case bool:
|
||||
if typed {
|
||||
output.WriteString("true")
|
||||
} else {
|
||||
output.WriteString("false")
|
||||
}
|
||||
case string:
|
||||
return appendString(output, typed)
|
||||
case json.Number:
|
||||
integer, err := strconv.ParseInt(string(typed), 10, 64)
|
||||
if err != nil {
|
||||
return fmt.Errorf("canonical JSON number %q is not a signed integer", typed)
|
||||
}
|
||||
output.WriteString(strconv.FormatInt(integer, 10))
|
||||
case []any:
|
||||
output.WriteByte('[')
|
||||
for index, item := range typed {
|
||||
if index > 0 {
|
||||
output.WriteByte(',')
|
||||
}
|
||||
if err := appendCanonical(output, item); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
output.WriteByte(']')
|
||||
case map[string]any:
|
||||
keys := make([]string, 0, len(typed))
|
||||
for key := range typed {
|
||||
if !utf8.ValidString(key) {
|
||||
return fmt.Errorf("canonical JSON object key is not valid UTF-8")
|
||||
}
|
||||
keys = append(keys, key)
|
||||
}
|
||||
sort.Slice(keys, func(left, right int) bool {
|
||||
return compareUTF16(keys[left], keys[right]) < 0
|
||||
})
|
||||
output.WriteByte('{')
|
||||
for index, key := range keys {
|
||||
if index > 0 {
|
||||
output.WriteByte(',')
|
||||
}
|
||||
if err := appendString(output, key); err != nil {
|
||||
return err
|
||||
}
|
||||
output.WriteByte(':')
|
||||
if err := appendCanonical(output, typed[key]); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
output.WriteByte('}')
|
||||
default:
|
||||
return fmt.Errorf("canonical JSON decoded unsupported type %T", value)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func appendString(output *bytes.Buffer, value string) error {
|
||||
if !utf8.ValidString(value) {
|
||||
return fmt.Errorf("canonical JSON string is not valid UTF-8")
|
||||
}
|
||||
output.WriteByte('"')
|
||||
for _, character := range value {
|
||||
switch character {
|
||||
case '"':
|
||||
output.WriteString(`\"`)
|
||||
case '\\':
|
||||
output.WriteString(`\\`)
|
||||
case '\b':
|
||||
output.WriteString(`\b`)
|
||||
case '\t':
|
||||
output.WriteString(`\t`)
|
||||
case '\n':
|
||||
output.WriteString(`\n`)
|
||||
case '\f':
|
||||
output.WriteString(`\f`)
|
||||
case '\r':
|
||||
output.WriteString(`\r`)
|
||||
default:
|
||||
if character < 0x20 {
|
||||
output.WriteString(`\u00`)
|
||||
output.WriteString(strconv.FormatInt(int64(character>>4), 16))
|
||||
output.WriteString(strconv.FormatInt(int64(character&0xf), 16))
|
||||
continue
|
||||
}
|
||||
output.WriteRune(character)
|
||||
}
|
||||
}
|
||||
output.WriteByte('"')
|
||||
return nil
|
||||
}
|
||||
|
||||
func compareUTF16(left, right string) int {
|
||||
leftUnits := utf16.Encode([]rune(left))
|
||||
rightUnits := utf16.Encode([]rune(right))
|
||||
limit := len(leftUnits)
|
||||
if len(rightUnits) < limit {
|
||||
limit = len(rightUnits)
|
||||
}
|
||||
for index := 0; index < limit; index++ {
|
||||
if leftUnits[index] < rightUnits[index] {
|
||||
return -1
|
||||
}
|
||||
if leftUnits[index] > rightUnits[index] {
|
||||
return 1
|
||||
}
|
||||
}
|
||||
switch {
|
||||
case len(leftUnits) < len(rightUnits):
|
||||
return -1
|
||||
case len(leftUnits) > len(rightUnits):
|
||||
return 1
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,61 @@
|
||||
package canonicaljson
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestMarshalUsesJCSObjectOrderAndPreservesArrayOrder(t *testing.T) {
|
||||
input := map[string]any{
|
||||
"\ue000": "private-use",
|
||||
"😀": "supplementary",
|
||||
"z": "<&>",
|
||||
"a": []any{int64(3), true, nil, "é"},
|
||||
}
|
||||
|
||||
got, err := Marshal(input)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal() error = %v", err)
|
||||
}
|
||||
want := `{"a":[3,true,null,"é"],"z":"<&>","😀":"supplementary","":"private-use"}`
|
||||
if string(got) != want {
|
||||
t.Fatalf("Marshal() = %s, want %s", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCanonicalizeNormalizesIntegersAndRejectsFloats(t *testing.T) {
|
||||
got, err := Canonicalize([]byte(`{"positive":1,"negative_zero":-0,"ordered":[3,2,1]}`))
|
||||
if err != nil {
|
||||
t.Fatalf("Canonicalize(integer input) error = %v", err)
|
||||
}
|
||||
want := `{"negative_zero":0,"ordered":[3,2,1],"positive":1}`
|
||||
if string(got) != want {
|
||||
t.Fatalf("Canonicalize(integer input) = %s, want %s", got, want)
|
||||
}
|
||||
|
||||
for _, raw := range []string{
|
||||
`{"value":1.5}`,
|
||||
`{"value":1e3}`,
|
||||
`{"value":9223372036854775808}`,
|
||||
`{"value":1,"value":2}`,
|
||||
`{"nested":{"value":1,"value":2}}`,
|
||||
`{"value":1}{"trailing":2}`,
|
||||
} {
|
||||
t.Run(raw, func(t *testing.T) {
|
||||
if _, err := Canonicalize([]byte(raw)); err == nil {
|
||||
t.Fatalf("Canonicalize(%s) error = nil, want rejection", raw)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestMarshalRejectsUnsupportedTypedValues(t *testing.T) {
|
||||
for _, value := range []any{
|
||||
float64(1),
|
||||
map[int]string{1: "unsupported-key"},
|
||||
func() {},
|
||||
} {
|
||||
if _, err := Marshal(value); err == nil {
|
||||
t.Fatalf("Marshal(%T) error = nil, want rejection", value)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -38,6 +38,14 @@ var (
|
||||
ErrUpstreamURLChangeConfirmationRequired = &APIError{HTTPStatus: http.StatusConflict, Code: "UPSTREAM_URL_CHANGE_CONFIRMATION_REQUIRED", Message: "Upstream URL change requires explicit confirmation"}
|
||||
ErrNoActiveUpstreamKey = &APIError{HTTPStatus: http.StatusConflict, Code: "NO_ACTIVE_UPSTREAM_KEY", Message: "No active upstream key available for this group"}
|
||||
ErrBadGateway = &APIError{HTTPStatus: http.StatusBadGateway, Code: "BAD_GATEWAY", Message: "Upstream service error"}
|
||||
ErrIdempotencyKeyRequired = &APIError{HTTPStatus: http.StatusPreconditionRequired, Code: "IDEMPOTENCY_KEY_REQUIRED", Message: "Idempotency-Key is required"}
|
||||
ErrInvalidIdempotencyKey = &APIError{HTTPStatus: http.StatusBadRequest, Code: "INVALID_IDEMPOTENCY_KEY", Message: "Idempotency-Key must be a canonical UUID v4"}
|
||||
ErrIdempotencyKeyReused = &APIError{HTTPStatus: http.StatusConflict, Code: "IDEMPOTENCY_KEY_REUSED", Message: "Idempotency-Key was already used for another request"}
|
||||
ErrIdempotencyResultExpired = &APIError{HTTPStatus: http.StatusGone, Code: "IDEMPOTENCY_RESULT_EXPIRED", Message: "The idempotent result retention period expired"}
|
||||
ErrControlOperationIncomplete = &APIError{HTTPStatus: http.StatusServiceUnavailable, Code: "CONTROL_OPERATION_INCOMPLETE", Message: "The resource was committed but runtime recovery is incomplete"}
|
||||
ErrControlRecoveryPending = &APIError{HTTPStatus: http.StatusServiceUnavailable, Code: "CONTROL_RECOVERY_PENDING", Message: "An earlier committed operation is still recovering"}
|
||||
ErrSettingsPreconditionRequired = &APIError{HTTPStatus: http.StatusPreconditionRequired, Code: "SETTINGS_PRECONDITION_REQUIRED", Message: "If-Match is required"}
|
||||
ErrSettingsVersionConflict = &APIError{HTTPStatus: http.StatusPreconditionFailed, Code: "SETTINGS_VERSION_CONFLICT", Message: "Settings changed since they were loaded"}
|
||||
)
|
||||
|
||||
// NewAPIErrorWithData creates a copy of an APIError with response data.
|
||||
|
||||
@@ -58,6 +58,16 @@ func GetLocalizer(acceptLang string) *i18n.Localizer {
|
||||
return i18n.NewLocalizer(bundle, langs...)
|
||||
}
|
||||
|
||||
// ResolveLanguage returns the canonical supported language selected from
|
||||
// Accept-Language.
|
||||
func ResolveLanguage(acceptLang string) string {
|
||||
languages := parseAcceptLanguage(acceptLang)
|
||||
if len(languages) == 0 {
|
||||
return "zh-CN"
|
||||
}
|
||||
return languages[0]
|
||||
}
|
||||
|
||||
// parseAcceptLanguage 解析 Accept-Language 头
|
||||
func parseAcceptLanguage(acceptLang string) []string {
|
||||
if acceptLang == "" {
|
||||
|
||||
@@ -2,17 +2,25 @@ package locales
|
||||
|
||||
// MessagesEnUS contains English (US) control-plane translations.
|
||||
var MessagesEnUS = map[string]string{
|
||||
"common.success": "Success",
|
||||
"bad_request": "Bad request",
|
||||
"request_too_large": "Request body is too large",
|
||||
"bad_gateway": "Upstream service error",
|
||||
"internal_error": "Internal error",
|
||||
"auth.invalid_key": "Invalid authorization key",
|
||||
"auth.locked": "Too many authentication attempts; try again later",
|
||||
"group.not_found": "Group not found",
|
||||
"group.name_exists": "Group name already exists",
|
||||
"group.in_use": "The group is still referenced by access keys",
|
||||
"group.upstream_url_conflict": "An existing group already uses this upstream URL",
|
||||
"common.success": "Success",
|
||||
"bad_request": "Bad request",
|
||||
"request_too_large": "Request body is too large",
|
||||
"bad_gateway": "Upstream service error",
|
||||
"internal_error": "Internal error",
|
||||
"idempotency.required": "Idempotency-Key is required",
|
||||
"idempotency.invalid": "Idempotency-Key must be a canonical UUID v4",
|
||||
"idempotency.reused": "Idempotency-Key was already used for another request",
|
||||
"idempotency.expired": "The idempotent result retention period expired",
|
||||
"control.operation_incomplete": "The resource was committed but runtime recovery is incomplete",
|
||||
"control.recovery_pending": "An earlier committed operation is still recovering",
|
||||
"settings.precondition_required": "If-Match is required",
|
||||
"settings.version_conflict": "Settings changed since they were loaded",
|
||||
"auth.invalid_key": "Invalid authorization key",
|
||||
"auth.locked": "Too many authentication attempts; try again later",
|
||||
"group.not_found": "Group not found",
|
||||
"group.name_exists": "Group name already exists",
|
||||
"group.in_use": "The group is still referenced by access keys",
|
||||
"group.upstream_url_conflict": "An existing group already uses this upstream URL",
|
||||
"group.upstream_url_change_confirmation_required": "Changing the upstream URL requires explicit confirmation",
|
||||
"group.no_active_upstream_key": "No active upstream key is available for this group",
|
||||
"key.not_found": "Key not found",
|
||||
|
||||
@@ -2,17 +2,25 @@ package locales
|
||||
|
||||
// MessagesJaJP contains Japanese control-plane translations.
|
||||
var MessagesJaJP = map[string]string{
|
||||
"common.success": "成功",
|
||||
"bad_request": "不正なリクエスト",
|
||||
"request_too_large": "リクエストボディが大きすぎます",
|
||||
"bad_gateway": "上流サービスエラー",
|
||||
"internal_error": "内部エラー",
|
||||
"auth.invalid_key": "無効な認証キー",
|
||||
"auth.locked": "認証試行回数が多すぎます。しばらくしてから再試行してください",
|
||||
"group.not_found": "グループが存在しません",
|
||||
"group.name_exists": "グループ名が既に存在します",
|
||||
"group.in_use": "グループはアクセスキーから参照されています",
|
||||
"group.upstream_url_conflict": "このアップストリームURLは既存のグループで使用されています",
|
||||
"common.success": "成功",
|
||||
"bad_request": "不正なリクエスト",
|
||||
"request_too_large": "リクエストボディが大きすぎます",
|
||||
"bad_gateway": "上流サービスエラー",
|
||||
"internal_error": "内部エラー",
|
||||
"idempotency.required": "Idempotency-Key が必要です",
|
||||
"idempotency.invalid": "Idempotency-Key は正規 UUID v4 である必要があります",
|
||||
"idempotency.reused": "Idempotency-Key は別のリクエストで使用済みです",
|
||||
"idempotency.expired": "冪等結果の保持期間が終了しました",
|
||||
"control.operation_incomplete": "リソースは確定しましたが、実行時の復旧が未完了です",
|
||||
"control.recovery_pending": "以前に確定した操作を復旧中です",
|
||||
"settings.precondition_required": "If-Match が必要です",
|
||||
"settings.version_conflict": "読み込み後に設定が変更されました",
|
||||
"auth.invalid_key": "無効な認証キー",
|
||||
"auth.locked": "認証試行回数が多すぎます。しばらくしてから再試行してください",
|
||||
"group.not_found": "グループが存在しません",
|
||||
"group.name_exists": "グループ名が既に存在します",
|
||||
"group.in_use": "グループはアクセスキーから参照されています",
|
||||
"group.upstream_url_conflict": "このアップストリームURLは既存のグループで使用されています",
|
||||
"group.upstream_url_change_confirmation_required": "アップストリームURLの変更には明示的な確認が必要です",
|
||||
"group.no_active_upstream_key": "このグループには利用可能な有効なアップストリームキーがありません",
|
||||
"key.not_found": "キーが存在しません",
|
||||
|
||||
@@ -2,17 +2,25 @@ package locales
|
||||
|
||||
// MessagesZhCN contains Simplified Chinese control-plane translations.
|
||||
var MessagesZhCN = map[string]string{
|
||||
"common.success": "操作成功",
|
||||
"bad_request": "请求错误",
|
||||
"request_too_large": "请求体过大",
|
||||
"bad_gateway": "上游服务错误",
|
||||
"internal_error": "内部错误",
|
||||
"auth.invalid_key": "无效的授权密钥",
|
||||
"auth.locked": "认证尝试过多,请稍后重试",
|
||||
"group.not_found": "分组不存在",
|
||||
"group.name_exists": "分组名称已存在",
|
||||
"group.in_use": "分组仍被访问密钥引用",
|
||||
"group.upstream_url_conflict": "已有分组使用该上游地址",
|
||||
"common.success": "操作成功",
|
||||
"bad_request": "请求错误",
|
||||
"request_too_large": "请求体过大",
|
||||
"bad_gateway": "上游服务错误",
|
||||
"internal_error": "内部错误",
|
||||
"idempotency.required": "必须提供 Idempotency-Key",
|
||||
"idempotency.invalid": "Idempotency-Key 必须是规范 UUID v4",
|
||||
"idempotency.reused": "Idempotency-Key 已用于其它请求",
|
||||
"idempotency.expired": "幂等结果保留期已过",
|
||||
"control.operation_incomplete": "资源已提交,但运行时恢复尚未完成",
|
||||
"control.recovery_pending": "较早提交的操作仍在恢复",
|
||||
"settings.precondition_required": "必须提供 If-Match",
|
||||
"settings.version_conflict": "设置已在加载后发生变化",
|
||||
"auth.invalid_key": "无效的授权密钥",
|
||||
"auth.locked": "认证尝试过多,请稍后重试",
|
||||
"group.not_found": "分组不存在",
|
||||
"group.name_exists": "分组名称已存在",
|
||||
"group.in_use": "分组仍被访问密钥引用",
|
||||
"group.upstream_url_conflict": "已有分组使用该上游地址",
|
||||
"group.upstream_url_change_confirmation_required": "修改上游地址需要明确确认",
|
||||
"group.no_active_upstream_key": "该分组没有可用的活跃上游密钥",
|
||||
"key.not_found": "密钥不存在",
|
||||
|
||||
@@ -8,6 +8,8 @@ import (
|
||||
const (
|
||||
// LocalizerKey 是 gin.Context 中存储 Localizer 的键
|
||||
LocalizerKey = "localizer"
|
||||
// LanguageKey 是 gin.Context 中存储规范语言代码的键
|
||||
LanguageKey = "language"
|
||||
)
|
||||
|
||||
// Middleware i18n 中间件
|
||||
@@ -18,14 +20,26 @@ func Middleware() gin.HandlerFunc {
|
||||
|
||||
// 获取 Localizer
|
||||
localizer := GetLocalizer(acceptLang)
|
||||
language := ResolveLanguage(acceptLang)
|
||||
|
||||
// 将 Localizer 存储到 Context 中
|
||||
c.Set(LocalizerKey, localizer)
|
||||
c.Set(LanguageKey, language)
|
||||
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
// GetLanguageFromContext 返回当前响应应使用的规范语言代码。
|
||||
func GetLanguageFromContext(c *gin.Context) string {
|
||||
if value, exists := c.Get(LanguageKey); exists {
|
||||
if language, ok := value.(string); ok && language != "" {
|
||||
return language
|
||||
}
|
||||
}
|
||||
return ResolveLanguage(c.GetHeader("Accept-Language"))
|
||||
}
|
||||
|
||||
// GetLocalizerFromContext 从 gin.Context 获取 Localizer
|
||||
func GetLocalizerFromContext(c *gin.Context) *i18n.Localizer {
|
||||
if localizer, exists := c.Get(LocalizerKey); exists {
|
||||
|
||||
@@ -403,14 +403,14 @@ func TestServiceListBatchLoadsCurrentAccessKeyNames(t *testing.T) {
|
||||
db := openRequestLogQueryDB(t)
|
||||
current := models.AccessKey{
|
||||
Name: "before-rename", KeyValue: "cipher-current", KeyHash: "hash-current",
|
||||
Status: "active", Filters: models.JSON(`{}`),
|
||||
KeySuffix: "0001", Status: "active", Filters: models.JSON(`{}`),
|
||||
}
|
||||
if err := db.Create(¤t).Error; err != nil {
|
||||
t.Fatalf("create current AccessKey: %v", err)
|
||||
}
|
||||
deleted := models.AccessKey{
|
||||
Name: "deleted", KeyValue: "cipher-deleted", KeyHash: "hash-deleted",
|
||||
Status: "active", Filters: models.JSON(`{}`),
|
||||
KeySuffix: "0002", Status: "active", Filters: models.JSON(`{}`),
|
||||
}
|
||||
if err := db.Create(&deleted).Error; err != nil {
|
||||
t.Fatalf("create deleted AccessKey: %v", err)
|
||||
|
||||
@@ -94,7 +94,10 @@ func queryCompileRows(ctx context.Context, db *gorm.DB) (compileRows, error) {
|
||||
if err := db.Order("id ASC").Find(&rows.groups).Error; err != nil {
|
||||
return compileRows{}, fmt.Errorf("query groups: %w", err)
|
||||
}
|
||||
if err := db.Order("id ASC").Find(&rows.accessKeys).Error; err != nil {
|
||||
if err := db.
|
||||
Select("id", "name", "key_hash", "status", "filters", "rpm_limit").
|
||||
Order("id ASC").
|
||||
Find(&rows.accessKeys).Error; err != nil {
|
||||
return compileRows{}, fmt.Errorf("query access keys: %w", err)
|
||||
}
|
||||
return rows, nil
|
||||
@@ -125,6 +128,38 @@ func BuildCompileInput(ctx context.Context, db *gorm.DB) (state.CompileInput, er
|
||||
return input, nil
|
||||
}
|
||||
|
||||
// BuildGroupKeyEntries maps the persisted key configuration for one group.
|
||||
// Runtime health fields are initialized to their durable baseline.
|
||||
func BuildGroupKeyEntries(
|
||||
ctx context.Context,
|
||||
db *gorm.DB,
|
||||
groupID uint,
|
||||
) ([]state.KeyEntry, error) {
|
||||
if groupID == 0 {
|
||||
return nil, fmt.Errorf("group id is required")
|
||||
}
|
||||
var group struct{ ID uint }
|
||||
if err := db.WithContext(ctx).
|
||||
Model(&models.Group{}).
|
||||
Select("id").
|
||||
Where("id = ?", groupID).
|
||||
Take(&group).Error; err != nil {
|
||||
return nil, fmt.Errorf("query group %d: %w", groupID, err)
|
||||
}
|
||||
var rows []models.UpstreamKey
|
||||
if err := db.WithContext(ctx).
|
||||
Where("group_id = ?", groupID).
|
||||
Order("id ASC").
|
||||
Find(&rows).Error; err != nil {
|
||||
return nil, fmt.Errorf("query group %d upstream keys: %w", groupID, err)
|
||||
}
|
||||
entries := mapUpstreamKeys(rows)
|
||||
if err := state.ValidateKeyEntries(entries); err != nil {
|
||||
return nil, fmt.Errorf("validate group %d upstream keys: %w", groupID, err)
|
||||
}
|
||||
return entries, nil
|
||||
}
|
||||
|
||||
func (l *Loader) read(ctx context.Context) (state.CompileInput, []state.KeyEntry, error) {
|
||||
input, err := BuildCompileInput(ctx, l.db)
|
||||
if err != nil {
|
||||
|
||||
@@ -277,7 +277,8 @@ func TestBuildCompileInputReadsUncommittedTransactionState(t *testing.T) {
|
||||
mustCreate(t, tx, &group)
|
||||
mustCreate(t, tx, &models.AccessKey{
|
||||
Name: "pending", KeyValue: "cipher", KeyHash: "pending-hash",
|
||||
Status: "active", Filters: models.JSON(fmt.Sprintf(`{"groups":[%d]}`, group.ID)),
|
||||
KeySuffix: "0000", Status: "active",
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d]}`, group.ID)),
|
||||
})
|
||||
|
||||
input, err := loader.BuildCompileInput(context.Background(), tx)
|
||||
@@ -298,7 +299,8 @@ func TestBuildCompileInputReturnsIndependentData(t *testing.T) {
|
||||
}
|
||||
mustCreate(t, db, &models.AccessKey{
|
||||
Name: "owned", KeyValue: "cipher", KeyHash: "owned-hash", Status: "active",
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d],"protocols":["openai"],"models":["gpt-owned"]}`, group.ID)),
|
||||
KeySuffix: "0001",
|
||||
Filters: models.JSON(fmt.Sprintf(`{"groups":[%d],"protocols":["openai"],"models":["gpt-owned"]}`, group.ID)),
|
||||
})
|
||||
|
||||
first, err := loader.BuildCompileInput(context.Background(), db)
|
||||
@@ -337,7 +339,7 @@ func TestBuildCompileInputDoesNotQueryUpstreamKeys(t *testing.T) {
|
||||
group := createRuntimeGroup(t, db, "query-boundary", protocol.OpenAI, "gpt-query")
|
||||
mustCreate(t, db, &models.AccessKey{
|
||||
Name: "query-boundary", KeyValue: "cipher", KeyHash: "query-boundary-hash",
|
||||
Status: "active", Filters: models.JSON(`{}`),
|
||||
KeySuffix: "0002", Status: "active", Filters: models.JSON(`{}`),
|
||||
})
|
||||
mustCreate(t, db, &models.UpstreamKey{
|
||||
GroupID: group.ID, KeyValue: "upstream-cipher", KeyHash: "upstream-query-hash",
|
||||
@@ -572,7 +574,7 @@ func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) {
|
||||
|
||||
activeAccess := models.AccessKey{
|
||||
Name: "active access", KeyValue: "access-cipher-active", KeyHash: "active-hash",
|
||||
Status: "active",
|
||||
KeySuffix: "0003", Status: "active",
|
||||
Filters: models.JSON(fmt.Sprintf(
|
||||
`{"groups":[%d,9999],"protocols":["openai"],"models":["Primary"]}`,
|
||||
firstGroup.ID,
|
||||
@@ -580,7 +582,7 @@ func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) {
|
||||
}
|
||||
disabledAccess := models.AccessKey{
|
||||
Name: "disabled access", KeyValue: "access-cipher-disabled", KeyHash: "disabled-hash",
|
||||
Status: "disabled", Filters: models.JSON(`{}`),
|
||||
KeySuffix: "0004", Status: "disabled", Filters: models.JSON(`{}`),
|
||||
}
|
||||
mustCreate(t, db, &activeAccess)
|
||||
mustCreate(t, db, &disabledAccess)
|
||||
@@ -671,11 +673,56 @@ func TestLoaderMapsAccessAndUpstreamKeys(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildGroupKeyEntriesReadsOnlyRequestedGroupInStableOrder(t *testing.T) {
|
||||
db := openMigratedDatabase(t)
|
||||
firstGroup := createRuntimeGroup(t, db, "first-entries", protocol.OpenAI, "gpt-4o")
|
||||
secondGroup := createRuntimeGroup(t, db, "second-entries", protocol.Anthropic, "claude")
|
||||
weight := 9
|
||||
keys := []models.UpstreamKey{
|
||||
{
|
||||
GroupID: secondGroup.ID, KeyValue: "other-cipher",
|
||||
KeyHash: "other-hash", Status: models.UpstreamKeyStatusActive,
|
||||
},
|
||||
{
|
||||
GroupID: firstGroup.ID, KeyValue: "first-cipher",
|
||||
KeyHash: "first-hash", Status: models.UpstreamKeyStatusActive,
|
||||
},
|
||||
{
|
||||
GroupID: firstGroup.ID, KeyValue: "second-cipher",
|
||||
KeyHash: "second-hash", Status: models.UpstreamKeyStatusDisabled,
|
||||
WeightManual: &weight,
|
||||
},
|
||||
}
|
||||
for index := range keys {
|
||||
mustCreate(t, db, &keys[index])
|
||||
}
|
||||
|
||||
got, err := loader.BuildGroupKeyEntries(t.Context(), db, firstGroup.ID)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildGroupKeyEntries() error = %v", err)
|
||||
}
|
||||
if len(got) != 2 || got[0].ID != keys[1].ID || got[1].ID != keys[2].ID {
|
||||
t.Fatalf("BuildGroupKeyEntries() = %#v", got)
|
||||
}
|
||||
if got[1].WeightManual == nil || *got[1].WeightManual != weight ||
|
||||
got[1].Status != state.KeyStatusDisabled ||
|
||||
got[1].EncryptedValue != "second-cipher" {
|
||||
t.Fatalf("second entry = %#v", got[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildGroupKeyEntriesRejectsMissingGroup(t *testing.T) {
|
||||
db := openMigratedDatabase(t)
|
||||
if _, err := loader.BuildGroupKeyEntries(t.Context(), db, 999); err == nil {
|
||||
t.Fatal("BuildGroupKeyEntries(missing group) error = nil")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoaderMapsAccessKeyRPMLimit(t *testing.T) {
|
||||
db := openMigratedDatabase(t)
|
||||
accessKey := models.AccessKey{
|
||||
Name: "rate-limited", KeyValue: "access-cipher", KeyHash: "rate-limited-hash",
|
||||
Status: "active", Filters: models.JSON(`{}`), RPMLimit: 27,
|
||||
KeySuffix: "0005", Status: "active", Filters: models.JSON(`{}`), RPMLimit: 27,
|
||||
}
|
||||
mustCreate(t, db, &accessKey)
|
||||
|
||||
@@ -708,7 +755,7 @@ func TestLoaderRejectsInvalidCredentialRowsWithoutPublishing(t *testing.T) {
|
||||
}()
|
||||
mustCreate(t, db, &models.AccessKey{
|
||||
Name: "invalid", KeyValue: "access-cipher", KeyHash: "invalid-status-hash",
|
||||
Status: "revoked", Filters: models.JSON(`{}`),
|
||||
KeySuffix: "0006", Status: "revoked", Filters: models.JSON(`{}`),
|
||||
})
|
||||
},
|
||||
wantError: "invalid status",
|
||||
@@ -718,7 +765,8 @@ func TestLoaderRejectsInvalidCredentialRowsWithoutPublishing(t *testing.T) {
|
||||
insert: func(t *testing.T, db *gorm.DB, _ models.Group) {
|
||||
mustCreate(t, db, &models.AccessKey{
|
||||
Name: "invalid", KeyValue: "access-cipher", KeyHash: "invalid-protocol-hash",
|
||||
Status: "active", Filters: models.JSON(`{"protocols":["unknown"]}`),
|
||||
KeySuffix: "0007", Status: "active",
|
||||
Filters: models.JSON(`{"protocols":["unknown"]}`),
|
||||
})
|
||||
},
|
||||
wantError: "invalid protocol",
|
||||
@@ -728,7 +776,8 @@ func TestLoaderRejectsInvalidCredentialRowsWithoutPublishing(t *testing.T) {
|
||||
insert: func(t *testing.T, db *gorm.DB, _ models.Group) {
|
||||
mustCreate(t, db, &models.AccessKey{
|
||||
Name: "invalid", KeyValue: "access-cipher", KeyHash: "blank-model-hash",
|
||||
Status: "active", Filters: models.JSON(`{"models":[" "]}`),
|
||||
KeySuffix: "0008", Status: "active",
|
||||
Filters: models.JSON(`{"models":[" "]}`),
|
||||
})
|
||||
},
|
||||
wantError: "filter model is required",
|
||||
@@ -738,7 +787,8 @@ func TestLoaderRejectsInvalidCredentialRowsWithoutPublishing(t *testing.T) {
|
||||
insert: func(t *testing.T, db *gorm.DB, _ models.Group) {
|
||||
mustCreate(t, db, &models.AccessKey{
|
||||
Name: "invalid", KeyValue: "access-cipher", KeyHash: "unknown-filter-field-hash",
|
||||
Status: "active", Filters: models.JSON(`{"protcols":["openai"]}`),
|
||||
KeySuffix: "0009", Status: "active",
|
||||
Filters: models.JSON(`{"protcols":["openai"]}`),
|
||||
})
|
||||
},
|
||||
wantError: "unknown field",
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package state
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
@@ -40,6 +41,24 @@ func (m *Manager) Publish(input CompileInput) (*ConfigSnapshot, error) {
|
||||
return m.publishCompiled(next, nil), nil
|
||||
}
|
||||
|
||||
// Matches reports whether input compiles to the currently published runtime
|
||||
// configuration. Snapshot revisions are ordering metadata and are ignored.
|
||||
func (m *Manager) Matches(input CompileInput) (bool, error) {
|
||||
next, err := Compile(input)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
m.publishMu.Lock()
|
||||
defer m.publishMu.Unlock()
|
||||
current := m.current.Load()
|
||||
if current == nil {
|
||||
return false, nil
|
||||
}
|
||||
currentValue := *current
|
||||
currentValue.Revision = 0
|
||||
return reflect.DeepEqual(¤tValue, next), nil
|
||||
}
|
||||
|
||||
func (m *Manager) publishCompiled(next *ConfigSnapshot, beforeLock func()) *ConfigSnapshot {
|
||||
if beforeLock != nil {
|
||||
beforeLock()
|
||||
|
||||
@@ -65,6 +65,41 @@ func TestManagerPublishFailureKeepsCurrentSnapshot(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerMatchesCompiledInputWithoutChangingRevision(t *testing.T) {
|
||||
manager := NewManager()
|
||||
input := managerCompileInput(1)
|
||||
published, err := manager.Publish(input)
|
||||
if err != nil {
|
||||
t.Fatalf("Publish() error = %v", err)
|
||||
}
|
||||
|
||||
matches, err := manager.Matches(input)
|
||||
if err != nil {
|
||||
t.Fatalf("Matches(same) error = %v", err)
|
||||
}
|
||||
if !matches {
|
||||
t.Fatal("Matches(same) = false, want true")
|
||||
}
|
||||
if current := manager.Current(); current != published || current.Revision != 1 {
|
||||
t.Fatalf("Matches() changed current snapshot: %#v", current)
|
||||
}
|
||||
|
||||
changed := managerCompileInput(2)
|
||||
matches, err = manager.Matches(changed)
|
||||
if err != nil {
|
||||
t.Fatalf("Matches(changed) error = %v", err)
|
||||
}
|
||||
if matches {
|
||||
t.Fatal("Matches(changed) = true, want false")
|
||||
}
|
||||
|
||||
invalid := managerCompileInput(1)
|
||||
invalid.Groups[0].Protocols = []protocol.Protocol{"reserved"}
|
||||
if _, err := manager.Matches(invalid); err == nil {
|
||||
t.Fatal("Matches(invalid) error = nil, want compile rejection")
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerConcurrentPublishAndCurrent(t *testing.T) {
|
||||
const goroutines = 32
|
||||
|
||||
|
||||
@@ -137,6 +137,111 @@ func (r *KeyRegistry) ApplyImport(groupID uint, entries []KeyEntry) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// MatchesGroup compares the persisted configuration-owned fields for one
|
||||
// group. Runtime health state is intentionally excluded.
|
||||
func (r *KeyRegistry) MatchesGroup(groupID uint, entries []KeyEntry) bool {
|
||||
if groupID == 0 || ValidateKeyEntries(entries) != nil {
|
||||
return false
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if entry.GroupID != groupID {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
r.mu.RLock()
|
||||
defer r.mu.RUnlock()
|
||||
return r.matchesGroupLocked(groupID, entries)
|
||||
}
|
||||
|
||||
// ReconcileGroup makes one group match DB truth while preserving runtime
|
||||
// health state for entries whose persisted configuration is already equal.
|
||||
func (r *KeyRegistry) ReconcileGroup(groupID uint, entries []KeyEntry) (bool, error) {
|
||||
if groupID == 0 {
|
||||
return false, fmt.Errorf("group id is required")
|
||||
}
|
||||
if err := ValidateKeyEntries(entries); err != nil {
|
||||
return false, err
|
||||
}
|
||||
for _, entry := range entries {
|
||||
if entry.GroupID != groupID {
|
||||
return false, fmt.Errorf(
|
||||
"key %d belongs to group %d, want %d",
|
||||
entry.ID,
|
||||
entry.GroupID,
|
||||
groupID,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, entry := range entries {
|
||||
if existingGroupID, exists := r.keyGroups[entry.ID]; exists &&
|
||||
existingGroupID != groupID {
|
||||
return false, fmt.Errorf(
|
||||
"key %d already belongs to group %d",
|
||||
entry.ID,
|
||||
existingGroupID,
|
||||
)
|
||||
}
|
||||
}
|
||||
if r.matchesGroupLocked(groupID, entries) {
|
||||
return false, nil
|
||||
}
|
||||
|
||||
previous := r.buckets[groupID]
|
||||
next := make(map[uint]*KeyEntry, len(entries))
|
||||
for _, desired := range entries {
|
||||
if existing := previous[desired.ID]; existing != nil &&
|
||||
samePersistedKeyConfig(*existing, desired) {
|
||||
next[desired.ID] = existing
|
||||
continue
|
||||
}
|
||||
cloned := cloneKeyEntry(desired)
|
||||
next[desired.ID] = &cloned
|
||||
}
|
||||
for keyID := range previous {
|
||||
delete(r.keyGroups, keyID)
|
||||
}
|
||||
if len(next) == 0 {
|
||||
delete(r.buckets, groupID)
|
||||
} else {
|
||||
r.buckets[groupID] = next
|
||||
for keyID := range next {
|
||||
r.keyGroups[keyID] = groupID
|
||||
}
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *KeyRegistry) matchesGroupLocked(groupID uint, entries []KeyEntry) bool {
|
||||
current := r.buckets[groupID]
|
||||
if len(current) != len(entries) {
|
||||
return false
|
||||
}
|
||||
for _, desired := range entries {
|
||||
existing := current[desired.ID]
|
||||
if existing == nil || !samePersistedKeyConfig(*existing, desired) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func samePersistedKeyConfig(left, right KeyEntry) bool {
|
||||
if left.ID != right.ID ||
|
||||
left.GroupID != right.GroupID ||
|
||||
left.Status != right.Status ||
|
||||
left.EncryptedValue != right.EncryptedValue {
|
||||
return false
|
||||
}
|
||||
if left.WeightManual == nil || right.WeightManual == nil {
|
||||
return left.WeightManual == nil && right.WeightManual == nil
|
||||
}
|
||||
return *left.WeightManual == *right.WeightManual
|
||||
}
|
||||
|
||||
func (r *KeyRegistry) RemoveKey(keyID uint) bool {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
|
||||
@@ -445,6 +445,102 @@ func TestKeyRegistryApplyImportUpsertsOnlyRequestedGroup(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyRegistryReconcileGroupUsesDBTruthAndPreservesSatisfiedRuntimeState(t *testing.T) {
|
||||
registry := NewKeyRegistry()
|
||||
mustReplaceKeyEntries(t, registry, []KeyEntry{
|
||||
{
|
||||
ID: 1, GroupID: 10, Status: KeyStatusActive,
|
||||
EncryptedValue: "cipher-one",
|
||||
},
|
||||
{
|
||||
ID: 2, GroupID: 10, Status: KeyStatusActive,
|
||||
EncryptedValue: "stale-cipher",
|
||||
},
|
||||
{
|
||||
ID: 3, GroupID: 10, Status: KeyStatusDisabled,
|
||||
EncryptedValue: "removed-from-db",
|
||||
},
|
||||
{
|
||||
ID: 4, GroupID: 20, Status: KeyStatusActive,
|
||||
EncryptedValue: "other-group",
|
||||
},
|
||||
})
|
||||
if !registry.SetCooldown(1, time.Now().Add(time.Hour)) {
|
||||
t.Fatal("SetCooldown(1) = false")
|
||||
}
|
||||
if _, ok := registry.IncrFailure(1); !ok {
|
||||
t.Fatal("IncrFailure(1) = false")
|
||||
}
|
||||
|
||||
wantGroup := []KeyEntry{
|
||||
{
|
||||
ID: 1, GroupID: 10, Status: KeyStatusActive,
|
||||
EncryptedValue: "cipher-one",
|
||||
},
|
||||
{
|
||||
ID: 2, GroupID: 10, Status: KeyStatusDisabled,
|
||||
EncryptedValue: "cipher-two",
|
||||
},
|
||||
}
|
||||
if registry.MatchesGroup(10, wantGroup) {
|
||||
t.Fatal("MatchesGroup(stale) = true, want false")
|
||||
}
|
||||
changed, err := registry.ReconcileGroup(10, wantGroup)
|
||||
if err != nil {
|
||||
t.Fatalf("ReconcileGroup() error = %v", err)
|
||||
}
|
||||
if !changed {
|
||||
t.Fatal("ReconcileGroup() changed = false, want true")
|
||||
}
|
||||
if !registry.MatchesGroup(10, wantGroup) {
|
||||
t.Fatal("MatchesGroup(reconciled) = false, want true")
|
||||
}
|
||||
if got := registryEntry(t, registry, 1); got.FailureCount != 1 ||
|
||||
got.CooldownUntil.IsZero() {
|
||||
t.Fatalf("satisfied key runtime state was reset: %#v", got)
|
||||
}
|
||||
if got := registryEntry(t, registry, 2); got.Status != KeyStatusDisabled ||
|
||||
got.EncryptedValue != "cipher-two" || got.FailureCount != 0 {
|
||||
t.Fatalf("stale key was not rebuilt from DB truth: %#v", got)
|
||||
}
|
||||
if _, ok := registry.EncryptedValue(3); ok {
|
||||
t.Fatal("key absent from DB truth remains in reconciled group")
|
||||
}
|
||||
if got, ok := registry.EncryptedValue(4); !ok || got != "other-group" {
|
||||
t.Fatalf("other group changed: %q, %t", got, ok)
|
||||
}
|
||||
|
||||
changed, err = registry.ReconcileGroup(10, wantGroup)
|
||||
if err != nil {
|
||||
t.Fatalf("second ReconcileGroup() error = %v", err)
|
||||
}
|
||||
if changed {
|
||||
t.Fatal("second ReconcileGroup() changed = true, want false")
|
||||
}
|
||||
if got := registryEntry(t, registry, 1); got.FailureCount != 1 {
|
||||
t.Fatalf("idempotent reconciliation reset runtime state: %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyRegistryReconcileGroupRejectsCrossGroupIdentityWithoutMutation(t *testing.T) {
|
||||
registry := NewKeyRegistry()
|
||||
mustReplaceKeyEntries(t, registry, []KeyEntry{{
|
||||
ID: 1, GroupID: 20, Status: KeyStatusActive,
|
||||
EncryptedValue: "other-group",
|
||||
}})
|
||||
|
||||
_, err := registry.ReconcileGroup(10, []KeyEntry{{
|
||||
ID: 1, GroupID: 10, Status: KeyStatusActive,
|
||||
EncryptedValue: "attempted-move",
|
||||
}})
|
||||
if err == nil {
|
||||
t.Fatal("ReconcileGroup(cross-group key) error = nil")
|
||||
}
|
||||
if got, ok := registry.EncryptedValue(1); !ok || got != "other-group" {
|
||||
t.Fatalf("failed reconciliation mutated registry: %q, %t", got, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyRegistryRemoveKey(t *testing.T) {
|
||||
registry := NewKeyRegistry()
|
||||
mustReplaceKeyEntries(t, registry, []KeyEntry{
|
||||
|
||||
+294
-21
@@ -1,7 +1,10 @@
|
||||
package storage
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/url"
|
||||
"os"
|
||||
@@ -16,12 +19,13 @@ import (
|
||||
"gorm.io/gorm/logger"
|
||||
|
||||
"gpt-load/internal/platform/config"
|
||||
"gpt-load/internal/platform/encryption"
|
||||
"gpt-load/internal/platform/securefile"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
// CurrentSchemaVersion identifies the SQLite schema supported by this binary.
|
||||
const CurrentSchemaVersion uint = 1
|
||||
const CurrentSchemaVersion uint = 2
|
||||
|
||||
const sqliteBusyTimeoutMS = 5000
|
||||
|
||||
@@ -107,19 +111,30 @@ func OpenWithSource(dsn string, source config.DatabaseSource) (*gorm.DB, error)
|
||||
|
||||
// AutoMigrate creates the current persistence schema and initializes schema_info.
|
||||
func AutoMigrate(db *gorm.DB) error {
|
||||
return autoMigrate(db, nil, "")
|
||||
}
|
||||
|
||||
// AutoMigrateWithEncryption upgrades credential-bearing legacy schemas before
|
||||
// runtime loading. dataDir is used only to back up a managed encryption.key.
|
||||
func AutoMigrateWithEncryption(
|
||||
db *gorm.DB,
|
||||
encryptionService encryption.Service,
|
||||
dataDir string,
|
||||
) error {
|
||||
return autoMigrate(db, encryptionService, dataDir)
|
||||
}
|
||||
|
||||
func autoMigrate(
|
||||
db *gorm.DB,
|
||||
encryptionService encryption.Service,
|
||||
dataDir string,
|
||||
) error {
|
||||
if db == nil {
|
||||
return fmt.Errorf("auto-migrate SQLite database: db is nil")
|
||||
}
|
||||
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
if tx.Migrator().HasTable(&schemaInfo{}) {
|
||||
if err := validateSchemaVersion(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
return migrateCurrentSchema(tx)
|
||||
}
|
||||
|
||||
tables, err := tx.Migrator().GetTables()
|
||||
if !db.Migrator().HasTable(&schemaInfo{}) {
|
||||
tables, err := db.Migrator().GetTables()
|
||||
if err != nil {
|
||||
return fmt.Errorf("list SQLite tables: %w", err)
|
||||
}
|
||||
@@ -128,15 +143,64 @@ func AutoMigrate(db *gorm.DB) error {
|
||||
return fmt.Errorf("initialize SQLite schema: non-empty database without schema_info")
|
||||
}
|
||||
}
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := migrateCurrentSchema(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Create(&schemaInfo{Version: CurrentSchemaVersion}).Error; err != nil {
|
||||
return fmt.Errorf("initialize schema_info: %w", err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
if err := migrateCurrentSchema(tx); err != nil {
|
||||
version, err := readSchemaVersion(db)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
switch version {
|
||||
case CurrentSchemaVersion:
|
||||
return db.Transaction(migrateCurrentSchema)
|
||||
case 1:
|
||||
var accessKeyCount int64
|
||||
if db.Migrator().HasTable(&models.AccessKey{}) {
|
||||
if err := db.Table("access_keys").Count(&accessKeyCount).Error; err != nil {
|
||||
return fmt.Errorf("count access_keys before schema upgrade: %w", err)
|
||||
}
|
||||
}
|
||||
if accessKeyCount > 0 && encryptionService == nil {
|
||||
return fmt.Errorf(
|
||||
"upgrade SQLite schema version 1: encryption service is required",
|
||||
)
|
||||
}
|
||||
if err := backupSchemaV1(db, dataDir); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := tx.Create(&schemaInfo{Version: CurrentSchemaVersion}).Error; err != nil {
|
||||
return fmt.Errorf("initialize schema_info: %w", err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
return db.Transaction(func(tx *gorm.DB) error {
|
||||
if err := migrateAccessKeysV1ToV2(tx, encryptionService); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := migrateCurrentSchema(tx); err != nil {
|
||||
return err
|
||||
}
|
||||
result := tx.Model(&schemaInfo{}).
|
||||
Where("version = ?", uint(1)).
|
||||
Update("version", CurrentSchemaVersion)
|
||||
if result.Error != nil {
|
||||
return fmt.Errorf("update schema_info: %w", result.Error)
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return fmt.Errorf("update schema_info: version changed concurrently")
|
||||
}
|
||||
return nil
|
||||
})
|
||||
default:
|
||||
return fmt.Errorf(
|
||||
"unsupported schema version %d, want %d",
|
||||
version,
|
||||
CurrentSchemaVersion,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func migrateCurrentSchema(db *gorm.DB) error {
|
||||
@@ -149,6 +213,7 @@ func migrateCurrentSchema(db *gorm.DB) error {
|
||||
&models.ModelPrice{},
|
||||
&models.SystemSetting{},
|
||||
&models.Job{},
|
||||
&models.ControlOperation{},
|
||||
&schemaInfo{},
|
||||
); err != nil {
|
||||
return fmt.Errorf("auto-migrate SQLite schema: %w", err)
|
||||
@@ -219,21 +284,229 @@ FROM model_prices`).Error; err != nil {
|
||||
}
|
||||
|
||||
func validateSchemaVersion(db *gorm.DB) error {
|
||||
version, err := readSchemaVersion(db)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if version != CurrentSchemaVersion {
|
||||
return fmt.Errorf("unsupported schema version %d, want %d", version, CurrentSchemaVersion)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func readSchemaVersion(db *gorm.DB) (uint, error) {
|
||||
var count int64
|
||||
if err := db.Model(&schemaInfo{}).Count(&count).Error; err != nil {
|
||||
return fmt.Errorf("count schema_info rows: %w", err)
|
||||
return 0, fmt.Errorf("count schema_info rows: %w", err)
|
||||
}
|
||||
if count != 1 {
|
||||
return fmt.Errorf("schema_info contains %d rows, want exactly one", count)
|
||||
return 0, fmt.Errorf("schema_info contains %d rows, want exactly one", count)
|
||||
}
|
||||
|
||||
var info schemaInfo
|
||||
if err := db.First(&info).Error; err != nil {
|
||||
return fmt.Errorf("read schema_info: %w", err)
|
||||
return 0, fmt.Errorf("read schema_info: %w", err)
|
||||
}
|
||||
if info.Version != CurrentSchemaVersion {
|
||||
return fmt.Errorf("unsupported schema version %d, want %d", info.Version, CurrentSchemaVersion)
|
||||
return info.Version, nil
|
||||
}
|
||||
|
||||
func migrateAccessKeysV1ToV2(
|
||||
db *gorm.DB,
|
||||
encryptionService encryption.Service,
|
||||
) error {
|
||||
if !db.Migrator().HasTable(&models.AccessKey{}) {
|
||||
return nil
|
||||
}
|
||||
if db.Migrator().HasColumn(&models.AccessKey{}, "key_suffix") {
|
||||
return fmt.Errorf("upgrade access_keys schema: unexpected key_suffix column")
|
||||
}
|
||||
if err := db.Exec("ALTER TABLE access_keys ADD COLUMN key_suffix char(4)").Error; err != nil {
|
||||
return fmt.Errorf("add access_keys key_suffix: %w", err)
|
||||
}
|
||||
|
||||
type encryptedAccessKey struct {
|
||||
ID uint
|
||||
KeyValue string `gorm:"column:key_value"`
|
||||
}
|
||||
var rows []encryptedAccessKey
|
||||
if err := db.Table("access_keys").
|
||||
Select("id", "key_value").
|
||||
Order("id ASC").
|
||||
Find(&rows).Error; err != nil {
|
||||
return fmt.Errorf("read access_keys for schema upgrade: %w", err)
|
||||
}
|
||||
for _, row := range rows {
|
||||
if encryptionService == nil {
|
||||
return fmt.Errorf("upgrade access_keys schema: encryption service is required")
|
||||
}
|
||||
plaintext, err := encryptionService.Decrypt(row.KeyValue)
|
||||
if err != nil || !validAccessKeyPlaintext(plaintext) {
|
||||
return fmt.Errorf("upgrade access_keys row %d: credential validation failed", row.ID)
|
||||
}
|
||||
suffix := plaintext[len(plaintext)-4:]
|
||||
result := db.Table("access_keys").
|
||||
Where("id = ? AND key_suffix IS NULL", row.ID).
|
||||
Update("key_suffix", suffix)
|
||||
if result.Error != nil {
|
||||
return fmt.Errorf("backfill access_keys key_suffix: %w", result.Error)
|
||||
}
|
||||
if result.RowsAffected != 1 {
|
||||
return fmt.Errorf("backfill access_keys key_suffix: row changed concurrently")
|
||||
}
|
||||
}
|
||||
|
||||
if err := db.Exec(`CREATE TABLE access_keys__v2 (
|
||||
id integer PRIMARY KEY AUTOINCREMENT,
|
||||
name varchar(255) NOT NULL,
|
||||
key_value text NOT NULL,
|
||||
key_hash varchar(128) NOT NULL,
|
||||
key_suffix char(4) NOT NULL,
|
||||
status varchar(32) NOT NULL DEFAULT 'active',
|
||||
filters json,
|
||||
rpm_limit integer NOT NULL DEFAULT 0,
|
||||
daily_cost_limit real NOT NULL DEFAULT 0,
|
||||
monthly_cost_limit real NOT NULL DEFAULT 0,
|
||||
created_at datetime,
|
||||
updated_at datetime,
|
||||
CONSTRAINT chk_access_key_suffix
|
||||
CHECK (key_suffix GLOB '[0-9a-f][0-9a-f][0-9a-f][0-9a-f]'),
|
||||
CONSTRAINT chk_access_key_status
|
||||
CHECK (status IN ('active','disabled'))
|
||||
)`).Error; err != nil {
|
||||
return fmt.Errorf("create rebuilt access_keys table: %w", err)
|
||||
}
|
||||
if err := db.Exec(`INSERT INTO access_keys__v2 (
|
||||
id,name,key_value,key_hash,key_suffix,status,filters,rpm_limit,
|
||||
daily_cost_limit,monthly_cost_limit,created_at,updated_at
|
||||
)
|
||||
SELECT
|
||||
id,name,key_value,key_hash,key_suffix,status,filters,rpm_limit,
|
||||
daily_cost_limit,monthly_cost_limit,created_at,updated_at
|
||||
FROM access_keys`).Error; err != nil {
|
||||
return fmt.Errorf("copy rebuilt access_keys table: %w", err)
|
||||
}
|
||||
if err := db.Exec("DROP TABLE access_keys").Error; err != nil {
|
||||
return fmt.Errorf("drop previous access_keys table: %w", err)
|
||||
}
|
||||
if err := db.Exec("ALTER TABLE access_keys__v2 RENAME TO access_keys").Error; err != nil {
|
||||
return fmt.Errorf("rename rebuilt access_keys table: %w", err)
|
||||
}
|
||||
if err := db.Exec(
|
||||
"CREATE UNIQUE INDEX idx_access_keys_key_hash ON access_keys(key_hash)",
|
||||
).Error; err != nil {
|
||||
return fmt.Errorf("create rebuilt access_keys key hash index: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validAccessKeyPlaintext(value string) bool {
|
||||
const accessKeyPrefix = "sk-gl-"
|
||||
if len(value) != len(accessKeyPrefix)+32 || !strings.HasPrefix(value, accessKeyPrefix) {
|
||||
return false
|
||||
}
|
||||
for _, character := range value[len(accessKeyPrefix):] {
|
||||
if character < '0' || character > '9' && character < 'a' || character > 'f' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func backupSchemaV1(db *gorm.DB, dataDir string) error {
|
||||
type databaseInfo struct {
|
||||
Name string
|
||||
File string
|
||||
}
|
||||
var databases []databaseInfo
|
||||
if err := db.Raw("PRAGMA database_list").Scan(&databases).Error; err != nil {
|
||||
return fmt.Errorf("inspect SQLite database for schema backup: %w", err)
|
||||
}
|
||||
var databasePath string
|
||||
for _, database := range databases {
|
||||
if database.Name == "main" {
|
||||
databasePath = database.File
|
||||
break
|
||||
}
|
||||
}
|
||||
if databasePath == "" {
|
||||
return nil
|
||||
}
|
||||
|
||||
suffix, err := schemaBackupSuffix()
|
||||
if err != nil {
|
||||
return fmt.Errorf("prepare schema backup identity: %w", err)
|
||||
}
|
||||
databaseBackup := databasePath + ".schema-v1-backup-" + suffix
|
||||
escapedBackup := strings.ReplaceAll(databaseBackup, "'", "''")
|
||||
if err := db.Exec("VACUUM INTO '" + escapedBackup + "'").Error; err != nil {
|
||||
return fmt.Errorf("back up SQLite schema version 1: %w", err)
|
||||
}
|
||||
if err := securefile.HardenManagedFileIfExists(databaseBackup); err != nil {
|
||||
return fmt.Errorf("secure SQLite schema backup: %w", err)
|
||||
}
|
||||
|
||||
if dataDir == "" {
|
||||
return nil
|
||||
}
|
||||
keyPath := filepath.Join(dataDir, encryption.KeyFileName)
|
||||
if _, err := os.Lstat(keyPath); os.IsNotExist(err) {
|
||||
return nil
|
||||
} else if err != nil {
|
||||
return fmt.Errorf("inspect encryption key for schema backup: %w", err)
|
||||
}
|
||||
keyBackup := keyPath + ".schema-v1-backup-" + suffix
|
||||
if err := copySchemaBackup(keyPath, keyBackup); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func schemaBackupSuffix() (string, error) {
|
||||
randomBytes := make([]byte, 8)
|
||||
if _, err := rand.Read(randomBytes); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return time.Now().UTC().Format("20060102T150405.000000000Z") +
|
||||
"-" + hex.EncodeToString(randomBytes), nil
|
||||
}
|
||||
|
||||
func copySchemaBackup(sourcePath, targetPath string) error {
|
||||
source, err := os.Open(sourcePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("open encryption key for schema backup: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
_ = source.Close()
|
||||
}()
|
||||
info, err := source.Stat()
|
||||
if err != nil || !info.Mode().IsRegular() {
|
||||
return fmt.Errorf("inspect encryption key for schema backup")
|
||||
}
|
||||
|
||||
target, err := os.OpenFile(targetPath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0o600)
|
||||
if err != nil {
|
||||
return fmt.Errorf("create encryption key schema backup: %w", err)
|
||||
}
|
||||
success := false
|
||||
defer func() {
|
||||
_ = target.Close()
|
||||
if !success {
|
||||
_ = os.Remove(targetPath)
|
||||
}
|
||||
}()
|
||||
if _, err := io.Copy(target, source); err != nil {
|
||||
return fmt.Errorf("copy encryption key schema backup: %w", err)
|
||||
}
|
||||
if err := target.Sync(); err != nil {
|
||||
return fmt.Errorf("sync encryption key schema backup: %w", err)
|
||||
}
|
||||
if err := target.Close(); err != nil {
|
||||
return fmt.Errorf("close encryption key schema backup: %w", err)
|
||||
}
|
||||
if err := securefile.HardenManagedFileIfExists(targetPath); err != nil {
|
||||
return fmt.Errorf("secure encryption key schema backup: %w", err)
|
||||
}
|
||||
success = true
|
||||
return nil
|
||||
}
|
||||
|
||||
|
||||
+21
-18
@@ -91,11 +91,12 @@ func TestAccessKeyStatusAcceptsOnlyDurableOperatorStates(t *testing.T) {
|
||||
db := openMigratedDatabase(t)
|
||||
for index, status := range []string{"active", "disabled"} {
|
||||
key := models.AccessKey{
|
||||
Name: "allowed-" + string(rune('a'+index)),
|
||||
KeyValue: "ciphertext",
|
||||
KeyHash: "allowed-status-" + string(rune('a'+index)),
|
||||
Status: status,
|
||||
Filters: models.JSON(`{}`),
|
||||
Name: "allowed-" + string(rune('a'+index)),
|
||||
KeyValue: "ciphertext",
|
||||
KeyHash: "allowed-status-" + string(rune('a'+index)),
|
||||
KeySuffix: "7f2a",
|
||||
Status: status,
|
||||
Filters: models.JSON(`{}`),
|
||||
}
|
||||
if err := db.Create(&key).Error; err != nil {
|
||||
t.Fatalf("create access key with status %q: %v", status, err)
|
||||
@@ -103,11 +104,12 @@ func TestAccessKeyStatusAcceptsOnlyDurableOperatorStates(t *testing.T) {
|
||||
}
|
||||
|
||||
invalid := models.AccessKey{
|
||||
Name: "invalid",
|
||||
KeyValue: "ciphertext",
|
||||
KeyHash: "invalid-status",
|
||||
Status: "blacklisted",
|
||||
Filters: models.JSON(`{}`),
|
||||
Name: "invalid",
|
||||
KeyValue: "ciphertext",
|
||||
KeyHash: "invalid-status",
|
||||
KeySuffix: "7f2a",
|
||||
Status: "blacklisted",
|
||||
Filters: models.JSON(`{}`),
|
||||
}
|
||||
if err := db.Create(&invalid).Error; err == nil {
|
||||
t.Fatal("blacklisted status error = nil, want CHECK constraint error")
|
||||
@@ -662,8 +664,8 @@ func TestAutoMigrateCreatesRequestLogV1FieldsAndCompositeIndexes(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
if storage.CurrentSchemaVersion != 1 {
|
||||
t.Fatalf("CurrentSchemaVersion = %d, want 1", storage.CurrentSchemaVersion)
|
||||
if storage.CurrentSchemaVersion != 2 {
|
||||
t.Fatalf("CurrentSchemaVersion = %d, want 2", storage.CurrentSchemaVersion)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -803,8 +805,8 @@ func TestAutoMigrateCreatesM4UsagePricingColumnsAndConstraints(t *testing.T) {
|
||||
t.Error("builtin ModelPrice Source was accepted")
|
||||
}
|
||||
|
||||
if storage.CurrentSchemaVersion != 1 {
|
||||
t.Errorf("CurrentSchemaVersion = %d, want 1", storage.CurrentSchemaVersion)
|
||||
if storage.CurrentSchemaVersion != 2 {
|
||||
t.Errorf("CurrentSchemaVersion = %d, want 2", storage.CurrentSchemaVersion)
|
||||
}
|
||||
if err := storage.AutoMigrate(db); err != nil {
|
||||
t.Fatalf("second AutoMigrate() error = %v", err)
|
||||
@@ -1079,10 +1081,11 @@ func TestAutoMigrateCreatesCriticalUniqueConstraints(t *testing.T) {
|
||||
|
||||
t.Run("access key hash", func(t *testing.T) {
|
||||
first := models.AccessKey{
|
||||
Name: "access-one",
|
||||
KeyValue: "ciphertext-one",
|
||||
KeyHash: "same-access-key-hash",
|
||||
Filters: models.JSON(`{}`),
|
||||
Name: "access-one",
|
||||
KeyValue: "ciphertext-one",
|
||||
KeyHash: "same-access-key-hash",
|
||||
KeySuffix: "7f2a",
|
||||
Filters: models.JSON(`{}`),
|
||||
}
|
||||
second := first
|
||||
second.ID = 0
|
||||
|
||||
@@ -8,6 +8,7 @@ type AccessKey struct {
|
||||
Name string `gorm:"type:varchar(255);not null"`
|
||||
KeyValue string `gorm:"type:text;not null"`
|
||||
KeyHash string `gorm:"type:varchar(128);not null;uniqueIndex"`
|
||||
KeySuffix string `gorm:"type:char(4);not null;check:chk_access_key_suffix,key_suffix GLOB '[0-9a-f][0-9a-f][0-9a-f][0-9a-f]'"`
|
||||
Status string `gorm:"type:varchar(32);not null;default:'active';check:chk_access_key_status,status IN ('active','disabled')"`
|
||||
Filters JSON `gorm:"type:json"`
|
||||
RPMLimit int64 `gorm:"not null;default:0"`
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
package models
|
||||
|
||||
import "time"
|
||||
|
||||
// ControlOperation durably identifies a create/import mutation and its
|
||||
// post-commit side-effect progress. It never stores plaintext credentials.
|
||||
type ControlOperation struct {
|
||||
CommitSequence uint64 `gorm:"primaryKey;autoIncrement"`
|
||||
OperationID string `gorm:"type:char(36);not null;uniqueIndex"`
|
||||
IdempotencyKey string `gorm:"type:char(36);not null;uniqueIndex"`
|
||||
DigestVersion uint `gorm:"not null;check:chk_control_operation_digest_version,digest_version > 0"`
|
||||
RequestDigest []byte `gorm:"type:blob;not null;check:chk_control_operation_digest,length(request_digest) = 32"`
|
||||
OperationKind string `gorm:"type:varchar(32);not null"`
|
||||
ResourceIdentity string `gorm:"type:varchar(64);not null"`
|
||||
CanonicalResult []byte `gorm:"type:blob"`
|
||||
RequiredStages JSON `gorm:"type:json"`
|
||||
LastCompletedStage string `gorm:"type:varchar(32)"`
|
||||
FailedStage string `gorm:"type:varchar(32)"`
|
||||
CompletedAt *time.Time `gorm:"index"`
|
||||
CompactedAt *time.Time
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
}
|
||||
@@ -0,0 +1,222 @@
|
||||
package storage_test
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"gpt-load/internal/platform/config"
|
||||
"gpt-load/internal/platform/encryption"
|
||||
"gpt-load/internal/storage"
|
||||
"gpt-load/internal/storage/models"
|
||||
)
|
||||
|
||||
func TestAutoMigrateCreatesSchemaV2SecurityTables(t *testing.T) {
|
||||
db := openMigratedDatabase(t)
|
||||
if storage.CurrentSchemaVersion != 2 {
|
||||
t.Fatalf("CurrentSchemaVersion = %d, want 2", storage.CurrentSchemaVersion)
|
||||
}
|
||||
|
||||
type columnInfo struct {
|
||||
Name string
|
||||
NotNull int `gorm:"column:notnull"`
|
||||
}
|
||||
var columns []columnInfo
|
||||
if err := db.Raw("PRAGMA table_info('access_keys')").Scan(&columns).Error; err != nil {
|
||||
t.Fatalf("inspect access_keys: %v", err)
|
||||
}
|
||||
var suffixNotNull bool
|
||||
for _, column := range columns {
|
||||
if column.Name == "key_suffix" {
|
||||
suffixNotNull = column.NotNull == 1
|
||||
}
|
||||
}
|
||||
if !suffixNotNull {
|
||||
t.Fatal("access_keys.key_suffix is missing or nullable")
|
||||
}
|
||||
if !db.Migrator().HasTable(&models.ControlOperation{}) {
|
||||
t.Fatal("control_operations table is missing")
|
||||
}
|
||||
|
||||
valid := models.AccessKey{
|
||||
Name: "valid", KeyValue: "ciphertext", KeyHash: "hash-valid", KeySuffix: "7f2a",
|
||||
Status: "active", Filters: models.JSON(`{"groups":[],"protocols":[],"models":[]}`),
|
||||
}
|
||||
if err := db.Create(&valid).Error; err != nil {
|
||||
t.Fatalf("create valid AccessKey: %v", err)
|
||||
}
|
||||
invalid := valid
|
||||
invalid.ID = 0
|
||||
invalid.Name = "invalid"
|
||||
invalid.KeyHash = "hash-invalid"
|
||||
invalid.KeySuffix = "7F2A"
|
||||
if err := db.Create(&invalid).Error; err == nil {
|
||||
t.Fatal("create invalid AccessKey suffix error = nil, want CHECK rejection")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAutoMigrateUpgradesV1AccessKeyWithBackups(t *testing.T) {
|
||||
dataDir := t.TempDir()
|
||||
keyMaterial := strings.Repeat("ab", 32)
|
||||
keyPath := filepath.Join(dataDir, encryption.KeyFileName)
|
||||
if err := os.WriteFile(keyPath, []byte(keyMaterial+"\n"), 0o600); err != nil {
|
||||
t.Fatalf("write encryption key fixture: %v", err)
|
||||
}
|
||||
keyService, err := encryption.NewServiceWithKeyFile("", dataDir)
|
||||
if err != nil {
|
||||
t.Fatalf("NewServiceWithKeyFile() error = %v", err)
|
||||
}
|
||||
dbPath := filepath.Join(dataDir, "gpt-load.db")
|
||||
db, err := storage.OpenWithSource(dbPath, config.DatabaseSourceManaged)
|
||||
if err != nil {
|
||||
t.Fatalf("OpenWithSource() error = %v", err)
|
||||
}
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
t.Fatalf("db.DB() error = %v", err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := sqlDB.Close(); err != nil {
|
||||
t.Errorf("close database: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
createLegacySchemaV1(t, db)
|
||||
plaintext := "sk-gl-00112233445566778899aabbccdd7f2a"
|
||||
ciphertext, err := keyService.Encrypt(plaintext)
|
||||
if err != nil {
|
||||
t.Fatalf("Encrypt() error = %v", err)
|
||||
}
|
||||
if err := db.Exec(`INSERT INTO access_keys
|
||||
(name,key_value,key_hash,status,filters,rpm_limit,daily_cost_limit,monthly_cost_limit)
|
||||
VALUES (?,?,?,?,?,?,?,?)`,
|
||||
"legacy", ciphertext, keyService.Hash(plaintext), "active",
|
||||
`{"groups":[],"protocols":[],"models":[]}`, 0, 0, 0,
|
||||
).Error; err != nil {
|
||||
t.Fatalf("insert legacy AccessKey: %v", err)
|
||||
}
|
||||
|
||||
if err := storage.AutoMigrateWithEncryption(db, keyService, dataDir); err != nil {
|
||||
t.Fatalf("AutoMigrateWithEncryption() error = %v", err)
|
||||
}
|
||||
|
||||
var suffix string
|
||||
if err := db.Raw("SELECT key_suffix FROM access_keys WHERE name = 'legacy'").Scan(&suffix).Error; err != nil {
|
||||
t.Fatalf("read migrated suffix: %v", err)
|
||||
}
|
||||
if suffix != "7f2a" {
|
||||
t.Fatalf("migrated suffix = %q, want 7f2a", suffix)
|
||||
}
|
||||
var version uint
|
||||
if err := db.Raw("SELECT version FROM schema_info").Scan(&version).Error; err != nil {
|
||||
t.Fatalf("read schema version: %v", err)
|
||||
}
|
||||
if version != 2 {
|
||||
t.Fatalf("schema version = %d, want 2", version)
|
||||
}
|
||||
assertSingleSecureBackup(t, dbPath+".schema-v1-backup-*")
|
||||
assertSingleSecureBackup(t, keyPath+".schema-v1-backup-*")
|
||||
}
|
||||
|
||||
func TestAutoMigrateV1AccessKeyFailureRollsBackWithoutLeakingCredential(t *testing.T) {
|
||||
db, err := storage.Open(":memory:")
|
||||
if err != nil {
|
||||
t.Fatalf("Open(:memory:) error = %v", err)
|
||||
}
|
||||
sqlDB, err := db.DB()
|
||||
if err != nil {
|
||||
t.Fatalf("db.DB() error = %v", err)
|
||||
}
|
||||
t.Cleanup(func() { _ = sqlDB.Close() })
|
||||
createLegacySchemaV1(t, db)
|
||||
|
||||
correct, err := encryption.NewService("correct-migration-key")
|
||||
if err != nil {
|
||||
t.Fatalf("NewService(correct) error = %v", err)
|
||||
}
|
||||
wrong, err := encryption.NewService("wrong-migration-key")
|
||||
if err != nil {
|
||||
t.Fatalf("NewService(wrong) error = %v", err)
|
||||
}
|
||||
const plaintext = "sk-gl-00112233445566778899aabbccddeeff"
|
||||
ciphertext, err := correct.Encrypt(plaintext)
|
||||
if err != nil {
|
||||
t.Fatalf("Encrypt() error = %v", err)
|
||||
}
|
||||
if err := db.Exec(`INSERT INTO access_keys
|
||||
(name,key_value,key_hash,status,filters,rpm_limit,daily_cost_limit,monthly_cost_limit)
|
||||
VALUES (?,?,?,?,?,?,?,?)`,
|
||||
"legacy", ciphertext, correct.Hash(plaintext), "active",
|
||||
`{"groups":[],"protocols":[],"models":[]}`, 0, 0, 0,
|
||||
).Error; err != nil {
|
||||
t.Fatalf("insert legacy AccessKey: %v", err)
|
||||
}
|
||||
|
||||
migrationErr := storage.AutoMigrateWithEncryption(db, wrong, "")
|
||||
if migrationErr == nil {
|
||||
t.Fatal("AutoMigrateWithEncryption(wrong key) error = nil, want failure")
|
||||
}
|
||||
for _, forbidden := range []string{plaintext, ciphertext} {
|
||||
if strings.Contains(migrationErr.Error(), forbidden) {
|
||||
t.Fatalf("migration error leaked credential material: %v", migrationErr)
|
||||
}
|
||||
}
|
||||
var version uint
|
||||
if err := db.Raw("SELECT version FROM schema_info").Scan(&version).Error; err != nil {
|
||||
t.Fatalf("read schema version: %v", err)
|
||||
}
|
||||
if version != 1 {
|
||||
t.Fatalf("schema version after failure = %d, want 1", version)
|
||||
}
|
||||
if db.Migrator().HasColumn("access_keys", "key_suffix") {
|
||||
t.Fatal("failed migration left key_suffix column behind")
|
||||
}
|
||||
}
|
||||
|
||||
func createLegacySchemaV1(t *testing.T, db *gorm.DB) {
|
||||
t.Helper()
|
||||
statements := []string{
|
||||
`CREATE TABLE schema_info (version integer PRIMARY KEY)`,
|
||||
`INSERT INTO schema_info(version) VALUES (1)`,
|
||||
`CREATE TABLE access_keys (
|
||||
id integer PRIMARY KEY AUTOINCREMENT,
|
||||
name varchar(255) NOT NULL,
|
||||
key_value text NOT NULL,
|
||||
key_hash varchar(128) NOT NULL UNIQUE,
|
||||
status varchar(32) NOT NULL DEFAULT 'active'
|
||||
CHECK (status IN ('active','disabled')),
|
||||
filters json,
|
||||
rpm_limit integer NOT NULL DEFAULT 0,
|
||||
daily_cost_limit real NOT NULL DEFAULT 0,
|
||||
monthly_cost_limit real NOT NULL DEFAULT 0,
|
||||
created_at datetime,
|
||||
updated_at datetime
|
||||
)`,
|
||||
}
|
||||
for _, statement := range statements {
|
||||
if err := db.Exec(statement).Error; err != nil {
|
||||
t.Fatalf("create schema v1 fixture: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func assertSingleSecureBackup(t *testing.T, pattern string) {
|
||||
t.Helper()
|
||||
matches, err := filepath.Glob(pattern)
|
||||
if err != nil {
|
||||
t.Fatalf("glob backup: %v", err)
|
||||
}
|
||||
if len(matches) != 1 {
|
||||
t.Fatalf("backup matches = %v, want exactly one", matches)
|
||||
}
|
||||
info, err := os.Stat(matches[0])
|
||||
if err != nil {
|
||||
t.Fatalf("stat backup: %v", err)
|
||||
}
|
||||
if info.Mode().Perm() != 0o600 {
|
||||
t.Fatalf("backup permissions = %o, want 600", info.Mode().Perm())
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user