From d717725d1019beaf8d1345708105c54ead993a7b Mon Sep 17 00:00:00 2001 From: tbphp Date: Sat, 18 Jul 2026 14:44:28 +0800 Subject: [PATCH] =?UTF-8?q?feat(control):=20=E5=AE=9E=E7=8E=B0M1=20S6?= =?UTF-8?q?=E6=8E=A7=E5=88=B6=E9=9D=A2API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- internal/container/container.go | 18 +- internal/container/container_test.go | 67 +- internal/control/access_keys.go | 311 +++++++++ internal/control/access_keys_test.go | 509 ++++++++++++++ internal/control/groups.go | 65 ++ internal/control/groups_test.go | 77 ++ internal/control/import.go | 536 ++++++++++++++ internal/control/import_test.go | 698 +++++++++++++++++++ internal/control/server.go | 203 ++++++ internal/control/server_test.go | 263 +++++++ internal/control/service.go | 81 +++ internal/control/service_integration_test.go | 135 ++++ internal/control/service_test.go | 159 +++++ internal/control/smoke_test.go | 260 +++++++ internal/control/test_helpers_test.go | 89 +++ internal/dialect/set.go | 16 + internal/dialect/set_test.go | 22 + internal/gateway/handler.go | 15 +- internal/gateway/handler_test.go | 6 +- internal/gateway/stream_integration_test.go | 2 +- internal/platform/config/config.go | 5 + internal/platform/config/config_test.go | 3 + internal/state/loader/loader.go | 60 +- internal/state/loader/loader_test.go | 106 +++ internal/storage/db.go | 12 +- internal/storage/db_test.go | 33 + scripts/m1-smoke.sh | 56 ++ 27 files changed, 3759 insertions(+), 48 deletions(-) create mode 100644 internal/control/access_keys.go create mode 100644 internal/control/access_keys_test.go create mode 100644 internal/control/groups.go create mode 100644 internal/control/groups_test.go create mode 100644 internal/control/import.go create mode 100644 internal/control/import_test.go create mode 100644 internal/control/server.go create mode 100644 internal/control/server_test.go create mode 100644 internal/control/service.go create mode 100644 internal/control/service_integration_test.go create mode 100644 internal/control/service_test.go create mode 100644 internal/control/smoke_test.go create mode 100644 internal/control/test_helpers_test.go create mode 100644 internal/dialect/set.go create mode 100644 internal/dialect/set_test.go create mode 100755 scripts/m1-smoke.sh diff --git a/internal/container/container.go b/internal/container/container.go index 07f8188a..b059a83a 100644 --- a/internal/container/container.go +++ b/internal/container/container.go @@ -11,6 +11,7 @@ import ( "gorm.io/gorm" "gpt-load/internal/app" + "gpt-load/internal/control" "gpt-load/internal/dialect" "gpt-load/internal/gateway" "gpt-load/internal/platform/config" @@ -59,10 +60,14 @@ func BuildContainer() (*dig.Container, error) { }) }, dialect.NewOpenAI, - gateway.NewDialectSet, + func(openAI *dialect.OpenAI) dialect.Set { + return dialect.NewSet(openAI) + }, gateway.NewForwarder, func(forwarder *gateway.Forwarder) gateway.AttemptForwarder { return forwarder }, gateway.NewHandler, + control.NewService, + control.NewServer, func(db *gorm.DB, manager *state.Manager, registry *state.KeyRegistry) app.RuntimeStateLoader { return stateloader.New(db, manager, registry) }, @@ -74,10 +79,15 @@ func BuildContainer() (*dig.Container, error) { return nil, err } } - if err := dependencyContainer.Invoke(func(engine *gin.Engine, handler *gateway.Handler) { - handler.RegisterRoutes(engine) + if err := dependencyContainer.Invoke(func( + engine *gin.Engine, + gatewayHandler *gateway.Handler, + controlServer *control.Server, + ) { + gatewayHandler.RegisterRoutes(engine) + controlServer.RegisterRoutes(engine) }); err != nil { - return nil, fmt.Errorf("register gateway routes: %w", err) + return nil, fmt.Errorf("register HTTP routes: %w", err) } return dependencyContainer, nil } diff --git a/internal/container/container_test.go b/internal/container/container_test.go index dae84a1f..44d0de18 100644 --- a/internal/container/container_test.go +++ b/internal/container/container_test.go @@ -14,18 +14,20 @@ import ( "gorm.io/gorm" "gpt-load/internal/app" + "gpt-load/internal/control" "gpt-load/internal/dialect" "gpt-load/internal/gateway" "gpt-load/internal/platform/config" "gpt-load/internal/platform/encryption" "gpt-load/internal/platform/httpclient" + "gpt-load/internal/platform/i18n" "gpt-load/internal/platform/redact" "gpt-load/internal/state" "gpt-load/internal/storage" "gpt-load/internal/storage/store" ) -func TestBuildContainerResolvesS5DependencyGraph(t *testing.T) { +func TestBuildContainerResolvesS6DependencyGraph(t *testing.T) { dataDir := t.TempDir() t.Setenv("AUTH_KEY", "test-auth-key") t.Setenv("DATA_DIR", dataDir) @@ -51,7 +53,9 @@ func TestBuildContainerResolvesS5DependencyGraph(t *testing.T) { runtimeState app.RuntimeStateLoader, _ *gateway.Handler, attemptForwarder gateway.AttemptForwarder, - _ gateway.DialectSet, + _ dialect.Set, + _ *control.Service, + _ *control.Server, _ *httpclient.HTTPClientManager, _ *redact.Redactor, _ *dialect.OpenAI, @@ -82,16 +86,71 @@ func TestBuildContainerResolvesS5DependencyGraph(t *testing.T) { resolved = true }) if err != nil { - t.Fatalf("resolve S5 dependency graph: %v", err) + t.Fatalf("resolve S6 dependency graph: %v", err) } if !resolved { - t.Fatal("S5 dependency graph was not invoked") + t.Fatal("S6 dependency graph was not invoked") } if _, err := os.Stat(filepath.Join(dataDir, encryption.KeyFileName)); err != nil { t.Fatalf("container did not initialize encryption keyfile: %v", err) } } +func TestBuildContainerRegistersS6ControlRoutesWithoutAffectingGateway(t *testing.T) { + t.Setenv("AUTH_KEY", "test-auth-key") + t.Setenv("DATA_DIR", t.TempDir()) + t.Setenv("DATABASE_DSN", ":memory:") + t.Setenv("ENCRYPTION_KEY", "test-master-key-long") + t.Setenv("REDIS_DSN", "") + if err := i18n.Init(); err != nil { + t.Fatalf("i18n.Init() error = %v", err) + } + + dependencyContainer, err := BuildContainer() + if err != nil { + t.Fatalf("BuildContainer() error = %v", err) + } + err = dependencyContainer.Invoke(func( + engine *gin.Engine, + db *gorm.DB, + runtimeState app.RuntimeStateLoader, + ) { + if err := storage.AutoMigrate(db); err != nil { + t.Fatalf("AutoMigrate() error = %v", err) + } + if err := runtimeState.Load(context.Background()); err != nil { + t.Fatalf("runtimeState.Load() error = %v", err) + } + + groupsRecorder := httptest.NewRecorder() + groupsRequest := httptest.NewRequest(http.MethodGet, "/api/groups", nil) + groupsRequest.Header.Set("Authorization", "Bearer test-auth-key") + engine.ServeHTTP(groupsRecorder, groupsRequest) + if groupsRecorder.Code != http.StatusOK { + t.Fatalf("groups status = %d, want 200; body=%s", groupsRecorder.Code, groupsRecorder.Body.String()) + } + + gatewayRecorder := httptest.NewRecorder() + gatewayRequest := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"gpt-4o"}`)) + gatewayRequest.Header.Set("Authorization", "Bearer test-auth-key") + engine.ServeHTTP(gatewayRecorder, gatewayRequest) + if gatewayRecorder.Code != http.StatusUnauthorized || !strings.Contains(gatewayRecorder.Body.String(), "invalid_access_key") { + t.Fatalf("gateway response = %d %s, want data-plane 401", gatewayRecorder.Code, gatewayRecorder.Body.String()) + } + + unknownRecorder := httptest.NewRecorder() + unknownRequest := httptest.NewRequest(http.MethodGet, "/api/unknown", nil) + unknownRequest.Header.Set("Authorization", "Bearer test-auth-key") + engine.ServeHTTP(unknownRecorder, unknownRequest) + if unknownRecorder.Code != http.StatusUnauthorized || !strings.Contains(unknownRecorder.Body.String(), "invalid_access_key") { + t.Fatalf("unknown /api response = %d %s, want documented gateway NoRoute 401", unknownRecorder.Code, unknownRecorder.Body.String()) + } + }) + if err != nil { + t.Fatalf("resolve S6 engine: %v", err) + } +} + func TestBuildContainerRegistersS5GatewayRoute(t *testing.T) { dataDir := t.TempDir() t.Setenv("AUTH_KEY", "test-auth-key") diff --git a/internal/control/access_keys.go b/internal/control/access_keys.go new file mode 100644 index 00000000..5c55586d --- /dev/null +++ b/internal/control/access_keys.go @@ -0,0 +1,311 @@ +package control + +import ( + "bytes" + "context" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "strings" + + "gorm.io/gorm" + + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/protocol" + "gpt-load/internal/state" + "gpt-load/internal/storage/models" +) + +type AccessKeyFilters struct { + Groups []uint `json:"groups"` + Protocols []protocol.Protocol `json:"protocols"` + Models []string `json:"models"` +} + +type AccessKeyCreateRequest struct { + Name string `json:"name"` + Filters *AccessKeyFilters `json:"filters"` +} + +type AccessKeyUpdateRequest struct { + Name *string `json:"name"` + Status *state.AccessKeyStatus `json:"status"` + Filters *AccessKeyFilters `json:"filters"` +} + +type AccessKeyResponse struct { + ID uint `json:"id"` + Name string `json:"name"` + Key string `json:"key"` + Status state.AccessKeyStatus `json:"status"` + Filters AccessKeyFilters `json:"filters"` +} + +func (s *Service) CreateAccessKey( + ctx context.Context, + request AccessKeyCreateRequest, +) (AccessKeyResponse, error) { + name, err := normalizeAccessKeyName(request.Name) + if err != nil { + return AccessKeyResponse{}, err + } + filters, err := normalizeAccessKeyFilters(request.Filters) + if err != nil { + return AccessKeyResponse{}, err + } + encodedFilters, err := json.Marshal(filters) + if err != nil { + return AccessKeyResponse{}, fmt.Errorf("encode access key filters: %w", err) + } + + var result AccessKeyResponse + _, err = s.writeConfig(ctx, func(tx *gorm.DB) error { + if err := validateFilterGroupReferences(tx, filters.Groups); err != nil { + return err + } + randomBytes := make([]byte, 16) + if _, err := io.ReadFull(s.random, randomBytes); err != nil { + return fmt.Errorf("generate access key: %w", err) + } + plaintext := "gl-" + hex.EncodeToString(randomBytes) + ciphertext, err := s.encryption.Encrypt(plaintext) + if err != nil { + return fmt.Errorf("encrypt access key: %w", err) + } + row := models.AccessKey{ + Name: name, KeyValue: ciphertext, KeyHash: s.encryption.Hash(plaintext), + Status: string(state.AccessKeyStatusActive), Filters: models.JSON(encodedFilters), + } + 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, + } + return nil + }, nil) + if err != nil { + return AccessKeyResponse{}, 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 { + return nil, app_errors.ParseDBError(err) + } + + result := make([]AccessKeyResponse, 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) + if err != nil { + return nil, fmt.Errorf("decode access key %d filters: %w", row.ID, 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, + }) + } + return result, nil +} + +func (s *Service) UpdateAccessKey( + ctx context.Context, + id uint, + request AccessKeyUpdateRequest, +) (AccessKeyResponse, error) { + if id == 0 || (request.Name == nil && request.Status == nil && request.Filters == nil) { + return AccessKeyResponse{}, app_errors.ErrBadRequest + } + + var name *string + if request.Name != nil { + normalized, err := normalizeAccessKeyName(*request.Name) + if err != nil { + return AccessKeyResponse{}, err + } + name = &normalized + } + if request.Status != nil && + *request.Status != state.AccessKeyStatusActive && + *request.Status != state.AccessKeyStatusDisabled { + return AccessKeyResponse{}, app_errors.ErrValidation + } + var filters *AccessKeyFilters + var encodedFilters []byte + if request.Filters != nil { + normalized, err := normalizeAccessKeyFilters(request.Filters) + if err != nil { + return AccessKeyResponse{}, err + } + encoded, err := json.Marshal(normalized) + if err != nil { + return AccessKeyResponse{}, fmt.Errorf("encode access key filters: %w", err) + } + filters = &normalized + encodedFilters = encoded + } + + var result AccessKeyResponse + _, err := s.writeConfig(ctx, func(tx *gorm.DB) error { + var row models.AccessKey + if err := tx.First(&row, id).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) + } + currentFilters, err := decodeStoredAccessKeyFilters(row.Filters) + if err != nil { + return fmt.Errorf("decode access key %d filters: %w", row.ID, err) + } + status := state.AccessKeyStatus(row.Status) + if status != state.AccessKeyStatusActive && status != state.AccessKeyStatusDisabled { + return fmt.Errorf("access key %d has invalid status", row.ID) + } + + updates := make(map[string]any, 3) + if name != nil { + row.Name = *name + updates["name"] = row.Name + } + if request.Status != nil { + status = *request.Status + updates["status"] = string(status) + } + if filters != nil { + if err := validateFilterGroupReferences(tx, filters.Groups); err != nil { + return err + } + currentFilters = *filters + updates["filters"] = models.JSON(encodedFilters) + } + if err := tx.Model(&row).Updates(updates).Error; err != nil { + return app_errors.ParseDBError(err) + } + result = AccessKeyResponse{ + ID: row.ID, Name: row.Name, Key: plaintext, Status: status, Filters: currentFilters, + } + return nil + }, nil) + if err != nil { + return AccessKeyResponse{}, err + } + return result, nil +} + +func (s *Service) DeleteAccessKey(ctx context.Context, id uint) error { + if id == 0 { + return app_errors.ErrBadRequest + } + _, err := s.writeConfig(ctx, func(tx *gorm.DB) error { + var row models.AccessKey + if err := tx.Select("id").First(&row, id).Error; err != nil { + return app_errors.ParseDBError(err) + } + if err := tx.Delete(&row).Error; err != nil { + return app_errors.ParseDBError(err) + } + return nil + }, nil) + return err +} + +func normalizeAccessKeyName(raw string) (string, error) { + normalized, err := normalizeImportName(&raw) + if err != nil { + return "", err + } + return *normalized, nil +} + +func normalizeAccessKeyFilters(input *AccessKeyFilters) (AccessKeyFilters, error) { + result := AccessKeyFilters{ + Groups: make([]uint, 0), Protocols: make([]protocol.Protocol, 0), Models: make([]string, 0), + } + if input == nil { + return result, nil + } + + seenGroups := make(map[uint]struct{}, len(input.Groups)) + for _, groupID := range input.Groups { + if groupID == 0 { + return AccessKeyFilters{}, app_errors.ErrValidation + } + if _, duplicate := seenGroups[groupID]; duplicate { + continue + } + seenGroups[groupID] = struct{}{} + result.Groups = append(result.Groups, groupID) + } + seenProtocols := make(map[protocol.Protocol]struct{}, len(input.Protocols)) + for _, value := range input.Protocols { + if !value.Valid() { + return AccessKeyFilters{}, app_errors.ErrValidation + } + if _, duplicate := seenProtocols[value]; duplicate { + continue + } + seenProtocols[value] = struct{}{} + result.Protocols = append(result.Protocols, value) + } + seenModels := make(map[string]struct{}, len(input.Models)) + for _, value := range input.Models { + normalized := strings.TrimSpace(value) + if normalized == "" { + return AccessKeyFilters{}, app_errors.ErrValidation + } + if _, duplicate := seenModels[normalized]; duplicate { + continue + } + seenModels[normalized] = struct{}{} + result.Models = append(result.Models, normalized) + } + return result, nil +} + +func validateFilterGroupReferences(tx *gorm.DB, groupIDs []uint) error { + if len(groupIDs) == 0 { + return nil + } + var count int64 + if err := tx.Model(&models.Group{}).Where("id IN ?", groupIDs).Count(&count).Error; err != nil { + return app_errors.ParseDBError(err) + } + if count != int64(len(groupIDs)) { + return app_errors.ErrValidation + } + return nil +} + +func decodeStoredAccessKeyFilters(raw models.JSON) (AccessKeyFilters, error) { + if len(raw) == 0 { + return normalizeAccessKeyFilters(nil) + } + var decoded AccessKeyFilters + decoder := json.NewDecoder(bytes.NewReader(raw)) + decoder.DisallowUnknownFields() + if err := decoder.Decode(&decoded); err != nil { + return AccessKeyFilters{}, err + } + var extra any + if err := decoder.Decode(&extra); !errors.Is(err, io.EOF) { + if err == nil { + return AccessKeyFilters{}, fmt.Errorf("multiple JSON values") + } + return AccessKeyFilters{}, err + } + return normalizeAccessKeyFilters(&decoded) +} diff --git a/internal/control/access_keys_test.go b/internal/control/access_keys_test.go new file mode 100644 index 00000000..6dab628f --- /dev/null +++ b/internal/control/access_keys_test.go @@ -0,0 +1,509 @@ +package control + +import ( + "bytes" + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "reflect" + "strings" + "sync" + "testing" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" + + "gpt-load/internal/platform/config" + "gpt-load/internal/protocol" + "gpt-load/internal/state" + stateloader "gpt-load/internal/state/loader" + "gpt-load/internal/storage/models" +) + +func TestCreateAccessKeyGeneratesEncryptedGLToken(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.random = bytes.NewReader([]byte{ + 0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, + 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, + }) + if err := fixture.registry.ApplyImport(77, []state.KeyEntry{{ + ID: 88, GroupID: 77, Status: state.KeyStatusActive, EncryptedValue: "existing-upstream-cipher", + }}); err != nil { + t.Fatalf("seed Registry: %v", err) + } + if count, ok := fixture.registry.IncrFailure(88); !ok || count != 1 { + t.Fatalf("seed failure count = %d, %t", count, ok) + } + before := fixture.manager.Current().Revision + + result, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: " client "}) + if err != nil { + t.Fatalf("CreateAccessKey() error = %v", err) + } + const plaintext = "gl-000102030405060708090a0b0c0d0e0f" + if result.Key != plaintext || result.Name != "client" || result.Status != state.AccessKeyStatusActive { + t.Fatalf("CreateAccessKey() = %#v", result) + } + if got := fixture.manager.Current().Revision; got != before+1 { + t.Fatalf("revision = %d, want %d", got, before+1) + } + + var row models.AccessKey + if err := fixture.db.First(&row, result.ID).Error; err != nil { + t.Fatalf("query AccessKey: %v", err) + } + if row.KeyValue == "" || row.KeyHash == "" || row.KeyValue == plaintext || strings.Contains(row.KeyValue, plaintext) { + t.Fatalf("stored AccessKey exposes plaintext: %#v", row) + } + decrypted, err := fixture.encryption.Decrypt(row.KeyValue) + if err != nil || decrypted != plaintext { + t.Fatalf("Decrypt() = %q, %v, want plaintext", decrypted, err) + } + if row.KeyHash != fixture.encryption.Hash(plaintext) { + t.Fatalf("KeyHash = %q, want stable HMAC", row.KeyHash) + } + if _, ok := fixture.manager.Current().AccessKeysByHash[row.KeyHash]; !ok { + t.Fatalf("published snapshot lacks hash %q", row.KeyHash) + } + if got, ok := fixture.registry.EncryptedValue(88); !ok || got != "existing-upstream-cipher" { + t.Fatalf("Registry value = %q, %t, want unchanged", got, ok) + } + if count, ok := fixture.registry.IncrFailure(88); !ok || count != 2 { + t.Fatalf("Registry failure count = %d, %t, want preserved count 2", count, ok) + } +} + +func TestAccessKeyFiltersNormalizeAndAcceptExistingDisabledGroups(t *testing.T) { + fixture := newServiceFixture(t) + enabled := validControlGroup("filter-enabled") + if err := fixture.db.Create(enabled).Error; err != nil { + t.Fatalf("create enabled group: %v", err) + } + disabled := validControlGroup("filter-disabled") + if err := fixture.db.Create(disabled).Error; err != nil { + t.Fatalf("create disabled group: %v", err) + } + if err := fixture.db.Model(disabled).Update("enabled", false).Error; err != nil { + t.Fatalf("disable group: %v", err) + } + fixture.service.random = bytes.NewReader(make([]byte, 16)) + + result, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{ + Name: "filtered", + Filters: &AccessKeyFilters{ + Groups: []uint{enabled.ID, disabled.ID, enabled.ID}, + Protocols: []protocol.Protocol{protocol.OpenAI, protocol.OpenAIResponse, protocol.OpenAI}, + Models: []string{" gpt-4o ", "gpt-4o", "claude"}, + }, + }) + if err != nil { + t.Fatalf("CreateAccessKey() error = %v", err) + } + if len(result.Filters.Groups) != 2 || result.Filters.Groups[0] != enabled.ID || result.Filters.Groups[1] != disabled.ID { + t.Fatalf("normalized groups = %#v", result.Filters.Groups) + } + if len(result.Filters.Protocols) != 2 || result.Filters.Protocols[1] != protocol.OpenAIResponse { + t.Fatalf("normalized protocols = %#v", result.Filters.Protocols) + } + if len(result.Filters.Models) != 2 || result.Filters.Models[0] != "gpt-4o" || result.Filters.Models[1] != "claude" { + t.Fatalf("normalized models = %#v", result.Filters.Models) + } + + var row models.AccessKey + if err := fixture.db.First(&row, result.ID).Error; err != nil { + t.Fatalf("query AccessKey: %v", err) + } + var document map[string]json.RawMessage + if err := json.Unmarshal(row.Filters, &document); err != nil { + t.Fatalf("decode stored filters: %v", err) + } + if len(document) != 3 || document["groups"] == nil || document["protocols"] == nil || document["models"] == nil { + t.Fatalf("stored filter document = %#v", document) + } + view := fixture.manager.Current().AccessKeysByHash[row.KeyHash] + if _, ok := view.Filters.Groups[enabled.ID]; !ok { + t.Fatalf("published filters = %#v, want enabled group", view.Filters) + } + if _, ok := view.Filters.Groups[disabled.ID]; !ok { + t.Fatalf("published filters = %#v, want disabled group reference retained", view.Filters) + } +} + +func TestAccessKeyFiltersRejectInvalidCurrentInputWithoutPublishing(t *testing.T) { + blank := " " + tooLong := strings.Repeat("名", 86) + controlName := "bad\nname" + tests := []struct { + name string + request AccessKeyCreateRequest + }{ + {name: "blank name", request: AccessKeyCreateRequest{Name: blank}}, + {name: "long name", request: AccessKeyCreateRequest{Name: tooLong}}, + {name: "control name", request: AccessKeyCreateRequest{Name: controlName}}, + {name: "zero group", request: AccessKeyCreateRequest{Name: "client", Filters: &AccessKeyFilters{Groups: []uint{0}}}}, + {name: "missing group", request: AccessKeyCreateRequest{Name: "client", Filters: &AccessKeyFilters{Groups: []uint{999}}}}, + {name: "unknown protocol", request: AccessKeyCreateRequest{Name: "client", Filters: &AccessKeyFilters{Protocols: []protocol.Protocol{"unknown"}}}}, + {name: "blank model", request: AccessKeyCreateRequest{Name: "client", Filters: &AccessKeyFilters{Models: []string{" "}}}}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.random = bytes.NewReader(make([]byte, 16)) + before := fixture.manager.Current().Revision + if _, err := fixture.service.CreateAccessKey(context.Background(), test.request); err == nil { + t.Fatal("CreateAccessKey() error = nil") + } + var count int64 + if err := fixture.db.Model(&models.AccessKey{}).Count(&count).Error; err != nil { + t.Fatalf("count AccessKeys: %v", err) + } + if count != 0 || fixture.manager.Current().Revision != before { + t.Fatalf("count/revision = %d/%d, want 0/%d", count, fixture.manager.Current().Revision, before) + } + }) + } +} + +func TestListAccessKeysReturnsPlaintextInIDOrderAndFailsClosed(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + randomBytes := make([]byte, 32) + for index := range randomBytes { + randomBytes[index] = byte(index) + } + fixture.service.random = bytes.NewReader(randomBytes) + first, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "first"}) + if err != nil { + t.Fatalf("create first: %v", err) + } + second, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "second"}) + if err != nil { + t.Fatalf("create second: %v", err) + } + + listed, err := fixture.service.ListAccessKeys(context.Background()) + 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 { + t.Fatalf("ListAccessKeys() = %#v", listed) + } + + const corruptCiphertext = "known-corrupt-ciphertext" + 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) + } + + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + recorder := httptest.NewRecorder() + 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") { + t.Fatalf("GET access keys = %d %s", recorder.Code, recorder.Body.String()) + } + for _, forbidden := range []string{first.Key, corruptCiphertext} { + if strings.Contains(recorder.Body.String(), forbidden) { + t.Fatalf("fail-closed response exposes %q: %s", forbidden, recorder.Body.String()) + } + } +} + +func TestUpdateAccessKeyPreservesCredentialAcrossPointerPatches(t *testing.T) { + fixture := newServiceFixture(t) + group := validControlGroup("access-key-update") + if err := fixture.db.Create(group).Error; err != nil { + t.Fatalf("create group: %v", err) + } + fixture.service.random = bytes.NewReader(make([]byte, 16)) + created, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{ + Name: "before", + Filters: &AccessKeyFilters{ + Groups: []uint{group.ID}, Protocols: []protocol.Protocol{protocol.OpenAI}, Models: []string{"gpt-4o"}, + }, + }) + if err != nil { + t.Fatalf("CreateAccessKey() error = %v", err) + } + original := loadAccessKeyRow(t, fixture.db, created.ID) + + name := " after " + updated, err := fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{Name: &name}) + if err != nil { + t.Fatalf("name-only UpdateAccessKey() error = %v", err) + } + if updated.Name != "after" || updated.Status != state.AccessKeyStatusActive || updated.Key != created.Key { + t.Fatalf("name-only update = %#v", updated) + } + assertAccessKeyCredentialUnchanged(t, fixture.db, created.ID, original) + + status := state.AccessKeyStatusDisabled + updated, err = fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{Status: &status}) + if err != nil { + t.Fatalf("status-only UpdateAccessKey() error = %v", err) + } + if updated.Name != "after" || updated.Status != status || updated.Filters.Models[0] != "gpt-4o" { + t.Fatalf("status-only update = %#v", updated) + } + assertAccessKeyCredentialUnchanged(t, fixture.db, created.ID, original) + + emptyFilters := AccessKeyFilters{} + updated, err = fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{Filters: &emptyFilters}) + if err != nil { + t.Fatalf("clear-filters UpdateAccessKey() error = %v", err) + } + if len(updated.Filters.Groups) != 0 || len(updated.Filters.Protocols) != 0 || len(updated.Filters.Models) != 0 { + t.Fatalf("cleared filters = %#v", updated.Filters) + } + assertAccessKeyCredentialUnchanged(t, fixture.db, created.ID, original) + + finalName := "final" + active := state.AccessKeyStatusActive + finalFilters := AccessKeyFilters{ + Groups: []uint{group.ID}, Protocols: []protocol.Protocol{protocol.Anthropic}, Models: []string{"claude-3-5"}, + } + updated, err = fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{ + Name: &finalName, Status: &active, Filters: &finalFilters, + }) + if err != nil { + t.Fatalf("all-fields UpdateAccessKey() error = %v", err) + } + if updated.Name != finalName || updated.Status != active || len(updated.Filters.Groups) != 1 || updated.Filters.Protocols[0] != protocol.Anthropic { + t.Fatalf("all-fields update = %#v", updated) + } + assertAccessKeyCredentialUnchanged(t, fixture.db, created.ID, original) + + before := fixture.manager.Current().Revision + if _, err := fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{}); err == nil { + t.Fatal("empty UpdateAccessKey() error = nil") + } + if got := fixture.manager.Current().Revision; got != before { + t.Fatalf("revision after empty update = %d, want %d", got, before) + } +} + +func TestUpdateAccessKeyStatusAndDeletePublishWithoutMutatingRegistry(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.random = bytes.NewReader(make([]byte, 16)) + created, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "toggle"}) + if err != nil { + t.Fatalf("CreateAccessKey() error = %v", err) + } + row := loadAccessKeyRow(t, fixture.db, created.ID) + if err := fixture.registry.ApplyImport(77, []state.KeyEntry{{ + ID: 88, GroupID: 77, Status: state.KeyStatusActive, EncryptedValue: "registry-cipher", + }}); err != nil { + t.Fatalf("seed Registry: %v", err) + } + if count, ok := fixture.registry.IncrFailure(88); !ok || count != 1 { + t.Fatalf("seed failure count = %d, %t", count, ok) + } + + before := fixture.manager.Current().Revision + disabled := state.AccessKeyStatusDisabled + if _, err := fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{Status: &disabled}); err != nil { + t.Fatalf("disable UpdateAccessKey() error = %v", err) + } + if got := fixture.manager.Current().Revision; got != before+1 { + t.Fatalf("revision after disable = %d, want %d", got, before+1) + } + if _, ok := fixture.manager.Current().AccessKeysByHash[row.KeyHash]; ok { + t.Fatal("disabled AccessKey remains in authentication snapshot") + } + + active := state.AccessKeyStatusActive + if _, err := fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{Status: &active}); err != nil { + t.Fatalf("enable UpdateAccessKey() error = %v", err) + } + if _, ok := fixture.manager.Current().AccessKeysByHash[row.KeyHash]; !ok { + t.Fatal("active AccessKey missing from authentication snapshot") + } + + invalid := state.AccessKeyStatus("paused") + before = fixture.manager.Current().Revision + if _, err := fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{Status: &invalid}); err == nil { + t.Fatal("invalid status UpdateAccessKey() error = nil") + } + if got := fixture.manager.Current().Revision; got != before { + t.Fatalf("revision after invalid status = %d, want %d", got, before) + } + + if err := fixture.service.DeleteAccessKey(context.Background(), created.ID); err != nil { + t.Fatalf("DeleteAccessKey() error = %v", err) + } + if got := fixture.manager.Current().Revision; got != before+1 { + t.Fatalf("revision after delete = %d, want %d", got, before+1) + } + if _, ok := fixture.manager.Current().AccessKeysByHash[row.KeyHash]; ok { + t.Fatal("deleted AccessKey remains in authentication snapshot") + } + var count int64 + if err := fixture.db.Model(&models.AccessKey{}).Where("id = ?", created.ID).Count(&count).Error; err != nil || count != 0 { + t.Fatalf("deleted DB row count = %d, error=%v", count, err) + } + + before = fixture.manager.Current().Revision + if err := fixture.service.DeleteAccessKey(context.Background(), created.ID); err == nil { + t.Fatal("repeated DeleteAccessKey() error = nil") + } + if _, err := fixture.service.UpdateAccessKey(context.Background(), created.ID, AccessKeyUpdateRequest{Name: stringPointer("missing")}); err == nil { + t.Fatal("unknown UpdateAccessKey() error = nil") + } + if got := fixture.manager.Current().Revision; got != before { + t.Fatalf("revision after missing mutations = %d, want %d", got, before) + } + if value, ok := fixture.registry.EncryptedValue(88); !ok || value != "registry-cipher" { + t.Fatalf("Registry value = %q, %t, want unchanged", value, ok) + } + if failures, ok := fixture.registry.IncrFailure(88); !ok || failures != 2 { + t.Fatalf("Registry failure count = %d, %t, want preserved count 2", failures, ok) + } +} + +func TestUpdateAccessKeyRollsBackWhenCredentialCannotDecrypt(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.random = bytes.NewReader(make([]byte, 16)) + created, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "before"}) + if err != nil { + t.Fatalf("CreateAccessKey() error = %v", err) + } + if err := fixture.db.Model(&models.AccessKey{}).Where("id = ?", created.ID).UpdateColumn("key_value", "corrupt").Error; err != nil { + 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") + } + 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) + } +} + +func TestAccessKeyDanglingFiltersDoNotBlockUnrelatedUpdate(t *testing.T) { + fixture := newServiceFixture(t) + danglingPlaintext := "gl-dangling-test-value" + danglingCiphertext, err := fixture.encryption.Encrypt(danglingPlaintext) + if err != nil { + t.Fatalf("encrypt dangling key: %v", err) + } + 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":[]}`), + }).Error; err != nil { + t.Fatalf("create dangling AccessKey: %v", err) + } + fixture.service.random = bytes.NewReader(make([]byte, 16)) + current, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "current"}) + if err != nil { + t.Fatalf("create current AccessKey: %v", err) + } + + if _, err := fixture.service.UpdateAccessKey(context.Background(), current.ID, AccessKeyUpdateRequest{Name: stringPointer("renamed")}); err != nil { + t.Fatalf("unrelated UpdateAccessKey() error = %v", err) + } + before := fixture.manager.Current().Revision + invalidFilters := AccessKeyFilters{Groups: []uint{999}} + if _, err := fixture.service.UpdateAccessKey(context.Background(), current.ID, AccessKeyUpdateRequest{Filters: &invalidFilters}); err == nil { + t.Fatal("current dangling filter UpdateAccessKey() error = nil") + } + if got := fixture.manager.Current().Revision; got != before { + t.Fatalf("revision after rejected dangling filter = %d, want %d", got, before) + } +} + +func TestConcurrentAccessKeyCRUDPublishesDatabaseTruth(t *testing.T) { + fixture := newServiceFixture(t) + rows := make([]AccessKeyResponse, 4) + for index := range rows { + created, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "seed-" + string(rune('a'+index))}) + if err != nil { + t.Fatalf("seed AccessKey %d: %v", index, err) + } + rows[index] = created + } + before := fixture.manager.Current().Revision + start := make(chan struct{}) + errors := make(chan error, 6) + var ready sync.WaitGroup + ready.Add(6) + operations := []func() error{ + func() error { + _, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "concurrent-a"}) + return err + }, + func() error { + _, err := fixture.service.CreateAccessKey(context.Background(), AccessKeyCreateRequest{Name: "concurrent-b"}) + return err + }, + func() error { + _, err := fixture.service.UpdateAccessKey(context.Background(), rows[0].ID, AccessKeyUpdateRequest{Name: stringPointer("updated-a")}) + return err + }, + func() error { + status := state.AccessKeyStatusDisabled + _, err := fixture.service.UpdateAccessKey(context.Background(), rows[1].ID, AccessKeyUpdateRequest{Status: &status}) + return err + }, + func() error { return fixture.service.DeleteAccessKey(context.Background(), rows[2].ID) }, + func() error { return fixture.service.DeleteAccessKey(context.Background(), rows[3].ID) }, + } + for _, operation := range operations { + operation := operation + go func() { + ready.Done() + <-start + errors <- operation() + }() + } + ready.Wait() + close(start) + for range operations { + if err := <-errors; err != nil { + t.Fatalf("concurrent AccessKey mutation error = %v", err) + } + } + + input, err := stateloader.BuildCompileInput(context.Background(), fixture.db) + if err != nil { + t.Fatalf("BuildCompileInput() error = %v", err) + } + want, err := state.Compile(input) + if err != nil { + t.Fatalf("Compile(DB input) error = %v", err) + } + got := fixture.manager.Current() + want.Revision = got.Revision + if !reflect.DeepEqual(got, want) { + t.Fatalf("published snapshot differs from DB compile\ngot=%#v\nwant=%#v", got, want) + } + if got.Revision != before+uint64(len(operations)) { + t.Fatalf("revision = %d, want %d", got.Revision, before+uint64(len(operations))) + } +} + +func loadAccessKeyRow(t *testing.T, db *gorm.DB, id uint) models.AccessKey { + t.Helper() + var row models.AccessKey + if err := db.First(&row, id).Error; err != nil { + t.Fatalf("query AccessKey %d: %v", id, err) + } + return row +} + +func assertAccessKeyCredentialUnchanged(t *testing.T, db *gorm.DB, id uint, original models.AccessKey) { + t.Helper() + current := loadAccessKeyRow(t, db, id) + if current.KeyValue != original.KeyValue || current.KeyHash != original.KeyHash { + t.Fatalf("credential changed: before=%q/%q after=%q/%q", original.KeyValue, original.KeyHash, current.KeyValue, current.KeyHash) + } +} + +func stringPointer(value string) *string { + return &value +} diff --git a/internal/control/groups.go b/internal/control/groups.go new file mode 100644 index 00000000..2efbf92f --- /dev/null +++ b/internal/control/groups.go @@ -0,0 +1,65 @@ +package control + +import ( + "context" + "encoding/json" + "fmt" + + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/protocol" + "gpt-load/internal/storage/models" +) + +type GroupModel struct { + ID string `json:"id"` + Alias string `json:"alias"` +} + +type GroupResponse struct { + ID uint `json:"id"` + Name string `json:"name"` + UpstreamURL string `json:"upstream_url"` + Protocols []protocol.Protocol `json:"protocols"` + Models []GroupModel `json:"models"` + Enabled bool `json:"enabled"` + KeyCount int64 `json:"key_count"` +} + +func (s *Service) ListGroups(ctx context.Context) ([]GroupResponse, error) { + var groups []models.Group + if err := s.db.WithContext(ctx).Order("id ASC").Find(&groups).Error; err != nil { + return nil, app_errors.ParseDBError(err) + } + + type keyCountRow struct { + GroupID uint + Count int64 + } + var countRows []keyCountRow + if err := s.db.WithContext(ctx).Model(&models.UpstreamKey{}). + Select("group_id, COUNT(*) AS count").Group("group_id").Find(&countRows).Error; err != nil { + return nil, app_errors.ParseDBError(err) + } + counts := make(map[uint]int64, len(countRows)) + for _, row := range countRows { + counts[row.GroupID] = row.Count + } + + result := make([]GroupResponse, 0, len(groups)) + for _, group := range groups { + var protocols []protocol.Protocol + if err := json.Unmarshal(group.Protocols, &protocols); err != nil { + return nil, fmt.Errorf("decode group %d protocols: %w", group.ID, err) + } + var groupModels []GroupModel + if err := json.Unmarshal(group.Models, &groupModels); err != nil { + return nil, fmt.Errorf("decode group %d models: %w", group.ID, err) + } + result = append(result, GroupResponse{ + ID: group.ID, Name: group.Name, UpstreamURL: group.UpstreamURL, + Protocols: protocols, Models: groupModels, Enabled: group.Enabled, + KeyCount: counts[group.ID], + }) + } + return result, nil +} diff --git a/internal/control/groups_test.go b/internal/control/groups_test.go new file mode 100644 index 00000000..722ac188 --- /dev/null +++ b/internal/control/groups_test.go @@ -0,0 +1,77 @@ +package control + +import ( + "context" + "encoding/json" + "strings" + "testing" + + "gorm.io/gorm" + + "gpt-load/internal/storage/models" +) + +func TestListGroupsReturnsAllGroupsAndTotalKeyCountsInTwoQueries(t *testing.T) { + fixture := newServiceFixture(t) + enabled := validControlGroup("enabled") + enabled.Models = models.JSON(`[{"id":"gpt-4o","alias":"Primary"}]`) + if err := fixture.db.Create(enabled).Error; err != nil { + t.Fatalf("create enabled group: %v", err) + } + disabled := validControlGroup("disabled") + if err := fixture.db.Create(disabled).Error; err != nil { + t.Fatalf("create disabled group: %v", err) + } + if err := fixture.db.Model(disabled).Update("enabled", false).Error; err != nil { + t.Fatalf("disable group: %v", err) + } + for _, key := range []models.UpstreamKey{ + {GroupID: enabled.ID, KeyValue: "cipher-active", KeyHash: "hash-active", Status: models.UpstreamKeyStatusActive}, + {GroupID: enabled.ID, KeyValue: "cipher-disabled", KeyHash: "hash-disabled", Status: models.UpstreamKeyStatusDisabled}, + } { + key := key + if err := fixture.db.Create(&key).Error; err != nil { + t.Fatalf("create upstream key: %v", err) + } + } + + queryCount := 0 + const callbackName = "test:list_groups_query_count" + if err := fixture.db.Callback().Query().After("gorm:query").Register(callbackName, func(*gorm.DB) { + queryCount++ + }); err != nil { + t.Fatalf("register query callback: %v", err) + } + + got, err := fixture.service.ListGroups(context.Background()) + if err != nil { + t.Fatalf("ListGroups() error = %v", err) + } + if queryCount != 2 { + t.Fatalf("ListGroups() query count = %d, want 2", queryCount) + } + if len(got) != 2 { + t.Fatalf("ListGroups() = %#v, want two groups", got) + } + if got[0].ID != enabled.ID || !got[0].Enabled || got[0].KeyCount != 2 { + t.Fatalf("enabled response = %#v, want enabled with two total keys", got[0]) + } + if len(got[0].Protocols) != 1 || len(got[0].Models) != 1 || got[0].Models[0].Alias != "Primary" { + t.Fatalf("enabled protocols/models = %#v/%#v", got[0].Protocols, got[0].Models) + } + if got[1].ID != disabled.ID || got[1].Enabled || got[1].KeyCount != 0 { + t.Fatalf("disabled response = %#v, want disabled with zero keys", got[1]) + } + + encoded, err := json.Marshal(got) + if err != nil { + t.Fatalf("json.Marshal(groups) error = %v", err) + } + for _, forbidden := range []string{ + "key_value", "key_hash", "cipher-active", "cipher-disabled", "hash-active", "hash-disabled", + } { + if strings.Contains(string(encoded), forbidden) { + t.Fatalf("groups response exposes %q: %s", forbidden, encoded) + } + } +} diff --git a/internal/control/import.go b/internal/control/import.go new file mode 100644 index 00000000..7ff6336e --- /dev/null +++ b/internal/control/import.go @@ -0,0 +1,536 @@ +package control + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "net" + "net/url" + "strings" + "unicode" + + "github.com/sirupsen/logrus" + "gorm.io/gorm" + + "gpt-load/internal/dialect" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/protocol" + "gpt-load/internal/state" + "gpt-load/internal/storage/models" +) + +const maxImportKeyLines = 1000 + +type ImportRequest struct { + UpstreamURL string `json:"upstream_url"` + Protocols []protocol.Protocol `json:"protocols"` + Name *string `json:"name"` + Keys string `json:"keys"` +} + +type ImportResult struct { + GroupID uint `json:"group_id"` + GroupName string `json:"group_name"` + Created bool `json:"created"` + KeysAdded int `json:"keys_added"` + KeysDuplicated int `json:"keys_duplicated"` + ModelsFetched bool `json:"models_fetched"` + Models []string `json:"models"` +} + +type importKeyCandidate struct { + plaintext string + hash string +} + +type normalizedImport struct { + upstreamURL string + signature string + hostname string + protocols []protocol.Protocol + explicitName *string + keys []importKeyCandidate + firstKey string + duplicateLines int +} + +type modelFetchResult struct { + fetched bool + models []string +} + +func normalizeUpstreamURL(raw string) (normalized, signature, hostname string, err error) { + parsed, parseErr := url.Parse(strings.TrimSpace(raw)) + if parseErr != nil || parsed.Opaque != "" || parsed.Host == "" { + return "", "", "", app_errors.ErrValidation + } + parsed.Scheme = strings.ToLower(parsed.Scheme) + if parsed.Scheme != "http" && parsed.Scheme != "https" { + return "", "", "", app_errors.ErrValidation + } + if parsed.User != nil || parsed.Fragment != "" { + return "", "", "", app_errors.ErrValidation + } + + hostname = strings.ToLower(parsed.Hostname()) + if hostname == "" { + return "", "", "", app_errors.ErrValidation + } + port := parsed.Port() + if port != "" { + parsed.Host = net.JoinHostPort(hostname, port) + } else if strings.Contains(hostname, ":") { + parsed.Host = "[" + hostname + "]" + } else { + parsed.Host = hostname + } + parsed.Path = strings.TrimRight(parsed.Path, "/") + parsed.RawPath = "" + normalized = parsed.String() + sum := sha256.Sum256([]byte(normalized)) + return normalized, hex.EncodeToString(sum[:]), hostname, nil +} + +func (s *Service) normalizeImportInput(request ImportRequest) (normalizedImport, error) { + upstreamURL, signature, hostname, err := normalizeUpstreamURL(request.UpstreamURL) + if err != nil { + return normalizedImport{}, err + } + + protocols, err := normalizeImportProtocols(request.Protocols) + if err != nil { + return normalizedImport{}, err + } + explicitName, err := normalizeImportName(request.Name) + if err != nil { + return normalizedImport{}, err + } + keys, firstKey, duplicateLines, err := s.normalizeImportKeys(request.Keys) + if err != nil { + return normalizedImport{}, err + } + + return normalizedImport{ + upstreamURL: upstreamURL, signature: signature, hostname: hostname, + protocols: protocols, explicitName: explicitName, keys: keys, + firstKey: firstKey, duplicateLines: duplicateLines, + }, nil +} + +func normalizeImportProtocols(values []protocol.Protocol) ([]protocol.Protocol, error) { + result := make([]protocol.Protocol, 0, len(values)) + seen := make(map[protocol.Protocol]struct{}, len(values)) + for _, value := range values { + if !value.Valid() || value == protocol.OpenAIResponse { + return nil, app_errors.ErrValidation + } + if _, duplicate := seen[value]; duplicate { + continue + } + seen[value] = struct{}{} + result = append(result, value) + } + if len(result) == 0 { + return nil, app_errors.ErrValidation + } + return result, nil +} + +func normalizeImportName(value *string) (*string, error) { + if value == nil { + return nil, nil + } + normalized := strings.TrimSpace(*value) + if normalized == "" || len([]byte(normalized)) > 255 { + return nil, app_errors.ErrValidation + } + for _, character := range normalized { + if unicode.IsControl(character) { + return nil, app_errors.ErrValidation + } + } + return &normalized, nil +} + +func (s *Service) normalizeImportKeys(raw string) ([]importKeyCandidate, string, int, error) { + result := make([]importKeyCandidate, 0) + seen := make(map[string]struct{}) + firstKey := "" + duplicateLines := 0 + nonEmptyLines := 0 + for _, line := range strings.Split(raw, "\n") { + plaintext := strings.TrimSpace(line) + if plaintext == "" { + continue + } + nonEmptyLines++ + if nonEmptyLines > maxImportKeyLines { + return nil, "", 0, app_errors.ErrValidation + } + if firstKey == "" { + firstKey = plaintext + } + hash := s.encryption.Hash(plaintext) + if _, duplicate := seen[hash]; duplicate { + duplicateLines++ + continue + } + seen[hash] = struct{}{} + result = append(result, importKeyCandidate{plaintext: plaintext, hash: hash}) + } + if len(result) == 0 { + return nil, "", 0, app_errors.ErrValidation + } + return result, firstKey, duplicateLines, nil +} + +func (s *Service) Import(ctx context.Context, request ImportRequest) (ImportResult, error) { + normalized, err := s.normalizeImportInput(request) + if err != nil { + return ImportResult{}, err + } + if isLiteralPrivateHost(normalized.hostname) { + logrus.WithFields(logrus.Fields{ + "host": normalized.hostname, + "signature": normalized.signature, + }).Warn("Importing upstream with a private or local host") + } + modelFetch, err := s.prefetchModels(ctx, normalized) + if err != nil { + return ImportResult{}, err + } + + result := ImportResult{ModelsFetched: modelFetch.fetched, Models: make([]string, 0)} + var requestedEntries []state.KeyEntry + _, err = s.writeConfig(ctx, func(tx *gorm.DB) error { + group, created, err := s.persistImportGroup(tx, normalized, modelFetch) + if err != nil { + return err + } + result.GroupID = group.ID + result.GroupName = group.Name + result.Created = created + result.Models, err = decodeImportModelIDs(group.Models) + if err != nil { + return err + } + + requestedEntries, result.KeysAdded, result.KeysDuplicated, err = + s.persistImportKeys(tx, group.ID, normalized) + if err != nil { + return err + } + if err := state.ValidateKeyEntries(requestedEntries); err != nil { + return fmt.Errorf("validate imported keys: %w", err) + } + return nil + }, func() error { + batch := make([]state.KeyEntry, 0, len(requestedEntries)) + for _, entry := range requestedEntries { + if _, exists := s.registry.EncryptedValue(entry.ID); !exists { + batch = append(batch, entry) + } + } + if len(batch) == 0 { + return nil + } + return s.registry.ApplyImport(result.GroupID, batch) + }) + if err != nil { + return ImportResult{}, err + } + return result, nil +} + +func (s *Service) prefetchModels(ctx context.Context, normalized normalizedImport) (modelFetchResult, error) { + if err := ctx.Err(); err != nil { + return modelFetchResult{}, err + } + + protocols := append([]protocol.Protocol(nil), normalized.protocols...) + rules := state.HeaderRules{} + var group models.Group + query := s.db.WithContext(ctx).Where("signature = ?", normalized.signature).Limit(1).Find(&group) + if query.Error != nil { + if err := ctx.Err(); err != nil { + return modelFetchResult{}, err + } + return modelFetchResult{}, app_errors.ParseDBError(query.Error) + } + if query.RowsAffected == 1 { + var existing []protocol.Protocol + if err := json.Unmarshal(group.Protocols, &existing); err != nil { + return modelFetchResult{}, fmt.Errorf("decode group %d protocols for model fetch: %w", group.ID, err) + } + protocols = stableProtocolUnion(existing, normalized.protocols) + if snapshot := s.manager.Current(); snapshot != nil { + if view, ok := snapshot.Groups[group.ID]; ok { + rules = view.HeaderRules + } + } + } + + var selected dialect.Dialect + for _, value := range protocols { + if candidate, ok := s.dialects[value]; ok { + selected = candidate + break + } + } + if selected == nil { + return modelFetchResult{}, nil + } + + fetchContext, cancel := context.WithTimeout(ctx, s.modelFetchTimeout) + models, err := selected.ListModels( + fetchContext, + normalized.upstreamURL, + normalized.firstKey, + rules, + ) + cancel() + if parentErr := ctx.Err(); parentErr != nil { + return modelFetchResult{}, parentErr + } + if err != nil { + return modelFetchResult{}, nil + } + validated, ok := normalizeFetchedModels(models) + if !ok { + return modelFetchResult{}, nil + } + return modelFetchResult{fetched: true, models: validated}, nil +} + +func normalizeFetchedModels(values []string) ([]string, bool) { + result := make([]string, 0, len(values)) + seen := make(map[string]struct{}, len(values)) + for _, value := range values { + normalized := strings.TrimSpace(value) + if normalized == "" || normalized != value { + return nil, false + } + if _, duplicate := seen[normalized]; duplicate { + continue + } + seen[normalized] = struct{}{} + result = append(result, normalized) + } + return result, true +} + +func (s *Service) persistImportGroup( + tx *gorm.DB, + normalized normalizedImport, + modelFetch modelFetchResult, +) (models.Group, bool, error) { + var group models.Group + query := tx.Where("signature = ?", normalized.signature).Limit(1).Find(&group) + if query.Error != nil { + return models.Group{}, false, app_errors.ParseDBError(query.Error) + } + if query.RowsAffected == 1 { + var existing []protocol.Protocol + if err := json.Unmarshal(group.Protocols, &existing); err != nil { + return models.Group{}, false, fmt.Errorf("decode group %d protocols: %w", group.ID, err) + } + merged := stableProtocolUnion(existing, normalized.protocols) + encoded, err := json.Marshal(merged) + if err != nil { + return models.Group{}, false, fmt.Errorf("encode group protocols: %w", err) + } + updates := map[string]any{"protocols": models.JSON(encoded)} + if modelFetch.fetched { + var storedModels []GroupModel + if err := json.Unmarshal(group.Models, &storedModels); err != nil { + return models.Group{}, false, fmt.Errorf("decode group %d models: %w", group.ID, err) + } + mergedModels := appendMissingGroupModels(storedModels, modelFetch.models) + encodedModels, err := json.Marshal(mergedModels) + if err != nil { + return models.Group{}, false, fmt.Errorf("encode group models: %w", err) + } + updates["models"] = models.JSON(encodedModels) + group.Models = models.JSON(encodedModels) + } + if err := tx.Model(&group).Updates(updates).Error; err != nil { + return models.Group{}, false, app_errors.ParseDBError(err) + } + group.Protocols = models.JSON(encoded) + return group, false, nil + } + + name, err := resolveImportGroupName(tx, normalized) + if err != nil { + return models.Group{}, false, err + } + protocols, err := json.Marshal(normalized.protocols) + if err != nil { + return models.Group{}, false, fmt.Errorf("encode group protocols: %w", err) + } + storedModels := make([]GroupModel, 0, len(modelFetch.models)) + if modelFetch.fetched { + for _, modelID := range modelFetch.models { + storedModels = append(storedModels, GroupModel{ID: modelID}) + } + } + encodedModels, err := json.Marshal(storedModels) + if err != nil { + return models.Group{}, false, fmt.Errorf("encode group models: %w", err) + } + group = models.Group{ + Name: name, UpstreamURL: normalized.upstreamURL, Signature: normalized.signature, + Protocols: models.JSON(protocols), Models: models.JSON(encodedModels), Config: models.JSON(`{}`), + Enabled: true, + } + if err := tx.Create(&group).Error; err != nil { + return models.Group{}, false, app_errors.ParseDBError(err) + } + return group, true, nil +} + +func appendMissingGroupModels(existing []GroupModel, fetched []string) []GroupModel { + result := append([]GroupModel(nil), existing...) + seen := make(map[string]struct{}, len(existing)+len(fetched)) + for _, model := range existing { + seen[model.ID] = struct{}{} + } + for _, modelID := range fetched { + if _, exists := seen[modelID]; exists { + continue + } + seen[modelID] = struct{}{} + result = append(result, GroupModel{ID: modelID}) + } + return result +} + +func resolveImportGroupName(tx *gorm.DB, normalized normalizedImport) (string, error) { + if normalized.explicitName != nil { + available, err := importGroupNameAvailable(tx, *normalized.explicitName) + if err != nil { + return "", err + } + if !available { + return "", app_errors.ErrDuplicateResource + } + return *normalized.explicitName, nil + } + + for suffix := 1; ; suffix++ { + candidate := normalized.hostname + if suffix > 1 { + candidate = fmt.Sprintf("%s-%d", normalized.hostname, suffix) + } + available, err := importGroupNameAvailable(tx, candidate) + if err != nil { + return "", err + } + if available { + return candidate, nil + } + } +} + +func importGroupNameAvailable(tx *gorm.DB, name string) (bool, error) { + var count int64 + if err := tx.Model(&models.Group{}).Where("name = ?", name).Count(&count).Error; err != nil { + return false, app_errors.ParseDBError(err) + } + return count == 0, nil +} + +func stableProtocolUnion(existing, requested []protocol.Protocol) []protocol.Protocol { + result := append([]protocol.Protocol(nil), existing...) + seen := make(map[protocol.Protocol]struct{}, len(existing)+len(requested)) + for _, value := range existing { + seen[value] = struct{}{} + } + for _, value := range requested { + if _, exists := seen[value]; exists { + continue + } + seen[value] = struct{}{} + result = append(result, value) + } + return result +} + +func (s *Service) persistImportKeys( + tx *gorm.DB, + groupID uint, + normalized normalizedImport, +) ([]state.KeyEntry, int, int, error) { + hashes := make([]string, 0, len(normalized.keys)) + for _, candidate := range normalized.keys { + hashes = append(hashes, candidate.hash) + } + var existingRows []models.UpstreamKey + if err := tx.Where("group_id = ? AND key_hash IN ?", groupID, hashes).Find(&existingRows).Error; err != nil { + return nil, 0, 0, app_errors.ParseDBError(err) + } + existingByHash := make(map[string]models.UpstreamKey, len(existingRows)) + for _, row := range existingRows { + existingByHash[row.KeyHash] = row + } + + entries := make([]state.KeyEntry, 0, len(normalized.keys)) + added := 0 + duplicated := normalized.duplicateLines + for _, candidate := range normalized.keys { + row, exists := existingByHash[candidate.hash] + if exists { + duplicated++ + } else { + ciphertext, err := s.encryption.Encrypt(candidate.plaintext) + if err != nil { + return nil, 0, 0, fmt.Errorf("encrypt imported key: %w", err) + } + row = models.UpstreamKey{ + GroupID: groupID, KeyValue: ciphertext, KeyHash: candidate.hash, + Status: models.UpstreamKeyStatusActive, + } + if err := tx.Create(&row).Error; err != nil { + return nil, 0, 0, app_errors.ParseDBError(err) + } + added++ + } + entries = append(entries, state.KeyEntry{ + ID: row.ID, GroupID: row.GroupID, WeightManual: row.WeightManual, + Status: state.KeyStatus(row.Status), EncryptedValue: row.KeyValue, + }) + } + return entries, added, duplicated, nil +} + +func decodeImportModelIDs(raw models.JSON) ([]string, error) { + var stored []GroupModel + if err := json.Unmarshal(raw, &stored); err != nil { + return nil, fmt.Errorf("decode imported group models: %w", err) + } + result := make([]string, 0, len(stored)) + seen := make(map[string]struct{}, len(stored)) + for _, model := range stored { + if _, exists := seen[model.ID]; exists { + continue + } + seen[model.ID] = struct{}{} + result = append(result, model.ID) + } + return result, nil +} + +func isLiteralPrivateHost(hostname string) bool { + normalized := strings.TrimSuffix(strings.ToLower(strings.TrimSpace(hostname)), ".") + if normalized == "localhost" { + return true + } + ip := net.ParseIP(strings.Trim(normalized, "[]")) + if ip == nil { + return false + } + return ip.IsPrivate() || ip.IsLoopback() || ip.IsLinkLocalUnicast() || + ip.IsLinkLocalMulticast() || ip.IsUnspecified() +} diff --git a/internal/control/import_test.go b/internal/control/import_test.go new file mode 100644 index 00000000..1ad6d723 --- /dev/null +++ b/internal/control/import_test.go @@ -0,0 +1,698 @@ +package control + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" + + "gpt-load/internal/dialect" + "gpt-load/internal/platform/config" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/protocol" + stateloader "gpt-load/internal/state/loader" + "gpt-load/internal/storage/models" +) + +func TestNormalizeUpstreamURL(t *testing.T) { + tests := []struct { + name string + input string + want string + wantError bool + }{ + {name: "trim lower host and slash", input: " HTTPS://API.Example.COM/v1/ ", want: "https://api.example.com/v1"}, + {name: "preserve fixed query", input: "https://API.Example.com/base/?api-version=2026-01", want: "https://api.example.com/base?api-version=2026-01"}, + {name: "preserve internal double slash", input: "https://api.example.com/a//b/", want: "https://api.example.com/a//b"}, + {name: "ipv6 without port", input: "https://[::1]/v1/", want: "https://[::1]/v1"}, + {name: "ipv6 with port", input: "http://[2001:db8::1]:8080/v1/", want: "http://[2001:db8::1]:8080/v1"}, + {name: "reject relative", input: "/v1", wantError: true}, + {name: "reject ftp", input: "ftp://api.example.com", wantError: true}, + {name: "reject userinfo", input: "https://user:pass@api.example.com", wantError: true}, + {name: "reject fragment", input: "https://api.example.com/#secret", wantError: true}, + {name: "reject missing host", input: "https:///v1", wantError: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, signature, _, err := normalizeUpstreamURL(test.input) + if test.wantError { + if err == nil { + t.Fatalf("normalizeUpstreamURL(%q) error = nil", test.input) + } + return + } + if err != nil { + t.Fatalf("normalizeUpstreamURL(%q) error = %v", test.input, err) + } + if got != test.want { + t.Fatalf("normalizeUpstreamURL(%q) = %q, want %q", test.input, got, test.want) + } + wantHash := sha256.Sum256([]byte(test.want)) + if signature != hex.EncodeToString(wantHash[:]) || len(signature) != 64 { + t.Fatalf("signature = %q, want SHA-256 of %q", signature, test.want) + } + }) + } + + _, firstSignature, _, err := normalizeUpstreamURL("https://API.Example.com/v1/") + if err != nil { + t.Fatal(err) + } + _, secondSignature, _, err := normalizeUpstreamURL(" https://api.example.com/v1 ") + if err != nil { + t.Fatal(err) + } + if firstSignature != secondSignature { + t.Fatalf("equivalent signatures differ: %q != %q", firstSignature, secondSignature) + } + _, otherQuerySignature, _, err := normalizeUpstreamURL("https://api.example.com/v1?tenant=other") + if err != nil { + t.Fatal(err) + } + if firstSignature == otherQuerySignature { + t.Fatal("different fixed query produced the same signature") + } +} + +func TestNormalizeImportInput(t *testing.T) { + fixture := newServiceFixture(t) + request := ImportRequest{ + UpstreamURL: "https://api.example.com/", + Protocols: []protocol.Protocol{protocol.OpenAI, protocol.OpenAI, protocol.Anthropic}, + Keys: " sk-a \n\n sk-a\nsk-b\n", + } + got, err := fixture.service.normalizeImportInput(request) + if err != nil { + t.Fatalf("normalizeImportInput() error = %v", err) + } + if len(got.protocols) != 2 || got.protocols[0] != protocol.OpenAI || got.protocols[1] != protocol.Anthropic { + t.Fatalf("protocols = %#v, want stable unique openai/anthropic", got.protocols) + } + if len(got.keys) != 2 || got.keys[0].plaintext != "sk-a" || got.keys[1].plaintext != "sk-b" { + t.Fatalf("keys = %#v, want two stable unique candidates", got.keys) + } + if got.keys[0].hash == "" || got.keys[0].hash == got.keys[1].hash { + t.Fatalf("key hashes = %#v, want distinct non-empty HMACs", got.keys) + } + if got.firstKey != "sk-a" || got.duplicateLines != 1 { + t.Fatalf("firstKey/duplicates = %q/%d, want sk-a/1", got.firstKey, got.duplicateLines) + } +} + +func TestNormalizeImportInputRejectsInvalidRequests(t *testing.T) { + fixture := newServiceFixture(t) + valid := ImportRequest{ + UpstreamURL: "https://api.example.com", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Keys: "sk-valid", + } + blank := " " + tooLong := strings.Repeat("名", 86) + controlName := "bad\nname" + tests := []struct { + name string + mutate func(*ImportRequest) + }{ + {name: "empty protocols", mutate: func(request *ImportRequest) { request.Protocols = nil }}, + {name: "unknown protocol", mutate: func(request *ImportRequest) { request.Protocols = []protocol.Protocol{"unknown"} }}, + {name: "reserved response protocol", mutate: func(request *ImportRequest) { request.Protocols = []protocol.Protocol{protocol.OpenAIResponse} }}, + {name: "empty keys", mutate: func(request *ImportRequest) { request.Keys = " \n\t" }}, + {name: "too many non-empty lines", mutate: func(request *ImportRequest) { request.Keys = importLines(1001) }}, + {name: "blank explicit name", mutate: func(request *ImportRequest) { request.Name = &blank }}, + {name: "name over 255 bytes", mutate: func(request *ImportRequest) { request.Name = &tooLong }}, + {name: "control character in name", mutate: func(request *ImportRequest) { request.Name = &controlName }}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + request := valid + test.mutate(&request) + if _, err := fixture.service.normalizeImportInput(request); err == nil { + t.Fatal("normalizeImportInput() error = nil") + } + }) + } +} + +func TestNormalizeImportInputAllowsExactlyOneThousandNonEmptyLines(t *testing.T) { + fixture := newServiceFixture(t) + got, err := fixture.service.normalizeImportInput(ImportRequest{ + UpstreamURL: "https://api.example.com", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Keys: importLines(1000), + }) + if err != nil { + t.Fatalf("normalizeImportInput() error = %v", err) + } + if len(got.keys)+got.duplicateLines != 1000 { + t.Fatalf("non-empty line count = %d, want 1000", len(got.keys)+got.duplicateLines) + } +} + +func TestPrivateImportHostUsesOnlyLiteralHostClassification(t *testing.T) { + tests := []struct { + host string + want bool + }{ + {host: "localhost", want: true}, + {host: "LOCALHOST.", want: true}, + {host: "127.0.0.1", want: true}, + {host: "10.1.2.3", want: true}, + {host: "172.16.0.1", want: true}, + {host: "192.168.1.1", want: true}, + {host: "::1", want: true}, + {host: "fe80::1", want: true}, + {host: "0.0.0.0", want: true}, + {host: "8.8.8.8", want: false}, + {host: "api.example.com", want: false}, + } + for _, test := range tests { + t.Run(test.host, func(t *testing.T) { + if got := isLiteralPrivateHost(test.host); got != test.want { + t.Fatalf("isLiteralPrivateHost(%q) = %t, want %t", test.host, got, test.want) + } + }) + } +} + +func TestImportNormalizesGroupsAndCountsDuplicates(t *testing.T) { + fixture := newServiceFixture(t) + before := fixture.manager.Current().Revision + first, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: "HTTPS://API.Example.COM/v1/", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Keys: "sk-a\nsk-a\nsk-b", + }) + if err != nil { + t.Fatalf("first Import() error = %v", err) + } + if !first.Created || first.GroupName != "api.example.com" || first.KeysAdded != 2 || first.KeysDuplicated != 1 { + t.Fatalf("first result = %#v", first) + } + if first.ModelsFetched || len(first.Models) != 0 { + t.Fatalf("first model result = %#v, want unfetched empty models", first) + } + if got := fixture.manager.Current().Revision; got != before+1 { + t.Fatalf("first revision = %d, want %d", got, before+1) + } + + rename := "ignored-new-name" + second, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: "https://api.example.com/v1", + Protocols: []protocol.Protocol{protocol.Anthropic}, + Name: &rename, + Keys: "sk-a\nsk-b", + }) + if err != nil { + t.Fatalf("second Import() error = %v", err) + } + if second.Created || second.GroupID != first.GroupID || second.GroupName != first.GroupName || + second.KeysAdded != 0 || second.KeysDuplicated != 2 { + t.Fatalf("second result = %#v", second) + } + if got := fixture.manager.Current().Revision; got != before+2 { + t.Fatalf("second revision = %d, want %d", got, before+2) + } + assertGroupCount(t, fixture.db, 1) + + var group models.Group + if err := fixture.db.First(&group, first.GroupID).Error; err != nil { + t.Fatalf("query imported group: %v", err) + } + var protocols []protocol.Protocol + if err := json.Unmarshal(group.Protocols, &protocols); err != nil { + t.Fatalf("decode protocols: %v", err) + } + if len(protocols) != 2 || protocols[0] != protocol.OpenAI || protocols[1] != protocol.Anthropic { + t.Fatalf("stored protocols = %#v, want stable union", protocols) + } + var keyCount int64 + if err := fixture.db.Model(&models.UpstreamKey{}).Count(&keyCount).Error; err != nil { + t.Fatalf("count upstream keys: %v", err) + } + if keyCount != 2 { + t.Fatalf("upstream key count = %d, want 2", keyCount) + } +} + +func TestImportPersistsOnlyEncryptedKeysAndPublishesRegistryEntries(t *testing.T) { + fixture := newServiceFixture(t) + result, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: "https://secure.example.com/v1/", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Keys: "sk-secret-a\nsk-secret-b", + }) + if err != nil { + t.Fatalf("Import() error = %v", err) + } + + var group models.Group + if err := fixture.db.First(&group, result.GroupID).Error; err != nil { + t.Fatalf("query group: %v", err) + } + if group.UpstreamURL != "https://secure.example.com/v1" || string(group.Config) != "{}" || string(group.Models) != "[]" || !group.Enabled { + t.Fatalf("stored group = %#v", group) + } + var rows []models.UpstreamKey + if err := fixture.db.Order("id ASC").Find(&rows).Error; err != nil { + t.Fatalf("query upstream keys: %v", err) + } + if len(rows) != 2 { + t.Fatalf("stored keys = %#v, want two", rows) + } + for index, plaintext := range []string{"sk-secret-a", "sk-secret-b"} { + row := rows[index] + if row.KeyValue == "" || row.KeyHash == "" || strings.Contains(row.KeyValue, plaintext) || strings.Contains(row.KeyHash, plaintext) { + t.Fatalf("stored key %d exposes plaintext: %#v", row.ID, row) + } + decrypted, err := fixture.encryption.Decrypt(row.KeyValue) + if err != nil || decrypted != plaintext { + t.Fatalf("Decrypt(key %d) = %q, %v, want %q", row.ID, decrypted, err, plaintext) + } + if got, ok := fixture.registry.EncryptedValue(row.ID); !ok || got != row.KeyValue { + t.Fatalf("Registry key %d = %q, %t, want persisted cipher", row.ID, got, ok) + } + } +} + +func TestImportRepairsMissingRegistryKeyWithoutResettingExistingRuntime(t *testing.T) { + fixture := newServiceFixture(t) + request := ImportRequest{ + UpstreamURL: "https://repair.example.com", + Protocols: []protocol.Protocol{protocol.OpenAI}, + Keys: "sk-repair", + } + if _, err := fixture.service.Import(context.Background(), request); err != nil { + t.Fatalf("initial Import() error = %v", err) + } + var row models.UpstreamKey + if err := fixture.db.First(&row).Error; err != nil { + t.Fatalf("query imported key: %v", err) + } + if count, ok := fixture.registry.IncrFailure(row.ID); !ok || count != 1 { + t.Fatalf("first failure count = %d, %t", count, ok) + } + if _, err := fixture.service.Import(context.Background(), request); err != nil { + t.Fatalf("present-key reimport error = %v", err) + } + if count, ok := fixture.registry.IncrFailure(row.ID); !ok || count != 2 { + t.Fatalf("second failure count = %d, %t, want preserved runtime count 2", count, ok) + } + + if !fixture.registry.RemoveKey(row.ID) { + t.Fatal("RemoveKey() = false, want simulated missing entry") + } + result, err := fixture.service.Import(context.Background(), request) + if err != nil { + t.Fatalf("repair Import() error = %v", err) + } + if result.KeysAdded != 0 || result.KeysDuplicated != 1 { + t.Fatalf("repair result = %#v, want added=0 duplicated=1", result) + } + if got, ok := fixture.registry.EncryptedValue(row.ID); !ok || got != row.KeyValue { + t.Fatalf("repaired Registry value = %q, %t, want %q", got, ok, row.KeyValue) + } +} + +func TestImportRejectsExplicitNameCollisionAndPreservesExistingName(t *testing.T) { + fixture := newServiceFixture(t) + name := "shared-name" + first, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: "https://first.example.com", Protocols: []protocol.Protocol{protocol.OpenAI}, + Name: &name, Keys: "sk-first", + }) + if err != nil { + t.Fatalf("first Import() error = %v", err) + } + beforeConflict := fixture.manager.Current().Revision + _, err = fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: "https://second.example.com", Protocols: []protocol.Protocol{protocol.OpenAI}, + Name: &name, Keys: "sk-second", + }) + var apiErr *app_errors.APIError + if !errors.As(err, &apiErr) || apiErr.HTTPStatus != http.StatusConflict { + t.Fatalf("name collision error = %#v, want typed 409", err) + } + assertGroupCount(t, fixture.db, 1) + if got := fixture.manager.Current().Revision; got != beforeConflict { + t.Fatalf("revision after collision = %d, want %d", got, beforeConflict) + } + + rename := "ignored-rename" + second, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: "https://first.example.com", Protocols: []protocol.Protocol{protocol.Anthropic}, + Name: &rename, Keys: "sk-first", + }) + if err != nil { + t.Fatalf("same-signature Import() error = %v", err) + } + if second.GroupID != first.GroupID || second.GroupName != name { + t.Fatalf("same-signature result = %#v, want original name %q", second, name) + } +} + +func TestImportHandlerRejectsInvalidInputWithoutMutation(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().Revision + + tests := []struct { + name string + body []byte + }{ + {name: "invalid json", body: []byte(`{"upstream_url":`)}, + {name: "reserved protocol", body: mustJSON(t, ImportRequest{ + UpstreamURL: "https://api.example.com", Protocols: []protocol.Protocol{protocol.OpenAIResponse}, Keys: "sk-one", + })}, + {name: "too many keys", body: mustJSON(t, ImportRequest{ + UpstreamURL: "https://api.example.com", Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: importLines(1001), + })}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/api/import", bytes.NewReader(test.body)) + request.Header.Set("Authorization", "Bearer test-auth-key") + request.Header.Set("Content-Type", "application/json") + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusBadRequest { + t.Fatalf("status = %d, want 400; body=%s", recorder.Code, recorder.Body.String()) + } + }) + } + assertGroupCount(t, fixture.db, 0) + if got := fixture.manager.Current().Revision; got != before { + t.Fatalf("revision = %d, want %d", got, before) + } +} + +func TestImportPrivateHostWarningDoesNotExposeURLOrCredentials(t *testing.T) { + fixture := newServiceFixture(t) + const ( + rawURL = "http://127.0.0.1/base?tenant_secret=known-query-secret" + plaintext = "known-upstream-key" + ) + var logs bytes.Buffer + previousOutput := logrus.StandardLogger().Out + logrus.SetOutput(&logs) + t.Cleanup(func() { logrus.SetOutput(previousOutput) }) + + result, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: rawURL, Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: plaintext, + }) + if err != nil { + t.Fatalf("Import() error = %v", err) + } + var row models.UpstreamKey + if err := fixture.db.First(&row).Error; err != nil { + t.Fatalf("query imported key: %v", err) + } + var group models.Group + if err := fixture.db.First(&group, result.GroupID).Error; err != nil { + t.Fatalf("query imported group: %v", err) + } + if !strings.Contains(logs.String(), "127.0.0.1") || !strings.Contains(logs.String(), group.Signature) { + t.Fatalf("private warning missing safe host/signature metadata: %s", logs.String()) + } + for _, forbidden := range []string{rawURL, "known-query-secret", plaintext, row.KeyValue, row.KeyHash} { + if strings.Contains(logs.String(), forbidden) { + t.Fatalf("private warning exposes %q: %s", forbidden, logs.String()) + } + } +} + +func TestImportFetchesModelsWithEmptyRules(t *testing.T) { + requests := make(chan http.Header, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + requests <- request.Header.Clone() + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"object":"list","data":[{"id":"gpt-4o"},{"id":"gpt-4.1"}]}`)) + })) + t.Cleanup(upstream.Close) + + fixture := newServiceFixture(t) + fixture.service.dialects = dialectSetForOpenAI(upstream.Client()) + result, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: upstream.URL, Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-model-list", + }) + if err != nil { + t.Fatalf("Import() error = %v", err) + } + if !result.ModelsFetched || fmt.Sprint(result.Models) != "[gpt-4o gpt-4.1]" { + t.Fatalf("model result = %#v", result) + } + headers := <-requests + if got := headers.Get("Authorization"); got != "Bearer sk-model-list" { + t.Fatalf("Authorization = %q, want default Bearer credential", got) + } + + var group models.Group + if err := fixture.db.First(&group, result.GroupID).Error; err != nil { + t.Fatalf("query imported group: %v", err) + } + if string(group.Models) != `[{"id":"gpt-4o","alias":""},{"id":"gpt-4.1","alias":""}]` { + t.Fatalf("stored models = %s", group.Models) + } +} + +func TestImportFetchesModelsWithExistingHeaderRulesAndStableMerge(t *testing.T) { + headersSeen := make(chan http.Header, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + headersSeen <- request.Header.Clone() + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"object":"list","data":[{"id":"existing"},{"id":"new"},{"id":"new"}]}`)) + })) + t.Cleanup(upstream.Close) + + fixture := newServiceFixture(t) + initial, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: upstream.URL, Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-initial", + }) + if err != nil { + t.Fatalf("initial Import() error = %v", err) + } + if err := fixture.db.Model(&models.Group{}).Where("id = ?", initial.GroupID).Updates(map[string]any{ + "models": models.JSON(`[{"id":"existing","alias":"Keep"}]`), + "config": models.JSON(`{"header_rules":{"set":{"Authorization":"Token ${API_KEY}","X-Custom":"key=${API_KEY}"}}}`), + }).Error; err != nil { + t.Fatalf("seed group runtime config: %v", err) + } + input, err := stateloader.BuildCompileInput(context.Background(), fixture.db) + if err != nil { + t.Fatalf("BuildCompileInput() error = %v", err) + } + if _, err := fixture.manager.Publish(input); err != nil { + t.Fatalf("Publish() error = %v", err) + } + fixture.service.dialects = dialectSetForOpenAI(upstream.Client()) + + result, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: upstream.URL, Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-new", + }) + if err != nil { + t.Fatalf("merged Import() error = %v", err) + } + if !result.ModelsFetched || fmt.Sprint(result.Models) != "[existing new]" { + t.Fatalf("merged model result = %#v", result) + } + headers := <-headersSeen + if headers.Get("Authorization") != "Token sk-new" || headers.Get("X-Custom") != "key=sk-new" { + t.Fatalf("model-list headers = %#v", headers) + } + + var group models.Group + if err := fixture.db.First(&group, initial.GroupID).Error; err != nil { + t.Fatalf("query merged group: %v", err) + } + var stored []GroupModel + if err := json.Unmarshal(group.Models, &stored); err != nil { + t.Fatalf("decode merged models: %v", err) + } + if len(stored) != 2 || stored[0].ID != "existing" || stored[0].Alias != "Keep" || stored[1].ID != "new" { + t.Fatalf("stored merged models = %#v", stored) + } +} + +func TestImportFetchesModelsTreatsUpstreamFailuresAsBestEffort(t *testing.T) { + tests := []struct { + name string + status int + body string + waitForCancel bool + wantFetched bool + }{ + {name: "empty list", status: http.StatusOK, body: `{"object":"list","data":[]}`, wantFetched: true}, + {name: "upstream error", status: http.StatusInternalServerError, body: `{"error":"failed"}`}, + {name: "malformed json", status: http.StatusOK, body: `{"data":`}, + {name: "blank model", status: http.StatusOK, body: `{"data":[{"id":""}]}`}, + {name: "boundary whitespace", status: http.StatusOK, body: `{"data":[{"id":" gpt-4o "}]}`}, + {name: "timeout", waitForCancel: true}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + if test.waitForCancel { + <-request.Context().Done() + return + } + writer.WriteHeader(test.status) + _, _ = writer.Write([]byte(test.body)) + })) + t.Cleanup(upstream.Close) + + fixture := newServiceFixture(t) + fixture.service.dialects = dialectSetForOpenAI(upstream.Client()) + if test.waitForCancel { + fixture.service.modelFetchTimeout = 20 * time.Millisecond + } + result, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: upstream.URL, Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-best-effort", + }) + if err != nil { + t.Fatalf("Import() error = %v", err) + } + if result.ModelsFetched != test.wantFetched || len(result.Models) != 0 { + t.Fatalf("result = %#v, want fetched=%t and empty models", result, test.wantFetched) + } + assertGroupCount(t, fixture.db, 1) + }) + } +} + +func TestImportFetchesModelsSkipsUnregisteredProtocols(t *testing.T) { + called := make(chan struct{}, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + called <- struct{}{} + })) + t.Cleanup(upstream.Close) + fixture := newServiceFixture(t) + fixture.service.dialects = dialectSetForOpenAI(upstream.Client()) + + result, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: upstream.URL, Protocols: []protocol.Protocol{protocol.Anthropic, protocol.Gemini}, Keys: "sk-no-lister", + }) + if err != nil { + t.Fatalf("Import() error = %v", err) + } + if result.ModelsFetched { + t.Fatalf("result = %#v, want models_fetched=false", result) + } + select { + case <-called: + t.Fatal("unregistered protocol triggered an HTTP model-list call") + default: + } +} + +func TestImportFetchesModelsStopsOnParentCancellation(t *testing.T) { + called := make(chan struct{}, 1) + upstream := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) { + called <- struct{}{} + })) + t.Cleanup(upstream.Close) + fixture := newServiceFixture(t) + fixture.service.dialects = dialectSetForOpenAI(upstream.Client()) + before := fixture.manager.Current().Revision + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := fixture.service.Import(ctx, ImportRequest{ + UpstreamURL: upstream.URL, Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-canceled", + }) + if !errors.Is(err, context.Canceled) { + t.Fatalf("Import() error = %v, want context.Canceled", err) + } + assertGroupCount(t, fixture.db, 0) + if got := fixture.manager.Current().Revision; got != before { + t.Fatalf("revision = %d, want %d", got, before) + } + select { + case <-called: + t.Fatal("canceled parent triggered a model-list call") + default: + } +} + +func TestImportModelFetchDoesNotHoldWriteLock(t *testing.T) { + entered := make(chan struct{}) + release := make(chan struct{}) + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, _ *http.Request) { + close(entered) + <-release + _, _ = writer.Write([]byte(`{"object":"list","data":[{"id":"gpt-4o"}]}`)) + })) + t.Cleanup(upstream.Close) + + fixture := newServiceFixture(t) + fixture.service.dialects = dialectSetForOpenAI(upstream.Client()) + firstDone := make(chan error, 1) + go func() { + _, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: upstream.URL, Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-first", + }) + firstDone <- err + }() + select { + case <-entered: + case <-time.After(2 * time.Second): + t.Fatal("first import did not enter ListModels") + } + + secondDone := make(chan error, 1) + go func() { + _, err := fixture.service.Import(context.Background(), ImportRequest{ + UpstreamURL: "https://second.example.com", Protocols: []protocol.Protocol{protocol.Anthropic}, Keys: "sk-second", + }) + secondDone <- err + }() + select { + case err := <-secondDone: + if err != nil { + t.Fatalf("second Import() error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("second control write blocked behind model fetch") + } + close(release) + select { + case err := <-firstDone: + if err != nil { + t.Fatalf("first Import() error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("first import did not finish after model fetch release") + } +} + +func dialectSetForOpenAI(client *http.Client) dialect.Set { + return dialect.NewSet(dialect.NewOpenAI(client)) +} + +func mustJSON(t *testing.T, value any) []byte { + t.Helper() + encoded, err := json.Marshal(value) + if err != nil { + t.Fatalf("json.Marshal(%T) error = %v", value, err) + } + return encoded +} + +func importLines(count int) string { + var builder strings.Builder + for index := range count { + fmt.Fprintf(&builder, "sk-%d\n", index) + } + return builder.String() +} diff --git a/internal/control/server.go b/internal/control/server.go new file mode 100644 index 00000000..8b8846fc --- /dev/null +++ b/internal/control/server.go @@ -0,0 +1,203 @@ +package control + +import ( + "crypto/sha256" + "crypto/subtle" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "strconv" + "strings" + + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" + + "gpt-load/internal/platform/config" + app_errors "gpt-load/internal/platform/errors" + "gpt-load/internal/platform/i18n" + "gpt-load/internal/platform/response" +) + +type Server struct { + authDigest [sha256.Size]byte + service *Service +} + +func NewServer(cfg *config.Config, service *Service) *Server { + return &Server{authDigest: sha256.Sum256([]byte(cfg.AuthKey)), service: service} +} + +func (s *Server) RegisterRoutes(engine *gin.Engine) { + api := engine.Group("/api") + api.Use(i18n.Middleware(), s.authenticate()) + api.GET("/groups", s.handleListGroups) + api.POST("/import", s.handleImport) + api.POST("/access-keys", s.handleCreateAccessKey) + api.GET("/access-keys", s.handleListAccessKeys) + api.PUT("/access-keys/:id", s.handleUpdateAccessKey) + api.DELETE("/access-keys/:id", s.handleDeleteAccessKey) +} + +func (s *Server) authenticate() gin.HandlerFunc { + return func(c *gin.Context) { + fields := strings.Fields(c.GetHeader("Authorization")) + token := "" + if len(fields) == 2 { + token = fields[1] + } + requestDigest := sha256.Sum256([]byte(token)) + formatValid := len(fields) == 2 && strings.EqualFold(fields[0], "Bearer") + matches := subtle.ConstantTimeCompare(requestDigest[:], s.authDigest[:]) == 1 + if !formatValid || !matches { + response.ErrorI18nFromAPIError(c, app_errors.ErrUnauthorized, "auth.invalid_key") + c.Abort() + return + } + c.Next() + } +} + +func (s *Server) handleListGroups(c *gin.Context) { + groups, err := s.service.ListGroups(c.Request.Context()) + if err != nil { + writeServiceError(c, "list_groups", err) + return + } + response.SuccessI18n(c, "common.success", groups) +} + +func (s *Server) handleImport(c *gin.Context) { + var request ImportRequest + if err := bindStrictJSON(c, &request); err != nil { + writeServiceError(c, "import", app_errors.ErrInvalidJSON) + return + } + result, err := s.service.Import(c.Request.Context(), request) + if err != nil { + writeServiceError(c, "import", err) + return + } + response.SuccessI18n(c, "common.success", result) +} + +func (s *Server) handleCreateAccessKey(c *gin.Context) { + var request AccessKeyCreateRequest + if err := bindStrictJSON(c, &request); err != nil { + writeServiceError(c, "create_access_key", app_errors.ErrInvalidJSON) + return + } + result, err := s.service.CreateAccessKey(c.Request.Context(), request) + if err != nil { + writeServiceError(c, "create_access_key", err) + return + } + response.SuccessI18n(c, "common.success", result) +} + +func (s *Server) handleListAccessKeys(c *gin.Context) { + result, err := s.service.ListAccessKeys(c.Request.Context()) + if err != nil { + writeServiceError(c, "list_access_keys", err) + return + } + response.SuccessI18n(c, "common.success", result) +} + +func (s *Server) handleUpdateAccessKey(c *gin.Context) { + id, ok := accessKeyID(c) + if !ok { + return + } + var request AccessKeyUpdateRequest + if err := bindStrictJSON(c, &request); err != nil { + writeServiceError(c, "update_access_key", app_errors.ErrInvalidJSON) + return + } + result, err := s.service.UpdateAccessKey(c.Request.Context(), id, request) + if err != nil { + writeServiceError(c, "update_access_key", err) + return + } + response.SuccessI18n(c, "common.success", result) +} + +func (s *Server) handleDeleteAccessKey(c *gin.Context) { + id, ok := accessKeyID(c) + if !ok { + return + } + if err := s.service.DeleteAccessKey(c.Request.Context(), id); err != nil { + writeServiceError(c, "delete_access_key", err) + return + } + response.SuccessI18n(c, "common.success", nil) +} + +func bindStrictJSON(c *gin.Context, target any) error { + decoder := json.NewDecoder(c.Request.Body) + decoder.DisallowUnknownFields() + if err := decoder.Decode(target); err != nil { + return err + } + + var extra json.RawMessage + if err := decoder.Decode(&extra); err != io.EOF { + if err == nil { + return fmt.Errorf("decode JSON request: multiple values") + } + return err + } + return nil +} + +func accessKeyID(c *gin.Context) (uint, bool) { + parsed, err := strconv.ParseUint(c.Param("id"), 10, strconv.IntSize) + if err != nil || parsed == 0 { + writeServiceError(c, "access_key_id", app_errors.ErrBadRequest) + return 0, false + } + return uint(parsed), true +} + +func writeServiceError(c *gin.Context, operation string, err error) { + var apiErr *app_errors.APIError + if errors.As(err, &apiErr) { + if apiErr.HTTPStatus >= http.StatusInternalServerError { + logServiceError(operation, err, apiErr.Code) + } + response.ErrorI18nFromAPIError(c, apiErr, serviceErrorMessageID(operation, apiErr)) + return + } + + logServiceError(operation, err, app_errors.ErrInternalServer.Code) + response.ErrorI18nFromAPIError(c, app_errors.ErrInternalServer, "internal_error") +} + +func serviceErrorMessageID(operation string, apiErr *app_errors.APIError) string { + switch apiErr.Code { + case app_errors.ErrBadRequest.Code, app_errors.ErrInvalidJSON.Code, app_errors.ErrValidation.Code: + return "bad_request" + case app_errors.ErrResourceNotFound.Code: + if operation == "list_groups" { + return "group.not_found" + } + return "key.not_found" + case app_errors.ErrDuplicateResource.Code: + if operation == "import" { + return "group.name_exists" + } + return "bad_request" + default: + return "internal_error" + } +} + +func logServiceError(operation string, err error, code string) { + logrus.WithFields(logrus.Fields{ + "operation": operation, + "error_code": code, + "error_type": fmt.Sprintf("%T", err), + }).Error("Control operation failed") +} diff --git a/internal/control/server_test.go b/internal/control/server_test.go new file mode 100644 index 00000000..654f4173 --- /dev/null +++ b/internal/control/server_test.go @@ -0,0 +1,263 @@ +package control + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "net/http" + "net/http/httptest" + "strconv" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" + + "gpt-load/internal/platform/config" + "gpt-load/internal/state" +) + +func TestManagementAuthRequiresConstantShapeBearerToken(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + + tests := []struct { + name string + header string + wantStatus int + }{ + {name: "missing", wantStatus: http.StatusUnauthorized}, + {name: "empty bearer", header: "Bearer", wantStatus: http.StatusUnauthorized}, + {name: "wrong different length", header: "Bearer x", wantStatus: http.StatusUnauthorized}, + {name: "extra field", header: "Bearer test-auth-key extra", wantStatus: http.StatusUnauthorized}, + {name: "wrong scheme", header: "Basic test-auth-key", wantStatus: http.StatusUnauthorized}, + {name: "case insensitive", header: "bEaReR test-auth-key", wantStatus: http.StatusOK}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/api/groups", nil) + if test.header != "" { + request.Header.Set("Authorization", test.header) + } + engine.ServeHTTP(recorder, request) + if recorder.Code != test.wantStatus { + t.Fatalf("status = %d, want %d; body=%s", recorder.Code, test.wantStatus, recorder.Body.String()) + } + if test.wantStatus == http.StatusUnauthorized { + var body struct { + Code string `json:"code"` + } + if err := json.Unmarshal(recorder.Body.Bytes(), &body); err != nil || body.Code != "UNAUTHORIZED" { + t.Fatalf("unauthorized body = %s, error=%v", recorder.Body.String(), err) + } + } + }) + } +} + +func TestManagementAuthLocalizesUnauthorized(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + + for _, test := range []struct { + language string + message string + }{ + {language: "en-US", message: "Invalid authorization key"}, + {language: "zh-CN", message: "无效的授权密钥"}, + } { + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/api/groups", nil) + request.Header.Set("Accept-Language", test.language) + engine.ServeHTTP(recorder, request) + if !strings.Contains(recorder.Body.String(), test.message) { + t.Fatalf("%s body = %s, want message %q", test.language, recorder.Body.String(), test.message) + } + } +} + +func TestManagementAuthDoesNotLogSecretOrDigest(t *testing.T) { + initControlI18n(t) + const authKey = "distinctive-control-auth-key" + fixture := newServiceFixture(t) + engine := gin.New() + NewServer(&config.Config{AuthKey: authKey}, fixture.service).RegisterRoutes(engine) + + var logs bytes.Buffer + previousOutput := logrus.StandardLogger().Out + logrus.SetOutput(&logs) + t.Cleanup(func() { logrus.SetOutput(previousOutput) }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/api/groups", nil) + request.Header.Set("Authorization", "Bearer "+authKey+"-wrong") + engine.ServeHTTP(recorder, request) + + digest := sha256.Sum256([]byte(authKey)) + for _, forbidden := range []string{authKey, hex.EncodeToString(digest[:])} { + if strings.Contains(logs.String(), forbidden) { + t.Fatalf("auth logs expose %q: %s", forbidden, logs.String()) + } + } +} + +func TestManagementWritesRejectUnknownFieldsAndMultipleJSONValues(t *testing.T) { + initControlI18n(t) + + t.Run("import rejects unknown top-level field", func(t *testing.T) { + fixture := newServiceFixture(t) + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + beforeRevision := fixture.manager.Current().Revision + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/api/import", strings.NewReader( + `{"upstream_url":"https://api.example.com","protocols":["openai"],"keys":"sk-test","unexpected":true}`, + )) + request.Header.Set("Authorization", "Bearer test-auth-key") + request.Header.Set("Content-Type", "application/json") + engine.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("POST /api/import = %d %s, want 400", recorder.Code, recorder.Body.String()) + } + assertGroupCount(t, fixture.db, 0) + if got := fixture.manager.Current().Revision; got != beforeRevision { + t.Fatalf("snapshot revision = %d, want unchanged %d", got, beforeRevision) + } + }) + + t.Run("access key rejects unknown nested filter field", func(t *testing.T) { + fixture := newServiceFixture(t) + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + beforeRevision := fixture.manager.Current().Revision + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/api/access-keys", strings.NewReader( + `{"name":"client","filters":{"gropus":[1]}}`, + )) + request.Header.Set("Authorization", "Bearer test-auth-key") + request.Header.Set("Content-Type", "application/json") + engine.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("POST /api/access-keys = %d %s, want 400", recorder.Code, recorder.Body.String()) + } + var count int64 + if err := fixture.db.Table("access_keys").Count(&count).Error; err != nil { + t.Fatalf("count access keys: %v", err) + } + if count != 0 { + t.Fatalf("access key count = %d, want 0", count) + } + if got := fixture.manager.Current().Revision; got != beforeRevision { + t.Fatalf("snapshot revision = %d, want unchanged %d", got, beforeRevision) + } + }) + + t.Run("access key update rejects a second JSON value", func(t *testing.T) { + fixture := newServiceFixture(t) + fixture.service.random = bytes.NewReader(make([]byte, 16)) + created, err := fixture.service.CreateAccessKey(t.Context(), AccessKeyCreateRequest{Name: "original"}) + if err != nil { + t.Fatalf("CreateAccessKey() error = %v", err) + } + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + beforeRevision := fixture.manager.Current().Revision + + recorder := httptest.NewRecorder() + request := httptest.NewRequest( + http.MethodPut, + "/api/access-keys/"+strconv.FormatUint(uint64(created.ID), 10), + strings.NewReader(`{"name":"changed"}{"status":"disabled"}`), + ) + request.Header.Set("Authorization", "Bearer test-auth-key") + request.Header.Set("Content-Type", "application/json") + engine.ServeHTTP(recorder, request) + + if recorder.Code != http.StatusBadRequest { + t.Fatalf("PUT /api/access-keys/:id = %d %s, want 400", recorder.Code, recorder.Body.String()) + } + row := loadAccessKeyRow(t, fixture.db, created.ID) + if row.Name != "original" || row.Status != string(state.AccessKeyStatusActive) { + t.Fatalf("access key row = %#v, want original active row", row) + } + if got := fixture.manager.Current().Revision; got != beforeRevision { + t.Fatalf("snapshot revision = %d, want unchanged %d", got, beforeRevision) + } + }) +} + +func TestUpdateAccessKeyRoutesParseIDsAndPreservePointerSemantics(t *testing.T) { + initControlI18n(t) + fixture := newServiceFixture(t) + fixture.service.random = bytes.NewReader(make([]byte, 16)) + created, err := fixture.service.CreateAccessKey(t.Context(), AccessKeyCreateRequest{Name: "route-key"}) + if err != nil { + t.Fatalf("CreateAccessKey() error = %v", err) + } + engine := gin.New() + NewServer(&config.Config{AuthKey: "test-auth-key"}, fixture.service).RegisterRoutes(engine) + + request := httptest.NewRequest(http.MethodPut, "/api/access-keys/"+strconv.FormatUint(uint64(created.ID), 10), strings.NewReader(`{"status":"disabled"}`)) + request.Header.Set("Authorization", "Bearer test-auth-key") + 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) { + t.Fatalf("PUT access key = %d %s", recorder.Code, recorder.Body.String()) + } + row := loadAccessKeyRow(t, fixture.db, created.ID) + if row.Name != "route-key" || row.Status != string(state.AccessKeyStatusDisabled) { + t.Fatalf("PUT row = %#v", row) + } + + for _, path := range []string{ + "/api/access-keys/0", + "/api/access-keys/-1", + "/api/access-keys/not-a-number", + "/api/access-keys/18446744073709551616", + } { + request = httptest.NewRequest(http.MethodPut, path, strings.NewReader(`{"name":"ignored"}`)) + request.Header.Set("Authorization", "Bearer test-auth-key") + request.Header.Set("Content-Type", "application/json") + recorder = httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusBadRequest { + t.Fatalf("PUT %s = %d %s, want 400", path, recorder.Code, recorder.Body.String()) + } + } + + request = httptest.NewRequest(http.MethodPut, "/api/access-keys/"+strconv.FormatUint(uint64(created.ID), 10), strings.NewReader(`{}`)) + request.Header.Set("Authorization", "Bearer test-auth-key") + request.Header.Set("Content-Type", "application/json") + recorder = httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusBadRequest { + t.Fatalf("empty PUT = %d %s, want 400", recorder.Code, recorder.Body.String()) + } + + request = httptest.NewRequest(http.MethodDelete, "/api/access-keys/"+strconv.FormatUint(uint64(created.ID), 10), nil) + request.Header.Set("Authorization", "Bearer test-auth-key") + recorder = httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusOK { + t.Fatalf("DELETE access key = %d %s", recorder.Code, recorder.Body.String()) + } + + request = httptest.NewRequest(http.MethodDelete, "/api/access-keys/"+strconv.FormatUint(uint64(created.ID), 10), nil) + request.Header.Set("Authorization", "Bearer test-auth-key") + recorder = httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + if recorder.Code != http.StatusNotFound { + t.Fatalf("repeated DELETE = %d %s, want 404", recorder.Code, recorder.Body.String()) + } +} diff --git a/internal/control/service.go b/internal/control/service.go new file mode 100644 index 00000000..5bbdc1c6 --- /dev/null +++ b/internal/control/service.go @@ -0,0 +1,81 @@ +package control + +import ( + "context" + "crypto/rand" + "fmt" + "io" + "sync" + "time" + + "gorm.io/gorm" + + "gpt-load/internal/dialect" + "gpt-load/internal/platform/encryption" + "gpt-load/internal/state" + stateloader "gpt-load/internal/state/loader" +) + +type Service struct { + db *gorm.DB + manager *state.Manager + registry *state.KeyRegistry + encryption encryption.Service + dialects dialect.Set + modelFetchTimeout time.Duration + random io.Reader + writeMu sync.Mutex +} + +func NewService( + db *gorm.DB, + manager *state.Manager, + registry *state.KeyRegistry, + encryptionService encryption.Service, + dialects dialect.Set, +) *Service { + return &Service{ + db: db, manager: manager, registry: registry, + encryption: encryptionService, dialects: dialects, + modelFetchTimeout: 3 * time.Second, random: rand.Reader, + } +} + +func (s *Service) writeConfig( + ctx context.Context, + mutate func(*gorm.DB) error, + afterPublish func() error, +) (*state.ConfigSnapshot, error) { + s.writeMu.Lock() + defer s.writeMu.Unlock() + + var input state.CompileInput + err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + if err := mutate(tx); err != nil { + return err + } + var err error + input, err = stateloader.BuildCompileInput(ctx, tx) + if err != nil { + return err + } + if _, err := state.Compile(input); err != nil { + return err + } + return nil + }) + if err != nil { + return nil, err + } + + snapshot, err := s.manager.Publish(input) + if err != nil { + return nil, fmt.Errorf("publish committed config: %w", err) + } + if afterPublish != nil { + if err := afterPublish(); err != nil { + return nil, fmt.Errorf("apply committed runtime update: %w", err) + } + } + return snapshot, nil +} diff --git a/internal/control/service_integration_test.go b/internal/control/service_integration_test.go new file mode 100644 index 00000000..8014bef1 --- /dev/null +++ b/internal/control/service_integration_test.go @@ -0,0 +1,135 @@ +package control + +import ( + "bytes" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gin-gonic/gin" + + "gpt-load/internal/dialect" + "gpt-load/internal/gateway" + platformhttp "gpt-load/internal/platform/httpclient" + "gpt-load/internal/platform/redact" + "gpt-load/internal/protocol" + "gpt-load/internal/state" +) + +func TestControlWriteLockDoesNotBlockDataPlane(t *testing.T) { + requestReached := make(chan bool, 2) + upstream := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + body, _ := io.ReadAll(request.Body) + streaming := bytes.Contains(body, []byte(`"stream":true`)) + requestReached <- streaming + if streaming { + writer.Header().Set("Content-Type", "text/event-stream") + _, _ = writer.Write([]byte("data: {\"id\":\"chunk\"}\n\n")) + writer.(http.Flusher).Flush() + return + } + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{"id":"response"}`)) + })) + t.Cleanup(upstream.Close) + + fixture := newServiceFixture(t) + const ( + groupID = uint(10) + upstreamKeyID = uint(20) + accessKey = "gl-data-plane-client" + upstreamKey = "sk-data-plane-upstream" + ) + if _, err := fixture.manager.Publish(state.CompileInput{ + Groups: []state.GroupConfig{{ + ID: groupID, Name: "data-plane", UpstreamURL: upstream.URL, + Protocols: []protocol.Protocol{protocol.OpenAI}, + Models: []state.ModelConfig{{ID: "gpt-4o"}}, Enabled: true, + }}, + AccessKeys: []state.AccessKeyConfig{{ + ID: 1, Name: "client", KeyHash: fixture.encryption.Hash(accessKey), + Status: state.AccessKeyStatusActive, + }}, + }); err != nil { + t.Fatalf("Publish(data plane) error = %v", err) + } + ciphertext, err := fixture.encryption.Encrypt(upstreamKey) + if err != nil { + t.Fatalf("Encrypt(upstream key) error = %v", err) + } + if err := fixture.registry.ApplyImport(groupID, []state.KeyEntry{{ + ID: upstreamKeyID, GroupID: groupID, Status: state.KeyStatusActive, EncryptedValue: ciphertext, + }}); err != nil { + t.Fatalf("ApplyImport() error = %v", err) + } + + openAI := dialect.NewOpenAI(http.DefaultClient) + handler := gateway.NewHandler( + fixture.manager, + fixture.registry, + fixture.encryption, + gateway.NewForwarder(platformhttp.NewHTTPClientManager(), redact.New()), + dialect.NewSet(openAI), + ) + gin.SetMode(gin.ReleaseMode) + engine := gin.New() + handler.RegisterRoutes(engine) + + fixture.service.writeMu.Lock() + locked := true + defer func() { + if locked { + fixture.service.writeMu.Unlock() + } + }() + + type response struct { + streaming bool + recorder *httptest.ResponseRecorder + } + responses := make(chan response, 2) + for _, streaming := range []bool{false, true} { + streaming := streaming + go func() { + body := `{"model":"gpt-4o"}` + if streaming { + body = `{"model":"gpt-4o","stream":true}` + } + request := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body)) + request.Header.Set("Authorization", "Bearer "+accessKey) + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + responses <- response{streaming: streaming, recorder: recorder} + }() + } + + for range 2 { + select { + case <-requestReached: + case <-time.After(2 * time.Second): + t.Fatal("data-plane request did not reach upstream while control writeMu was held") + } + } + for range 2 { + select { + case got := <-responses: + if got.recorder.Code != http.StatusOK { + t.Fatalf("streaming=%t response = %d %s", got.streaming, got.recorder.Code, got.recorder.Body.String()) + } + if got.streaming && !bytes.Contains(got.recorder.Body.Bytes(), []byte("data:")) { + t.Fatalf("streaming body = %q, want SSE data", got.recorder.Body.Bytes()) + } + case <-time.After(2 * time.Second): + t.Fatal("data-plane request did not complete while control writeMu was held") + } + } + if got, ok := fixture.registry.ActiveEncryptedValue(upstreamKeyID, groupID); !ok || got != ciphertext { + t.Fatalf("Registry value = %q, %t, want active ciphertext", got, ok) + } + + fixture.service.writeMu.Unlock() + locked = false +} diff --git a/internal/control/service_test.go b/internal/control/service_test.go new file mode 100644 index 00000000..f81d259c --- /dev/null +++ b/internal/control/service_test.go @@ -0,0 +1,159 @@ +package control + +import ( + "context" + "fmt" + "reflect" + "sync" + "testing" + + "gorm.io/gorm" + + "gpt-load/internal/protocol" + "gpt-load/internal/state" + stateloader "gpt-load/internal/state/loader" + "gpt-load/internal/storage/models" +) + +func TestWriteConfigRollsBackWhenCompileRejectsCandidate(t *testing.T) { + fixture := newServiceFixture(t) + before := fixture.manager.Current().Revision + + _, err := fixture.service.writeConfig(context.Background(), func(tx *gorm.DB) error { + return tx.Create(&models.Group{ + Name: "invalid", UpstreamURL: "https://invalid.example", + Signature: "invalid", Protocols: models.JSON(`[]`), + Models: models.JSON(`[]`), Config: models.JSON(`{}`), Enabled: true, + }).Error + }, nil) + if err == nil { + t.Fatal("writeConfig() error = nil, want Compile rejection") + } + assertGroupCount(t, fixture.db, 0) + if got := fixture.manager.Current().Revision; got != before { + t.Fatalf("revision = %d, want %d", got, before) + } +} + +func TestWriteConfigPublishesOnceBeforeCallbackWhileHoldingLock(t *testing.T) { + fixture := newServiceFixture(t) + before := fixture.manager.Current().Revision + callbackRevision := make(chan uint64, 1) + + snapshot, err := fixture.service.writeConfig(context.Background(), func(tx *gorm.DB) error { + return tx.Create(validControlGroup("published")).Error + }, func() error { + if fixture.service.writeMu.TryLock() { + fixture.service.writeMu.Unlock() + return fmt.Errorf("writeMu was not held during afterPublish") + } + callbackRevision <- fixture.manager.Current().Revision + return nil + }) + if err != nil { + t.Fatalf("writeConfig() error = %v", err) + } + if snapshot.Revision != before+1 { + t.Fatalf("snapshot revision = %d, want %d", snapshot.Revision, before+1) + } + if got := <-callbackRevision; got != snapshot.Revision { + t.Fatalf("callback revision = %d, want %d", got, snapshot.Revision) + } + assertGroupCount(t, fixture.db, 1) +} + +func TestWriteConfigSerializesConcurrentDatabaseAndSnapshotPublication(t *testing.T) { + fixture := newServiceFixture(t) + before := fixture.manager.Current().Revision + start := make(chan struct{}) + errors := make(chan error, 2) + var ready sync.WaitGroup + ready.Add(2) + + for _, name := range []string{"first", "second"} { + name := name + go func() { + ready.Done() + <-start + _, err := fixture.service.writeConfig(context.Background(), func(tx *gorm.DB) error { + return tx.Create(validControlGroup(name)).Error + }, nil) + errors <- err + }() + } + ready.Wait() + close(start) + for range 2 { + if err := <-errors; err != nil { + t.Fatalf("concurrent writeConfig() error = %v", err) + } + } + + assertGroupCount(t, fixture.db, 2) + snapshot := fixture.manager.Current() + if snapshot.Revision != before+2 { + t.Fatalf("revision = %d, want %d", snapshot.Revision, before+2) + } + if len(snapshot.Groups) != 2 { + t.Fatalf("snapshot groups = %#v, want two", snapshot.Groups) + } +} + +func TestConcurrentImportsPublishDatabaseTruth(t *testing.T) { + fixture := newServiceFixture(t) + before := fixture.manager.Current().Revision + requests := []ImportRequest{ + {UpstreamURL: "https://shared.example.com/v1", Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-shared-a"}, + {UpstreamURL: "https://shared.example.com/v1", Protocols: []protocol.Protocol{protocol.Anthropic}, Keys: "sk-shared-b"}, + {UpstreamURL: "https://one.example.com/v1", Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-one"}, + {UpstreamURL: "https://two.example.com/v1", Protocols: []protocol.Protocol{protocol.Gemini}, Keys: "sk-two"}, + {UpstreamURL: "https://three.example.com/v1", Protocols: []protocol.Protocol{protocol.OpenAI}, Keys: "sk-three"}, + {UpstreamURL: "https://four.example.com/v1", Protocols: []protocol.Protocol{protocol.Anthropic}, Keys: "sk-four"}, + } + + start := make(chan struct{}) + errors := make(chan error, len(requests)) + var ready sync.WaitGroup + ready.Add(len(requests)) + for _, request := range requests { + request := request + go func() { + ready.Done() + <-start + _, err := fixture.service.Import(context.Background(), request) + errors <- err + }() + } + ready.Wait() + close(start) + for range requests { + if err := <-errors; err != nil { + t.Fatalf("concurrent Import() error = %v", err) + } + } + + input, err := stateloader.BuildCompileInput(context.Background(), fixture.db) + if err != nil { + t.Fatalf("BuildCompileInput() error = %v", err) + } + want, err := state.Compile(input) + if err != nil { + t.Fatalf("Compile(DB input) error = %v", err) + } + got := fixture.manager.Current() + want.Revision = got.Revision + if !reflect.DeepEqual(got, want) { + t.Fatalf("published snapshot differs from DB compile\ngot=%#v\nwant=%#v", got, want) + } + if got.Revision != before+uint64(len(requests)) { + t.Fatalf("revision = %d, want %d", got.Revision, before+uint64(len(requests))) + } +} + +func validControlGroup(name string) *models.Group { + return &models.Group{ + Name: name, UpstreamURL: "https://" + name + ".example/v1", + Signature: name + "-signature", Protocols: models.JSON(`["openai"]`), + Models: models.JSON(`[{"id":"gpt-4o"}]`), Config: models.JSON(`{}`), Enabled: true, + } +} diff --git a/internal/control/smoke_test.go b/internal/control/smoke_test.go new file mode 100644 index 00000000..bcd37db8 --- /dev/null +++ b/internal/control/smoke_test.go @@ -0,0 +1,260 @@ +package control_test + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "regexp" + "strings" + "testing" + + "github.com/gin-gonic/gin" + "github.com/sirupsen/logrus" + "gorm.io/gorm" + + "gpt-load/internal/app" + "gpt-load/internal/container" + "gpt-load/internal/platform/i18n" + "gpt-load/internal/platform/redact" + "gpt-load/internal/state" + "gpt-load/internal/storage" + "gpt-load/internal/storage/models" + "gpt-load/internal/storage/store" + "gpt-load/internal/testutil/fakeupstream" +) + +func TestM1S6NoUISmoke(t *testing.T) { + const ( + authKey = "s6-control-auth-distinctive-value" + upstreamKey = "sk-s6-upstream-distinctive-credential" + model = "gpt-4o" + ) + upstream := fakeupstream.New( + fakeupstream.Step{Status: http.StatusOK, Fixture: "openai/models.json"}, + fakeupstream.Step{Status: http.StatusOK, Fixture: "openai/success.json"}, + ) + defer upstream.Close() + + t.Setenv("AUTH_KEY", authKey) + t.Setenv("DATA_DIR", t.TempDir()) + t.Setenv("DATABASE_DSN", ":memory:") + t.Setenv("ENCRYPTION_KEY", "s6-smoke-encryption-master-key") + t.Setenv("REDIS_DSN", "") + if err := i18n.Init(); err != nil { + t.Fatalf("i18n.Init() error = %v", err) + } + + dependencyContainer, err := container.BuildContainer() + if err != nil { + t.Fatalf("BuildContainer() error = %v", err) + } + + logger := logrus.StandardLogger() + previousOutput := logger.Out + previousHooks := logger.ReplaceHooks(make(logrus.LevelHooks)) + var logs bytes.Buffer + logger.SetOutput(&logs) + t.Cleanup(func() { + logger.SetOutput(previousOutput) + logger.ReplaceHooks(previousHooks) + }) + + err = dependencyContainer.Invoke(func( + engine *gin.Engine, + db *gorm.DB, + runtimeState app.RuntimeStateLoader, + registry *state.KeyRegistry, + storageStore store.Store, + runtimeRedactor *redact.Redactor, + ) { + logger.AddHook(redact.NewHook(runtimeRedactor)) + t.Cleanup(func() { + _ = storageStore.Close() + if sqlDB, dbErr := db.DB(); dbErr == nil { + _ = sqlDB.Close() + } + }) + if err := storage.AutoMigrate(db); err != nil { + t.Fatalf("AutoMigrate() error = %v", err) + } + if err := runtimeState.Load(context.Background()); err != nil { + t.Fatalf("runtimeState.Load() error = %v", err) + } + + responses := make(map[string]string) + invalid := serveJSON(t, engine, http.MethodPost, "/api/import", authKey, map[string]any{ + "upstream_url": "not-a-url", "protocols": []string{"openai"}, "keys": upstreamKey, + }) + responses["error"] = invalid.Body.String() + if invalid.Code != http.StatusBadRequest { + t.Fatalf("invalid import = %d %s, want 400", invalid.Code, invalid.Body.String()) + } + + importRecorder := serveJSON(t, engine, http.MethodPost, "/api/import", authKey, map[string]any{ + "upstream_url": upstream.URL, "protocols": []string{"openai"}, "name": "s6-smoke", "keys": upstreamKey, + }) + responses["import"] = importRecorder.Body.String() + if importRecorder.Code != http.StatusOK { + t.Fatalf("import = %d %s", importRecorder.Code, importRecorder.Body.String()) + } + var imported struct { + Data struct { + GroupID uint `json:"group_id"` + KeysAdded int `json:"keys_added"` + ModelsFetched bool `json:"models_fetched"` + Models []string `json:"models"` + } `json:"data"` + } + decodeSmokeResponse(t, importRecorder, &imported) + if imported.Data.GroupID == 0 || imported.Data.KeysAdded != 1 || !imported.Data.ModelsFetched || + len(imported.Data.Models) != 2 || imported.Data.Models[0] != model { + t.Fatalf("import data = %#v", imported.Data) + } + + groupsRecorder := serveJSON(t, engine, http.MethodGet, "/api/groups", authKey, nil) + responses["groups"] = groupsRecorder.Body.String() + if groupsRecorder.Code != http.StatusOK { + t.Fatalf("groups = %d %s", groupsRecorder.Code, groupsRecorder.Body.String()) + } + var groups struct { + Data []struct { + ID uint `json:"id"` + KeyCount int `json:"key_count"` + } `json:"data"` + } + decodeSmokeResponse(t, groupsRecorder, &groups) + if len(groups.Data) != 1 || groups.Data[0].ID != imported.Data.GroupID || groups.Data[0].KeyCount != 1 { + t.Fatalf("groups data = %#v", groups.Data) + } + + createRecorder := serveJSON(t, engine, http.MethodPost, "/api/access-keys", authKey, map[string]any{ + "name": "s6-client", + "filters": map[string]any{ + "groups": []uint{imported.Data.GroupID}, "protocols": []string{"openai"}, "models": []string{model}, + }, + }) + responses["create_access_key"] = createRecorder.Body.String() + if createRecorder.Code != http.StatusOK { + t.Fatalf("create AccessKey = %d %s", createRecorder.Code, createRecorder.Body.String()) + } + var created struct { + Data struct { + Key string `json:"key"` + } `json:"data"` + } + decodeSmokeResponse(t, createRecorder, &created) + if !regexp.MustCompile(`^gl-[0-9a-f]{32}$`).MatchString(created.Data.Key) { + t.Fatalf("generated AccessKey has invalid format") + } + + listRecorder := serveJSON(t, engine, http.MethodGet, "/api/access-keys", authKey, nil) + responses["list_access_keys"] = listRecorder.Body.String() + if listRecorder.Code != http.StatusOK || !strings.Contains(listRecorder.Body.String(), created.Data.Key) { + t.Fatalf("list AccessKeys = %d %s", listRecorder.Code, listRecorder.Body.String()) + } + + chatRecorder := serveJSON(t, engine, http.MethodPost, "/v1/chat/completions", created.Data.Key, map[string]any{ + "model": model, "messages": []map[string]string{{"role": "user", "content": "ping"}}, + }) + responses["gateway"] = chatRecorder.Body.String() + if chatRecorder.Code != http.StatusOK || !strings.Contains(chatRecorder.Body.String(), `"content":"pong"`) { + t.Fatalf("gateway = %d %s", chatRecorder.Code, chatRecorder.Body.String()) + } + if chatRecorder.Header().Get("X-GPTLoad-Attempts") != "1" { + t.Fatalf("X-GPTLoad-Attempts = %q, want 1", chatRecorder.Header().Get("X-GPTLoad-Attempts")) + } + + requests := upstream.Requests() + if len(requests) != 2 || requests[0].Method != http.MethodGet || requests[0].Path != "/v1/models" || + requests[1].Method != http.MethodPost || requests[1].Path != "/v1/chat/completions" { + t.Fatalf("upstream requests = %#v", requests) + } + for index, request := range requests { + if request.Headers.Get("Authorization") != "Bearer "+upstreamKey { + t.Fatalf("upstream request %d credential was not injected", index) + } + } + if strings.Contains(string(requests[1].Body), created.Data.Key) { + t.Fatal("gateway forwarded downstream AccessKey in request body") + } + + var upstreamRow models.UpstreamKey + if err := db.First(&upstreamRow).Error; err != nil { + t.Fatalf("query UpstreamKey: %v", err) + } + var accessRow models.AccessKey + if err := db.First(&accessRow).Error; err != nil { + t.Fatalf("query AccessKey: %v", err) + } + if upstreamRow.KeyValue == upstreamKey || accessRow.KeyValue == created.Data.Key { + t.Fatal("database stored a plaintext credential") + } + if _, ok := registry.EncryptedValue(upstreamRow.ID); !ok { + t.Fatal("imported upstream key is absent from KeyRegistry") + } + + secrets := []string{authKey, upstreamKey, created.Data.Key, upstreamRow.KeyValue, upstreamRow.KeyHash, accessRow.KeyValue, accessRow.KeyHash} + for name, body := range responses { + allowedAccessKey := name == "create_access_key" || name == "list_access_keys" + for secretIndex, secret := range secrets { + if secret == created.Data.Key && allowedAccessKey { + continue + } + if secret != "" && strings.Contains(body, secret) { + t.Fatalf("%s response exposes secret index %d", name, secretIndex) + } + } + } + for secretIndex, secret := range secrets { + if secret != "" && strings.Contains(logs.String(), secret) { + t.Fatalf("logs expose secret index %d", secretIndex) + } + } + databaseText := fmt.Sprintf("%#v %#v", upstreamRow, accessRow) + for _, plaintext := range []string{authKey, upstreamKey, created.Data.Key} { + if strings.Contains(databaseText, plaintext) { + t.Fatalf("database contains plaintext credential") + } + } + }) + if err != nil { + t.Fatalf("resolve S6 smoke dependencies: %v", err) + } +} + +func serveJSON( + t *testing.T, + engine *gin.Engine, + method, path, bearer string, + payload any, +) *httptest.ResponseRecorder { + t.Helper() + var body *bytes.Reader + if payload == nil { + body = bytes.NewReader(nil) + } else { + encoded, err := json.Marshal(payload) + if err != nil { + t.Fatalf("marshal request payload: %v", err) + } + body = bytes.NewReader(encoded) + } + request := httptest.NewRequest(method, path, body) + request.Header.Set("Authorization", "Bearer "+bearer) + if payload != nil { + request.Header.Set("Content-Type", "application/json") + } + recorder := httptest.NewRecorder() + engine.ServeHTTP(recorder, request) + return recorder +} + +func decodeSmokeResponse(t *testing.T, recorder *httptest.ResponseRecorder, target any) { + t.Helper() + if err := json.Unmarshal(recorder.Body.Bytes(), target); err != nil { + t.Fatalf("decode response %d %s: %v", recorder.Code, recorder.Body.String(), err) + } +} diff --git a/internal/control/test_helpers_test.go b/internal/control/test_helpers_test.go new file mode 100644 index 00000000..ec7b5cfb --- /dev/null +++ b/internal/control/test_helpers_test.go @@ -0,0 +1,89 @@ +package control + +import ( + "sync" + "testing" + + "github.com/gin-gonic/gin" + "gorm.io/gorm" + + "gpt-load/internal/dialect" + "gpt-load/internal/platform/encryption" + "gpt-load/internal/platform/i18n" + "gpt-load/internal/state" + "gpt-load/internal/storage" +) + +var ( + controlI18nOnce sync.Once + controlI18nErr error +) + +type serviceFixture struct { + db *gorm.DB + manager *state.Manager + registry *state.KeyRegistry + encryption encryption.Service + service *Service +} + +func initControlI18n(t *testing.T) { + t.Helper() + controlI18nOnce.Do(func() { + gin.SetMode(gin.ReleaseMode) + controlI18nErr = i18n.Init() + }) + if controlI18nErr != nil { + t.Fatalf("i18n.Init() error = %v", controlI18nErr) + } +} + +func newServiceFixture(t *testing.T) serviceFixture { + t.Helper() + db := openControlTestDB(t) + manager := state.NewManager() + registry := state.NewKeyRegistry() + keyService, err := encryption.NewService("control-test-master-key-material-2026") + if err != nil { + t.Fatalf("encryption.NewService() error = %v", err) + } + if _, err := manager.Publish(state.CompileInput{}); err != nil { + t.Fatalf("manager.Publish(empty) error = %v", err) + } + return serviceFixture{ + db: db, manager: manager, registry: registry, encryption: keyService, + service: NewService(db, manager, registry, keyService, dialect.NewSet()), + } +} + +func openControlTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := storage.Open(":memory:") + if err != nil { + t.Fatalf("storage.Open(:memory:) 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 control test database: %v", err) + } + }) + if err := storage.AutoMigrate(db); err != nil { + t.Fatalf("storage.AutoMigrate() error = %v", err) + } + return db +} + +func assertGroupCount(t *testing.T, db *gorm.DB, want int64) { + t.Helper() + var got int64 + if err := db.Table("groups").Count(&got).Error; err != nil { + t.Fatalf("count groups: %v", err) + } + if got != want { + t.Fatalf("group count = %d, want %d", got, want) + } +} diff --git a/internal/dialect/set.go b/internal/dialect/set.go new file mode 100644 index 00000000..179c3231 --- /dev/null +++ b/internal/dialect/set.go @@ -0,0 +1,16 @@ +package dialect + +import "gpt-load/internal/protocol" + +type Set map[protocol.Protocol]Dialect + +func NewSet(values ...Dialect) Set { + result := make(Set, len(values)) + for _, value := range values { + if value == nil { + continue + } + result[value.Protocol()] = value + } + return result +} diff --git a/internal/dialect/set_test.go b/internal/dialect/set_test.go new file mode 100644 index 00000000..9743b157 --- /dev/null +++ b/internal/dialect/set_test.go @@ -0,0 +1,22 @@ +package dialect + +import ( + "net/http" + "testing" + + "gpt-load/internal/protocol" +) + +func TestNewSetIndexesDialectsByProtocol(t *testing.T) { + openAI := NewOpenAI(http.DefaultClient) + got := NewSet(openAI) + if len(got) != 1 || got[protocol.OpenAI] != openAI { + t.Fatalf("NewSet(OpenAI) = %#v, want OpenAI indexed", got) + } +} + +func TestNewSetAllowsAnEmptyRuntimeSet(t *testing.T) { + if got := NewSet(); len(got) != 0 { + t.Fatalf("NewSet() = %#v, want empty", got) + } +} diff --git a/internal/gateway/handler.go b/internal/gateway/handler.go index cedd2700..21c450b6 100644 --- a/internal/gateway/handler.go +++ b/internal/gateway/handler.go @@ -15,7 +15,6 @@ import ( "gpt-load/internal/health" "gpt-load/internal/platform/encryption" "gpt-load/internal/platform/utils" - "gpt-load/internal/protocol" "gpt-load/internal/scheduler" "gpt-load/internal/state" ) @@ -42,22 +41,12 @@ type runtimeKeyRegistry interface { ActiveEncryptedValue(keyID, expectedGroupID uint) (string, bool) } -type DialectSet map[protocol.Protocol]dialect.Dialect - -func NewDialectSet(openAI *dialect.OpenAI) DialectSet { - dialects := make(DialectSet) - if openAI != nil { - dialects[protocol.OpenAI] = openAI - } - return dialects -} - type Handler struct { manager *state.Manager registry runtimeKeyRegistry encryption encryption.Service forwarder AttemptForwarder - dialects DialectSet + dialects dialect.Set newRandom func() *rand.Rand } @@ -66,7 +55,7 @@ func NewHandler( registry *state.KeyRegistry, encryptionService encryption.Service, forwarder AttemptForwarder, - dialects DialectSet, + dialects dialect.Set, ) *Handler { return &Handler{ manager: manager, registry: registry, encryption: encryptionService, diff --git a/internal/gateway/handler_test.go b/internal/gateway/handler_test.go index 8949fe27..9e660fee 100644 --- a/internal/gateway/handler_test.go +++ b/internal/gateway/handler_test.go @@ -719,7 +719,7 @@ func TestHandlerSkipsCandidateChangedAfterCollection(t *testing.T) { mutate: func() { tt.mutate(t, registry, keyService) }, } openAI := dialect.NewOpenAI(http.DefaultClient) - handler := NewHandler(manager, registry, keyService, forwarder, NewDialectSet(openAI)) + handler := NewHandler(manager, registry, keyService, forwarder, dialect.NewSet(openAI)) handler.registry = runtimeRegistry handler.newRandom = func() *rand.Rand { return rand.New(rand.NewSource(1)) } engine := gin.New() @@ -789,7 +789,7 @@ func newRealGatewayEngine(t *testing.T, upstreamURL string, upstreamKeys ...stri registry, keyService, NewForwarder(clients, redact.New()), - NewDialectSet(openAI), + dialect.NewSet(openAI), ) handler.newRandom = func() *rand.Rand { return rand.New(rand.NewSource(1)) } engine := gin.New() @@ -847,7 +847,7 @@ func newHandlerTestRuntime( } openAI := dialect.NewOpenAI(http.DefaultClient) - handler := NewHandler(manager, registry, keyService, forwarder, NewDialectSet(openAI)) + handler := NewHandler(manager, registry, keyService, forwarder, dialect.NewSet(openAI)) handler.newRandom = func() *rand.Rand { return rand.New(rand.NewSource(1)) } engine := gin.New() handler.RegisterRoutes(engine) diff --git a/internal/gateway/stream_integration_test.go b/internal/gateway/stream_integration_test.go index 0d49677b..b796a571 100644 --- a/internal/gateway/stream_integration_test.go +++ b/internal/gateway/stream_integration_test.go @@ -487,7 +487,7 @@ func newStreamingGatewayEngine(t *testing.T, groups ...streamGatewayGroup) (*gin registry, keyService, NewForwarder(clients, redact.New()), - NewDialectSet(dialect.NewOpenAI(http.DefaultClient)), + dialect.NewSet(dialect.NewOpenAI(http.DefaultClient)), ) handler.newRandom = func() *rand.Rand { return rand.New(zeroSource{}) } engine := gin.New() diff --git a/internal/platform/config/config.go b/internal/platform/config/config.go index 10d855be..c3bcc41e 100644 --- a/internal/platform/config/config.go +++ b/internal/platform/config/config.go @@ -7,6 +7,7 @@ import ( "os" "path/filepath" "strconv" + "strings" "github.com/joho/godotenv" ) @@ -82,6 +83,10 @@ func Load() (*Config, error) { if authKey == "" { return nil, fmt.Errorf("AUTH_KEY is required") } + fields := strings.Fields(authKey) + if len(fields) != 1 || fields[0] != authKey { + return nil, fmt.Errorf("AUTH_KEY must not contain whitespace") + } dataDir := valueOrDefault("DATA_DIR", defaultDataDir) databaseDSN := os.Getenv("DATABASE_DSN") diff --git a/internal/platform/config/config_test.go b/internal/platform/config/config_test.go index 641bb702..618622a4 100644 --- a/internal/platform/config/config_test.go +++ b/internal/platform/config/config_test.go @@ -101,6 +101,9 @@ func TestLoadRejectsInvalidRequiredAndNumericValues(t *testing.T) { env map[string]string }{ {name: "missing auth key", env: map[string]string{}}, + {name: "whitespace-only auth key", env: map[string]string{"AUTH_KEY": " "}}, + {name: "auth key with internal space", env: map[string]string{"AUTH_KEY": "admin key"}}, + {name: "auth key with tab", env: map[string]string{"AUTH_KEY": "admin\tkey"}}, {name: "invalid port", env: map[string]string{"AUTH_KEY": "x", "PORT": "nope"}}, {name: "port out of range", env: map[string]string{"AUTH_KEY": "x", "PORT": "70000"}}, {name: "invalid shutdown timeout", env: map[string]string{"AUTH_KEY": "x", "GRACEFUL_SHUTDOWN_TIMEOUT": "0"}}, diff --git a/internal/state/loader/loader.go b/internal/state/loader/loader.go index df344210..5dfb3a26 100644 --- a/internal/state/loader/loader.go +++ b/internal/state/loader/loader.go @@ -24,11 +24,10 @@ type Loader struct { registry *state.KeyRegistry } -type storedRows struct { - settings []models.SystemSetting - groups []models.Group - accessKeys []models.AccessKey - upstreamKeys []models.UpstreamKey +type compileRows struct { + settings []models.SystemSetting + groups []models.Group + accessKeys []models.AccessKey } type modelDTO struct { @@ -86,39 +85,56 @@ func (l *Loader) Load(ctx context.Context) error { return nil } -func (l *Loader) queryRows(ctx context.Context) (storedRows, error) { - db := l.db.WithContext(ctx) - var rows storedRows +func queryCompileRows(ctx context.Context, db *gorm.DB) (compileRows, error) { + db = db.WithContext(ctx) + var rows compileRows if err := db.Order("key ASC").Find(&rows.settings).Error; err != nil { - return storedRows{}, fmt.Errorf("query system settings: %w", err) + return compileRows{}, fmt.Errorf("query system settings: %w", err) } if err := db.Order("id ASC").Find(&rows.groups).Error; err != nil { - return storedRows{}, fmt.Errorf("query groups: %w", err) + return compileRows{}, fmt.Errorf("query groups: %w", err) } if err := db.Order("id ASC").Find(&rows.accessKeys).Error; err != nil { - return storedRows{}, fmt.Errorf("query access keys: %w", err) - } - if err := db.Order("id ASC").Find(&rows.upstreamKeys).Error; err != nil { - return storedRows{}, fmt.Errorf("query upstream keys: %w", err) + return compileRows{}, fmt.Errorf("query access keys: %w", err) } return rows, nil } -func (l *Loader) read(ctx context.Context) (state.CompileInput, []state.KeyEntry, error) { - rows, err := l.queryRows(ctx) +func queryUpstreamKeys(ctx context.Context, db *gorm.DB) ([]models.UpstreamKey, error) { + var rows []models.UpstreamKey + if err := db.WithContext(ctx).Order("id ASC").Find(&rows).Error; err != nil { + return nil, fmt.Errorf("query upstream keys: %w", err) + } + return rows, nil +} + +// BuildCompileInput maps persisted configuration rows into a runtime compiler input. +func BuildCompileInput(ctx context.Context, db *gorm.DB) (state.CompileInput, error) { + rows, err := queryCompileRows(ctx, db) if err != nil { - return state.CompileInput{}, nil, err + return state.CompileInput{}, err } input, err := mapSystemAndGroups(rows) if err != nil { - return state.CompileInput{}, nil, err + return state.CompileInput{}, err } - accessKeys, err := mapAccessKeys(rows.accessKeys) + input.AccessKeys, err = mapAccessKeys(rows.accessKeys) + if err != nil { + return state.CompileInput{}, err + } + return input, nil +} + +func (l *Loader) read(ctx context.Context) (state.CompileInput, []state.KeyEntry, error) { + input, err := BuildCompileInput(ctx, l.db) if err != nil { return state.CompileInput{}, nil, err } - input.AccessKeys = accessKeys - return input, mapUpstreamKeys(rows.upstreamKeys), nil + upstreamKeys, err := queryUpstreamKeys(ctx, l.db) + if err != nil { + return state.CompileInput{}, nil, err + } + return input, mapUpstreamKeys(upstreamKeys), nil } func decodeJSON(raw models.JSON, target any) error { @@ -164,7 +180,7 @@ func decodeSettingValue(raw string) (any, error) { return value, nil } -func mapSystemAndGroups(rows storedRows) (state.CompileInput, error) { +func mapSystemAndGroups(rows compileRows) (state.CompileInput, error) { input := state.CompileInput{ SystemSettings: make(config.Settings, len(rows.settings)), Groups: make([]state.GroupConfig, 0, len(rows.groups)), diff --git a/internal/state/loader/loader_test.go b/internal/state/loader/loader_test.go index 9bb2cd7e..85d548cd 100644 --- a/internal/state/loader/loader_test.go +++ b/internal/state/loader/loader_test.go @@ -46,6 +46,112 @@ func TestLoaderLoadsEmptyMigratedDatabase(t *testing.T) { } } +func TestBuildCompileInputReadsUncommittedTransactionState(t *testing.T) { + db := openMigratedDatabase(t) + tx := db.Begin() + if tx.Error != nil { + t.Fatalf("Begin() error = %v", tx.Error) + } + t.Cleanup(func() { + _ = tx.Rollback().Error + }) + + group := models.Group{ + Name: "pending", UpstreamURL: "https://pending.example", + Signature: "pending-signature", Protocols: models.JSON(`["openai"]`), + Models: models.JSON(`[{"id":"gpt-pending"}]`), Config: models.JSON(`{}`), Enabled: true, + } + 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)), + }) + + input, err := loader.BuildCompileInput(context.Background(), tx) + if err != nil { + t.Fatalf("BuildCompileInput() error = %v", err) + } + if len(input.Groups) != 1 || len(input.AccessKeys) != 1 { + t.Fatalf("input = %#v, want one pending group and access key", input) + } +} + +func TestBuildCompileInputReturnsIndependentData(t *testing.T) { + db := openMigratedDatabase(t) + mustCreate(t, db, &models.SystemSetting{Key: "connect_timeout", Value: "20"}) + group := createRuntimeGroup(t, db, "owned", protocol.OpenAI, "gpt-owned") + if err := db.Model(&group).Update("config", models.JSON(`{"request_timeout":30}`)).Error; err != nil { + t.Fatalf("update group config: %v", err) + } + 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)), + }) + + first, err := loader.BuildCompileInput(context.Background(), db) + if err != nil { + t.Fatalf("first BuildCompileInput() error = %v", err) + } + first.SystemSettings["connect_timeout"] = 99 + first.Groups[0].Protocols[0] = protocol.Anthropic + first.Groups[0].Settings["request_timeout"] = 99 + first.AccessKeys[0].Filters.Groups[999] = struct{}{} + first.AccessKeys[0].Filters.Models["mutated"] = struct{}{} + + second, err := loader.BuildCompileInput(context.Background(), db) + if err != nil { + t.Fatalf("second BuildCompileInput() error = %v", err) + } + if got := fmt.Sprint(second.SystemSettings["connect_timeout"]); got != "20" { + t.Fatalf("connect_timeout = %q, want 20", got) + } + if second.Groups[0].Protocols[0] != protocol.OpenAI { + t.Fatalf("protocol = %q, want openai", second.Groups[0].Protocols[0]) + } + if got := fmt.Sprint(second.Groups[0].Settings["request_timeout"]); got != "30" { + t.Fatalf("request_timeout = %q, want 30", got) + } + if _, ok := second.AccessKeys[0].Filters.Groups[999]; ok { + t.Fatal("second input retained mutated group filter") + } + if _, ok := second.AccessKeys[0].Filters.Models["mutated"]; ok { + t.Fatal("second input retained mutated model filter") + } +} + +func TestBuildCompileInputDoesNotQueryUpstreamKeys(t *testing.T) { + db := openMigratedDatabase(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(`{}`), + }) + mustCreate(t, db, &models.UpstreamKey{ + GroupID: group.ID, KeyValue: "upstream-cipher", KeyHash: "upstream-query-hash", + Status: models.UpstreamKeyStatusActive, + }) + + const callbackName = "test:build_compile_input_tables" + seen := make(map[string]int) + if err := db.Callback().Query().After("gorm:query").Register(callbackName, func(tx *gorm.DB) { + seen[tx.Statement.Table]++ + }); err != nil { + t.Fatalf("register query callback: %v", err) + } + + if _, err := loader.BuildCompileInput(context.Background(), db); err != nil { + t.Fatalf("BuildCompileInput() error = %v", err) + } + if seen["upstream_keys"] != 0 { + t.Fatalf("upstream_keys query count = %d, want 0", seen["upstream_keys"]) + } + for _, table := range []string{"system_settings", "groups", "access_keys"} { + if seen[table] != 1 { + t.Errorf("%s query count = %d, want 1; all=%#v", table, seen[table], seen) + } + } +} + func TestLoaderMapsSystemAndGroupRows(t *testing.T) { db := openMigratedDatabase(t) mustCreate(t, db, &models.SystemSetting{Key: "connect_timeout", Value: "20"}) diff --git a/internal/storage/db.go b/internal/storage/db.go index 94496e4c..f9c1b491 100644 --- a/internal/storage/db.go +++ b/internal/storage/db.go @@ -2,13 +2,16 @@ package storage import ( "fmt" + "log" "net/url" "os" "path/filepath" "strings" + "time" "github.com/glebarez/sqlite" "gorm.io/gorm" + "gorm.io/gorm/logger" "gpt-load/internal/storage/models" ) @@ -20,6 +23,13 @@ type schemaInfo struct { Version uint `gorm:"primaryKey;autoIncrement:false"` } +var databaseLogger = logger.New(log.New(os.Stdout, "\r\n", log.LstdFlags), logger.Config{ + SlowThreshold: 200 * time.Millisecond, + LogLevel: logger.Warn, + ParameterizedQueries: true, + Colorful: true, +}) + func (schemaInfo) TableName() string { return "schema_info" } @@ -35,7 +45,7 @@ func Open(dsn string) (*gorm.DB, error) { return nil, err } - db, err := gorm.Open(sqlite.Open(withForeignKeysEnabled(dsn)), &gorm.Config{}) + db, err := gorm.Open(sqlite.Open(withForeignKeysEnabled(dsn)), &gorm.Config{Logger: databaseLogger}) if err != nil { return nil, fmt.Errorf("open SQLite database: %w", err) } diff --git a/internal/storage/db_test.go b/internal/storage/db_test.go index aebd42d7..bb65f2ea 100644 --- a/internal/storage/db_test.go +++ b/internal/storage/db_test.go @@ -1,6 +1,7 @@ package storage_test import ( + "context" "path/filepath" "strings" "testing" @@ -126,6 +127,38 @@ func TestOpenCreatesSQLiteDatabase(t *testing.T) { } } +func TestOpenConfiguresParameterizedSQLLogging(t *testing.T) { + t.Parallel() + + 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() { + if err := sqlDB.Close(); err != nil { + t.Errorf("close database: %v", err) + } + }) + + filter, ok := db.Logger.(gorm.ParamsFilter) + if !ok { + t.Fatalf("database logger %T does not implement gorm.ParamsFilter", db.Logger) + } + const query = "INSERT INTO secrets(value) VALUES (?)" + const secret = "known-plaintext-secret" + filteredQuery, params := filter.ParamsFilter(context.Background(), query, secret) + if filteredQuery != query { + t.Fatalf("filtered query = %q, want %q", filteredQuery, query) + } + if len(params) != 0 { + t.Fatalf("parameterized SQL logger retained %d parameter(s), want 0", len(params)) + } +} + func TestAutoMigrateCreatesNineTablesAndSchemaVersion(t *testing.T) { t.Parallel() diff --git a/scripts/m1-smoke.sh b/scripts/m1-smoke.sh new file mode 100755 index 00000000..fae4dea5 --- /dev/null +++ b/scripts/m1-smoke.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +set -euo pipefail + +for command_name in curl jq; do + command -v "${command_name}" >/dev/null || { + echo "missing command: ${command_name}" >&2 + exit 1 + } +done + +for name in GPT_LOAD_URL AUTH_KEY UPSTREAM_URL UPSTREAM_KEY MODEL; do + [[ -n "${!name:-}" ]] || { + echo "missing environment variable: ${name}" >&2 + exit 1 + } +done + +base_url="${GPT_LOAD_URL%/}" +import_payload="$(jq -n \ + --arg url "${UPSTREAM_URL}" \ + --arg key "${UPSTREAM_KEY}" \ + '{upstream_url:$url,protocols:["openai"],keys:$key}')" +import_response="$(curl -fsS \ + -H "Authorization: Bearer ${AUTH_KEY}" \ + -H 'Content-Type: application/json' \ + -d "${import_payload}" \ + "${base_url}/api/import")" +group_id="$(jq -er '.data.group_id' <<<"${import_response}")" + +curl -fsS \ + -H "Authorization: Bearer ${AUTH_KEY}" \ + "${base_url}/api/groups" \ + | jq -e --argjson id "${group_id}" '.data[] | select(.id == $id)' >/dev/null + +access_payload="$(jq -n \ + --arg model "${MODEL}" \ + --argjson group "${group_id}" \ + '{name:"m1-smoke",filters:{groups:[$group],protocols:["openai"],models:[$model]}}')" +access_response="$(curl -fsS \ + -H "Authorization: Bearer ${AUTH_KEY}" \ + -H 'Content-Type: application/json' \ + -d "${access_payload}" \ + "${base_url}/api/access-keys")" +access_key="$(jq -er '.data.key' <<<"${access_response}")" + +chat_payload="$(jq -n \ + --arg model "${MODEL}" \ + '{model:$model,messages:[{role:"user",content:"Reply with OK"}]}')" +curl -fsS \ + -H "Authorization: Bearer ${access_key}" \ + -H 'Content-Type: application/json' \ + -d "${chat_payload}" \ + "${base_url}/v1/chat/completions" \ + | jq -e . >/dev/null + +printf 'M1 S6 smoke passed for group_id=%s\n' "${group_id}"