easyai-ai-gateway/apps/api/internal/store/security_events_integration_test.go

package store

import (
	"context"
	"crypto/sha256"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

func TestSecurityEventWatermarkVerificationAndFallbackLifecycle(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 security event 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)
	}
	defer db.Close()
	issuer := "https://auth.test.example/ssf"
	subjectIssuer := "https://auth.test/realms/tenant"
	audience := "urn:easyai:ssf:receiver:" + uuid.NewString()
	streamID, tenantID, subject := uuid.NewString(), uuid.NewString(), uuid.NewString()
	t.Cleanup(func() {
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_receipts WHERE issuer=$1`, issuer)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_oidc_revocation_watermarks WHERE issuer=$1`, subjectIssuer)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_security_event_stream_state WHERE issuer=$1`, issuer)
		_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_audit_logs WHERE actor_source='ssf' AND target_id=ANY($1)`, []string{shortSecurityEventHash(subject), shortSecurityEventHash(streamID)})
	})
	if err := db.EnsureSecurityEventStreamState(ctx, issuer, audience, streamID); err != nil {
		t.Fatal(err)
	}
	now := time.Now().UTC().Truncate(time.Second)
	if err := db.AdvanceSecurityEventStreamState(ctx, issuer, audience, now, 180*time.Second); err != nil {
		t.Fatalf("advance fresh stream state: %v", err)
	}
	evaluation, err := db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, tenantID, subject, now, now, 180*time.Second)
	if err != nil || !evaluation.RequireIntrospection || evaluation.Mode != "bootstrap" {
		t.Fatalf("bootstrap evaluation=%#v error=%v", evaluation, err)
	}

	confirmAt := func(label string, at time.Time) {
		t.Helper()
		stateHash := sha256.Sum256([]byte(label))
		if err := db.BeginSecurityEventVerification(ctx, issuer, audience, stateHash[:], at); err != nil {
			t.Fatal(err)
		}
		matched, err := db.ConfirmSecurityEventVerification(ctx, issuer, audience, streamID, uuid.NewString(), stateHash[:], at)
		if err != nil || !matched {
			t.Fatalf("confirm matched=%v error=%v", matched, err)
		}
	}
	inserted, err := db.ApplySecurityEventStreamUpdated(ctx, issuer, audience, uuid.NewString(), streamID, "enabled", "onboarding", now)
	if err != nil || !inserted {
		t.Fatalf("enable stream inserted=%v error=%v", inserted, err)
	}
	confirmAt("first-verification-state", now)
	confirmAt("second-verification-state", now.Add(time.Minute))
	evaluation, err = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, tenantID, subject, now, now.Add(time.Minute), 180*time.Second)
	if err != nil || !evaluation.RequireIntrospection || evaluation.Mode != "introspection_fallback" {
		t.Fatalf("bootstrap overlap evaluation=%#v error=%v", evaluation, err)
	}
	confirmAt("post-bootstrap-verification-state", now.Add(361*time.Second))
	evaluation, err = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, tenantID, subject, now, now.Add(361*time.Second), 180*time.Second)
	if err != nil || evaluation.RequireIntrospection || evaluation.Mode != "push_healthy" {
		t.Fatalf("healthy evaluation=%#v error=%v", evaluation, err)
	}
	streamUpdateJTI := uuid.NewString()
	inserted, err = db.ApplySecurityEventStreamUpdated(ctx, issuer, audience, streamUpdateJTI, streamID, "paused", "maintenance", now.Add(362*time.Second))
	if err != nil || !inserted {
		t.Fatalf("stream update inserted=%v error=%v", inserted, err)
	}
	inserted, err = db.ApplySecurityEventStreamUpdated(ctx, issuer, audience, streamUpdateJTI, streamID, "paused", "maintenance", now)
	if err != nil || inserted {
		t.Fatalf("duplicate stream update inserted=%v error=%v", inserted, err)
	}
	confirmAt("paused-verification-one", now.Add(363*time.Second))
	confirmAt("paused-verification-two", now.Add(364*time.Second))
	evaluation, err = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, tenantID, subject, now, now.Add(364*time.Second), 180*time.Second)
	if err != nil || !evaluation.RequireIntrospection || evaluation.Mode != "introspection_fallback" {
		t.Fatalf("stream update fallback=%#v error=%v", evaluation, err)
	}
	_, _ = db.pool.Exec(ctx, `UPDATE gateway_security_event_stream_state SET stream_status='enabled',mode='push_healthy',fallback_since=NULL,
		consecutive_verifications=0,last_verification_at=$3 WHERE issuer=$1 AND audience=$2`, issuer, audience, now)

	revokedAt := now.Add(-10 * time.Second)
	input := ApplySessionRevokedInput{Issuer: issuer, Audience: audience, JTI: uuid.NewString(), TransactionID: uuid.NewString(), SubjectIssuer: subjectIssuer, TenantID: tenantID, Subject: subject, EventTimestamp: revokedAt, InitiatingEntity: "admin"}
	result, err := db.ApplySessionRevoked(ctx, input)
	if err != nil || !result.WatermarkMoved || result.AuditID == "" {
		t.Fatalf("apply result=%#v error=%v", result, err)
	}
	duplicate, err := db.ApplySessionRevoked(ctx, input)
	if err != nil || !duplicate.Duplicate {
		t.Fatalf("duplicate result=%#v error=%v", duplicate, err)
	}
	evaluation, _ = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, tenantID, subject, revokedAt, now, 180*time.Second)
	if !evaluation.Revoked {
		t.Fatal("token at watermark was accepted")
	}
	evaluation, _ = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, "https://other.test/issuer", tenantID, subject, revokedAt, now, 180*time.Second)
	if evaluation.Revoked {
		t.Fatal("watermark crossed OIDC issuer boundary")
	}
	evaluation, _ = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, tenantID, subject, now, now, 180*time.Second)
	if evaluation.Revoked {
		t.Fatal("token after watermark was rejected")
	}

	_, _ = db.pool.Exec(ctx, `UPDATE gateway_security_event_stream_state SET last_verification_at=$3
		WHERE issuer=$1 AND audience=$2`, issuer, audience, now.Add(-181*time.Second))
	evaluation, err = db.EvaluateOIDCSecurityEvent(ctx, issuer, audience, subjectIssuer, tenantID, subject, now, now, 180*time.Second)
	if err != nil || !evaluation.RequireIntrospection || evaluation.Mode != "introspection_fallback" {
		t.Fatalf("fallback evaluation=%#v error=%v", evaluation, err)
	}
}