From 15ae2bc4cae35bcfb42fb0dd59c372a592f5544d Mon Sep 17 00:00:00 2001 From: chengcheng Date: Tue, 21 Jul 2026 10:57:33 +0800 Subject: [PATCH] =?UTF-8?q?fix(auth):=20=E6=B6=88=E9=99=A4=20API=20Key=20?= =?UTF-8?q?=E6=A0=A1=E9=AA=8C=E8=BF=9E=E6=8E=A5=E6=B1=A0=E6=AD=BB=E9=94=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 先完整收集同前缀候选项并关闭查询结果，再执行 bcrypt 比对和 last_used_at 更新，避免小连接池下查询与更新相互等待。 新增 Rows 关闭顺序、前缀碰撞、MaxConns=1 和 8 并发真实 PostgreSQL 回归测试。 --- .../store/api_key_auth_integration_test.go | 99 ++++++++++ apps/api/internal/store/api_key_auth_test.go | 169 ++++++++++++++++++ .../identity_pairing_integration_test.go | 2 + apps/api/internal/store/postgres.go | 96 ++++++---- 4 files changed, 331 insertions(+), 35 deletions(-) create mode 100644 apps/api/internal/store/api_key_auth_integration_test.go create mode 100644 apps/api/internal/store/api_key_auth_test.go diff --git a/apps/api/internal/store/api_key_auth_integration_test.go b/apps/api/internal/store/api_key_auth_integration_test.go new file mode 100644 index 0000000..8247c5e --- /dev/null +++ b/apps/api/internal/store/api_key_auth_integration_test.go @@ -0,0 +1,99 @@ +package store + +import ( + "context" + "sync" + "testing" + "time" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" + "github.com/jackc/pgx/v5/pgxpool" +) + +func TestVerifyLocalAPIKeyWorksWithSingleConnectionPool(t *testing.T) { + db, verificationStore, created, user := newLocalAPIKeyVerificationFixture(t, 1) + ctx := context.Background() + + verifyCtx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + verified, err := verificationStore.VerifyLocalAPIKey(verifyCtx, created.Secret) + if err != nil { + t.Fatalf("verify local API key with one connection: %v", err) + } + if verified.APIKeyID != created.APIKey.ID || verified.GatewayUserID != user.ID { + t.Fatalf("verified identity = %+v, want API key %q and user %q", verified, created.APIKey.ID, user.ID) + } + + var lastUsedAt *time.Time + if err := db.pool.QueryRow(ctx, `SELECT last_used_at FROM gateway_api_keys WHERE id=$1::uuid`, created.APIKey.ID).Scan(&lastUsedAt); err != nil { + t.Fatalf("read API key last_used_at: %v", err) + } + if lastUsedAt == nil { + t.Fatal("successful API key verification did not update last_used_at") + } +} + +func TestVerifyLocalAPIKeyHandlesEightConcurrentRequestsWithFourConnections(t *testing.T) { + _, verificationStore, created, _ := newLocalAPIKeyVerificationFixture(t, 4) + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + const requestCount = 8 + start := make(chan struct{}) + errorsByRequest := make(chan error, requestCount) + var requests sync.WaitGroup + requests.Add(requestCount) + for range requestCount { + go func() { + defer requests.Done() + <-start + _, err := verificationStore.VerifyLocalAPIKey(ctx, created.Secret) + errorsByRequest <- err + }() + } + close(start) + requests.Wait() + close(errorsByRequest) + + for err := range errorsByRequest { + if err != nil { + t.Fatalf("concurrent API key verification failed: %v", err) + } + } + if acquired := verificationStore.pool.Stat().AcquiredConns(); acquired != 0 { + t.Fatalf("API key verification left %d connections acquired", acquired) + } +} + +func newLocalAPIKeyVerificationFixture(t *testing.T, maxConnections int32) (*Store, *Store, CreatedAPIKey, GatewayUser) { + t.Helper() + db := newIdentityPairingPostgresTestStore(t) + ctx := context.Background() + user, err := db.RegisterLocalUser(ctx, LocalRegisterInput{ + Username: "api-key-verification-user", + Password: "api-key-verification-password", + }) + if err != nil { + t.Fatalf("register API key verification user: %v", err) + } + created, err := db.CreateAPIKey(ctx, CreateAPIKeyInput{Name: "API key verification fixture"}, &auth.User{ + ID: user.ID, + GatewayUserID: user.ID, + GatewayTenantID: user.GatewayTenantID, + TenantID: user.TenantID, + TenantKey: user.TenantKey, + }) + if err != nil { + t.Fatalf("create API key verification fixture: %v", err) + } + + config := db.pool.Config() + config.MaxConns = maxConnections + config.MinConns = 0 + pool, err := pgxpool.NewWithConfig(ctx, config) + if err != nil { + t.Fatalf("create verification pool with %d connections: %v", maxConnections, err) + } + t.Cleanup(pool.Close) + return db, &Store{pool: pool}, created, user +} diff --git a/apps/api/internal/store/api_key_auth_test.go b/apps/api/internal/store/api_key_auth_test.go new file mode 100644 index 0000000..5b8e299 --- /dev/null +++ b/apps/api/internal/store/api_key_auth_test.go @@ -0,0 +1,169 @@ +package store + +import ( + "context" + "errors" + "testing" + + "github.com/easyai/easyai-ai-gateway/apps/api/internal/auth" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" + "golang.org/x/crypto/bcrypt" +) + +func TestVerifyLocalAPIKeyClosesCandidateRowsBeforeUpdatingUsage(t *testing.T) { + secret := "sk-gw-matching-secret" + wrongHash, err := bcrypt.GenerateFromPassword([]byte("sk-gw-different-secret"), bcrypt.MinCost) + if err != nil { + t.Fatalf("hash non-matching API key: %v", err) + } + matchingHash, err := bcrypt.GenerateFromPassword([]byte(secret), bcrypt.MinCost) + if err != nil { + t.Fatalf("hash matching API key: %v", err) + } + rows := &fakeLocalAPIKeyRows{candidates: []localAPIKeyCandidate{ + {apiKeyID: "wrong-key", hash: string(wrongHash), keyPrefix: apiKeyPrefix(secret)}, + { + apiKeyID: "matching-key", + hash: string(matchingHash), + keyPrefix: apiKeyPrefix(secret), + keyName: "Matching key", + scopesBytes: []byte(`["chat"]`), + userGroupID: "group-id", + gatewayUserID: "user-id", + username: "api-key-user", + rolesBytes: []byte(`["user"]`), + gatewayTenantID: "gateway-tenant-id", + tenantID: "tenant-id", + tenantKey: "tenant-key", + }, + }} + database := &fakeLocalAPIKeyDatabase{rows: rows} + + user, err := verifyLocalAPIKey(context.Background(), database, secret) + if err != nil { + t.Fatalf("verify local API key: %v", err) + } + if !rows.closed { + t.Fatal("candidate rows remained open after API key verification") + } + if database.updatedAPIKeyID != "matching-key" { + t.Fatalf("updated API key = %q, want matching-key", database.updatedAPIKeyID) + } + if user.APIKeyID != "matching-key" || user.GatewayUserID != "user-id" { + t.Fatalf("verified user = %+v", user) + } +} + +func TestVerifyLocalAPIKeyReturnsUnauthorizedAfterClosingCandidateRows(t *testing.T) { + secret := "sk-gw-unknown-secret" + wrongHash, err := bcrypt.GenerateFromPassword([]byte("sk-gw-different-secret"), bcrypt.MinCost) + if err != nil { + t.Fatalf("hash non-matching API key: %v", err) + } + rows := &fakeLocalAPIKeyRows{candidates: []localAPIKeyCandidate{{ + apiKeyID: "wrong-key", + hash: string(wrongHash), + }}} + database := &fakeLocalAPIKeyDatabase{rows: rows} + + _, err = verifyLocalAPIKey(context.Background(), database, secret) + if !errors.Is(err, auth.ErrUnauthorized) { + t.Fatalf("verify error = %v, want unauthorized", err) + } + if !rows.closed { + t.Fatal("candidate rows remained open after unsuccessful API key verification") + } + if database.updatedAPIKeyID != "" { + t.Fatalf("unexpected API key usage update for %q", database.updatedAPIKeyID) + } +} + +type fakeLocalAPIKeyDatabase struct { + rows *fakeLocalAPIKeyRows + updatedAPIKeyID string +} + +func (database *fakeLocalAPIKeyDatabase) Query(context.Context, string, ...any) (pgx.Rows, error) { + return database.rows, nil +} + +func (database *fakeLocalAPIKeyDatabase) Exec(_ context.Context, _ string, arguments ...any) (pgconn.CommandTag, error) { + if !database.rows.closed { + return pgconn.CommandTag{}, errors.New("API key usage update started before candidate rows closed") + } + database.updatedAPIKeyID, _ = arguments[0].(string) + return pgconn.NewCommandTag("UPDATE 1"), nil +} + +type fakeLocalAPIKeyRows struct { + candidates []localAPIKeyCandidate + current int + closed bool +} + +func (rows *fakeLocalAPIKeyRows) Close() { + rows.closed = true +} + +func (rows *fakeLocalAPIKeyRows) Err() error { + return nil +} + +func (rows *fakeLocalAPIKeyRows) CommandTag() pgconn.CommandTag { + return pgconn.CommandTag{} +} + +func (rows *fakeLocalAPIKeyRows) FieldDescriptions() []pgconn.FieldDescription { + return nil +} + +func (rows *fakeLocalAPIKeyRows) Next() bool { + if rows.current >= len(rows.candidates) { + rows.Close() + return false + } + rows.current++ + return true +} + +func (rows *fakeLocalAPIKeyRows) Scan(destinations ...any) error { + candidate := rows.candidates[rows.current-1] + values := []any{ + candidate.apiKeyID, + candidate.hash, + candidate.keyPrefix, + candidate.keyName, + candidate.scopesBytes, + candidate.userGroupID, + candidate.gatewayUserID, + candidate.username, + candidate.rolesBytes, + candidate.gatewayTenantID, + candidate.tenantID, + candidate.tenantKey, + } + for index, value := range values { + switch destination := destinations[index].(type) { + case *string: + *destination = value.(string) + case *[]byte: + *destination = value.([]byte) + default: + return errors.New("unsupported fake row destination") + } + } + return nil +} + +func (rows *fakeLocalAPIKeyRows) Values() ([]any, error) { + return nil, errors.New("not implemented") +} + +func (rows *fakeLocalAPIKeyRows) RawValues() [][]byte { + return nil +} + +func (rows *fakeLocalAPIKeyRows) Conn() *pgx.Conn { + return nil +} diff --git a/apps/api/internal/store/identity_pairing_integration_test.go b/apps/api/internal/store/identity_pairing_integration_test.go index 65c8b16..6837ee2 100644 --- a/apps/api/internal/store/identity_pairing_integration_test.go +++ b/apps/api/internal/store/identity_pairing_integration_test.go @@ -685,11 +685,13 @@ func newIdentityPairingPostgresTestStore(t *testing.T) *Store { migrationDirectory := filepath.Join(filepath.Dir(filename), "..", "..", "migrations") for _, migrationName := range []string{ "0001_init.sql", + "0017_task_record_enrichment.sql", "0061_oidc_server_sessions.sql", "0065_identity_configuration_revisions.sql", "0066_identity_onboarding_exchanges.sql", "0067_identity_secret_cleanup_queue.sql", "0068_identity_pairing_start_reservation.sql", + "0069_billing_correctness_v2.sql", } { migration, err := os.ReadFile(filepath.Join(migrationDirectory, migrationName)) if err != nil { diff --git a/apps/api/internal/store/postgres.go b/apps/api/internal/store/postgres.go index 8de239e..08c5c21 100644 --- a/apps/api/internal/store/postgres.go +++ b/apps/api/internal/store/postgres.go @@ -1519,11 +1519,35 @@ WHERE subject_type = 'api_key' AND subject_id = $1::uuid`, apiKeyID); err != nil } func (s *Store) VerifyLocalAPIKey(ctx context.Context, secret string) (*auth.User, error) { + return verifyLocalAPIKey(ctx, s.pool, secret) +} + +type localAPIKeyDatabase interface { + Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) + Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) +} + +type localAPIKeyCandidate struct { + apiKeyID string + hash string + keyPrefix string + keyName string + scopesBytes []byte + userGroupID string + gatewayUserID string + username string + rolesBytes []byte + gatewayTenantID string + tenantID string + tenantKey string +} + +func verifyLocalAPIKey(ctx context.Context, database localAPIKeyDatabase, secret string) (*auth.User, error) { prefix := apiKeyPrefix(secret) if prefix == "" { return nil, auth.ErrUnauthorized } - rows, err := s.pool.Query(ctx, ` + rows, err := database.Query(ctx, ` SELECT k.id::text, k.key_hash, k.key_prefix, k.name, k.scopes, COALESCE(k.user_group_id::text, u.default_user_group_id::text, ''), u.id::text, u.username, u.roles, COALESCE(u.gateway_tenant_id::text, ''), COALESCE(u.tenant_id, ''), COALESCE(u.tenant_key, '') @@ -1538,49 +1562,51 @@ WHERE k.key_prefix = $1 if err != nil { return nil, err } - defer rows.Close() + candidates, err := pgx.CollectRows(rows, func(row pgx.CollectableRow) (localAPIKeyCandidate, error) { + var candidate localAPIKeyCandidate + err := row.Scan( + &candidate.apiKeyID, + &candidate.hash, + &candidate.keyPrefix, + &candidate.keyName, + &candidate.scopesBytes, + &candidate.userGroupID, + &candidate.gatewayUserID, + &candidate.username, + &candidate.rolesBytes, + &candidate.gatewayTenantID, + &candidate.tenantID, + &candidate.tenantKey, + ) + return candidate, err + }) + if err != nil { + return nil, err + } - for rows.Next() { - var apiKeyID string - var hash string - var keyPrefix string - var keyName string - var scopesBytes []byte - var userGroupID string - var gatewayUserID string - var username string - var rolesBytes []byte - var gatewayTenantID string - var tenantID string - var tenantKey string - if err := rows.Scan(&apiKeyID, &hash, &keyPrefix, &keyName, &scopesBytes, &userGroupID, &gatewayUserID, &username, &rolesBytes, &gatewayTenantID, &tenantID, &tenantKey); err != nil { - return nil, err - } - if bcrypt.CompareHashAndPassword([]byte(hash), []byte(secret)) != nil { + for _, candidate := range candidates { + if bcrypt.CompareHashAndPassword([]byte(candidate.hash), []byte(secret)) != nil { continue } - if _, err := s.pool.Exec(ctx, `UPDATE gateway_api_keys SET last_used_at = now(), updated_at = now() WHERE id = $1::uuid`, apiKeyID); err != nil { + if _, err := database.Exec(ctx, `UPDATE gateway_api_keys SET last_used_at = now(), updated_at = now() WHERE id = $1::uuid`, candidate.apiKeyID); err != nil { return nil, err } return &auth.User{ - ID: gatewayUserID, - Username: username, - Roles: decodeStringArray(rolesBytes), - TenantID: tenantID, - GatewayTenantID: gatewayTenantID, - TenantKey: tenantKey, + ID: candidate.gatewayUserID, + Username: candidate.username, + Roles: decodeStringArray(candidate.rolesBytes), + TenantID: candidate.tenantID, + GatewayTenantID: candidate.gatewayTenantID, + TenantKey: candidate.tenantKey, Source: "gateway", - GatewayUserID: gatewayUserID, - UserGroupID: userGroupID, - APIKeyID: apiKeyID, - APIKeyName: keyName, - APIKeyPrefix: keyPrefix, - APIKeyScopes: decodeStringArray(scopesBytes), + GatewayUserID: candidate.gatewayUserID, + UserGroupID: candidate.userGroupID, + APIKeyID: candidate.apiKeyID, + APIKeyName: candidate.keyName, + APIKeyPrefix: candidate.keyPrefix, + APIKeyScopes: decodeStringArray(candidate.scopesBytes), }, nil } - if err := rows.Err(); err != nil { - return nil, err - } return nil, auth.ErrUnauthorized }