diff --git a/apps/api/internal/httpapi/identity_pairing_coordinator.go b/apps/api/internal/httpapi/identity_pairing_coordinator.go index 3dfd16c..2e9421c 100644 --- a/apps/api/internal/httpapi/identity_pairing_coordinator.go +++ b/apps/api/internal/httpapi/identity_pairing_coordinator.go @@ -7,9 +7,11 @@ import ( "time" "github.com/easyai/easyai-ai-gateway/apps/api/internal/identity" + "github.com/easyai/easyai-ai-gateway/apps/api/internal/store" ) const identityPairingReconcileInterval = 5 * time.Second +const identityPairingLeadershipTimeout = 5 * time.Second type identityPairingReconciliationStore interface { PendingIdentityPairingExchanges(context.Context) ([]identity.PairingExchange, error) @@ -21,9 +23,9 @@ type identityPairingRestorer interface { RestoreCompletedSecurityEvents(context.Context, string) error } -func (s *Server) reconcileCanonicalIdentityPairing() { +func (s *Server) reconcileCanonicalIdentityPairing(ctx context.Context) { reconcileCanonicalIdentityPairing( - s.ctx, + ctx, s.store, s.identityPairing, &s.identityRestoredPairings, @@ -34,15 +36,57 @@ func (s *Server) reconcileCanonicalIdentityPairing() { } func (s *Server) runIdentityPairingCoordinator() { - ticker := time.NewTicker(identityPairingReconcileInterval) - defer ticker.Stop() + retry := time.NewTicker(identityPairingReconcileInterval) + defer retry.Stop() for { + acquireCtx, cancel := context.WithTimeout(s.ctx, identityPairingLeadershipTimeout) + leadership, acquired, err := s.store.TryAcquireIdentityCoordinatorLeadership(acquireCtx) + cancel() + if err != nil && s.logger != nil { + s.logger.Warn("identity pairing coordinator leader election failed", + "error_category", "identity_coordinator_leader_election_failed", + "error", err, + ) + } + if acquired { + if s.logger != nil { + s.logger.Info("identity pairing coordinator leadership acquired") + } + if !s.runIdentityPairingLeader(leadership) { + return + } + } select { case <-s.ctx.Done(): return + case <-retry.C: + } + } +} + +func (s *Server) runIdentityPairingLeader(leadership store.Leadership) bool { + defer leadership.Release() + ticker := time.NewTicker(identityPairingReconcileInterval) + defer ticker.Stop() + for { + cycleCtx, cancel := context.WithTimeout(s.ctx, identityPairingLeadershipTimeout) + if err := leadership.KeepAlive(cycleCtx); err != nil { + cancel() + if s.logger != nil { + s.logger.Warn("identity pairing coordinator leadership lost", + "error_category", "identity_coordinator_leadership_lost", + "error", err, + ) + } + return true + } + s.reconcileIdentityRuntime(cycleCtx) + s.reconcileCanonicalIdentityPairing(cycleCtx) + cancel() + select { + case <-s.ctx.Done(): + return false case <-ticker.C: - s.reconcileIdentityRuntime() - s.reconcileCanonicalIdentityPairing() } } } @@ -105,14 +149,15 @@ func reconcileCanonicalIdentityPairing( } } -func (s *Server) reconcileIdentityRuntime() { +func (s *Server) reconcileIdentityRuntime(ctx context.Context) { if s.identityRuntime == nil || !s.identityRuntime.ReconciliationRequired() { return } - ctx, cancel := context.WithTimeout(s.ctx, 5*time.Second) - defer cancel() if err := s.identityRuntime.ReconcileActive(ctx); err != nil && s.logger != nil { - s.logger.Warn("identity runtime reconciliation failed and will retry", "error_category", "identity_runtime_reconciliation_failed") + s.logger.Warn("identity runtime reconciliation failed and will retry", + "error_category", "identity_runtime_reconciliation_failed", + "error", err, + ) } } diff --git a/apps/api/internal/httpapi/server.go b/apps/api/internal/httpapi/server.go index ecd9c51..df288ed 100644 --- a/apps/api/internal/httpapi/server.go +++ b/apps/api/internal/httpapi/server.go @@ -81,43 +81,8 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor } server.auth.ServerMainInternalKey = cfg.ServerMainInternalKey server.auth.ServerMainInternalSecret = cfg.ServerMainInternalSecret - secretStore, err := identitySecretStore(cfg) - if err != nil { - panic("invalid identity SecretStore: " + err.Error()) - } - go ssfreceiver.RunSecurityEventRetirementWorker(ctx, db) - go ssfreceiver.RunIdentitySecretCleanupWorker(ctx, db, secretStore) - runtimeBuilder := identityruntime.NewRuntimeBuilder(ctx, db, secretStore, identityruntime.RuntimeBuilderConfig{ - AppEnv: cfg.AppEnv, JWKSCacheTTL: 5 * time.Minute, - HeartbeatInterval: time.Duration(cfg.IdentitySecurityEventHeartbeatIntervalSeconds) * time.Second, - StaleAfter: time.Duration(cfg.IdentitySecurityEventStaleAfterSeconds) * time.Second, - ClockSkew: time.Duration(cfg.IdentitySecurityEventClockSkewSeconds) * time.Second, - }, securityEventMetrics) - server.identityRuntime = identityruntime.NewManager(db, runtimeBuilder) - if err := server.identityRuntime.LoadActive(ctx); err != nil && logger != nil { - logger.Error("load active identity runtime failed; local management login remains available", "error_category", "identity_runtime_load_failed") - } - server.identityPairing = identity.NewPairingService(db, secretStore, func(baseURL string) (identity.OnboardingRemote, error) { - return identity.NewOnboardingClient(baseURL, nil, cfg.AppEnv) - }, server.identityRuntime, cfg.AppEnv) - server.reconcileCanonicalIdentityPairing() - go server.runIdentityPairingCoordinator() - server.auth.OIDCVerifierProvider = func() *auth.OIDCVerifier { - if runtime := server.identityRuntime.Current(); runtime != nil { - return runtime.Verifier - } - return nil - } - server.auth.OIDCSessionResolverProvider = func(ctx context.Context, sessionID string) (*auth.User, error) { - runtime := server.identityRuntime.Current() - if runtime == nil || runtime.Sessions == nil { - return nil, auth.ErrUnauthorized - } - user, resolveErr := runtime.Sessions.Resolve(ctx, sessionID) - return user, oidcSessionRequestError(resolveErr) - } - server.auth.LegacyJWTEnabledProvider = func() bool { - return server.identityRuntime.LegacyJWTEnabled() + if cfg.RunsPublicHTTP() { + server.initializeIdentityRuntime(ctx, db, securityEventMetrics) } server.auth.LocalAPIKeyVerifier = db.VerifyLocalAPIKey server.runner.StartAdmissionNotifier(ctx) @@ -359,6 +324,54 @@ func NewServerWithContext(ctx context.Context, cfg config.Config, db *store.Stor return server.recover(server.cors(server.protectOIDCSessionCookie(mux))) } +func (server *Server) initializeIdentityRuntime( + ctx context.Context, + db *store.Store, + securityEventMetrics *ssfreceiver.Metrics, +) { + secretStore, err := identitySecretStore(server.cfg) + if err != nil { + panic("invalid identity SecretStore: " + err.Error()) + } + go ssfreceiver.RunSecurityEventRetirementWorker(ctx, db) + go ssfreceiver.RunIdentitySecretCleanupWorker(ctx, db, secretStore) + runtimeBuilder := identityruntime.NewRuntimeBuilder(ctx, db, secretStore, identityruntime.RuntimeBuilderConfig{ + AppEnv: server.cfg.AppEnv, JWKSCacheTTL: 5 * time.Minute, + HeartbeatInterval: time.Duration(server.cfg.IdentitySecurityEventHeartbeatIntervalSeconds) * time.Second, + StaleAfter: time.Duration(server.cfg.IdentitySecurityEventStaleAfterSeconds) * time.Second, + ClockSkew: time.Duration(server.cfg.IdentitySecurityEventClockSkewSeconds) * time.Second, + }, securityEventMetrics) + server.identityRuntime = identityruntime.NewManager(db, runtimeBuilder) + if err := server.identityRuntime.LoadActive(ctx); err != nil && server.logger != nil { + server.logger.Error( + "load active identity runtime failed; local management login remains available", + "error_category", "identity_runtime_load_failed", + "error", err, + ) + } + server.identityPairing = identity.NewPairingService(db, secretStore, func(baseURL string) (identity.OnboardingRemote, error) { + return identity.NewOnboardingClient(baseURL, nil, server.cfg.AppEnv) + }, server.identityRuntime, server.cfg.AppEnv) + go server.runIdentityPairingCoordinator() + server.auth.OIDCVerifierProvider = func() *auth.OIDCVerifier { + if runtime := server.identityRuntime.Current(); runtime != nil { + return runtime.Verifier + } + return nil + } + server.auth.OIDCSessionResolverProvider = func(ctx context.Context, sessionID string) (*auth.User, error) { + runtime := server.identityRuntime.Current() + if runtime == nil || runtime.Sessions == nil { + return nil, auth.ErrUnauthorized + } + user, resolveErr := runtime.Sessions.Resolve(ctx, sessionID) + return user, oidcSessionRequestError(resolveErr) + } + server.auth.LegacyJWTEnabledProvider = func() bool { + return server.identityRuntime.LegacyJWTEnabled() + } +} + func identitySecretStore(cfg config.Config) (ssfreceiver.SecretStore, error) { switch strings.ToLower(strings.TrimSpace(cfg.IdentitySecretStore)) { case "", "file": diff --git a/apps/api/internal/runner/queue_worker.go b/apps/api/internal/runner/queue_worker.go index 0ed6ba8..13aaaa9 100644 --- a/apps/api/internal/runner/queue_worker.go +++ b/apps/api/internal/runner/queue_worker.go @@ -215,12 +215,13 @@ func (s *Service) loadAsyncWorkerCapacity(ctx context.Context) (store.AsyncWorke return store.AsyncWorkerCapacitySnapshot{}, err } allocation, err := s.store.RegisterWorkerInstance(ctx, store.WorkerRegistrationInput{ - InstanceID: s.workerInstanceID, - PodUID: strings.TrimSpace(os.Getenv("POD_UID")), - PodName: strings.TrimSpace(os.Getenv("POD_NAME")), - Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")), - Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")), - DesiredCapacity: snapshot.Capacity, + InstanceID: s.workerInstanceID, + PodUID: strings.TrimSpace(os.Getenv("POD_UID")), + PodName: strings.TrimSpace(os.Getenv("POD_NAME")), + Site: strings.TrimSpace(os.Getenv("EASYAI_SITE")), + Revision: strings.TrimSpace(os.Getenv("AI_GATEWAY_REVISION")), + DesiredCapacity: snapshot.Capacity, + HeartbeatStaleAfter: time.Duration(s.cfg.AsyncWorkerRefreshIntervalSeconds) * 6 * time.Second, }) if err != nil { return store.AsyncWorkerCapacitySnapshot{}, err diff --git a/apps/api/internal/securityevents/service.go b/apps/api/internal/securityevents/service.go index e3436d0..4d04268 100644 --- a/apps/api/internal/securityevents/service.go +++ b/apps/api/internal/securityevents/service.go @@ -52,6 +52,16 @@ type Service struct { metrics *Metrics } +type heartbeatLeadershipRepository interface { + TryAcquireSecurityEventHeartbeatLeadership(context.Context) (store.Leadership, bool, error) +} + +const ( + heartbeatLeadershipRetryInterval = 5 * time.Second + heartbeatLeadershipKeepAlive = 5 * time.Second + heartbeatExecutionTimeout = 30 * time.Second +) + func (s *Service) SetMetrics(metrics *Metrics) { s.metrics = metrics } func NewService(repository StateRepository, config ServiceConfig) (*Service, error) { @@ -108,7 +118,34 @@ func (s *Service) Evaluate(ctx context.Context, identity auth.OIDCSecurityEventI } func (s *Service) run(ctx context.Context) { - s.sendVerification(ctx) + elector, distributed := s.repository.(heartbeatLeadershipRepository) + if !distributed { + s.runHeartbeatLoop(ctx) + return + } + retry := time.NewTicker(heartbeatLeadershipRetryInterval) + defer retry.Stop() + for { + acquireCtx, cancel := context.WithTimeout(ctx, heartbeatLeadershipKeepAlive) + leadership, acquired, err := elector.TryAcquireSecurityEventHeartbeatLeadership(acquireCtx) + cancel() + if err != nil { + _ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "leader_election_failed") + } else if acquired { + if !s.runLeaderHeartbeatLoop(ctx, leadership) { + return + } + } + select { + case <-ctx.Done(): + return + case <-retry.C: + } + } +} + +func (s *Service) runHeartbeatLoop(ctx context.Context) { + s.sendScheduledVerification(ctx) ticker := time.NewTicker(s.config.HeartbeatInterval) defer ticker.Stop() for { @@ -116,11 +153,42 @@ func (s *Service) run(ctx context.Context) { case <-ctx.Done(): return case <-ticker.C: - s.sendVerification(ctx) + s.sendScheduledVerification(ctx) } } } +func (s *Service) runLeaderHeartbeatLoop(ctx context.Context, leadership store.Leadership) bool { + defer leadership.Release() + s.sendScheduledVerification(ctx) + heartbeat := time.NewTicker(s.config.HeartbeatInterval) + defer heartbeat.Stop() + keepalive := time.NewTicker(heartbeatLeadershipKeepAlive) + defer keepalive.Stop() + for { + select { + case <-ctx.Done(): + return false + case <-heartbeat.C: + s.sendScheduledVerification(ctx) + case <-keepalive.C: + keepaliveCtx, cancel := context.WithTimeout(ctx, heartbeatLeadershipKeepAlive) + err := leadership.KeepAlive(keepaliveCtx) + cancel() + if err != nil { + _ = s.repository.RecordSecurityEventHeartbeatFailure(ctx, s.config.Issuer, s.config.Audience, "leader_keepalive_failed") + return true + } + } + } +} + +func (s *Service) sendScheduledVerification(ctx context.Context) { + verificationCtx, cancel := context.WithTimeout(ctx, heartbeatExecutionTimeout) + defer cancel() + s.sendVerification(verificationCtx) +} + func (s *Service) sendVerification(ctx context.Context) { outcome := "failed" defer func() { diff --git a/apps/api/internal/securityevents/service_test.go b/apps/api/internal/securityevents/service_test.go index 36b5089..636f688 100644 --- a/apps/api/internal/securityevents/service_test.go +++ b/apps/api/internal/securityevents/service_test.go @@ -15,6 +15,8 @@ import ( type fakeStateRepository struct { mutex sync.Mutex stateHash []byte + beginCount int + beginSignal chan struct{} verification chan struct{} evaluation store.SecurityEventEvaluation } @@ -25,6 +27,13 @@ func (*fakeStateRepository) EnsureSecurityEventStreamState(context.Context, stri func (f *fakeStateRepository) BeginSecurityEventVerification(_ context.Context, _, _ string, hash []byte, _ time.Time) error { f.mutex.Lock() f.stateHash = append([]byte(nil), hash...) + f.beginCount++ + if f.beginSignal != nil { + select { + case f.beginSignal <- struct{}{}: + default: + } + } f.mutex.Unlock() return nil } @@ -38,6 +47,105 @@ func (f *fakeStateRepository) EvaluateOIDCSecurityEvent(context.Context, string, return f.evaluation, nil } +type deniedHeartbeatLeadershipRepository struct { + *fakeStateRepository +} + +func (*deniedHeartbeatLeadershipRepository) TryAcquireSecurityEventHeartbeatLeadership(context.Context) (store.Leadership, bool, error) { + return nil, false, nil +} + +type fakeHeartbeatLeadership struct { + released chan struct{} + once sync.Once +} + +func (*fakeHeartbeatLeadership) KeepAlive(context.Context) error { return nil } +func (leadership *fakeHeartbeatLeadership) Release() { + leadership.once.Do(func() { close(leadership.released) }) +} + +type grantedHeartbeatLeadershipRepository struct { + *fakeStateRepository + leadership *fakeHeartbeatLeadership +} + +func (repository *grantedHeartbeatLeadershipRepository) TryAcquireSecurityEventHeartbeatLeadership(context.Context) (store.Leadership, bool, error) { + return repository.leadership, true, nil +} + +func TestServiceFollowerDoesNotRunVerificationHeartbeat(t *testing.T) { + repository := &deniedHeartbeatLeadershipRepository{fakeStateRepository: &fakeStateRepository{}} + service, err := NewService(repository, ServiceConfig{ + Issuer: "https://auth.test.example/ssf", Audience: "urn:easyai:ssf:receiver:app", + StreamID: "00000000-0000-4000-8000-000000000001", + ManagementTokenURL: "https://auth.test.example/token", + ManagementClientID: "gateway-ssf", ManagementSecret: "management-secret", + HeartbeatInterval: time.Hour, StaleAfter: 2 * time.Hour, + }) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + service.Run(ctx) + close(done) + }() + time.Sleep(100 * time.Millisecond) + cancel() + <-done + repository.mutex.Lock() + defer repository.mutex.Unlock() + if repository.beginCount != 0 { + t.Fatalf("follower persisted %d verification heartbeats", repository.beginCount) + } +} + +func TestServiceLeaderRunsHeartbeatAndReleasesLeadership(t *testing.T) { + repository := &grantedHeartbeatLeadershipRepository{ + fakeStateRepository: &fakeStateRepository{beginSignal: make(chan struct{}, 1)}, + leadership: &fakeHeartbeatLeadership{released: make(chan struct{})}, + } + tokenServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/token" { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token":"management-token","token_type":"Bearer","expires_in":300}`)) + return + } + w.WriteHeader(http.StatusNoContent) + })) + defer tokenServer.Close() + service, err := NewService(repository, ServiceConfig{ + Issuer: tokenServer.URL, Audience: "urn:easyai:ssf:receiver:app", + StreamID: "00000000-0000-4000-8000-000000000001", + ManagementTokenURL: tokenServer.URL + "/token", + ManagementClientID: "gateway-ssf", ManagementSecret: "management-secret", + HeartbeatInterval: time.Hour, StaleAfter: 2 * time.Hour, HTTPClient: tokenServer.Client(), + }) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan struct{}) + go func() { + service.Run(ctx) + close(done) + }() + select { + case <-repository.beginSignal: + case <-time.After(2 * time.Second): + t.Fatal("leader did not persist a verification heartbeat") + } + cancel() + <-done + select { + case <-repository.leadership.released: + default: + t.Fatal("leader did not release PostgreSQL leadership") + } +} + func TestServiceRequestsVerificationWithMachineTokenAndPersistsStateFirst(t *testing.T) { repository := &fakeStateRepository{verification: make(chan struct{}, 1)} var server *httptest.Server diff --git a/apps/api/internal/store/admission_queue_integration_test.go b/apps/api/internal/store/admission_queue_integration_test.go index f7fa4bd..d3b5b1e 100644 --- a/apps/api/internal/store/admission_queue_integration_test.go +++ b/apps/api/internal/store/admission_queue_integration_test.go @@ -488,8 +488,8 @@ func TestWorkerCapacityAllocationAndFailover(t *testing.T) { } if _, err := db.pool.Exec(ctx, ` UPDATE gateway_worker_instances -SET heartbeat_at = now() - interval '16 seconds' -WHERE instance_id = $1`, secondID); err != nil { +SET heartbeat_at = now() - $2::interval +WHERE instance_id = $1`, secondID, (workerHeartbeatStaleAfter + time.Second).String()); err != nil { t.Fatalf("expire second heartbeat: %v", err) } first, err = db.RegisterWorkerInstance(ctx, WorkerRegistrationInput{InstanceID: firstID, DesiredCapacity: 5}) diff --git a/apps/api/internal/store/leadership.go b/apps/api/internal/store/leadership.go new file mode 100644 index 0000000..78d66e5 --- /dev/null +++ b/apps/api/internal/store/leadership.go @@ -0,0 +1,93 @@ +package store + +import ( + "context" + "errors" + "sync" + "time" + + "github.com/jackc/pgx/v5/pgxpool" +) + +const leadershipReleaseTimeout = 5 * time.Second + +const ( + identityCoordinatorLeadershipKey = "easyai-gateway-identity-coordinator" + securityEventHeartbeatLeadershipKey = "easyai-gateway-security-event-heartbeat" +) + +// Leadership holds a PostgreSQL session advisory lock. Callers must keep the +// lease alive and release it when leadership is no longer needed. +type Leadership interface { + KeepAlive(context.Context) error + Release() +} + +type postgresLeadership struct { + conn *pgxpool.Conn + key string + once sync.Once +} + +func (s *Store) TryAcquireIdentityCoordinatorLeadership(ctx context.Context) (Leadership, bool, error) { + return s.tryAcquireLeadership(ctx, identityCoordinatorLeadershipKey) +} + +func (s *Store) TryAcquireSecurityEventHeartbeatLeadership(ctx context.Context) (Leadership, bool, error) { + return s.tryAcquireLeadership(ctx, securityEventHeartbeatLeadershipKey) +} + +func (s *Store) tryAcquireLeadership(ctx context.Context, key string) (Leadership, bool, error) { + conn, err := s.pool.Acquire(ctx) + if err != nil { + return nil, false, err + } + var acquired bool + if err := conn.QueryRow(ctx, + `SELECT pg_try_advisory_lock(hashtextextended($1, 0))`, + key, + ).Scan(&acquired); err != nil { + conn.Release() + return nil, false, err + } + if !acquired { + conn.Release() + return nil, false, nil + } + return &postgresLeadership{conn: conn, key: key}, true, nil +} + +func (leadership *postgresLeadership) KeepAlive(ctx context.Context) error { + if leadership == nil || leadership.conn == nil { + return errors.New("leadership lease is closed") + } + var alive int + return leadership.conn.QueryRow(ctx, `SELECT 1`).Scan(&alive) +} + +func (leadership *postgresLeadership) Release() { + if leadership == nil { + return + } + leadership.once.Do(func() { + unlockCtx, unlockCancel := context.WithTimeout(context.Background(), leadershipReleaseTimeout) + var unlocked bool + err := leadership.conn.QueryRow(unlockCtx, + `SELECT pg_advisory_unlock(hashtextextended($1, 0))`, + leadership.key, + ).Scan(&unlocked) + unlockCancel() + if err != nil || !unlocked { + // A session-scoped advisory lock must never return to the pool when + // unlock is uncertain. Closing the connection lets PostgreSQL + // release the lock and forces the pool to replace the session. + conn := leadership.conn.Hijack() + closeCtx, closeCancel := context.WithTimeout(context.Background(), leadershipReleaseTimeout) + _ = conn.Close(closeCtx) + closeCancel() + } else { + leadership.conn.Release() + } + leadership.conn = nil + }) +} diff --git a/apps/api/internal/store/leadership_integration_test.go b/apps/api/internal/store/leadership_integration_test.go new file mode 100644 index 0000000..e1437d9 --- /dev/null +++ b/apps/api/internal/store/leadership_integration_test.go @@ -0,0 +1,52 @@ +package store + +import ( + "context" + "os" + "strings" + "testing" + "time" +) + +func TestPostgresLeadershipIsExclusiveAndRecoverable(t *testing.T) { + databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL")) + if databaseURL == "" { + t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run PostgreSQL leadership integration tests") + } + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + first, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatal(err) + } + defer first.Close() + second, err := Connect(ctx, databaseURL) + if err != nil { + t.Fatal(err) + } + defer second.Close() + + leader, acquired, err := first.TryAcquireSecurityEventHeartbeatLeadership(ctx) + if err != nil || !acquired { + t.Fatalf("first leadership acquired=%v err=%v", acquired, err) + } + defer leader.Release() + if err := leader.KeepAlive(ctx); err != nil { + t.Fatalf("leadership keepalive: %v", err) + } + blocked, acquired, err := second.TryAcquireSecurityEventHeartbeatLeadership(ctx) + if err != nil { + t.Fatal(err) + } + if acquired { + blocked.Release() + t.Fatal("second store acquired leadership while first store held it") + } + + leader.Release() + replacement, acquired, err := second.TryAcquireSecurityEventHeartbeatLeadership(ctx) + if err != nil || !acquired { + t.Fatalf("replacement leadership acquired=%v err=%v", acquired, err) + } + replacement.Release() +} diff --git a/apps/api/internal/store/security_events.go b/apps/api/internal/store/security_events.go index 7f00a28..209f3e0 100644 --- a/apps/api/internal/store/security_events.go +++ b/apps/api/internal/store/security_events.go @@ -10,6 +10,42 @@ import ( "github.com/jackc/pgx/v5" ) +const ( + securityEventTransactionLockTimeout = 5 * time.Second + securityEventTransactionIdleTimeout = 15 * time.Second + transactionRollbackTimeout = 5 * time.Second +) + +func (s *Store) beginSecurityEventTransaction(ctx context.Context) (pgx.Tx, error) { + tx, err := s.pool.Begin(ctx) + if err != nil { + return nil, err + } + settings := []struct { + name string + value string + }{ + {name: "idle_in_transaction_session_timeout", value: securityEventTransactionIdleTimeout.String()}, + {name: "lock_timeout", value: securityEventTransactionLockTimeout.String()}, + } + for _, setting := range settings { + if _, err := tx.Exec(ctx, `SELECT set_config($1, $2, true)`, setting.name, setting.value); err != nil { + rollbackSecurityEventTransaction(tx) + return nil, err + } + } + return tx, nil +} + +func rollbackSecurityEventTransaction(tx pgx.Tx) { + if tx == nil { + return + } + ctx, cancel := context.WithTimeout(context.Background(), transactionRollbackTimeout) + defer cancel() + _ = tx.Rollback(ctx) +} + type ApplySessionRevokedInput struct { Issuer string Audience string @@ -41,11 +77,11 @@ func (s *Store) ApplySessionRevoked(ctx context.Context, input ApplySessionRevok if input.SubjectType == "" { input.SubjectType = "principal" } - tx, err := s.pool.Begin(ctx) + tx, err := s.beginSecurityEventTransaction(ctx) if err != nil { return ApplySecurityEventResult{}, err } - defer func() { _ = tx.Rollback(ctx) }() + defer rollbackSecurityEventTransaction(tx) subjectHash := shortSecurityEventHash(input.Subject) tag, err := tx.Exec(ctx, ` INSERT INTO gateway_security_event_receipts ( @@ -146,11 +182,11 @@ WHERE session.gateway_user_id = gateway_user.id } func (s *Store) EnsureSecurityEventStreamState(ctx context.Context, issuer, audience, streamID string) error { - tx, err := s.pool.Begin(ctx) + tx, err := s.beginSecurityEventTransaction(ctx) if err != nil { return err } - defer func() { _ = tx.Rollback(ctx) }() + defer rollbackSecurityEventTransaction(tx) if _, err := tx.Exec(ctx, ` INSERT INTO gateway_security_event_stream_state(issuer,audience,stream_id,mode) VALUES($1,$2,$3::uuid,'bootstrap') @@ -184,11 +220,11 @@ func (s *Store) BeginSecurityEventVerification(ctx context.Context, issuer, audi if len(stateHash) != sha256.Size { return errors.New("verification state hash must be SHA-256") } - tx, err := s.pool.Begin(ctx) + tx, err := s.beginSecurityEventTransaction(ctx) if err != nil { return err } - defer func() { _ = tx.Rollback(ctx) }() + defer rollbackSecurityEventTransaction(tx) tag, err := tx.Exec(ctx, ` UPDATE gateway_security_event_stream_state SET pending_state_hash=$3,pending_state_created_at=$4,updated_at=now() @@ -212,11 +248,11 @@ ON CONFLICT(issuer,audience,state_hash) DO NOTHING`, issuer, audience, stateHash } func (s *Store) ConfirmSecurityEventVerification(ctx context.Context, issuer, audience, streamID, jti string, stateHash []byte, now time.Time) (bool, error) { - tx, err := s.pool.Begin(ctx) + tx, err := s.beginSecurityEventTransaction(ctx) if err != nil { return false, err } - defer func() { _ = tx.Rollback(ctx) }() + defer rollbackSecurityEventTransaction(tx) tag, err := tx.Exec(ctx, ` INSERT INTO gateway_security_event_receipts(issuer,jti,audience,event_type,subject_hash) VALUES($1,$2::uuid,$3,$4,$5) @@ -294,11 +330,11 @@ ON CONFLICT(issuer,jti) DO NOTHING`, issuer, jti, audience, eventType, shortSecu // status change and immediately moves OIDC traffic to introspection. Push is // trusted again only after the normal consecutive-verification recovery path. func (s *Store) ApplySecurityEventStreamUpdated(ctx context.Context, issuer, audience, jti, streamID, status, reason string, now time.Time) (bool, error) { - tx, err := s.pool.Begin(ctx) + tx, err := s.beginSecurityEventTransaction(ctx) if err != nil { return false, err } - defer func() { _ = tx.Rollback(ctx) }() + defer rollbackSecurityEventTransaction(tx) tag, err := tx.Exec(ctx, ` INSERT INTO gateway_security_event_receipts(issuer,jti,audience,event_type,subject_hash) VALUES($1,$2::uuid,$3,$4,$5) diff --git a/apps/api/internal/store/worker_registry.go b/apps/api/internal/store/worker_registry.go index 1a9e2b3..c8dda15 100644 --- a/apps/api/internal/store/worker_registry.go +++ b/apps/api/internal/store/worker_registry.go @@ -9,15 +9,16 @@ import ( "github.com/jackc/pgx/v5" ) -const workerHeartbeatStaleAfter = 15 * time.Second +const workerHeartbeatStaleAfter = 30 * time.Second type WorkerRegistrationInput struct { - InstanceID string - PodUID string - PodName string - Site string - Revision string - DesiredCapacity int + InstanceID string + PodUID string + PodName string + Site string + Revision string + DesiredCapacity int + HeartbeatStaleAfter time.Duration } type WorkerAllocation struct { @@ -36,6 +37,10 @@ func (s *Store) RegisterWorkerInstance(ctx context.Context, input WorkerRegistra if input.DesiredCapacity < 0 { return WorkerAllocation{}, errors.New("worker desired capacity cannot be negative") } + staleAfter := input.HeartbeatStaleAfter + if staleAfter < workerHeartbeatStaleAfter { + staleAfter = workerHeartbeatStaleAfter + } tx, err := s.pool.Begin(ctx) if err != nil { return WorkerAllocation{}, err @@ -73,7 +78,7 @@ UPDATE gateway_worker_instances SET allocated_capacity = 0, updated_at = now() WHERE status <> 'active' - OR heartbeat_at <= now() - $1::interval`, workerHeartbeatStaleAfter.String()); err != nil { + OR heartbeat_at <= now() - $1::interval`, staleAfter.String()); err != nil { return WorkerAllocation{}, err } @@ -83,7 +88,7 @@ FROM gateway_worker_instances WHERE status = 'active' AND heartbeat_at > now() - $1::interval ORDER BY instance_id ASC -FOR UPDATE`, workerHeartbeatStaleAfter.String()) +FOR UPDATE`, staleAfter.String()) if err != nil { return WorkerAllocation{}, err }