fix(auth): 消除 API Key 校验连接池死锁
先完整收集同前缀候选项并关闭查询结果,再执行 bcrypt 比对和 last_used_at 更新,避免小连接池下查询与更新相互等待。 新增 Rows 关闭顺序、前缀碰撞、MaxConns=1 和 8 并发真实 PostgreSQL 回归测试。
This commit is contained in:
@@ -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
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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) {
|
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)
|
prefix := apiKeyPrefix(secret)
|
||||||
if prefix == "" {
|
if prefix == "" {
|
||||||
return nil, auth.ErrUnauthorized
|
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, ''),
|
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, ''),
|
u.id::text, u.username, u.roles, COALESCE(u.gateway_tenant_id::text, ''),
|
||||||
COALESCE(u.tenant_id, ''), COALESCE(u.tenant_key, '')
|
COALESCE(u.tenant_id, ''), COALESCE(u.tenant_key, '')
|
||||||
@@ -1503,49 +1527,51 @@ WHERE k.key_prefix = $1
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return nil, err
|
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() {
|
for _, candidate := range candidates {
|
||||||
var apiKeyID string
|
if bcrypt.CompareHashAndPassword([]byte(candidate.hash), []byte(secret)) != nil {
|
||||||
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 {
|
|
||||||
continue
|
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 nil, err
|
||||||
}
|
}
|
||||||
return &auth.User{
|
return &auth.User{
|
||||||
ID: gatewayUserID,
|
ID: candidate.gatewayUserID,
|
||||||
Username: username,
|
Username: candidate.username,
|
||||||
Roles: decodeStringArray(rolesBytes),
|
Roles: decodeStringArray(candidate.rolesBytes),
|
||||||
TenantID: tenantID,
|
TenantID: candidate.tenantID,
|
||||||
GatewayTenantID: gatewayTenantID,
|
GatewayTenantID: candidate.gatewayTenantID,
|
||||||
TenantKey: tenantKey,
|
TenantKey: candidate.tenantKey,
|
||||||
Source: "gateway",
|
Source: "gateway",
|
||||||
GatewayUserID: gatewayUserID,
|
GatewayUserID: candidate.gatewayUserID,
|
||||||
UserGroupID: userGroupID,
|
UserGroupID: candidate.userGroupID,
|
||||||
APIKeyID: apiKeyID,
|
APIKeyID: candidate.apiKeyID,
|
||||||
APIKeyName: keyName,
|
APIKeyName: candidate.keyName,
|
||||||
APIKeyPrefix: keyPrefix,
|
APIKeyPrefix: candidate.keyPrefix,
|
||||||
APIKeyScopes: decodeStringArray(scopesBytes),
|
APIKeyScopes: decodeStringArray(candidate.scopesBytes),
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
if err := rows.Err(); err != nil {
|
|
||||||
return nil, err
|
|
||||||
}
|
|
||||||
return nil, auth.ErrUnauthorized
|
return nil, auth.ErrUnauthorized
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user