BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(auth): 消除 API Key 校验连接池死锁

Browse Source 
先完整收集同前缀候选项并关闭查询结果，再执行 bcrypt 比对和 last_used_at 更新，避免小连接池下查询与更新相互等待。

新增 Rows 关闭顺序、前缀碰撞、MaxConns=1 和 8 并发真实 PostgreSQL 回归测试。
This commit is contained in:
chengcheng
2026-07-21 10:57:33 +08:00
parent dcf5c4f340
commit 4f45f38b8e
3 changed files with 329 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/api_key_auth_integration_test.go
+99
View File
@@ -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
}
apps/api/internal/store/api_key_auth_test.go
+169
View File
@@ -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
}
apps/api/internal/store/postgres.go
+61 -35
View File
@@ -1484,11 +1484,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, '')
@@ -1503,49 +1527,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
}