fix(identity): 完善统一认证配对恢复与安全退役

修复 credentials_saved 状态无法恢复、配对与激活并发冲突,以及 SSF 和身份 Secret 生命周期不完整的问题。新增持久化协调器、取消与清理状态机、事务级并发门禁、受控 SSF 凭据交接、禁用后的延迟 Secret 清理,并对生产环境统一认证及 Discovery 端点强制 HTTPS。

验证:go test ./...;go test -race ./internal/auth ./internal/identity ./internal/identityruntime ./internal/securityevents ./internal/httpapi ./internal/store -count=1;go vet ./...;真实 PostgreSQL 并发及清理成功/冲突回滚测试;pnpm openapi。
This commit is contained in:
2026-07-17 18:31:12 +08:00
parent cdfca61304
commit a312ad880d
55 changed files with 9225 additions and 419 deletions
@@ -0,0 +1,141 @@
package store
import (
"context"
"errors"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/identity"
"github.com/google/uuid"
)
func TestDisableActiveIdentityRevisionQueuesSecretsAfterRuntimeRetirementAndClearsReferences(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
seedIdentityActivationPrerequisites(t, ctx, db, "default")
revisionID := seedIdentityLifecycleRevision(t, ctx, db, identity.RevisionActive, "default")
machineReference := "identity-machine-disable-" + uuid.NewString()
sessionReference := "identity-session-disable-" + uuid.NewString()
if _, err := db.pool.Exec(ctx, `UPDATE gateway_identity_configuration_revisions
SET machine_credential_ref=$2,session_encryption_key_ref=$3 WHERE id=$1::uuid`,
revisionID, machineReference, sessionReference); err != nil {
t.Fatalf("seed active identity Secret references: %v", err)
}
if err := db.QueueIdentitySecretCleanup(ctx, machineReference, time.Now().Add(-time.Minute)); err != nil {
t.Fatalf("seed prematurely due machine Secret cleanup: %v", err)
}
disableStartedAt := time.Now().UTC()
disabled, err := db.DisableActiveIdentityRevision(ctx, 1, "trace-disable-secrets", "audit-disable-secrets")
if err != nil {
t.Fatalf("disable active identity revision: %v", err)
}
if disabled.State != identity.RevisionSuperseded || disabled.Version != 2 ||
disabled.MachineCredentialRef != "" || disabled.SessionEncryptionKeyRef != "" {
t.Fatalf("disabled revision retained Secret references: %#v", disabled)
}
persisted, err := db.IdentityConfigurationRevision(ctx, revisionID)
if err != nil {
t.Fatalf("read disabled identity revision: %v", err)
}
if persisted.State != identity.RevisionSuperseded || persisted.MachineCredentialRef != "" ||
persisted.SessionEncryptionKeyRef != "" {
t.Fatalf("persisted disabled revision retained Secret references: %#v", persisted)
}
type queuedSecret struct {
reference string
status string
notBefore time.Time
}
rows, err := db.pool.Query(ctx, `SELECT secret_ref,status,not_before
FROM gateway_identity_secret_cleanup_queue ORDER BY secret_ref`)
if err != nil {
t.Fatalf("read disabled identity Secret cleanup queue: %v", err)
}
defer rows.Close()
queued := make(map[string]queuedSecret, 2)
for rows.Next() {
var secret queuedSecret
if err := rows.Scan(&secret.reference, &secret.status, &secret.notBefore); err != nil {
t.Fatalf("scan disabled identity Secret cleanup: %v", err)
}
queued[secret.reference] = secret
}
if err := rows.Err(); err != nil {
t.Fatalf("iterate disabled identity Secret cleanup queue: %v", err)
}
if len(queued) != 2 {
t.Fatalf("disabled identity cleanup queue entries=%v, want both Secret references", queued)
}
minimumSafeCleanup := disableStartedAt.Add(45 * time.Second)
for _, reference := range []string{machineReference, sessionReference} {
secret, ok := queued[reference]
if !ok || secret.status != "pending" || secret.notBefore.Before(minimumSafeCleanup) {
t.Fatalf("queued Secret %q=%#v present=%t, want pending cleanup no earlier than %s",
reference, secret, ok, minimumSafeCleanup)
}
}
claims, err := db.ClaimIdentitySecretCleanups(ctx, 10, time.Minute)
if err != nil {
t.Fatalf("claim not-yet-due disabled identity Secrets: %v", err)
}
if len(claims) != 0 {
t.Fatalf("disabled identity Secrets became claimable before Runtime retirement: %#v", claims)
}
}
func TestDisableActiveIdentityRevisionCleanupConflictRollsBackStateReferencesAndQueue(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
seedIdentityActivationPrerequisites(t, ctx, db, "default")
revisionID := seedIdentityLifecycleRevision(t, ctx, db, identity.RevisionActive, "default")
machineReference := "identity-machine-disable-conflict-" + uuid.NewString()
sessionReference := "identity-session-disable-conflict-" + uuid.NewString()
if _, err := db.pool.Exec(ctx, `UPDATE gateway_identity_configuration_revisions
SET machine_credential_ref=$2,session_encryption_key_ref=$3 WHERE id=$1::uuid`,
revisionID, machineReference, sessionReference); err != nil {
t.Fatalf("seed conflicting active identity Secret references: %v", err)
}
if err := db.QueueIdentitySecretCleanup(ctx, machineReference, time.Now().Add(-time.Minute)); err != nil {
t.Fatalf("queue machine Secret before conflict claim: %v", err)
}
claims, err := db.ClaimIdentitySecretCleanups(ctx, 1, time.Minute)
if err != nil || len(claims) != 1 || claims[0].Reference != machineReference {
t.Fatalf("claim machine Secret before disable: claims=%#v error=%v", claims, err)
}
_, err = db.DisableActiveIdentityRevision(ctx, 1, "trace-disable-conflict", "audit-disable-conflict")
if !errors.Is(err, ErrIdentitySecretCleanupConflict) {
t.Fatalf("disable with claimed Secret cleanup error=%v, want cleanup conflict", err)
}
persisted, err := db.IdentityConfigurationRevision(ctx, revisionID)
if err != nil {
t.Fatalf("read active identity revision after cleanup conflict: %v", err)
}
if persisted.State != identity.RevisionActive || persisted.Version != 1 ||
persisted.MachineCredentialRef != machineReference || persisted.SessionEncryptionKeyRef != sessionReference {
t.Fatalf("cleanup conflict partially disabled identity revision: %#v", persisted)
}
var status, claimToken string
if err := db.pool.QueryRow(ctx, `SELECT status,claim_token::text FROM gateway_identity_secret_cleanup_queue
WHERE secret_ref=$1`, machineReference).Scan(&status, &claimToken); err != nil {
t.Fatalf("read claimed machine Secret after disable rollback: %v", err)
}
if status != "claimed" || claimToken != claims[0].ClaimToken {
t.Fatalf("disable rollback changed claimed cleanup status=%q token=%q", status, claimToken)
}
var sessionQueueEntries int
if err := db.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_identity_secret_cleanup_queue
WHERE secret_ref=$1`, sessionReference).Scan(&sessionQueueEntries); err != nil {
t.Fatalf("count session Secret cleanup after disable rollback: %v", err)
}
if sessionQueueEntries != 0 {
t.Fatalf("disable cleanup conflict left %d partial session cleanup entries", sessionQueueEntries)
}
}
@@ -29,6 +29,55 @@ COALESCE(machine_credential_ref,''),COALESCE(session_encryption_key_ref,''),sess
session_refresh_seconds,version,COALESCE(last_error_category,''),COALESCE(last_trace_id,''),COALESCE(last_audit_id,''),
validated_at,activated_at,superseded_at,created_at,updated_at`
// identityConfigurationLifecycleLockID serializes the database-side boundary
// between onboarding reservations and Active Revision mutations. A transaction
// must hold this lock before deciding that no Active configuration or foreign
// pairing reservation exists.
const identityConfigurationLifecycleLockID int64 = 0x4541494944454e54
// identityDisabledSecretCleanupDelay is intentionally longer than the
// IdentityRuntimeManager's 30-second old-Runtime retirement window. Starting
// cleanup one minute after the database handoff leaves an additional 30
// seconds for commit, Runtime publication, timer scheduling, and worker jitter.
const identityDisabledSecretCleanupDelay = time.Minute
func lockIdentityConfigurationLifecycle(ctx context.Context, tx pgx.Tx) error {
_, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock($1)`, identityConfigurationLifecycleLockID)
return err
}
// beginIdentityConfigurationLifecycleTx deliberately uses READ COMMITTED.
// PostgreSQL fixes the transaction snapshot when a SERIALIZABLE transaction
// executes the advisory-lock SELECT, before that statement finishes waiting.
// A later state check can therefore miss the preceding lock holder's commit.
// READ COMMITTED gives every statement after lock acquisition a fresh snapshot,
// while the transaction-scoped advisory lock serializes all lifecycle decisions.
func (s *Store) beginIdentityConfigurationLifecycleTx(ctx context.Context) (pgx.Tx, error) {
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.ReadCommitted})
if err != nil {
return nil, err
}
if err := lockIdentityConfigurationLifecycle(ctx, tx); err != nil {
_ = tx.Rollback(ctx)
return nil, err
}
return tx, nil
}
func queueDisabledIdentitySecretCleanupTx(ctx context.Context, tx pgx.Tx, reference string, notBefore time.Time) error {
tag, err := tx.Exec(ctx, `INSERT INTO gateway_identity_secret_cleanup_queue(secret_ref,not_before)
VALUES($1,$2) ON CONFLICT(secret_ref) DO UPDATE
SET not_before=GREATEST(gateway_identity_secret_cleanup_queue.not_before,EXCLUDED.not_before),updated_at=now()
WHERE gateway_identity_secret_cleanup_queue.status='pending'`, reference, notBefore)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrIdentitySecretCleanupConflict
}
return nil
}
func (s *Store) CreateIdentityConfigurationRevision(ctx context.Context, revision identity.Revision) (identity.Revision, error) {
scopes, _ := json.Marshal(revision.Scopes)
capabilities, _ := json.Marshal(revision.Capabilities)
@@ -58,7 +107,16 @@ FROM gateway_identity_configuration_revisions WHERE state='active'`))
func (s *Store) LatestInactiveIdentityConfigurationRevision(ctx context.Context) (identity.Revision, error) {
revision, err := scanIdentityRevision(s.pool.QueryRow(ctx, `SELECT `+identityRevisionColumns+`
FROM gateway_identity_configuration_revisions WHERE state IN ('draft','validated','failed') ORDER BY created_at DESC LIMIT 1`))
FROM gateway_identity_configuration_revisions
WHERE id=(
SELECT revision.id
FROM gateway_identity_configuration_revisions revision
JOIN gateway_identity_onboarding_exchanges exchange ON exchange.revision_id=revision.id
WHERE revision.state IN ('draft','validated','failed')
ORDER BY NOT (exchange.status='cancelled' AND exchange.cleanup_status='completed') DESC,
revision.created_at DESC,exchange.created_at DESC
LIMIT 1
)`))
return revision, normalizeIdentityRevisionError(err)
}
@@ -108,10 +166,16 @@ VALUES($1,$2,$3,$4::jsonb) ON CONFLICT(operation,idempotency_key) DO NOTHING`,
}
func (s *Store) ApplyIdentityManifest(ctx context.Context, id string, expectedVersion int64, applied identity.ManifestApplication) (identity.Revision, error) {
current, err := s.IdentityConfigurationRevision(ctx, id)
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
if err != nil {
return identity.Revision{}, err
}
defer tx.Rollback(ctx)
current, err := scanIdentityRevision(tx.QueryRow(ctx, `SELECT `+identityRevisionColumns+`
FROM gateway_identity_configuration_revisions WHERE id=$1::uuid FOR UPDATE`, id))
if err != nil {
return identity.Revision{}, normalizeIdentityRevisionError(err)
}
if current.Version != expectedVersion {
return identity.Revision{}, identity.ErrRevisionConflict
}
@@ -119,9 +183,20 @@ func (s *Store) ApplyIdentityManifest(ctx context.Context, id string, expectedVe
if err != nil {
return identity.Revision{}, err
}
newReferences := make([]string, 0, 2)
for _, reference := range []string{updated.MachineCredentialRef, updated.SessionEncryptionKeyRef} {
if reference != "" {
newReferences = append(newReferences, reference)
}
}
if len(newReferences) > 0 {
if err := adoptPendingIdentitySecrets(ctx, tx, newReferences...); err != nil {
return identity.Revision{}, err
}
}
scopes, _ := json.Marshal(updated.Scopes)
capabilities, _ := json.Marshal(updated.Capabilities)
revision, err := scanIdentityRevision(s.pool.QueryRow(ctx, `
revision, err := scanIdentityRevision(tx.QueryRow(ctx, `
UPDATE gateway_identity_configuration_revisions SET
issuer=$3,tenant_id=$4,application_id=$5,audience=NULLIF($6,''),browser_client_id=NULLIF($7,''),machine_client_id=NULLIF($8,''),
scopes=$9::jsonb,capabilities=$10::jsonb,token_introspection=$11,session_revocation=$12,
@@ -138,14 +213,35 @@ RETURNING `+identityRevisionColumns,
if errors.Is(err, pgx.ErrNoRows) {
return identity.Revision{}, identity.ErrRevisionConflict
}
return revision, err
if err != nil {
return identity.Revision{}, err
}
activeReferences := map[string]struct{}{}
for _, reference := range newReferences {
activeReferences[reference] = struct{}{}
}
for _, reference := range []string{current.MachineCredentialRef, current.SessionEncryptionKeyRef} {
if reference == "" {
continue
}
if _, stillActive := activeReferences[reference]; stillActive {
continue
}
if err := queueIdentitySecretCleanupTx(ctx, tx, reference, time.Now().UTC()); err != nil {
return identity.Revision{}, err
}
}
if err := tx.Commit(ctx); err != nil {
return identity.Revision{}, err
}
return revision, nil
}
func (s *Store) MarkIdentityRevisionValidated(ctx context.Context, id string, expectedVersion int64, traceID, auditID string) (identity.Revision, error) {
revision, err := scanIdentityRevision(s.pool.QueryRow(ctx, `
UPDATE gateway_identity_configuration_revisions SET state='validated',validated_at=now(),last_error_category=NULL,
last_trace_id=NULLIF($3,''),last_audit_id=NULLIF($4,''),version=version+1,updated_at=now()
WHERE id=$1::uuid AND version=$2 AND state IN ('draft','superseded')
WHERE id=$1::uuid AND version=$2 AND state='draft'
RETURNING `+identityRevisionColumns, id, expectedVersion, traceID, auditID))
if errors.Is(err, pgx.ErrNoRows) {
return identity.Revision{}, identity.ErrRevisionConflict
@@ -178,11 +274,22 @@ RETURNING `+identityRevisionColumns, id, expectedVersion, category, traceID, aud
}
func (s *Store) ActivateIdentityRevision(ctx context.Context, id string, expectedVersion int64, traceID, auditID string) (identity.Revision, bool, error) {
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
tx, err := s.beginIdentityConfigurationLifecycleTx(ctx)
if err != nil {
return identity.Revision{}, false, err
}
defer tx.Rollback(ctx)
var foreignPairingReservation bool
if err := tx.QueryRow(ctx, `SELECT EXISTS(
SELECT 1 FROM gateway_identity_pairing_start_reservation
WHERE (state='starting' AND expires_at > now())
OR (state='paired' AND revision_id IS DISTINCT FROM $1::uuid)
)`, id).Scan(&foreignPairingReservation); err != nil {
return identity.Revision{}, false, err
}
if foreignPairingReservation {
return identity.Revision{}, false, identity.ErrPairingInProgress
}
if ok, err := hasBreakGlassManager(ctx, tx); err != nil {
return identity.Revision{}, false, err
} else if !ok {
@@ -197,9 +304,16 @@ SELECT local_tenant_key FROM gateway_identity_configuration_revisions WHERE id=$
return identity.Revision{}, false, identity.ErrLocalTenantInvalid
}
var previousID, previousIssuer, previousTenant, previousAudience, previousClient, previousSessionRef string
_ = tx.QueryRow(ctx, `SELECT id::text,COALESCE(issuer,''),COALESCE(tenant_id,''),COALESCE(audience,''),
if err := tx.QueryRow(ctx, `SELECT id::text,COALESCE(issuer,''),COALESCE(tenant_id,''),COALESCE(audience,''),
COALESCE(browser_client_id,''),COALESCE(session_encryption_key_ref,'') FROM gateway_identity_configuration_revisions
WHERE state='active' FOR UPDATE`).Scan(&previousID, &previousIssuer, &previousTenant, &previousAudience, &previousClient, &previousSessionRef)
WHERE state='active' FOR UPDATE`).Scan(
&previousID, &previousIssuer, &previousTenant, &previousAudience, &previousClient, &previousSessionRef,
); err != nil && !errors.Is(err, pgx.ErrNoRows) {
return identity.Revision{}, false, err
}
if previousID != "" && previousID != id {
return identity.Revision{}, false, identity.ErrActiveConfigurationHandoffRequired
}
var nextIssuer, nextTenant, nextAudience, nextClient, nextSessionRef string
if err := tx.QueryRow(ctx, `SELECT COALESCE(issuer,''),COALESCE(tenant_id,''),COALESCE(audience,''),
COALESCE(browser_client_id,''),COALESCE(session_encryption_key_ref,'') FROM gateway_identity_configuration_revisions
@@ -228,6 +342,9 @@ WHERE id=$1::uuid AND version=$2 AND state='validated' RETURNING `+identityRevis
return identity.Revision{}, false, err
}
}
if _, err := tx.Exec(ctx, `DELETE FROM gateway_identity_pairing_start_reservation WHERE revision_id=$1::uuid`, id); err != nil {
return identity.Revision{}, false, err
}
if err := tx.Commit(ctx); err != nil {
return identity.Revision{}, false, err
}
@@ -235,7 +352,7 @@ WHERE id=$1::uuid AND version=$2 AND state='validated' RETURNING `+identityRevis
}
func (s *Store) DisableActiveIdentityRevision(ctx context.Context, expectedVersion int64, traceID, auditID string) (identity.Revision, error) {
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
tx, err := s.beginIdentityConfigurationLifecycleTx(ctx)
if err != nil {
return identity.Revision{}, err
}
@@ -245,9 +362,33 @@ func (s *Store) DisableActiveIdentityRevision(ctx context.Context, expectedVersi
} else if !ok {
return identity.Revision{}, identity.ErrBreakGlassRequired
}
var activeID, machineReference, sessionReference string
if err := tx.QueryRow(ctx, `SELECT id::text,COALESCE(machine_credential_ref,''),COALESCE(session_encryption_key_ref,'')
FROM gateway_identity_configuration_revisions
WHERE state='active' AND version=$1 FOR UPDATE`, expectedVersion).Scan(
&activeID, &machineReference, &sessionReference,
); errors.Is(err, pgx.ErrNoRows) {
return identity.Revision{}, identity.ErrRevisionConflict
} else if err != nil {
return identity.Revision{}, err
}
cleanupNotBefore := time.Now().UTC().Add(identityDisabledSecretCleanupDelay)
references := make(map[string]struct{}, 2)
for _, reference := range []string{machineReference, sessionReference} {
if reference != "" {
references[reference] = struct{}{}
}
}
for reference := range references {
if err := queueDisabledIdentitySecretCleanupTx(ctx, tx, reference, cleanupNotBefore); err != nil {
return identity.Revision{}, err
}
}
revision, err := scanIdentityRevision(tx.QueryRow(ctx, `UPDATE gateway_identity_configuration_revisions SET state='superseded',
superseded_at=now(),last_trace_id=NULLIF($2,''),last_audit_id=NULLIF($3,''),version=version+1,updated_at=now()
WHERE state='active' AND version=$1 RETURNING `+identityRevisionColumns, expectedVersion, traceID, auditID))
superseded_at=now(),machine_credential_ref=NULL,session_encryption_key_ref=NULL,
last_trace_id=NULLIF($3,''),last_audit_id=NULLIF($4,''),version=version+1,updated_at=now()
WHERE id=$1::uuid AND state='active' AND version=$2
RETURNING `+identityRevisionColumns, activeID, expectedVersion, traceID, auditID))
if errors.Is(err, pgx.ErrNoRows) {
return identity.Revision{}, identity.ErrRevisionConflict
}
+327 -9
View File
@@ -2,24 +2,156 @@ package store
import (
"context"
"encoding/json"
"errors"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/identity"
"github.com/jackc/pgx/v5"
)
const identityPairingColumns = `
id::text,revision_id::text,remote_exchange_id::text,exchange_token_ref,status,remote_version,expires_at,version,
COALESCE(last_error_category,''),COALESCE(auth_center_audit_id,''),COALESCE(last_trace_id,''),created_at,updated_at`
id::text,revision_id::text,remote_exchange_id::text,exchange_token_ref,status,cleanup_status,remote_version,expires_at,version,
COALESCE(last_error_category,''),COALESCE(auth_center_audit_id,''),COALESCE(last_trace_id,''),cancelled_at,cleanup_completed_at,
created_at,updated_at`
func (s *Store) ReserveIdentityPairingStart(ctx context.Context, attemptID string, expiresAt time.Time) error {
tx, err := s.beginIdentityConfigurationLifecycleTx(ctx)
if err != nil {
return err
}
defer tx.Rollback(ctx)
if _, err := tx.Exec(ctx, `DELETE FROM gateway_identity_pairing_start_reservation
WHERE state='starting' AND expires_at <= now()`); err != nil {
return err
}
var activeConfigurationExists bool
if err := tx.QueryRow(ctx, `SELECT EXISTS(
SELECT 1 FROM gateway_identity_configuration_revisions
WHERE state='active'
)`).Scan(&activeConfigurationExists); err != nil {
return err
}
if activeConfigurationExists {
return identity.ErrActiveConfigurationHandoffRequired
}
var blocked bool
if err := tx.QueryRow(ctx, `SELECT EXISTS(
SELECT 1 FROM gateway_identity_configuration_revisions revision
JOIN gateway_identity_onboarding_exchanges exchange ON exchange.revision_id=revision.id
WHERE revision.state IN ('draft','validated','failed')
AND NOT (exchange.status='cancelled' AND exchange.cleanup_status='completed')
)`).Scan(&blocked); err != nil {
return err
}
if blocked {
return identity.ErrPairingInProgress
}
if _, err := tx.Exec(ctx, `INSERT INTO gateway_identity_pairing_start_reservation(singleton,attempt_id,state,expires_at)
VALUES(true,$1::uuid,'starting',$2)`, attemptID, expiresAt); err != nil {
if isUniqueViolation(err) {
return identity.ErrPairingInProgress
}
return err
}
return tx.Commit(ctx)
}
func (s *Store) ReleaseIdentityPairingStart(ctx context.Context, attemptID string) error {
_, err := s.pool.Exec(ctx, `DELETE FROM gateway_identity_pairing_start_reservation
WHERE attempt_id=$1::uuid AND state='starting'`, attemptID)
return err
}
func (s *Store) CommitIdentityPairingStart(ctx context.Context, revision identity.Revision, exchange identity.PairingExchange) (identity.PairingExchange, error) {
tx, err := s.beginIdentityConfigurationLifecycleTx(ctx)
if err != nil {
return identity.PairingExchange{}, err
}
defer tx.Rollback(ctx)
var reserved bool
if err := tx.QueryRow(ctx, `SELECT true FROM gateway_identity_pairing_start_reservation
WHERE singleton=true AND attempt_id=$1::uuid AND state='starting' FOR UPDATE`, exchange.ID).Scan(&reserved); errors.Is(err, pgx.ErrNoRows) {
return identity.PairingExchange{}, identity.ErrPairingInProgress
} else if err != nil {
return identity.PairingExchange{}, err
}
if !reserved {
return identity.PairingExchange{}, identity.ErrPairingInProgress
}
var activeConfigurationExists bool
if err := tx.QueryRow(ctx, `SELECT EXISTS(
SELECT 1 FROM gateway_identity_configuration_revisions WHERE state='active'
)`).Scan(&activeConfigurationExists); err != nil {
return identity.PairingExchange{}, err
}
if activeConfigurationExists {
return identity.PairingExchange{}, identity.ErrActiveConfigurationHandoffRequired
}
if err := adoptPendingIdentitySecrets(ctx, tx, exchange.ExchangeTokenRef); err != nil {
return identity.PairingExchange{}, err
}
scopes, _ := json.Marshal(revision.Scopes)
capabilities, _ := json.Marshal(revision.Capabilities)
if _, err := tx.Exec(ctx, `INSERT INTO gateway_identity_configuration_revisions (
id,state,schema_version,auth_center_url,role_prefix,local_tenant_key,public_base_url,web_base_url,
jit_enabled,legacy_jwt_enabled,scopes,capabilities,session_idle_seconds,session_absolute_seconds,session_refresh_seconds,last_trace_id
) VALUES ($1::uuid,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11::jsonb,$12::jsonb,$13,$14,$15,NULLIF($16,''))`,
revision.ID, revision.State, revision.SchemaVersion, revision.AuthCenterURL, revision.RolePrefix,
revision.LocalTenantKey, revision.PublicBaseURL, revision.WebBaseURL, revision.JITEnabled, revision.LegacyJWTEnabled,
string(scopes), string(capabilities), revision.SessionIdleSeconds, revision.SessionAbsoluteSeconds,
revision.SessionRefreshSeconds, revision.LastTraceID,
); err != nil {
return identity.PairingExchange{}, err
}
created, err := scanIdentityPairing(tx.QueryRow(ctx, `INSERT INTO gateway_identity_onboarding_exchanges (
id,revision_id,remote_exchange_id,exchange_token_ref,status,cleanup_status,remote_version,expires_at,last_trace_id
) VALUES ($1::uuid,$2::uuid,$3::uuid,$4,$5,$6,$7,$8,NULLIF($9,''))
RETURNING `+identityPairingColumns,
exchange.ID, revision.ID, exchange.RemoteExchangeID, exchange.ExchangeTokenRef,
exchange.Status, exchange.CleanupStatus, exchange.RemoteVersion, exchange.ExpiresAt, exchange.LastTraceID,
))
if err != nil {
return identity.PairingExchange{}, err
}
tag, err := tx.Exec(ctx, `UPDATE gateway_identity_pairing_start_reservation
SET state='paired',revision_id=$2::uuid,expires_at=$3,updated_at=now()
WHERE attempt_id=$1::uuid AND state='starting'`, exchange.ID, revision.ID, exchange.ExpiresAt)
if err != nil {
return identity.PairingExchange{}, err
}
if tag.RowsAffected() != 1 {
return identity.PairingExchange{}, identity.ErrPairingInProgress
}
if err := tx.Commit(ctx); err != nil {
return identity.PairingExchange{}, err
}
return created, nil
}
func (s *Store) IdentityPairingStartBlocked(ctx context.Context) (bool, error) {
var blocked bool
err := s.pool.QueryRow(ctx, `
SELECT EXISTS(
SELECT 1 FROM gateway_identity_configuration_revisions revision
JOIN gateway_identity_onboarding_exchanges exchange ON exchange.revision_id=revision.id
WHERE revision.state IN ('draft','validated','failed')
AND NOT (exchange.status='cancelled' AND exchange.cleanup_status='completed')
) OR EXISTS (
SELECT 1 FROM gateway_identity_pairing_start_reservation
WHERE state='paired' OR expires_at > now()
)`).Scan(&blocked)
return blocked, err
}
func (s *Store) CreateIdentityPairingExchange(ctx context.Context, exchange identity.PairingExchange) (identity.PairingExchange, error) {
return scanIdentityPairing(s.pool.QueryRow(ctx, `
INSERT INTO gateway_identity_onboarding_exchanges (
id,revision_id,remote_exchange_id,exchange_token_ref,status,remote_version,expires_at,last_trace_id
) VALUES ($1::uuid,$2::uuid,$3::uuid,$4,$5,$6,$7,NULLIF($8,''))
id,revision_id,remote_exchange_id,exchange_token_ref,status,cleanup_status,remote_version,expires_at,last_trace_id
) VALUES ($1::uuid,$2::uuid,$3::uuid,$4,$5,$6,$7,$8,NULLIF($9,''))
RETURNING `+identityPairingColumns,
exchange.ID, exchange.RevisionID, exchange.RemoteExchangeID, exchange.ExchangeTokenRef,
exchange.Status, exchange.RemoteVersion, exchange.ExpiresAt, exchange.LastTraceID,
exchange.Status, exchange.CleanupStatus, exchange.RemoteVersion, exchange.ExpiresAt, exchange.LastTraceID,
))
}
@@ -43,8 +175,67 @@ FROM gateway_identity_onboarding_exchanges WHERE revision_id=$1::uuid`, revision
func (s *Store) PendingIdentityPairingExchanges(ctx context.Context) ([]identity.PairingExchange, error) {
rows, err := s.pool.Query(ctx, `SELECT `+identityPairingColumns+`
FROM gateway_identity_onboarding_exchanges exchange
WHERE exchange.id=(
SELECT candidate_exchange.id
FROM gateway_identity_onboarding_exchanges candidate_exchange
JOIN gateway_identity_configuration_revisions revision ON revision.id=candidate_exchange.revision_id
WHERE revision.state IN ('draft','validated','failed')
ORDER BY NOT (candidate_exchange.status='cancelled' AND candidate_exchange.cleanup_status='completed') DESC,
revision.created_at DESC,candidate_exchange.created_at DESC
LIMIT 1
)
AND exchange.status IN ('metadata_pending','preparing','ready','credentials_saved')`)
if err != nil {
return nil, err
}
defer rows.Close()
exchanges := make([]identity.PairingExchange, 0)
for rows.Next() {
exchange, scanErr := scanIdentityPairing(rows)
if scanErr != nil {
return nil, scanErr
}
exchanges = append(exchanges, exchange)
}
return exchanges, rows.Err()
}
func (s *Store) PendingIdentityPairingCleanups(ctx context.Context) ([]identity.PairingExchange, error) {
rows, err := s.pool.Query(ctx, `SELECT `+identityPairingColumns+`
FROM gateway_identity_onboarding_exchanges
WHERE status IN ('metadata_pending','preparing','ready','credentials_saved') ORDER BY created_at`)
WHERE status='cancelled' AND cleanup_status='pending' ORDER BY updated_at`)
if err != nil {
return nil, err
}
defer rows.Close()
exchanges := make([]identity.PairingExchange, 0)
for rows.Next() {
exchange, scanErr := scanIdentityPairing(rows)
if scanErr != nil {
return nil, scanErr
}
exchanges = append(exchanges, exchange)
}
return exchanges, rows.Err()
}
func (s *Store) CompletedIdentityPairingsAwaitingActivation(ctx context.Context) ([]identity.PairingExchange, error) {
rows, err := s.pool.Query(ctx, `SELECT `+identityPairingColumns+`
FROM gateway_identity_onboarding_exchanges exchange
WHERE exchange.id=(
SELECT candidate_exchange.id
FROM gateway_identity_onboarding_exchanges candidate_exchange
JOIN gateway_identity_configuration_revisions revision ON revision.id=candidate_exchange.revision_id
WHERE revision.state IN ('draft','validated','failed')
ORDER BY NOT (candidate_exchange.status='cancelled' AND candidate_exchange.cleanup_status='completed') DESC,
revision.created_at DESC,candidate_exchange.created_at DESC
LIMIT 1
)
AND exchange.status='completed' AND EXISTS (
SELECT 1 FROM gateway_identity_configuration_revisions revision
WHERE revision.id=exchange.revision_id AND revision.state IN ('draft','validated')
)`)
if err != nil {
return nil, err
}
@@ -61,13 +252,139 @@ WHERE status IN ('metadata_pending','preparing','ready','credentials_saved') ORD
}
func (s *Store) UpdateIdentityPairingExchange(ctx context.Context, id string, expectedVersion int64, update identity.PairingExchangeUpdate) (identity.PairingExchange, error) {
exchange, err := scanIdentityPairing(s.pool.QueryRow(ctx, `
tx, err := s.pool.Begin(ctx)
if err != nil {
return identity.PairingExchange{}, err
}
defer tx.Rollback(ctx)
exchange, err := scanIdentityPairing(tx.QueryRow(ctx, `
UPDATE gateway_identity_onboarding_exchanges SET status=$3,remote_version=$4,last_error_category=NULLIF($5,''),
auth_center_audit_id=COALESCE(NULLIF($6,''),auth_center_audit_id),version=version+1,updated_at=now()
WHERE id=$1::uuid AND version=$2
RETURNING `+identityPairingColumns,
id, expectedVersion, update.Status, update.RemoteVersion, update.LastErrorCategory, update.AuthCenterAuditID,
))
if errors.Is(err, pgx.ErrNoRows) {
return identity.PairingExchange{}, identity.ErrRevisionConflict
}
if err != nil {
return identity.PairingExchange{}, err
}
if update.Status == identity.PairingCompleted || update.Status == identity.PairingFailed || update.Status == identity.PairingExpired {
if err := queueIdentitySecretCleanupTx(ctx, tx, exchange.ExchangeTokenRef, time.Now().UTC()); err != nil {
return identity.PairingExchange{}, err
}
}
if err := tx.Commit(ctx); err != nil {
return identity.PairingExchange{}, err
}
return exchange, nil
}
func (s *Store) RecordIdentityPairingRetryFailure(ctx context.Context, id string, status identity.PairingStatus, category string) (identity.PairingExchange, error) {
exchange, err := scanIdentityPairing(s.pool.QueryRow(ctx, `
UPDATE gateway_identity_onboarding_exchanges SET last_error_category=$3,updated_at=now()
WHERE id=$1::uuid AND status=$2 AND cleanup_status='none'
RETURNING `+identityPairingColumns, id, status, category))
if errors.Is(err, pgx.ErrNoRows) {
return identity.PairingExchange{}, identity.ErrRevisionConflict
}
return exchange, err
}
func (s *Store) CancelIdentityPairingExchange(ctx context.Context, id string, expectedVersion int64, traceID, auditID string) (identity.PairingExchange, error) {
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
if err != nil {
return identity.PairingExchange{}, err
}
defer tx.Rollback(ctx)
exchange, err := scanIdentityPairing(tx.QueryRow(ctx, `SELECT `+identityPairingColumns+`
FROM gateway_identity_onboarding_exchanges WHERE id=$1::uuid FOR UPDATE`, id))
if errors.Is(err, pgx.ErrNoRows) {
return identity.PairingExchange{}, identity.ErrRevisionNotFound
}
if err != nil {
return identity.PairingExchange{}, err
}
if exchange.Status == identity.PairingCancelled {
return exchange, tx.Commit(ctx)
}
if exchange.Version != expectedVersion {
return identity.PairingExchange{}, identity.ErrRevisionConflict
}
revision, err := scanIdentityRevision(tx.QueryRow(ctx, `SELECT `+identityRevisionColumns+`
FROM gateway_identity_configuration_revisions WHERE id=$1::uuid FOR UPDATE`, exchange.RevisionID))
if err != nil {
return identity.PairingExchange{}, normalizeIdentityRevisionError(err)
}
if revision.State != identity.RevisionDraft && revision.State != identity.RevisionValidated && revision.State != identity.RevisionFailed {
return identity.PairingExchange{}, identity.ErrPairingNotCancellable
}
if _, err := tx.Exec(ctx, `UPDATE gateway_identity_configuration_revisions SET state='failed',
last_error_category='pairing_cancelled',last_trace_id=NULLIF($2,''),last_audit_id=NULLIF($3,''),
version=version+1,updated_at=now() WHERE id=$1::uuid`, exchange.RevisionID, traceID, auditID); err != nil {
return identity.PairingExchange{}, err
}
exchange, err = scanIdentityPairing(tx.QueryRow(ctx, `
UPDATE gateway_identity_onboarding_exchanges SET status='cancelled',cleanup_status='pending',cancelled_at=now(),
cleanup_completed_at=NULL,version=version+1,updated_at=now() WHERE id=$1::uuid AND version=$2
RETURNING `+identityPairingColumns, id, expectedVersion))
if errors.Is(err, pgx.ErrNoRows) {
return identity.PairingExchange{}, identity.ErrRevisionConflict
}
if err != nil {
return identity.PairingExchange{}, err
}
if err := queueIdentitySecretCleanupTx(ctx, tx, exchange.ExchangeTokenRef, time.Now().UTC()); err != nil {
return identity.PairingExchange{}, err
}
if err := tx.Commit(ctx); err != nil {
return identity.PairingExchange{}, err
}
return exchange, nil
}
func (s *Store) CompleteIdentityPairingCleanup(ctx context.Context, id string, expectedVersion int64) (identity.PairingExchange, error) {
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
if err != nil {
return identity.PairingExchange{}, err
}
defer tx.Rollback(ctx)
var revisionID string
if err := tx.QueryRow(ctx, `SELECT revision_id::text FROM gateway_identity_onboarding_exchanges
WHERE id=$1::uuid AND version=$2 AND status='cancelled' AND cleanup_status='pending' FOR UPDATE`, id, expectedVersion).Scan(&revisionID); errors.Is(err, pgx.ErrNoRows) {
return identity.PairingExchange{}, identity.ErrRevisionConflict
} else if err != nil {
return identity.PairingExchange{}, err
}
tag, err := tx.Exec(ctx, `UPDATE gateway_identity_configuration_revisions SET machine_credential_ref=NULL,
session_encryption_key_ref=NULL,version=version+1,updated_at=now() WHERE id=$1::uuid AND state='failed'`, revisionID)
if err != nil {
return identity.PairingExchange{}, err
}
if tag.RowsAffected() != 1 {
return identity.PairingExchange{}, identity.ErrRevisionConflict
}
exchange, err := scanIdentityPairing(tx.QueryRow(ctx, `UPDATE gateway_identity_onboarding_exchanges
SET cleanup_status='completed',cleanup_completed_at=now(),version=version+1,updated_at=now()
WHERE id=$1::uuid AND version=$2 AND status='cancelled' AND cleanup_status='pending'
RETURNING `+identityPairingColumns, id, expectedVersion))
if err != nil {
return identity.PairingExchange{}, err
}
if _, err := tx.Exec(ctx, `DELETE FROM gateway_identity_pairing_start_reservation WHERE revision_id=$1::uuid`, revisionID); err != nil {
return identity.PairingExchange{}, err
}
if err := tx.Commit(ctx); err != nil {
return identity.PairingExchange{}, err
}
return exchange, nil
}
func (s *Store) RecordIdentityPairingCleanupFailure(ctx context.Context, id, category string) (identity.PairingExchange, error) {
exchange, err := scanIdentityPairing(s.pool.QueryRow(ctx, `UPDATE gateway_identity_onboarding_exchanges
SET last_error_category=$2,updated_at=now() WHERE id=$1::uuid AND status='cancelled' AND cleanup_status='pending'
RETURNING `+identityPairingColumns, id, category))
if errors.Is(err, pgx.ErrNoRows) {
return identity.PairingExchange{}, identity.ErrRevisionConflict
}
@@ -78,9 +395,10 @@ func scanIdentityPairing(row scanner) (identity.PairingExchange, error) {
var exchange identity.PairingExchange
var status string
if err := row.Scan(
&exchange.ID, &exchange.RevisionID, &exchange.RemoteExchangeID, &exchange.ExchangeTokenRef, &status,
&exchange.ID, &exchange.RevisionID, &exchange.RemoteExchangeID, &exchange.ExchangeTokenRef, &status, &exchange.CleanupStatus,
&exchange.RemoteVersion, &exchange.ExpiresAt, &exchange.Version, &exchange.LastErrorCategory,
&exchange.AuthCenterAuditID, &exchange.LastTraceID, &exchange.CreatedAt, &exchange.UpdatedAt,
&exchange.AuthCenterAuditID, &exchange.LastTraceID, &exchange.CancelledAt, &exchange.CleanupCompletedAt,
&exchange.CreatedAt, &exchange.UpdatedAt,
); err != nil {
return identity.PairingExchange{}, err
}
@@ -0,0 +1,715 @@
package store
import (
"context"
"errors"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/identity"
"github.com/google/uuid"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestIdentityPairingCancellationKeepsOlderInactivePairingBlockedUntilEveryCleanupCompletes(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx := context.Background()
type pairingFixture struct {
revisionID string
pairingID string
machineRef string
sessionRef string
createdAt time.Time
displayName string
}
now := time.Now().UTC()
fixtures := []pairingFixture{
{
revisionID: uuid.NewString(), pairingID: uuid.NewString(),
machineRef: "identity-machine-older", sessionRef: "identity-session-older",
createdAt: now.Add(-time.Minute), displayName: "older",
},
{
revisionID: uuid.NewString(), pairingID: uuid.NewString(),
machineRef: "identity-machine-latest", sessionRef: "identity-session-latest",
createdAt: now, displayName: "latest",
},
}
for _, fixture := range fixtures {
if _, err := db.pool.Exec(ctx, `
INSERT INTO gateway_identity_configuration_revisions (
id,state,auth_center_url,local_tenant_key,public_base_url,web_base_url,
machine_credential_ref,session_encryption_key_ref,created_at,updated_at
) VALUES ($1::uuid,'draft','https://auth.test.example','default','https://gateway.test.example',
'https://gateway-web.test.example',$2,$3,$4,$4)`,
fixture.revisionID, fixture.machineRef, fixture.sessionRef, fixture.createdAt); err != nil {
t.Fatalf("seed %s inactive revision: %v", fixture.displayName, err)
}
if _, err := db.CreateIdentityPairingExchange(ctx, identity.PairingExchange{
ID: fixture.pairingID, RevisionID: fixture.revisionID, RemoteExchangeID: uuid.NewString(),
ExchangeTokenRef: "identity-exchange-" + fixture.pairingID,
Status: identity.PairingCredentialsSaved, CleanupStatus: identity.PairingCleanupNone,
RemoteVersion: 4, ExpiresAt: now.Add(time.Hour), LastTraceID: "trace-start-" + fixture.displayName,
}); err != nil {
t.Fatalf("seed %s pairing: %v", fixture.displayName, err)
}
}
orphanRevisionID := uuid.NewString()
if _, err := db.pool.Exec(ctx, `
INSERT INTO gateway_identity_configuration_revisions (
id,state,auth_center_url,local_tenant_key,public_base_url,web_base_url,created_at,updated_at
) VALUES ($1::uuid,'draft','https://auth.test.example','default','https://orphan-api.test.example',
'https://orphan-web.test.example',$2,$2)`, orphanRevisionID, now.Add(time.Minute)); err != nil {
t.Fatalf("seed orphan revision: %v", err)
}
canonical, err := db.LatestInactiveIdentityConfigurationRevision(ctx)
if err != nil || canonical.ID != fixtures[1].revisionID {
t.Fatalf("canonical inactive revision=%#v error=%v, want latest exchange-backed revision", canonical, err)
}
blocked, err := db.IdentityPairingStartBlocked(ctx)
if err != nil || !blocked {
t.Fatalf("two inactive pairings blocked=%t error=%v, want blocked", blocked, err)
}
latest := fixtures[1]
latestCancelled, err := db.CancelIdentityPairingExchange(ctx, latest.pairingID, 1, "trace-cancel-latest", "audit-cancel-latest")
if err != nil {
t.Fatalf("cancel latest pairing: %v", err)
}
latestReplay, err := db.CancelIdentityPairingExchange(ctx, latest.pairingID, 1, "trace-replayed", "audit-replayed")
if err != nil {
t.Fatalf("replay latest cancellation with original version: %v", err)
}
if latestReplay.Status != identity.PairingCancelled || latestReplay.CleanupStatus != identity.PairingCleanupPending ||
latestReplay.Version != latestCancelled.Version || latestCancelled.CancelledAt == nil || latestReplay.CancelledAt == nil ||
!latestReplay.CancelledAt.Equal(*latestCancelled.CancelledAt) {
t.Fatalf("cancellation replay changed persisted intent: first=%#v replay=%#v", latestCancelled, latestReplay)
}
latestCleaned, err := db.CompleteIdentityPairingCleanup(ctx, latest.pairingID, latestCancelled.Version)
if err != nil {
t.Fatalf("complete latest cleanup: %v", err)
}
if latestCleaned.CleanupStatus != identity.PairingCleanupCompleted || latestCleaned.CleanupCompletedAt == nil {
t.Fatalf("latest cleanup state=%#v", latestCleaned)
}
assertIdentityRevisionSecretRefs(t, ctx, db, latest.revisionID, "", "")
assertIdentityRevisionSecretRefs(t, ctx, db, fixtures[0].revisionID, fixtures[0].machineRef, fixtures[0].sessionRef)
blocked, err = db.IdentityPairingStartBlocked(ctx)
if err != nil || !blocked {
t.Fatalf("older pairing was ignored after latest cleanup: blocked=%t error=%v", blocked, err)
}
canonical, err = db.LatestInactiveIdentityConfigurationRevision(ctx)
if err != nil || canonical.ID != fixtures[0].revisionID {
t.Fatalf("canonical inactive revision after latest cleanup=%#v error=%v, want older blocker", canonical, err)
}
older := fixtures[0]
olderCancelled, err := db.CancelIdentityPairingExchange(ctx, older.pairingID, 1, "trace-cancel-older", "audit-cancel-older")
if err != nil {
t.Fatalf("cancel older pairing: %v", err)
}
olderReplay, err := db.CancelIdentityPairingExchange(ctx, older.pairingID, 1, "trace-replayed", "audit-replayed")
if err != nil || olderReplay.Version != olderCancelled.Version {
t.Fatalf("replay older cancellation: replay=%#v error=%v", olderReplay, err)
}
if _, err := db.CompleteIdentityPairingCleanup(ctx, older.pairingID, olderCancelled.Version); err != nil {
t.Fatalf("complete older cleanup: %v", err)
}
assertIdentityRevisionSecretRefs(t, ctx, db, older.revisionID, "", "")
blocked, err = db.IdentityPairingStartBlocked(ctx)
if err != nil || blocked {
t.Fatalf("all inactive pairing cleanup blocked=%t error=%v, want unblocked", blocked, err)
}
canonical, err = db.LatestInactiveIdentityConfigurationRevision(ctx)
if err != nil || canonical.ID != fixtures[1].revisionID || canonical.ID == orphanRevisionID {
t.Fatalf("canonical cleaned revision=%#v error=%v, orphan draft must stay hidden", canonical, err)
}
}
func TestIdentityPairingStartReservationSerializesAndRejectsStaleCommit(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx := context.Background()
firstAttempt := uuid.NewString()
secondAttempt := uuid.NewString()
if err := db.ReserveIdentityPairingStart(ctx, firstAttempt, time.Now().Add(time.Minute)); err != nil {
t.Fatalf("reserve first pairing start: %v", err)
}
if err := db.ReserveIdentityPairingStart(ctx, secondAttempt, time.Now().Add(time.Minute)); !errors.Is(err, identity.ErrPairingInProgress) {
t.Fatalf("second pairing reservation error=%v, want in-progress conflict", err)
}
if _, err := db.pool.Exec(ctx, `UPDATE gateway_identity_pairing_start_reservation SET expires_at=now()-interval '1 second'
WHERE attempt_id=$1::uuid`, firstAttempt); err != nil {
t.Fatalf("expire first pairing reservation: %v", err)
}
if err := db.ReserveIdentityPairingStart(ctx, secondAttempt, time.Now().Add(time.Minute)); err != nil {
t.Fatalf("take over expired pairing reservation: %v", err)
}
staleRevision := identity.Revision{
ID: uuid.NewString(), State: identity.RevisionDraft, SchemaVersion: 1, AuthCenterURL: "https://auth.test.example",
RolePrefix: "gateway.", LocalTenantKey: "default", PublicBaseURL: "https://gateway.test.example",
WebBaseURL: "https://gateway-web.test.example", Scopes: []string{}, Capabilities: []string{},
SessionIdleSeconds: 1800, SessionAbsoluteSeconds: 28800, SessionRefreshSeconds: 60,
}
if _, err := db.CommitIdentityPairingStart(ctx, staleRevision, identity.PairingExchange{
ID: firstAttempt, RevisionID: staleRevision.ID, RemoteExchangeID: uuid.NewString(),
ExchangeTokenRef: "identity-exchange-" + firstAttempt, Status: identity.PairingMetadataPending,
CleanupStatus: identity.PairingCleanupNone, RemoteVersion: 1, ExpiresAt: time.Now().Add(time.Minute),
}); !errors.Is(err, identity.ErrPairingInProgress) {
t.Fatalf("stale start commit error=%v, want in-progress conflict", err)
}
validRevision := staleRevision
validRevision.ID = uuid.NewString()
tokenReference := "identity-exchange-" + secondAttempt
if err := db.QueueIdentitySecretCleanup(ctx, tokenReference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatalf("stage replacement exchange token reference: %v", err)
}
committed, err := db.CommitIdentityPairingStart(ctx, validRevision, identity.PairingExchange{
ID: secondAttempt, RevisionID: validRevision.ID, RemoteExchangeID: uuid.NewString(),
ExchangeTokenRef: tokenReference, Status: identity.PairingMetadataPending,
CleanupStatus: identity.PairingCleanupNone, RemoteVersion: 1, ExpiresAt: time.Now().Add(time.Minute),
})
if err != nil {
t.Fatalf("commit replacement pairing: %v", err)
}
if err := db.ReleaseIdentityPairingStart(ctx, secondAttempt); err != nil {
t.Fatalf("release paired reservation: %v", err)
}
var reservationState string
if err := db.pool.QueryRow(ctx, `SELECT state FROM gateway_identity_pairing_start_reservation
WHERE attempt_id=$1::uuid`, secondAttempt).Scan(&reservationState); err != nil || reservationState != "paired" {
t.Fatalf("paired reservation state=%q error=%v", reservationState, err)
}
blocked, err := db.IdentityPairingStartBlocked(ctx)
if err != nil || !blocked {
t.Fatalf("committed pairing did not retain singleton reservation: blocked=%t error=%v", blocked, err)
}
cancelled, err := db.CancelIdentityPairingExchange(ctx, committed.ID, committed.Version, "trace-cancel", "audit-cancel")
if err != nil {
t.Fatalf("cancel replacement pairing: %v", err)
}
if _, err := db.CompleteIdentityPairingCleanup(ctx, committed.ID, cancelled.Version); err != nil {
t.Fatalf("complete replacement pairing cleanup: %v", err)
}
blocked, err = db.IdentityPairingStartBlocked(ctx)
if err != nil || blocked {
t.Fatalf("completed cleanup retained singleton reservation: blocked=%t error=%v", blocked, err)
}
}
func TestIdentityPairingStartRejectsAnyActiveConfigurationBeforeReservation(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx := context.Background()
if _, err := db.pool.Exec(ctx, `INSERT INTO gateway_identity_configuration_revisions (
id,state,auth_center_url,local_tenant_key,public_base_url,web_base_url
) VALUES ($1::uuid,'active','https://auth.test.example','default','https://gateway.test.example',
'https://gateway-web.test.example')`, uuid.NewString()); err != nil {
t.Fatalf("seed active identity revision: %v", err)
}
err := db.ReserveIdentityPairingStart(ctx, uuid.NewString(), time.Now().Add(time.Minute))
if !errors.Is(err, identity.ErrActiveConfigurationHandoffRequired) {
t.Fatalf("active identity configuration reservation error=%v", err)
}
var reservations int
if err := db.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_identity_pairing_start_reservation`).Scan(&reservations); err != nil {
t.Fatal(err)
}
if reservations != 0 {
t.Fatalf("active credential guard left %d start reservations", reservations)
}
}
func TestConcurrentIdentityPairingStartReservationsHaveSingleWinner(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
attempts := []string{uuid.NewString(), uuid.NewString()}
ready := make(chan struct{}, len(attempts))
start := make(chan struct{})
type reservationResult struct {
attemptID string
err error
}
results := make(chan reservationResult, len(attempts))
for _, attemptID := range attempts {
go func() {
ready <- struct{}{}
<-start
results <- reservationResult{
attemptID: attemptID,
err: db.ReserveIdentityPairingStart(ctx, attemptID, time.Now().Add(5*time.Minute)),
}
}()
}
for range attempts {
<-ready
}
close(start)
var winner string
conflicts := 0
for range attempts {
result := <-results
switch {
case result.err == nil:
if winner != "" {
t.Fatalf("multiple concurrent reservations succeeded: %s and %s", winner, result.attemptID)
}
winner = result.attemptID
case errors.Is(result.err, identity.ErrPairingInProgress):
conflicts++
default:
t.Fatalf("reserve concurrent pairing start %s: %v", result.attemptID, result.err)
}
}
if winner == "" || conflicts != 1 {
t.Fatalf("concurrent reservation winner=%q conflicts=%d, want one winner and one conflict", winner, conflicts)
}
var persistedAttemptID, state string
var reservations int
if err := db.pool.QueryRow(ctx, `SELECT count(*),COALESCE(max(attempt_id::text),''),COALESCE(max(state),'')
FROM gateway_identity_pairing_start_reservation`).Scan(&reservations, &persistedAttemptID, &state); err != nil {
t.Fatalf("read concurrent reservation result: %v", err)
}
if reservations != 1 || persistedAttemptID != winner || state != "starting" {
t.Fatalf("persisted reservations=%d attempt=%q state=%q, want winner %q in starting state",
reservations, persistedAttemptID, state, winner)
}
}
func TestConcurrentIdentityActivationAndPairingReservationRemainMutuallyExclusive(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
seedIdentityActivationPrerequisites(t, ctx, db, "default")
revisionID := seedIdentityLifecycleRevision(t, ctx, db, identity.RevisionValidated, "default")
attemptID := uuid.NewString()
gate, err := db.pool.Begin(ctx)
if err != nil {
t.Fatalf("begin identity lifecycle gate: %v", err)
}
defer gate.Rollback(context.Background())
if err := lockIdentityConfigurationLifecycle(ctx, gate); err != nil {
t.Fatalf("lock identity lifecycle gate: %v", err)
}
reserveResult := make(chan error, 1)
go func() {
reserveResult <- db.ReserveIdentityPairingStart(ctx, attemptID, time.Now().Add(5*time.Minute))
}()
waitForIdentityLifecycleLockWaiters(t, ctx, db, 1)
type activationResult struct {
revision identity.Revision
err error
}
activateResult := make(chan activationResult, 1)
go func() {
revision, _, activateErr := db.ActivateIdentityRevision(ctx, revisionID, 1, "trace-activate", "audit-activate")
activateResult <- activationResult{revision: revision, err: activateErr}
}()
waitForIdentityLifecycleLockWaiters(t, ctx, db, 2)
if err := gate.Commit(ctx); err != nil {
t.Fatalf("release identity lifecycle gate: %v", err)
}
reserveErr := <-reserveResult
activated := <-activateResult
if reserveErr != nil {
t.Fatalf("first queued pairing reservation failed: %v", reserveErr)
}
var activeRevisions, startingReservations int
if err := db.pool.QueryRow(ctx, `SELECT
(SELECT count(*) FROM gateway_identity_configuration_revisions WHERE state='active'),
(SELECT count(*) FROM gateway_identity_pairing_start_reservation WHERE state='starting')`).Scan(
&activeRevisions, &startingReservations,
); err != nil {
t.Fatalf("read activation and reservation invariant: %v", err)
}
if !errors.Is(activated.err, identity.ErrPairingInProgress) || activeRevisions != 0 || startingReservations != 1 {
t.Fatalf("racing activation error=%v revision_state=%q left active=%d starting reservations=%d; "+
"want pairing-in-progress, active=0, reservation=1",
activated.err, activated.revision.State, activeRevisions, startingReservations)
}
}
func TestConcurrentIdentityPairingCommitRejectsActiveRevisionCommittedAheadOfLifecycleLock(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
attemptID := uuid.NewString()
if err := db.ReserveIdentityPairingStart(ctx, attemptID, time.Now().Add(5*time.Minute)); err != nil {
t.Fatalf("reserve pairing start before commit race: %v", err)
}
tokenReference := "identity-exchange-" + attemptID
if err := db.QueueIdentitySecretCleanup(ctx, tokenReference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatalf("stage exchange token reference before commit race: %v", err)
}
draft := identity.Revision{
ID: uuid.NewString(), State: identity.RevisionDraft, SchemaVersion: 1, AuthCenterURL: "https://auth.test.example",
RolePrefix: "gateway.", LocalTenantKey: "default", PublicBaseURL: "https://gateway.test.example",
WebBaseURL: "https://gateway-web.test.example", Scopes: []string{}, Capabilities: []string{},
SessionIdleSeconds: 1800, SessionAbsoluteSeconds: 28800, SessionRefreshSeconds: 60,
}
exchange := identity.PairingExchange{
ID: attemptID, RevisionID: draft.ID, RemoteExchangeID: uuid.NewString(), ExchangeTokenRef: tokenReference,
Status: identity.PairingMetadataPending, CleanupStatus: identity.PairingCleanupNone,
RemoteVersion: 1, ExpiresAt: time.Now().Add(5 * time.Minute),
}
gate, err := db.pool.Begin(ctx)
if err != nil {
t.Fatalf("begin identity lifecycle gate: %v", err)
}
defer gate.Rollback(context.Background())
if err := lockIdentityConfigurationLifecycle(ctx, gate); err != nil {
t.Fatalf("lock identity lifecycle gate: %v", err)
}
commitResult := make(chan error, 1)
go func() {
_, commitErr := db.CommitIdentityPairingStart(ctx, draft, exchange)
commitResult <- commitErr
}()
waitForIdentityLifecycleLockWaiters(t, ctx, db, 1)
activeID := uuid.NewString()
if _, err := gate.Exec(ctx, `INSERT INTO gateway_identity_configuration_revisions (
id,state,auth_center_url,local_tenant_key,public_base_url,web_base_url
) VALUES ($1::uuid,'active','https://auth.test.example','default','https://active-gateway.test.example',
'https://active-gateway-web.test.example')`, activeID); err != nil {
t.Fatalf("commit active revision ahead of pairing commit: %v", err)
}
if err := gate.Commit(ctx); err != nil {
t.Fatalf("release identity lifecycle gate with active revision: %v", err)
}
commitErr := <-commitResult
if !errors.Is(commitErr, identity.ErrActiveConfigurationHandoffRequired) {
t.Fatalf("pairing commit after active revision error=%v, want active handoff required", commitErr)
}
var activeRevisions, draftRevisions, exchanges, startingReservations int
if err := db.pool.QueryRow(ctx, `SELECT
(SELECT count(*) FROM gateway_identity_configuration_revisions WHERE state='active'),
(SELECT count(*) FROM gateway_identity_configuration_revisions WHERE id=$1::uuid),
(SELECT count(*) FROM gateway_identity_onboarding_exchanges WHERE id=$2::uuid),
(SELECT count(*) FROM gateway_identity_pairing_start_reservation WHERE attempt_id=$2::uuid AND state='starting')`,
draft.ID, attemptID).Scan(&activeRevisions, &draftRevisions, &exchanges, &startingReservations); err != nil {
t.Fatalf("read pairing commit lifecycle invariant: %v", err)
}
if activeRevisions != 1 || draftRevisions != 0 || exchanges != 0 || startingReservations != 1 {
t.Fatalf("pairing commit race left active=%d draft=%d exchanges=%d starting reservations=%d; "+
"want active=1 draft=0 exchanges=0 reservation=1",
activeRevisions, draftRevisions, exchanges, startingReservations)
}
}
func TestConcurrentIdentityDisableAndActivationPreserveSingleActiveRevision(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
seedIdentityActivationPrerequisites(t, ctx, db, "default")
previousID := seedIdentityLifecycleRevision(t, ctx, db, identity.RevisionActive, "default")
nextID := seedIdentityLifecycleRevision(t, ctx, db, identity.RevisionValidated, "default")
gate, err := db.pool.Begin(ctx)
if err != nil {
t.Fatalf("begin identity lifecycle gate: %v", err)
}
defer gate.Rollback(context.Background())
if err := lockIdentityConfigurationLifecycle(ctx, gate); err != nil {
t.Fatalf("lock identity lifecycle gate: %v", err)
}
disableResult := make(chan error, 1)
go func() {
_, disableErr := db.DisableActiveIdentityRevision(ctx, 1, "trace-disable", "audit-disable")
disableResult <- disableErr
}()
waitForIdentityLifecycleLockWaiters(t, ctx, db, 1)
activateResult := make(chan error, 1)
go func() {
_, _, activateErr := db.ActivateIdentityRevision(ctx, nextID, 1, "trace-activate", "audit-activate")
activateResult <- activateErr
}()
waitForIdentityLifecycleLockWaiters(t, ctx, db, 2)
if err := gate.Commit(ctx); err != nil {
t.Fatalf("release identity lifecycle gate: %v", err)
}
if err := <-disableResult; err != nil {
t.Fatalf("disable first queued active revision: %v", err)
}
activateErr := <-activateResult
if activateErr != nil {
t.Fatalf("activate revision after first queued disable: %v", activateErr)
}
var activeIDs []string
rows, err := db.pool.Query(ctx, `SELECT id::text FROM gateway_identity_configuration_revisions
WHERE state='active' ORDER BY id`)
if err != nil {
t.Fatalf("read active revisions after disable/activate race: %v", err)
}
defer rows.Close()
for rows.Next() {
var id string
if err := rows.Scan(&id); err != nil {
t.Fatalf("scan active revision after disable/activate race: %v", err)
}
activeIDs = append(activeIDs, id)
}
if err := rows.Err(); err != nil {
t.Fatalf("iterate active revisions after disable/activate race: %v", err)
}
if len(activeIDs) != 1 || activeIDs[0] != nextID {
t.Fatalf("disable/activate race left active revisions=%v; previous=%q next=%q",
activeIDs, previousID, nextID)
}
}
func TestCompletedIdentityPairingRestoreUsesCanonicalVisibleRevision(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx := context.Background()
now := time.Now().UTC()
var latestRevisionID string
for index, createdAt := range []time.Time{now.Add(-time.Minute), now} {
revisionID, pairingID := uuid.NewString(), uuid.NewString()
if _, err := db.pool.Exec(ctx, `
INSERT INTO gateway_identity_configuration_revisions (
id,state,auth_center_url,local_tenant_key,public_base_url,web_base_url,created_at,updated_at
) VALUES ($1::uuid,'validated','https://auth.test.example','default','https://gateway.test.example',
'https://gateway-web.test.example',$2,$2)`, revisionID, createdAt); err != nil {
t.Fatalf("seed completed revision %d: %v", index, err)
}
if _, err := db.CreateIdentityPairingExchange(ctx, identity.PairingExchange{
ID: pairingID, RevisionID: revisionID, RemoteExchangeID: uuid.NewString(),
ExchangeTokenRef: "identity-exchange-" + pairingID, Status: identity.PairingCompleted,
CleanupStatus: identity.PairingCleanupNone, RemoteVersion: 4, ExpiresAt: now.Add(time.Hour),
}); err != nil {
t.Fatalf("seed completed pairing %d: %v", index, err)
}
latestRevisionID = revisionID
}
completed, err := db.CompletedIdentityPairingsAwaitingActivation(ctx)
if err != nil {
t.Fatal(err)
}
if len(completed) != 1 || completed[0].RevisionID != latestRevisionID {
t.Fatalf("restored completed pairings=%#v, want only canonical latest revision %s", completed, latestRevisionID)
}
visible, err := db.LatestInactiveIdentityConfigurationRevision(ctx)
if err != nil || visible.ID != latestRevisionID {
t.Fatalf("visible inactive revision=%#v error=%v, want same canonical revision", visible, err)
}
}
func TestIdentityPairingPersistsBindingErrorCategoryAllowedByLatestMigration(t *testing.T) {
db := newIdentityPairingPostgresTestStore(t)
ctx := context.Background()
revisionID := uuid.NewString()
if _, err := db.pool.Exec(ctx, `INSERT INTO gateway_identity_configuration_revisions (
id,state,auth_center_url,local_tenant_key,public_base_url,web_base_url
) VALUES ($1::uuid,'draft','https://auth.test.example','default','https://gateway.test.example',
'https://gateway-web.test.example')`, revisionID); err != nil {
t.Fatalf("seed revision for binding error category: %v", err)
}
pairingID := uuid.NewString()
created, err := db.CreateIdentityPairingExchange(ctx, identity.PairingExchange{
ID: pairingID, RevisionID: revisionID, RemoteExchangeID: uuid.NewString(),
ExchangeTokenRef: "identity-exchange-" + pairingID, Status: identity.PairingCredentialsSaved,
CleanupStatus: identity.PairingCleanupNone, RemoteVersion: 4, ExpiresAt: time.Now().Add(time.Hour),
})
if err != nil {
t.Fatalf("create pairing for binding error category: %v", err)
}
const category = "security_event_connection_binding_mismatch"
updated, err := db.RecordIdentityPairingRetryFailure(ctx, created.ID, created.Status, category)
if err != nil {
t.Fatalf("persist binding error category: %v", err)
}
if updated.LastErrorCategory != category {
t.Fatalf("binding error category=%q, want %q", updated.LastErrorCategory, category)
}
var persisted string
if err := db.pool.QueryRow(ctx, `SELECT last_error_category FROM gateway_identity_onboarding_exchanges
WHERE id=$1::uuid`, created.ID).Scan(&persisted); err != nil {
t.Fatalf("read persisted binding error category: %v", err)
}
if persisted != category {
t.Fatalf("persisted binding error category=%q, want %q", persisted, category)
}
}
func assertIdentityRevisionSecretRefs(t *testing.T, ctx context.Context, db *Store, revisionID, wantMachine, wantSession string) {
t.Helper()
revision, err := db.IdentityConfigurationRevision(ctx, revisionID)
if err != nil {
t.Fatalf("read revision %s: %v", revisionID, err)
}
if revision.MachineCredentialRef != wantMachine || revision.SessionEncryptionKeyRef != wantSession {
t.Fatalf("revision %s secret refs machine=%q session=%q, want machine=%q session=%q",
revisionID, revision.MachineCredentialRef, revision.SessionEncryptionKeyRef, wantMachine, wantSession)
}
}
func seedIdentityActivationPrerequisites(t *testing.T, ctx context.Context, db *Store, tenantKey string) {
t.Helper()
if _, err := db.pool.Exec(ctx, `INSERT INTO gateway_tenants(tenant_key,name,status)
VALUES($1,$2,'active')`, tenantKey, "Identity lifecycle test tenant"); err != nil {
t.Fatalf("seed identity lifecycle tenant: %v", err)
}
if _, err := db.pool.Exec(ctx, `INSERT INTO gateway_users(user_key,username,password_hash,roles,source,status)
VALUES($1,$2,$3,'["manager"]'::jsonb,'gateway','active')`,
"identity-lifecycle-manager", "identity-lifecycle-manager", "test-only-password-hash"); err != nil {
t.Fatalf("seed identity lifecycle break-glass manager: %v", err)
}
}
func seedIdentityLifecycleRevision(
t *testing.T,
ctx context.Context,
db *Store,
state identity.RevisionState,
tenantKey string,
) string {
t.Helper()
revisionID := uuid.NewString()
if _, err := db.pool.Exec(ctx, `INSERT INTO gateway_identity_configuration_revisions (
id,state,auth_center_url,issuer,tenant_id,application_id,audience,browser_client_id,
local_tenant_key,public_base_url,web_base_url,validated_at,activated_at
) VALUES (
$1::uuid,$2,'https://auth.test.example','https://issuer.test.example',$3,$4,$5,$6,
$7,'https://gateway.test.example','https://gateway-web.test.example',
CASE WHEN $2 IN ('validated','active') THEN now() ELSE NULL END,
CASE WHEN $2='active' THEN now() ELSE NULL END
)`, revisionID, state, uuid.NewString(), uuid.NewString(), "urn:easyai:resource:"+revisionID,
"browser-"+revisionID, tenantKey); err != nil {
t.Fatalf("seed %s identity lifecycle revision: %v", state, err)
}
return revisionID
}
func waitForIdentityLifecycleLockWaiters(t *testing.T, ctx context.Context, db *Store, want int) {
t.Helper()
classID := int64(uint64(identityConfigurationLifecycleLockID) >> 32)
objectID := int64(uint64(identityConfigurationLifecycleLockID) & 0xffffffff)
deadline := time.Now().Add(5 * time.Second)
for {
var waiters int
if err := db.pool.QueryRow(ctx, `SELECT count(*) FROM pg_locks
WHERE locktype='advisory' AND database=(SELECT oid FROM pg_database WHERE datname=current_database())
AND classid::bigint=$1 AND objid::bigint=$2 AND objsubid=1 AND NOT granted`, classID, objectID).Scan(&waiters); err != nil {
t.Fatalf("read identity lifecycle lock waiters: %v", err)
}
if waiters >= want {
return
}
if time.Now().After(deadline) {
t.Fatalf("identity lifecycle lock waiters=%d, want at least %d", waiters, want)
}
time.Sleep(10 * time.Millisecond)
}
}
func newIdentityPairingPostgresTestStore(t *testing.T) *Store {
t.Helper()
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run identity pairing PostgreSQL integration tests")
}
ctx := context.Background()
admin, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect identity pairing test database: %v", err)
}
var databaseName string
if err := admin.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
admin.Close()
t.Fatalf("read identity pairing test database name: %v", err)
}
if !strings.Contains(strings.ToLower(databaseName), "test") {
admin.Close()
t.Fatalf("refusing to use non-test database %q", databaseName)
}
schemaName := "gateway_identity_pairing_" + strings.ReplaceAll(uuid.NewString(), "-", "")
schemaIdentifier := pgx.Identifier{schemaName}.Sanitize()
if _, err := admin.Exec(ctx, `CREATE SCHEMA `+schemaIdentifier); err != nil {
admin.Close()
t.Fatalf("create identity pairing test schema: %v", err)
}
config, err := pgxpool.ParseConfig(databaseURL)
if err != nil {
_, _ = admin.Exec(ctx, `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`)
admin.Close()
t.Fatalf("parse identity pairing test database URL: %v", err)
}
config.ConnConfig.RuntimeParams["search_path"] = schemaName
pool, err := pgxpool.NewWithConfig(ctx, config)
if err != nil {
_, _ = admin.Exec(ctx, `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`)
admin.Close()
t.Fatalf("connect identity pairing test schema: %v", err)
}
t.Cleanup(func() {
pool.Close()
if _, err := admin.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`); err != nil {
t.Errorf("drop identity pairing test schema: %v", err)
}
admin.Close()
})
_, filename, _, ok := runtime.Caller(0)
if !ok {
t.Fatal("locate identity pairing integration test")
}
migrationDirectory := filepath.Join(filepath.Dir(filename), "..", "..", "migrations")
for _, migrationName := range []string{
"0001_init.sql",
"0061_oidc_server_sessions.sql",
"0067_identity_configuration_revisions.sql",
"0068_identity_onboarding_exchanges.sql",
"0069_identity_pairing_cancellation.sql",
"0070_identity_secret_cleanup_queue.sql",
"0071_identity_pairing_start_reservation.sql",
"0072_identity_secret_cleanup_claim_lifecycle.sql",
"0073_identity_pairing_start_reservation_upgrade.sql",
"0074_identity_pairing_error_categories.sql",
} {
migration, err := os.ReadFile(filepath.Join(migrationDirectory, migrationName))
if err != nil {
t.Fatalf("read %s: %v", migrationName, err)
}
tx, err := pool.Begin(ctx)
if err != nil {
t.Fatalf("begin %s: %v", migrationName, err)
}
if _, err := tx.Exec(ctx, string(migration)); err != nil {
_ = tx.Rollback(ctx)
t.Fatalf("execute %s: %v", migrationName, err)
}
if err := tx.Commit(ctx); err != nil {
t.Fatalf("commit %s: %v", migrationName, err)
}
}
return &Store{pool: pool}
}
@@ -0,0 +1,206 @@
package store
import (
"context"
"errors"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
"time"
"github.com/google/uuid"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestIdentitySecretCleanupClaimLifecycleUpgradeMigrationExists(t *testing.T) {
payload, err := os.ReadFile(identitySecretCleanupMigrationPath(t, "0072_identity_secret_cleanup_claim_lifecycle.sql"))
if err != nil {
t.Fatal(err)
}
content := strings.ToLower(string(payload))
for _, required := range []string{
"add column if not exists status text",
"add column if not exists claim_token uuid",
"add column if not exists lease_expires_at timestamptz",
"alter column status set default 'pending'",
"alter column status set not null",
"gateway_identity_secret_cleanup_status_check",
"gateway_identity_secret_cleanup_claim_check",
"drop index if exists idx_gateway_identity_secret_cleanup_due",
"create index idx_gateway_identity_secret_cleanup_due",
} {
if !strings.Contains(content, required) {
t.Fatalf("identity Secret cleanup lifecycle migration is missing %q", required)
}
}
for _, forbidden := range []string{"secret_value", "client_secret", "machine_secret", "token text"} {
if strings.Contains(content, forbidden) {
t.Fatalf("identity Secret cleanup lifecycle migration stores forbidden value field %q", forbidden)
}
}
}
func TestIdentitySecretCleanupClaimLifecycleUpgradeMigratesLegacyQueue(t *testing.T) {
db := newIdentitySecretCleanupMigrationPostgresStore(t)
ctx := context.Background()
reference := "identity-cleanup-legacy-" + uuid.NewString()
if _, err := db.pool.Exec(ctx, `
CREATE TABLE gateway_identity_secret_cleanup_queue (
secret_ref text PRIMARY KEY CHECK (secret_ref ~ '^[a-z0-9][a-z0-9-]{0,127}$'),
not_before timestamptz NOT NULL,
created_at timestamptz NOT NULL DEFAULT now(),
updated_at timestamptz NOT NULL DEFAULT now()
);`); err != nil {
t.Fatalf("create legacy identity Secret cleanup queue: %v", err)
}
if _, err := db.pool.Exec(ctx, `CREATE INDEX idx_gateway_identity_secret_cleanup_due
ON gateway_identity_secret_cleanup_queue(not_before,updated_at)`); err != nil {
t.Fatalf("create legacy identity Secret cleanup due index: %v", err)
}
if _, err := db.pool.Exec(ctx, `INSERT INTO gateway_identity_secret_cleanup_queue(secret_ref,not_before)
VALUES($1,now()-interval '1 minute')`, reference); err != nil {
t.Fatalf("seed legacy identity Secret cleanup row: %v", err)
}
applyIdentitySecretCleanupMigration(t, ctx, db.pool, "0072_identity_secret_cleanup_claim_lifecycle.sql")
var status string
var claimToken, leaseExpiresAt *string
if err := db.pool.QueryRow(ctx, `
SELECT status,claim_token::text,lease_expires_at::text
FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=$1`, reference).
Scan(&status, &claimToken, &leaseExpiresAt); err != nil {
t.Fatalf("read migrated cleanup row: %v", err)
}
if status != "pending" || claimToken != nil || leaseExpiresAt != nil {
t.Fatalf("legacy cleanup row was not backfilled to an unclaimed pending lifecycle")
}
claims, err := db.ClaimIdentitySecretCleanups(ctx, 1, time.Minute)
if err != nil {
t.Fatalf("claim migrated cleanup row: %v", err)
}
if len(claims) != 1 || claims[0].Reference != reference || claims[0].ClaimToken == "" {
t.Fatalf("claim migrated cleanup row returned unexpected claim metadata")
}
completed, err := db.CompleteIdentitySecretCleanup(ctx, reference, claims[0].ClaimToken)
if err != nil || !completed {
t.Fatalf("complete migrated cleanup claim: completed=%t err=%v", completed, err)
}
invalidReference := "identity-cleanup-invalid-" + uuid.NewString()
_, err = db.pool.Exec(ctx, `
INSERT INTO gateway_identity_secret_cleanup_queue(
secret_ref,not_before,status,claim_token,lease_expires_at
) VALUES($1,now(),'claimed',NULL,NULL)`, invalidReference)
requireIdentitySecretCleanupMigrationCheckViolation(t, err)
var indexDefinition string
if err := db.pool.QueryRow(ctx, `
SELECT pg_get_indexdef(indexrelid)
FROM pg_index
WHERE indexrelid='idx_gateway_identity_secret_cleanup_due'::regclass`).Scan(&indexDefinition); err != nil {
t.Fatalf("read upgraded cleanup due index: %v", err)
}
if !strings.Contains(strings.ToLower(indexDefinition), "(status, not_before, lease_expires_at, updated_at)") {
t.Fatalf("cleanup due index does not cover the claim lifecycle")
}
// The production migration runner applies each version once. Reapplying here
// proves that the upgrade also accepts the already-current 0070 schema shape.
applyIdentitySecretCleanupMigration(t, ctx, db.pool, "0072_identity_secret_cleanup_claim_lifecycle.sql")
}
func newIdentitySecretCleanupMigrationPostgresStore(t *testing.T) *Store {
t.Helper()
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run identity Secret cleanup migration integration tests")
}
ctx := context.Background()
admin, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatalf("connect identity Secret cleanup migration test database: %v", err)
}
var databaseName string
if err := admin.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
admin.Close()
t.Fatalf("read identity Secret cleanup migration test database name: %v", err)
}
if !strings.Contains(strings.ToLower(databaseName), "test") {
admin.Close()
t.Fatalf("refusing to use non-test database %q", databaseName)
}
schemaName := "gateway_identity_cleanup_migration_" + strings.ReplaceAll(uuid.NewString(), "-", "")
schemaIdentifier := pgx.Identifier{schemaName}.Sanitize()
if _, err := admin.Exec(ctx, `CREATE SCHEMA `+schemaIdentifier); err != nil {
admin.Close()
t.Fatalf("create identity Secret cleanup migration test schema: %v", err)
}
config, err := pgxpool.ParseConfig(databaseURL)
if err != nil {
_, _ = admin.Exec(ctx, `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`)
admin.Close()
t.Fatalf("parse identity Secret cleanup migration test database URL: %v", err)
}
config.ConnConfig.RuntimeParams["search_path"] = schemaName
pool, err := pgxpool.NewWithConfig(ctx, config)
if err != nil {
_, _ = admin.Exec(ctx, `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`)
admin.Close()
t.Fatalf("connect identity Secret cleanup migration test schema: %v", err)
}
t.Cleanup(func() {
pool.Close()
if _, err := admin.Exec(context.Background(), `DROP SCHEMA IF EXISTS `+schemaIdentifier+` CASCADE`); err != nil {
t.Errorf("drop identity Secret cleanup migration test schema: %v", err)
}
admin.Close()
})
return &Store{pool: pool}
}
func applyIdentitySecretCleanupMigration(t *testing.T, ctx context.Context, pool *pgxpool.Pool, name string) {
t.Helper()
payload, err := os.ReadFile(identitySecretCleanupMigrationPath(t, name))
if err != nil {
t.Fatalf("read identity Secret cleanup migration: %v", err)
}
tx, err := pool.Begin(ctx)
if err != nil {
t.Fatalf("begin identity Secret cleanup migration: %v", err)
}
if _, err := tx.Exec(ctx, string(payload)); err != nil {
_ = tx.Rollback(ctx)
t.Fatalf("execute identity Secret cleanup migration: %v", err)
}
if err := tx.Commit(ctx); err != nil {
t.Fatalf("commit identity Secret cleanup migration: %v", err)
}
}
func identitySecretCleanupMigrationPath(t *testing.T, name string) string {
t.Helper()
_, filename, _, ok := runtime.Caller(0)
if !ok {
t.Fatal("locate identity Secret cleanup migration test")
}
return filepath.Join(filepath.Dir(filename), "..", "..", "migrations", name)
}
func requireIdentitySecretCleanupMigrationCheckViolation(t *testing.T, err error) {
t.Helper()
if err == nil {
t.Fatal("invalid cleanup claim lifecycle unexpectedly satisfied migration constraints")
}
var postgresError *pgconn.PgError
if !errors.As(err, &postgresError) || postgresError.Code != "23514" {
t.Fatalf("invalid cleanup claim lifecycle error=%v, want PostgreSQL check violation", err)
}
}
@@ -2,6 +2,7 @@ package store
import (
"context"
"database/sql"
"encoding/json"
"errors"
"time"
@@ -13,6 +14,7 @@ import (
var (
ErrSecurityEventConnectionNotFound = errors.New("security event connection not found")
ErrSecurityEventConnectionConflict = errors.New("security event connection conflicts with current state")
ErrIdentitySecretCleanupConflict = errors.New("identity secret cleanup conflicts with current state")
)
type SecurityEventConnection struct {
@@ -44,6 +46,121 @@ type SecurityEventConnectionIdempotency struct {
Response json.RawMessage
}
type IdentitySecretCleanupClaim struct {
Reference string
ClaimToken string
LeaseExpiresAt time.Time
}
func (s *Store) QueueIdentitySecretCleanup(ctx context.Context, reference string, notBefore time.Time) error {
tag, err := s.pool.Exec(ctx, `INSERT INTO gateway_identity_secret_cleanup_queue(secret_ref,not_before)
VALUES($1,$2) ON CONFLICT(secret_ref) DO UPDATE
SET not_before=LEAST(gateway_identity_secret_cleanup_queue.not_before,EXCLUDED.not_before),updated_at=now()
WHERE gateway_identity_secret_cleanup_queue.status='pending'`, reference, notBefore)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrIdentitySecretCleanupConflict
}
return nil
}
func (s *Store) RemovePendingIdentitySecretCleanup(ctx context.Context, reference string) (bool, error) {
tag, err := s.pool.Exec(ctx, `DELETE FROM gateway_identity_secret_cleanup_queue
WHERE secret_ref=$1 AND status='pending'`, reference)
if err != nil {
return false, err
}
return tag.RowsAffected() == 1, nil
}
func (s *Store) ClaimIdentitySecretCleanups(ctx context.Context, limit int, lease time.Duration) ([]IdentitySecretCleanupClaim, error) {
if limit <= 0 || limit > 1000 {
limit = 100
}
if lease <= 0 || lease > time.Hour {
lease = 2 * time.Minute
}
leaseSeconds := int64((lease + time.Second - 1) / time.Second)
rows, err := s.pool.Query(ctx, `WITH candidates AS MATERIALIZED (
SELECT secret_ref FROM gateway_identity_secret_cleanup_queue
WHERE (status='pending' AND not_before <= now())
OR (status='claimed' AND lease_expires_at <= now())
ORDER BY COALESCE(lease_expires_at,not_before),created_at
FOR UPDATE SKIP LOCKED
LIMIT $1
)
UPDATE gateway_identity_secret_cleanup_queue queued
SET status='claimed',claim_token=gen_random_uuid(),
lease_expires_at=now()+make_interval(secs => $2::int),updated_at=now()
FROM candidates
WHERE queued.secret_ref=candidates.secret_ref
RETURNING queued.secret_ref,queued.claim_token::text,queued.lease_expires_at`, limit, leaseSeconds)
if err != nil {
return nil, err
}
defer rows.Close()
claims := make([]IdentitySecretCleanupClaim, 0)
for rows.Next() {
var claim IdentitySecretCleanupClaim
if err := rows.Scan(&claim.Reference, &claim.ClaimToken, &claim.LeaseExpiresAt); err != nil {
return nil, err
}
claims = append(claims, claim)
}
return claims, rows.Err()
}
func (s *Store) CompleteIdentitySecretCleanup(ctx context.Context, reference, claimToken string) (bool, error) {
if _, err := uuid.Parse(claimToken); err != nil {
return false, ErrIdentitySecretCleanupConflict
}
tag, err := s.pool.Exec(ctx, `DELETE FROM gateway_identity_secret_cleanup_queue
WHERE secret_ref=$1 AND status='claimed' AND claim_token=$2::uuid`, reference, claimToken)
if err != nil {
return false, err
}
return tag.RowsAffected() == 1, nil
}
func adoptPendingIdentitySecrets(ctx context.Context, tx pgx.Tx, references ...string) error {
unique := make(map[string]struct{}, len(references))
for _, reference := range references {
if reference == "" {
return ErrIdentitySecretCleanupConflict
}
unique[reference] = struct{}{}
}
values := make([]string, 0, len(unique))
for reference := range unique {
values = append(values, reference)
}
tag, err := tx.Exec(ctx, `DELETE FROM gateway_identity_secret_cleanup_queue
WHERE secret_ref=ANY($1) AND status='pending'`, values)
if err != nil {
return err
}
if tag.RowsAffected() != int64(len(values)) {
return ErrIdentitySecretCleanupConflict
}
return nil
}
func queueIdentitySecretCleanupTx(ctx context.Context, tx pgx.Tx, reference string, notBefore time.Time) error {
tag, err := tx.Exec(ctx, `INSERT INTO gateway_identity_secret_cleanup_queue(secret_ref,not_before)
VALUES($1,$2) ON CONFLICT(secret_ref) DO UPDATE
SET not_before=LEAST(gateway_identity_secret_cleanup_queue.not_before,EXCLUDED.not_before),updated_at=now()
WHERE gateway_identity_secret_cleanup_queue.status='pending'`, reference, notBefore)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrIdentitySecretCleanupConflict
}
return nil
}
func (s *Store) SecurityEventConnectionIdempotency(ctx context.Context, operation, key string) (SecurityEventConnectionIdempotency, error) {
var value SecurityEventConnectionIdempotency
err := s.pool.QueryRow(ctx, `SELECT request_hash,response FROM gateway_security_event_connection_idempotency
@@ -64,20 +181,22 @@ VALUES($1,$2,$3,$4::jsonb) ON CONFLICT(operation,idempotency_key) DO NOTHING`, o
}
func (s *Store) CreateSecurityEventConnection(ctx context.Context, input CreateSecurityEventConnectionInput) (SecurityEventConnection, error) {
var value SecurityEventConnection
err := s.pool.QueryRow(ctx, `
tx, err := s.pool.Begin(ctx)
if err != nil {
return SecurityEventConnection{}, err
}
defer tx.Rollback(ctx)
_, err = tx.Exec(ctx, `
INSERT INTO gateway_security_event_connections(
connection_id,transmitter_issuer,endpoint_url,credential_ref,management_client_id,management_credential_ref,lifecycle_status,idempotency_key
) VALUES($1::uuid,$2,$3,$4,$5,$6,'connecting',$7)
RETURNING connection_id::text,transmitter_issuer,endpoint_url,audience,stream_id::text,credential_ref,
next_credential_ref,management_client_id,management_credential_ref,lifecycle_status,version,last_error_category,idempotency_key,created_at,updated_at`,
`,
input.ConnectionID, input.TransmitterIssuer, input.EndpointURL, input.CredentialRef,
input.ManagementClientID, input.ManagementCredentialRef, input.IdempotencyKey,
).Scan(&value.ConnectionID, &value.TransmitterIssuer, &value.EndpointURL, &value.Audience, &value.StreamID,
&value.CredentialRef, &value.NextCredentialRef, &value.ManagementClientID, &value.ManagementCredentialRef, &value.LifecycleStatus, &value.Version,
&value.LastErrorCategory, &value.IdempotencyKey, &value.CreatedAt, &value.UpdatedAt)
)
if err != nil {
if isUniqueViolation(err) {
_ = tx.Rollback(ctx)
existing, getErr := s.SecurityEventConnection(ctx)
if getErr == nil && existing.IdempotencyKey == input.IdempotencyKey && existing.TransmitterIssuer == input.TransmitterIssuer {
return existing, nil
@@ -86,7 +205,13 @@ RETURNING connection_id::text,transmitter_issuer,endpoint_url,audience,stream_id
}
return SecurityEventConnection{}, err
}
return value, nil
if err := adoptPendingIdentitySecrets(ctx, tx, input.CredentialRef, input.ManagementCredentialRef); err != nil {
return SecurityEventConnection{}, err
}
if err := tx.Commit(ctx); err != nil {
return SecurityEventConnection{}, err
}
return s.SecurityEventConnection(ctx)
}
func (s *Store) SecurityEventConnection(ctx context.Context) (SecurityEventConnection, error) {
@@ -126,49 +251,56 @@ WHERE connection_id=$1::uuid AND version=$5`, connectionID, audience, streamID,
return s.SecurityEventConnection(ctx)
}
func (s *Store) SetSecurityEventManagementCredential(ctx context.Context, connectionID, clientID, reference string) (SecurityEventConnection, error) {
tag, err := s.pool.Exec(ctx, `UPDATE gateway_security_event_connections
SET management_client_id=$2,management_credential_ref=$3,last_error_category=NULL,updated_at=now()
WHERE connection_id=$1::uuid`, connectionID, clientID, reference)
func (s *Store) SetSecurityEventManagementCredential(ctx context.Context, connectionID, clientID, reference string, expectedVersion int64) (SecurityEventConnection, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return SecurityEventConnection{}, err
}
if tag.RowsAffected() == 0 {
defer tx.Rollback(ctx)
var oldReference sql.NullString
var lifecycle string
var version int64
if err := tx.QueryRow(ctx, `SELECT management_credential_ref,lifecycle_status,version FROM gateway_security_event_connections
WHERE connection_id=$1::uuid FOR UPDATE`, connectionID).Scan(&oldReference, &lifecycle, &version); errors.Is(err, pgx.ErrNoRows) {
return SecurityEventConnection{}, ErrSecurityEventConnectionNotFound
}
return s.SecurityEventConnection(ctx)
}
func (s *Store) UpdateSecurityEventConnectionLifecycle(ctx context.Context, connectionID, lifecycle string, errorCategory *string) (SecurityEventConnection, error) {
tag, err := s.pool.Exec(ctx, `UPDATE gateway_security_event_connections
SET lifecycle_status=$2,last_error_category=$3,version=version+1,updated_at=now() WHERE connection_id=$1::uuid`, connectionID, lifecycle, errorCategory)
if err != nil {
} else if err != nil {
return SecurityEventConnection{}, err
}
if tag.RowsAffected() == 0 {
return SecurityEventConnection{}, ErrSecurityEventConnectionNotFound
if version != expectedVersion || lifecycle == "retiring" || lifecycle == "disconnect_pending" {
return SecurityEventConnection{}, ErrSecurityEventConnectionConflict
}
return s.SecurityEventConnection(ctx)
}
func (s *Store) SetSecurityEventNextCredential(ctx context.Context, connectionID, reference string) (SecurityEventConnection, error) {
tag, err := s.pool.Exec(ctx, `UPDATE gateway_security_event_connections
SET next_credential_ref=$2,lifecycle_status='rotating',last_error_category=NULL,version=version+1,updated_at=now()
WHERE connection_id=$1::uuid`, connectionID, reference)
if err != nil {
return SecurityEventConnection{}, err
}
if tag.RowsAffected() == 0 {
return SecurityEventConnection{}, ErrSecurityEventConnectionNotFound
}
return s.SecurityEventConnection(ctx)
}
func (s *Store) PromoteSecurityEventCredential(ctx context.Context, connectionID, nextReference string) (SecurityEventConnection, error) {
tag, err := s.pool.Exec(ctx, `UPDATE gateway_security_event_connections
SET credential_ref=$2,next_credential_ref=NULL,lifecycle_status='bootstrap',last_error_category=NULL,
tag, err := tx.Exec(ctx, `UPDATE gateway_security_event_connections
SET management_client_id=$2,management_credential_ref=$3,last_error_category=NULL,
version=version+1,updated_at=now()
WHERE connection_id=$1::uuid AND next_credential_ref=$2`, connectionID, nextReference)
WHERE connection_id=$1::uuid AND version=$4
AND lifecycle_status NOT IN ('retiring','disconnect_pending')`, connectionID, clientID, reference, expectedVersion)
if err != nil {
return SecurityEventConnection{}, err
}
if tag.RowsAffected() == 0 {
return SecurityEventConnection{}, ErrSecurityEventConnectionConflict
}
if err := adoptPendingIdentitySecrets(ctx, tx, reference); err != nil {
return SecurityEventConnection{}, err
}
if oldReference.Valid && oldReference.String != reference {
if err := queueIdentitySecretCleanupTx(ctx, tx, oldReference.String, time.Now().UTC()); err != nil {
return SecurityEventConnection{}, err
}
}
if err := tx.Commit(ctx); err != nil {
return SecurityEventConnection{}, err
}
return s.SecurityEventConnection(ctx)
}
func (s *Store) TransitionSecurityEventConnectionLifecycle(ctx context.Context, connectionID, expectedLifecycle string, expectedVersion int64, lifecycle string, errorCategory *string) (SecurityEventConnection, error) {
tag, err := s.pool.Exec(ctx, `UPDATE gateway_security_event_connections
SET lifecycle_status=$4,last_error_category=$5,version=version+1,updated_at=now()
WHERE connection_id=$1::uuid AND lifecycle_status=$2 AND version=$3
AND (lifecycle_status NOT IN ('retiring','disconnect_pending')
OR $4 IN ('retiring','disconnect_pending'))`,
connectionID, expectedLifecycle, expectedVersion, lifecycle, errorCategory)
if err != nil {
return SecurityEventConnection{}, err
}
@@ -178,16 +310,168 @@ WHERE connection_id=$1::uuid AND next_credential_ref=$2`, connectionID, nextRefe
return s.SecurityEventConnection(ctx)
}
func (s *Store) DeleteSecurityEventConnection(ctx context.Context, connectionID string) error {
if _, err := uuid.Parse(connectionID); err != nil {
return ErrSecurityEventConnectionNotFound
func (s *Store) SetSecurityEventNextCredential(ctx context.Context, connectionID, reference string, expectedVersion int64) (SecurityEventConnection, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return SecurityEventConnection{}, err
}
tag, err := s.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
defer tx.Rollback(ctx)
tag, err := tx.Exec(ctx, `UPDATE gateway_security_event_connections
SET next_credential_ref=$2,lifecycle_status='rotating',last_error_category=NULL,version=version+1,updated_at=now()
WHERE connection_id=$1::uuid AND version=$3 AND next_credential_ref IS NULL
AND lifecycle_status IN ('connecting','verifying','bootstrap','enabled','degraded','error')`, connectionID, reference, expectedVersion)
if err != nil {
return SecurityEventConnection{}, err
}
if tag.RowsAffected() == 0 {
return SecurityEventConnection{}, ErrSecurityEventConnectionConflict
}
if err := adoptPendingIdentitySecrets(ctx, tx, reference); err != nil {
return SecurityEventConnection{}, err
}
if err := tx.Commit(ctx); err != nil {
return SecurityEventConnection{}, err
}
return s.SecurityEventConnection(ctx)
}
func (s *Store) PromoteSecurityEventCredential(ctx context.Context, connectionID, nextReference string) (SecurityEventConnection, error) {
tx, err := s.pool.Begin(ctx)
if err != nil {
return SecurityEventConnection{}, err
}
defer tx.Rollback(ctx)
var oldReference string
if err := tx.QueryRow(ctx, `SELECT credential_ref FROM gateway_security_event_connections
WHERE connection_id=$1::uuid AND next_credential_ref=$2 AND lifecycle_status='rotating' FOR UPDATE`, connectionID, nextReference).Scan(&oldReference); errors.Is(err, pgx.ErrNoRows) {
return SecurityEventConnection{}, ErrSecurityEventConnectionConflict
} else if err != nil {
return SecurityEventConnection{}, err
}
tag, err := tx.Exec(ctx, `UPDATE gateway_security_event_connections
SET credential_ref=$2,next_credential_ref=NULL,lifecycle_status='bootstrap',last_error_category=NULL,
version=version+1,updated_at=now()
WHERE connection_id=$1::uuid AND next_credential_ref=$2 AND lifecycle_status='rotating'`, connectionID, nextReference)
if err != nil {
return SecurityEventConnection{}, err
}
if tag.RowsAffected() == 0 {
return SecurityEventConnection{}, ErrSecurityEventConnectionConflict
}
if oldReference != nextReference {
if err := queueIdentitySecretCleanupTx(ctx, tx, oldReference, time.Now().UTC()); err != nil {
return SecurityEventConnection{}, err
}
}
if err := tx.Commit(ctx); err != nil {
return SecurityEventConnection{}, err
}
return s.SecurityEventConnection(ctx)
}
// DiscardPreparedSecurityEventConnection atomically retires the local Secrets
// and removes only the exact, still-unbound connection generation observed by
// the caller. A concurrent Bind or credential update changes the version and
// therefore preserves both the connection and its Secrets.
func (s *Store) DiscardPreparedSecurityEventConnection(ctx context.Context, connectionID, ownerKey string, expectedVersion int64) error {
if _, err := uuid.Parse(connectionID); err != nil || ownerKey == "" || expectedVersion <= 0 {
return ErrSecurityEventConnectionConflict
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return err
}
if tag.RowsAffected() == 0 {
return ErrSecurityEventConnectionNotFound
defer tx.Rollback(ctx)
var credentialReference string
var nextReference, managementReference sql.NullString
err = tx.QueryRow(ctx, `SELECT credential_ref,next_credential_ref,management_credential_ref
FROM gateway_security_event_connections
WHERE connection_id=$1::uuid AND idempotency_key=$2 AND version=$3 AND stream_id IS NULL
FOR UPDATE`, connectionID, ownerKey, expectedVersion).Scan(
&credentialReference, &nextReference, &managementReference,
)
if errors.Is(err, pgx.ErrNoRows) {
return ErrSecurityEventConnectionConflict
}
return nil
if err != nil {
return err
}
references := map[string]struct{}{credentialReference: {}}
if nextReference.Valid {
references[nextReference.String] = struct{}{}
}
if managementReference.Valid {
references[managementReference.String] = struct{}{}
}
now := time.Now().UTC()
for reference := range references {
if err := queueIdentitySecretCleanupTx(ctx, tx, reference, now); err != nil {
return err
}
}
tag, err := tx.Exec(ctx, `DELETE FROM gateway_security_event_connections
WHERE connection_id=$1::uuid AND idempotency_key=$2 AND version=$3 AND stream_id IS NULL`,
connectionID, ownerKey, expectedVersion)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrSecurityEventConnectionConflict
}
return tx.Commit(ctx)
}
// FinalizeRetiringSecurityEventConnection atomically transfers every Secret
// reference owned by the exact retiring generation to the durable cleanup
// queue before removing that generation. The SecretStore worker can therefore
// resume after a crash without the database losing its last cleanup intent.
func (s *Store) FinalizeRetiringSecurityEventConnection(ctx context.Context, connectionID string, expectedVersion int64) error {
if _, err := uuid.Parse(connectionID); err != nil || expectedVersion <= 0 {
return ErrSecurityEventConnectionConflict
}
tx, err := s.pool.Begin(ctx)
if err != nil {
return err
}
defer tx.Rollback(ctx)
var credentialReference string
var nextReference, managementReference sql.NullString
err = tx.QueryRow(ctx, `SELECT credential_ref,next_credential_ref,management_credential_ref
FROM gateway_security_event_connections
WHERE connection_id=$1::uuid AND lifecycle_status='retiring' AND version=$2
FOR UPDATE`, connectionID, expectedVersion).Scan(
&credentialReference, &nextReference, &managementReference,
)
if errors.Is(err, pgx.ErrNoRows) {
return ErrSecurityEventConnectionConflict
}
if err != nil {
return err
}
references := map[string]struct{}{credentialReference: {}}
if nextReference.Valid {
references[nextReference.String] = struct{}{}
}
if managementReference.Valid {
references[managementReference.String] = struct{}{}
}
now := time.Now().UTC()
for reference := range references {
if err := queueIdentitySecretCleanupTx(ctx, tx, reference, now); err != nil {
return err
}
}
tag, err := tx.Exec(ctx, `DELETE FROM gateway_security_event_connections
WHERE connection_id=$1::uuid AND lifecycle_status='retiring' AND version=$2`, connectionID, expectedVersion)
if err != nil {
return err
}
if tag.RowsAffected() != 1 {
return ErrSecurityEventConnectionConflict
}
return tx.Commit(ctx)
}
@@ -0,0 +1,665 @@
package store
import (
"context"
"errors"
"os"
"strings"
"sync"
"testing"
"time"
"github.com/google/uuid"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestSecurityEventCredentialReferenceSwitchesRemainCrashRecoverable(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
connectionID := uuid.NewString()
pushReference := "ssf-push-" + connectionID
managementReference := "ssf-management-" + connectionID
nextReference := "ssf-push-next-" + uuid.NewString()
newManagementReference := "ssf-management-next-" + uuid.NewString()
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`,
[]string{pushReference, managementReference, nextReference, newManagementReference})
})
if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
t.Fatal(err)
}
for _, reference := range []string{pushReference, managementReference} {
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
}
connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
})
if err != nil {
t.Fatal(err)
}
assertIdentitySecretQueueState(t, ctx, db, pushReference, false)
assertIdentitySecretQueueState(t, ctx, db, managementReference, false)
if err := db.QueueIdentitySecretCleanup(ctx, newManagementReference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
if _, err := db.SetSecurityEventManagementCredential(ctx, connectionID, "machine-client", newManagementReference, connection.Version); err != nil {
t.Fatal(err)
}
assertIdentitySecretQueueState(t, ctx, db, newManagementReference, false)
assertIdentitySecretQueueState(t, ctx, db, managementReference, true)
if err := db.QueueIdentitySecretCleanup(ctx, nextReference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
connection, err = db.SecurityEventConnection(ctx)
if err != nil {
t.Fatal(err)
}
connection, err = db.SetSecurityEventNextCredential(ctx, connectionID, nextReference, connection.Version)
if err != nil {
t.Fatal(err)
}
assertIdentitySecretQueueState(t, ctx, db, nextReference, false)
if _, err := db.PromoteSecurityEventCredential(ctx, connection.ConnectionID, nextReference); err != nil {
t.Fatal(err)
}
assertIdentitySecretQueueState(t, ctx, db, pushReference, true)
}
func TestDiscardPreparedSecurityEventConnectionQueuesAllSecretsAtomically(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
connectionID := uuid.NewString()
ownerKey := "identity-pairing-ssf-" + uuid.NewString()
pushReference := "ssf-push-" + connectionID
managementReference := "ssf-management-" + connectionID
nextReference := "ssf-push-next-" + uuid.NewString()
references := []string{pushReference, managementReference, nextReference}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
})
if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
t.Fatal(err)
}
for _, reference := range []string{pushReference, managementReference} {
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
}
connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
IdempotencyKey: ownerKey,
})
if err != nil {
t.Fatal(err)
}
if err := db.QueueIdentitySecretCleanup(ctx, nextReference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
connection, err = db.SetSecurityEventNextCredential(ctx, connectionID, nextReference, connection.Version)
if err != nil {
t.Fatal(err)
}
if err := db.DiscardPreparedSecurityEventConnection(ctx, connectionID, ownerKey, connection.Version); err != nil {
t.Fatal(err)
}
if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
t.Fatalf("discarded prepared connection still exists: %v", err)
}
for _, reference := range references {
assertIdentitySecretQueueStatus(t, ctx, db, reference, "pending")
}
}
func TestConcurrentPreparedDiscardAndBindPreserveTheWinningGeneration(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
connectionID := uuid.NewString()
ownerKey := "identity-pairing-ssf-" + uuid.NewString()
pushReference := "ssf-push-" + connectionID
managementReference := "ssf-management-" + connectionID
references := []string{pushReference, managementReference}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
})
if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
t.Fatal(err)
}
for _, reference := range references {
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
}
connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
IdempotencyKey: ownerKey,
})
if err != nil {
t.Fatal(err)
}
type result struct {
operation string
err error
}
start := make(chan struct{})
results := make(chan result, 2)
audience := "urn:easyai:ssf:receiver:" + uuid.NewString()
streamID := uuid.NewString()
go func() {
<-start
_, err := db.BindSecurityEventConnection(ctx, connectionID, audience, streamID, "verifying", connection.Version)
results <- result{operation: "bind", err: err}
}()
go func() {
<-start
err := db.DiscardPreparedSecurityEventConnection(ctx, connectionID, ownerKey, connection.Version)
results <- result{operation: "discard", err: err}
}()
close(start)
outcomes := map[string]error{}
for range 2 {
outcome := <-results
outcomes[outcome.operation] = outcome.err
}
bindWon := outcomes["bind"] == nil
discardWon := outcomes["discard"] == nil
if bindWon == discardWon {
t.Fatalf("concurrent outcomes bind=%v discard=%v, want exactly one winner", outcomes["bind"], outcomes["discard"])
}
loser := outcomes["bind"]
if bindWon {
loser = outcomes["discard"]
}
if !errors.Is(loser, ErrSecurityEventConnectionConflict) {
t.Fatalf("concurrent loser error=%v, want connection conflict", loser)
}
if bindWon {
persisted, err := db.SecurityEventConnection(ctx)
if err != nil || persisted.StreamID == nil || *persisted.StreamID != streamID {
t.Fatalf("winning bound generation was not preserved: connection=%+v err=%v", persisted, err)
}
for _, reference := range references {
assertIdentitySecretQueueState(t, ctx, db, reference, false)
}
return
}
if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
t.Fatalf("winning discard left a connection behind: %v", err)
}
for _, reference := range references {
assertIdentitySecretQueueStatus(t, ctx, db, reference, "pending")
}
}
func TestConcurrentSecurityEventNextCredentialsAdoptOnlyOneStagedSecret(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
connectionID := uuid.NewString()
pushReference := "ssf-push-" + connectionID
managementReference := "ssf-management-" + connectionID
candidates := []string{
"ssf-push-next-" + uuid.NewString(),
"ssf-push-next-" + uuid.NewString(),
}
allReferences := append([]string{pushReference, managementReference}, candidates...)
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, allReferences)
})
if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
t.Fatal(err)
}
for _, reference := range []string{pushReference, managementReference} {
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
}
connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
})
if err != nil {
t.Fatal(err)
}
for _, reference := range candidates {
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
}
type result struct {
reference string
err error
}
start := make(chan struct{})
results := make(chan result, len(candidates))
var workers sync.WaitGroup
for _, reference := range candidates {
workers.Add(1)
go func(reference string) {
defer workers.Done()
<-start
_, err := db.SetSecurityEventNextCredential(ctx, connectionID, reference, connection.Version)
results <- result{reference: reference, err: err}
}(reference)
}
close(start)
workers.Wait()
close(results)
var adopted, pending string
for result := range results {
switch {
case result.err == nil:
if adopted != "" {
t.Fatal("more than one concurrent next credential was adopted")
}
adopted = result.reference
case errors.Is(result.err, ErrSecurityEventConnectionConflict):
pending = result.reference
default:
t.Fatalf("concurrent next credential returned unexpected error: %v", result.err)
}
}
if adopted == "" || pending == "" {
t.Fatalf("concurrent next credential outcomes adopted=%t pending=%t", adopted != "", pending != "")
}
connection, err = db.SecurityEventConnection(ctx)
if err != nil || connection.NextCredentialRef == nil || *connection.NextCredentialRef != adopted {
t.Fatalf("connection did not retain the sole adopted next credential: %v", err)
}
assertIdentitySecretQueueState(t, ctx, db, adopted, false)
assertIdentitySecretQueueStatus(t, ctx, db, pending, "pending")
}
func TestRetiringSecurityEventConnectionCannotAdoptManagementSecret(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
connectionID := uuid.NewString()
pushReference := "ssf-push-" + connectionID
managementReference := "ssf-management-" + connectionID
newManagementReference := "ssf-management-next-" + uuid.NewString()
allReferences := []string{pushReference, managementReference, newManagementReference}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, allReferences)
})
if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
t.Fatal(err)
}
for _, reference := range []string{pushReference, managementReference} {
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
}
connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
})
if err != nil {
t.Fatal(err)
}
connection, err = db.TransitionSecurityEventConnectionLifecycle(
ctx, connectionID, connection.LifecycleStatus, connection.Version, "retiring", nil,
)
if err != nil {
t.Fatal(err)
}
if err := db.QueueIdentitySecretCleanup(ctx, newManagementReference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
if _, err := db.SetSecurityEventManagementCredential(ctx, connectionID, "machine-client", newManagementReference, connection.Version); !errors.Is(err, ErrSecurityEventConnectionConflict) {
t.Fatalf("retiring connection management adoption error=%v, want conflict", err)
}
connection, err = db.SecurityEventConnection(ctx)
if err != nil || connection.ManagementCredentialRef == nil || *connection.ManagementCredentialRef != managementReference {
t.Fatalf("retiring connection changed its management credential: %v", err)
}
assertIdentitySecretQueueStatus(t, ctx, db, newManagementReference, "pending")
}
func TestSecurityEventRetirementRequiresTheCurrentConnectionGeneration(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
connectionID := uuid.NewString()
pushReference := "ssf-push-" + connectionID
managementReference := "ssf-management-" + connectionID
references := []string{pushReference, managementReference}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
})
if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
t.Fatal(err)
}
for _, reference := range references {
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
}
connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
})
if err != nil {
t.Fatal(err)
}
if _, err := db.TransitionSecurityEventConnectionLifecycle(
ctx, connectionID, "connecting", connection.Version-1, "retiring", nil,
); !errors.Is(err, ErrSecurityEventConnectionConflict) {
t.Fatalf("stale retirement transition error=%v, want conflict", err)
}
current, err := db.SecurityEventConnection(ctx)
if err != nil || current.LifecycleStatus != "connecting" || current.Version != connection.Version {
t.Fatalf("stale transition changed connection generation: lifecycle=%q version=%d err=%v",
current.LifecycleStatus, current.Version, err)
}
current, err = db.TransitionSecurityEventConnectionLifecycle(
ctx, connectionID, "connecting", connection.Version, "retiring", nil,
)
if err != nil {
t.Fatal(err)
}
if _, err := db.TransitionSecurityEventConnectionLifecycle(
ctx, connectionID, "retiring", current.Version, "error", nil,
); !errors.Is(err, ErrSecurityEventConnectionConflict) {
t.Fatalf("retiring generation revival error=%v, want conflict", err)
}
if err := db.FinalizeRetiringSecurityEventConnection(ctx, connectionID, connection.Version); !errors.Is(err, ErrSecurityEventConnectionConflict) {
t.Fatalf("stale retirement deletion error=%v, want conflict", err)
}
if persisted, err := db.SecurityEventConnection(ctx); err != nil || persisted.LifecycleStatus != "retiring" || persisted.Version != current.Version {
t.Fatalf("stale retirement deletion removed current generation: lifecycle=%q version=%d err=%v",
persisted.LifecycleStatus, persisted.Version, err)
}
if err := db.FinalizeRetiringSecurityEventConnection(ctx, connectionID, current.Version); err != nil {
t.Fatal(err)
}
if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
t.Fatalf("current retiring generation was not deleted: %v", err)
}
for _, reference := range references {
assertIdentitySecretQueueStatus(t, ctx, db, reference, "pending")
}
}
func TestConcurrentSecurityEventRetirementFinalizationQueuesOneGenerationAtomically(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
connectionID := uuid.NewString()
pushReference := "ssf-push-retire-worker-" + connectionID
managementReference := "ssf-management-retire-worker-" + connectionID
nextReference := "ssf-push-next-retire-worker-" + uuid.NewString()
references := []string{pushReference, managementReference, nextReference}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
})
if _, err := db.pool.Exec(ctx, `DELETE FROM gateway_security_event_connections`); err != nil {
t.Fatal(err)
}
for _, reference := range []string{pushReference, managementReference} {
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
}
connection, err := db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
})
if err != nil {
t.Fatal(err)
}
if err := db.QueueIdentitySecretCleanup(ctx, nextReference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
connection, err = db.SetSecurityEventNextCredential(ctx, connectionID, nextReference, connection.Version)
if err != nil {
t.Fatal(err)
}
connection, err = db.TransitionSecurityEventConnectionLifecycle(
ctx, connectionID, connection.LifecycleStatus, connection.Version, "retiring", nil,
)
if err != nil {
t.Fatal(err)
}
start := make(chan struct{})
results := make(chan error, 2)
var workers sync.WaitGroup
for range 2 {
workers.Add(1)
go func() {
defer workers.Done()
<-start
results <- db.FinalizeRetiringSecurityEventConnection(ctx, connectionID, connection.Version)
}()
}
close(start)
workers.Wait()
close(results)
succeeded, conflicted := 0, 0
for err := range results {
switch {
case err == nil:
succeeded++
case errors.Is(err, ErrSecurityEventConnectionConflict):
conflicted++
default:
t.Fatalf("unexpected concurrent finalization error: %v", err)
}
}
if succeeded != 1 || conflicted != 1 {
t.Fatalf("concurrent finalization succeeded=%d conflicted=%d, want 1/1", succeeded, conflicted)
}
if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
t.Fatalf("finalized connection still exists: %v", err)
}
for _, reference := range references {
assertIdentitySecretQueueStatus(t, ctx, db, reference, "pending")
}
}
func TestClaimedStagingSecretsCannotBeAdoptedByAConnection(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
connectionID := uuid.NewString()
pushReference := "ssf-push-claimed-" + uuid.NewString()
managementReference := "ssf-management-claimed-" + uuid.NewString()
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_connections WHERE connection_id=$1::uuid`, connectionID)
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`,
[]string{pushReference, managementReference})
})
if err := db.QueueIdentitySecretCleanup(ctx, pushReference, time.Now().Add(10*time.Minute)); err != nil {
t.Fatal(err)
}
if err := db.QueueIdentitySecretCleanup(ctx, managementReference, time.Now().Add(-time.Minute)); err != nil {
t.Fatal(err)
}
claims, err := db.ClaimIdentitySecretCleanups(ctx, 1, time.Minute)
if err != nil || len(claims) != 1 || claims[0].Reference != managementReference {
t.Fatalf("claim staging management Secret: count=%d reference_matches=%t err=%v",
len(claims), len(claims) == 1 && claims[0].Reference == managementReference, err)
}
_, err = db.CreateSecurityEventConnection(ctx, CreateSecurityEventConnectionInput{
ConnectionID: connectionID, TransmitterIssuer: "https://auth.test.example/ssf",
EndpointURL: "https://gateway.test.example/api/v1/security-events/ssf", CredentialRef: pushReference,
ManagementClientID: "machine-client", ManagementCredentialRef: managementReference,
IdempotencyKey: "identity-pairing-ssf-" + uuid.NewString(),
})
if !errors.Is(err, ErrIdentitySecretCleanupConflict) {
t.Fatalf("claimed Secret adoption error=%v, want cleanup conflict", err)
}
if _, err := db.SecurityEventConnection(ctx); !errors.Is(err, ErrSecurityEventConnectionNotFound) {
t.Fatalf("failed adoption committed connection: %v", err)
}
assertIdentitySecretQueueStatus(t, ctx, db, pushReference, "pending")
assertIdentitySecretQueueStatus(t, ctx, db, managementReference, "claimed")
}
func TestIdentitySecretCleanupCompletionRejectsExpiredClaimABA(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
reference := "ssf-management-aba-" + uuid.NewString()
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=$1`, reference)
})
if err := db.QueueIdentitySecretCleanup(ctx, reference, time.Now().Add(-time.Minute)); err != nil {
t.Fatal(err)
}
first, err := db.ClaimIdentitySecretCleanups(ctx, 1, time.Minute)
if err != nil || len(first) != 1 {
t.Fatalf("first claim count=%d err=%v", len(first), err)
}
if _, err := db.pool.Exec(ctx, `UPDATE gateway_identity_secret_cleanup_queue SET lease_expires_at=now()-interval '1 second' WHERE secret_ref=$1`, reference); err != nil {
t.Fatal(err)
}
second, err := db.ClaimIdentitySecretCleanups(ctx, 1, time.Minute)
if err != nil || len(second) != 1 || second[0].ClaimToken == first[0].ClaimToken {
t.Fatalf("second claim count=%d token_rotated=%t err=%v",
len(second), len(second) == 1 && second[0].ClaimToken != first[0].ClaimToken, err)
}
completed, err := db.CompleteIdentitySecretCleanup(ctx, reference, first[0].ClaimToken)
if err != nil || completed {
t.Fatalf("stale completion completed=%t err=%v", completed, err)
}
assertIdentitySecretQueueStatus(t, ctx, db, reference, "claimed")
completed, err = db.CompleteIdentitySecretCleanup(ctx, reference, second[0].ClaimToken)
if err != nil || !completed {
t.Fatalf("current completion completed=%t err=%v", completed, err)
}
}
func TestConcurrentIdentitySecretCleanupClaimsAreDisjoint(t *testing.T) {
db := newSecurityEventSecretCleanupPostgresStore(t)
ctx := context.Background()
references := make([]string, 20)
for index := range references {
references[index] = "ssf-cleanup-concurrent-" + uuid.NewString()
if err := db.QueueIdentitySecretCleanup(ctx, references[index], time.Now().Add(-time.Minute)); err != nil {
t.Fatal(err)
}
}
t.Cleanup(func() {
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=ANY($1)`, references)
})
start := make(chan struct{})
results := make(chan []IdentitySecretCleanupClaim, 2)
errorsCh := make(chan error, 2)
var workers sync.WaitGroup
for range 2 {
workers.Add(1)
go func() {
defer workers.Done()
<-start
claims, err := db.ClaimIdentitySecretCleanups(ctx, len(references), time.Minute)
results <- claims
errorsCh <- err
}()
}
close(start)
workers.Wait()
close(results)
close(errorsCh)
for err := range errorsCh {
if err != nil {
t.Fatal(err)
}
}
claimed := map[string]string{}
for claims := range results {
for _, claim := range claims {
if previous := claimed[claim.Reference]; previous != "" {
t.Fatalf("Secret %q was claimed by more than one worker", claim.Reference)
}
claimed[claim.Reference] = claim.ClaimToken
}
}
if len(claimed) != len(references) {
t.Fatalf("claimed %d Secrets, want %d", len(claimed), len(references))
}
}
func newSecurityEventSecretCleanupPostgresStore(t *testing.T) *Store {
t.Helper()
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
if databaseURL == "" {
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run identity Secret cleanup PostgreSQL integration tests")
}
ctx := context.Background()
probe, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatal(err)
}
var databaseName string
if err := probe.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
probe.Close()
t.Fatal(err)
}
probe.Close()
if !strings.Contains(strings.ToLower(databaseName), "test") {
t.Fatalf("refusing to migrate non-test database %q", databaseName)
}
applyOIDCJITTestMigrations(t, ctx, databaseURL)
db, err := Connect(ctx, databaseURL)
if err != nil {
t.Fatal(err)
}
t.Cleanup(db.Close)
return db
}
func assertIdentitySecretQueueState(t *testing.T, ctx context.Context, db *Store, reference string, expected bool) {
t.Helper()
var exists bool
if err := db.pool.QueryRow(ctx, `SELECT EXISTS(
SELECT 1 FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=$1)`, reference).Scan(&exists); err != nil {
t.Fatal(err)
}
if exists != expected {
t.Fatalf("Secret cleanup reference %q exists=%t want=%t", reference, exists, expected)
}
}
func assertIdentitySecretQueueStatus(t *testing.T, ctx context.Context, db *Store, reference, expected string) {
t.Helper()
var status string
if err := db.pool.QueryRow(ctx, `SELECT status FROM gateway_identity_secret_cleanup_queue WHERE secret_ref=$1`, reference).Scan(&status); err != nil {
t.Fatal(err)
}
if status != expected {
t.Fatalf("Secret cleanup reference %q status=%q want=%q", reference, status, expected)
}
}