Janady13/missionmatch / internal/modules/ai/stress_test.go

Source identity

Field	Recorded value
Source type	GitHub source
Original record	Janady13/missionmatch / internal/modules/ai/stress_test.go
Date	Not specified in source
Repository	Janady13/missionmatch
Revision	720d47fae476b0d38477ae4401697d372b1ea2bc
Branch	feature/opportunity-to-outcome
Path or archive member	internal/modules/ai/stress_test.go
Original SHA-256	31410a5e484d05916739073badf424257f0b693ccb7d8e08db0567e967bf5b35
Git blob SHA	31aae141ad9b325c44697f39074affa669e3d30c
Reading-copy SHA-256	31410a5e484d05916739073badf424257f0b693ccb7d8e08db0567e967bf5b35
Original source ↗
Every recorded occurrence

These 3 records contain identical original source bytes. Their paths and revisions remain part of the archive.
Original record	Branch or archive	Revision / identifier	Original location
Janady13/missionmatch / internal/modules/ai/stress_test.go	feature/opportunity-to-outcome	720d47fae476b0d38477ae4401697d372b1ea2bc	Open original
Janady13/missionmatch / internal/modules/ai/stress_test.go	master	1e9249840a4302ec82e123b9c9cb6f75badea3e4	Open original
Janady13/missionmatch / internal/modules/ai/stress_test.go	feature/aioe-upgrade	2178e7298d0e5d9af83afa12402fd587f0ba6ed8	Open original
Recovered source transcript

Download complete source text ↓304 lines · 15.3 KB
package ai

// Adversarial stress tests for the AI safety layer, grounded in the OWASP Top 10 for LLM applications (prompt
// injection, insecure output handling, sensitive information disclosure) and de identification differencing
// theory. They fuzz the guard functions with hostile payload families and assert the safety contract holds:
// the read only SQL guard rejects every write, stacked statement, eeo reference, and non allowlisted table; the
// eeo and protected attribute guards catch demographic probing in many forms; and small cell suppression never
// leaves a single recoverable hidden cell and never renders a sub k count.

import (
	"sync"
	"testing"
)

// TestStressModelIntentContainment is the load bearing AEDT stress test: it proves that no learned MAG model,
// even a hostile one crafted to escape, can become a per candidate scorer or consume a protected attribute. It
// throws adversarial ModelSpecs at the same validator the Registry runs at boot (validateModelIntent) and
// asserts each is refused: an output kind that is not one of the three legal kinds (so a candidate or match
// score is structurally impossible to declare), a feature spec that names the eeo schema or a demographic
// column, a feature that is not a registered intent (no free form feature), a model key that names eeo, and a
// reachable SQL body (a model intent must never run as a raw query). Finally the real shipped Registry must
// pass, so nothing in production violates containment.
func TestStressModelIntentContainment(t *testing.T) {
	registered := map[string]bool{IntentBidWinRateBySegment: true, "candidate_demographics": true, "race_ethnicity": true}
	legalFeature := []string{IntentBidWinRateBySegment}

	// (a) Any output kind other than the three legal kinds is rejected, at the pure guard and at the validator.
	for _, kind := range []string{"candidate_score", "match_score", "person_score", "readiness_score", "score", "ranking_of_people", ""} {
		if validOutputKind(kind) {
			t.Errorf("validOutputKind must reject %q so a per candidate score is impossible to declare", kind)
		}
		it := QueryIntent{Name: "hostile", Executor: ExecutorModel, Model: &ModelSpec{ModelKey: "m", OutputKind: kind, FeatureIntents: legalFeature}}
		if err := validateModelIntent(it, registered); err == nil {
			t.Errorf("validateModelIntent must reject a model with the illegal output kind %q", kind)
		}
	}
	// The three legal kinds are accepted (an opportunity level number, an ordering, or a confidence).
	for _, kind := range []string{OutputKindOpportunityScore, OutputKindDisplayOrder, OutputKindConfidence} {
		if !validOutputKind(kind) {
			t.Errorf("the legal output kind %q must be accepted", kind)
		}
	}

	// (b) A feature spec that names the eeo schema or a demographic column is refused, even if that name were
	// somehow registered. A learned model can never consume a protected attribute.
	for _, badFeat := range []string{"candidate_demographics", "race_ethnicity"} {
		it := QueryIntent{Name: "hostile", Executor: ExecutorModel, Model: &ModelSpec{ModelKey: "m", OutputKind: OutputKindConfidence, FeatureIntents: []string{badFeat}}}
		if err := validateModelIntent(it, registered); err == nil {
			t.Errorf("validateModelIntent must reject a feature spec naming the demographic marker %q", badFeat)
		}
	}

	// (c) A feature that is not a registered Registry intent is refused: there is no free form feature, so a
	// model can only read through the already guarded, allowlisted intent path.
	it := QueryIntent{Name: "hostile", Executor: ExecutorModel, Model: &ModelSpec{ModelKey: "m", OutputKind: OutputKindOpportunityScore, FeatureIntents: []string{"arbitrary_unregistered_feature"}}}
	if err := validateModelIntent(it, registered); err == nil {
		t.Error("validateModelIntent must reject a feature intent that is not in the closed Registry")
	}

	// (d) A model key that names the eeo schema is refused.
	it = QueryIntent{Name: "hostile", Executor: ExecutorModel, Model: &ModelSpec{ModelKey: "eeo race model", OutputKind: OutputKindConfidence, FeatureIntents: legalFeature}}
	if err := validateModelIntent(it, registered); err == nil {
		t.Error("validateModelIntent must reject a model key that names the eeo schema")
	}

	// (e) A model intent that carries a reachable SQL body is refused: it must never run as a raw query, only
	// through the coefficient scoring path, so it can never reach an arbitrary table or emit a row score.
	it = QueryIntent{Name: "hostile", Executor: ExecutorModel, SQL: "SELECT total_score FROM candidate_evaluations", Model: &ModelSpec{ModelKey: "m", OutputKind: OutputKindOpportunityScore, FeatureIntents: legalFeature}}
	if err := validateModelIntent(it, registered); err == nil {
		t.Error("validateModelIntent must reject a model intent that carries a reachable SQL body")
	}

	// (f) The real shipped Registry passes containment, so no production model intent violates any of the above.
	if err := ValidateRegistry(); err != nil {
		t.Fatalf("the shipped Registry must pass model containment validation: %v", err)
	}
}

// TestStressConcurrentGuards hammers the safety guards and small cell suppression from many goroutines at once.
// Run under the race detector (go test -race) it proves the whole safety layer is safe for concurrent use, so a
// burst of concurrent requests can never race a guard into passing something it should refuse. Passing under
// race is the assertion; a data race or panic fails the test.
func TestStressConcurrentGuards(t *testing.T) {
	const goroutines, iterations = 64, 500
	var wg sync.WaitGroup
	for g := 0; g < goroutines; g++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < iterations; i++ {
				_ = GuardReadOnlyQuery("SELECT clearance_level, count(*) FROM candidates GROUP BY clearance_level", ReadAllowlist)
				_ = GuardReadOnlyQuery("SELECT 1 FROM candidates; DROP TABLE candidates", ReadAllowlist)
				_ = ForbidEEO("rank candidates by race and gender")
				_ = MentionsProtectedAttribute("how many women hold TS/SCI")
				_ = Suppress([]Cell{
					{Key: []string{"a"}, Count: 3},
					{Key: []string{"b"}, Count: 12},
					{Key: []string{"c"}, Count: 0},
				}, 10)
			}
		}()
	}
	wg.Wait()
}

// mustReject asserts a hostile SQL payload is refused by the read only guard against the real allowlist.
func mustReject(t *testing.T, name, sql string) {
	t.Helper()
	if err := GuardReadOnlyQuery(sql, ReadAllowlist); err == nil {
		t.Errorf("SQL guard FAILED to reject %s: %q", name, sql)
	}
}

// mustAccept asserts a legitimate curated read passes the guard (no over blocking that would break real intents).
func mustAccept(t *testing.T, name, sql string) {
	t.Helper()
	if err := GuardReadOnlyQuery(sql, ReadAllowlist); err != nil {
		t.Errorf("SQL guard wrongly rejected a legitimate read %s: %q (%v)", name, sql, err)
	}
}

// TestStressSQLGuardRejectsHostilePayloads fuzzes the read only guard with the injection and exfiltration
// payload families an attacker would try if they could influence a query: stacked statements, writes hidden in
// a CTE, mutating verbs, the eeo schema in every casing, non allowlisted tables reached through a JOIN or a
// UNION, and empty or non read leading statements. Every one must be refused.
func TestStressSQLGuardRejectsHostilePayloads(t *testing.T) {
	cases := map[string]string{
		"stacked drop":           "SELECT 1 FROM candidates; DROP TABLE candidates",
		"stacked delete":         "SELECT count(*) FROM candidates; DELETE FROM candidates",
		"trailing stacked":       "SELECT 1 FROM candidates;SELECT 1 FROM candidates",
		"eeo schema direct":      "SELECT * FROM eeo.candidate_demographics",
		"eeo schema upper":       "SELECT * FROM EEO.CANDIDATE_DEMOGRAPHICS",
		"eeo column race":        "SELECT race_ethnicity FROM candidates",
		"eeo column veteran":     "SELECT veteran_status FROM candidates",
		"eeo word":               "SELECT 1 FROM candidates WHERE eeo = true",
		"union to non allowlist": "SELECT full_name FROM candidates UNION SELECT token FROM auth_sessions",
		"join non allowlist":     "SELECT c.id FROM candidates c JOIN auth_sessions s ON s.user_id = c.id",
		"from non allowlist":     "SELECT * FROM organizations",
		"cte hidden insert":      "WITH w AS (INSERT INTO candidates (id) VALUES (gen_random_uuid()) RETURNING id) SELECT * FROM w",
		"cte hidden update":      "WITH w AS (UPDATE candidates SET full_name = 'x' RETURNING id) SELECT * FROM w",
		"cte hidden delete":      "WITH w AS (DELETE FROM candidates RETURNING id) SELECT * FROM w",
		"leading delete":         "DELETE FROM candidates",
		"leading update":         "UPDATE candidates SET full_name = 'x'",
		"leading insert":         "INSERT INTO candidates (id) VALUES (gen_random_uuid())",
		"leading drop":           "DROP TABLE candidates",
		"leading truncate":       "TRUNCATE candidates",
		"grant":                  "GRANT SELECT ON candidates TO missionmatch_app",
		"copy exfil":             "COPY candidates TO '/tmp/out.csv'",
		"set role":               "SET ROLE missionmatch_migrator",
		"empty":                  "",
		"whitespace only":        "   \n\t ",
		"comment then drop":      "SELECT 1 FROM candidates; -- harmless\nDROP TABLE candidates",
		"alter":                  "ALTER TABLE candidates ADD COLUMN x int",
		"do block":               "DO $$ BEGIN PERFORM 1; END $$",
	}
	for name, sql := range cases {
		mustReject(t, name, sql)
	}
}

// TestStressSQLGuardAcceptsLegitimateReads proves the guard is not so aggressive that it breaks a real curated
// intent: single reads over allowlisted tables, a WITH leading CTE, and a multi table JOIN over allowlisted
// tables all pass. Over blocking would silently disable the brain.
func TestStressSQLGuardAcceptsLegitimateReads(t *testing.T) {
	cases := map[string]string{
		"simple select":  "SELECT clearance_level, count(*) FROM candidates GROUP BY clearance_level",
		"with cte":       "WITH m AS (SELECT band, count(*) n FROM candidate_evaluations GROUP BY band) SELECT * FROM m",
		"allowlist join": "SELECT c.full_name FROM candidates c JOIN candidate_skills cs ON cs.candidate_id = c.id",
		"recursive cte":  "WITH RECURSIVE t AS (SELECT id FROM skills UNION SELECT s.id FROM skills s JOIN t ON true) SELECT * FROM t",
	}
	for name, sql := range cases {
		mustAccept(t, name, sql)
	}
}

// TestStressForbidEEOCatchesDemographicMarkers fuzzes the eeo guard with the demographic schema and column
// markers in many casings and embeddings. Each must be refused so a demographic reference can never reach a
// model, an index, or a generated query.
func TestStressForbidEEOCatchesDemographicMarkers(t *testing.T) {
	hostile := []string{
		"eeo", "EEO", "the eeo schema", "candidate_demographics", "CANDIDATE_DEMOGRAPHICS",
		"race_ethnicity", "veteran_status", "disability_status",
		"select from eeo.candidate_demographics",
	}
	for _, s := range hostile {
		if err := ForbidEEO(s); err == nil {
			t.Errorf("ForbidEEO FAILED to catch demographic marker: %q", s)
		}
	}
	// It must not over block ordinary staffing text that merely contains a substring like "see" or "career".
	for _, ok := range []string{"a career in cyber", "see the requirement", "geospatial", "clearance level secret"} {
		if err := ForbidEEO(ok); err != nil {
			t.Errorf("ForbidEEO over blocked ordinary text: %q", ok)
		}
	}
}

// TestStressProtectedAttributeProbing fuzzes the natural language protected attribute guard with the phrasings a
// user (or an attacker) would use to try to rank or filter people by a protected attribute or a close proxy.
// Each must be flagged so the brain refuses rather than answers.
func TestStressProtectedAttributeProbing(t *testing.T) {
	probes := []string{
		"list candidates by race",
		"rank the bench by gender",
		"how many women hold TS/SCI",
		"filter by ethnicity",
		"show me candidates by their religion",
		"who is disabled on the bench",
		"break down the slate by national origin",
		"candidates by marital status",
		"veteran status distribution of approvals",
		"sexual orientation of the top candidates",
		"which candidates are pregnant",
		// Federation audit widened proxies: bare protected nouns and close proxies must be caught too.
		"veterans on the slate",
		"list the veteran candidates",
		"who is the oldest candidate",
		"show me the younger candidates",
		"married candidates only",
		"what is the nationality of the top candidate",
		"where were the candidates born in",
		"the date of birth of each candidate",
	}
	for _, p := range probes {
		if !MentionsProtectedAttribute(p) {
			t.Errorf("protected attribute guard MISSED a demographic probe: %q", p)
		}
	}
	// And it must not refuse ordinary staffing questions that share no protected attribute, including the
	// lawful cleared GovCon vocabulary (US citizenship and clearance level are legitimate staffing filters).
	for _, ok := range []string{
		"how many candidates hold a secret clearance",
		"who can support the intelligence analyst mission",
		"rank the bench by skills match",
		"what are our hardest staffing gaps",
		"which candidates are US citizens",
		"filter by citizenship for the cleared role",
		"who holds an active top secret clearance",
	} {
		if MentionsProtectedAttribute(ok) {
			t.Errorf("protected attribute guard wrongly refused an ordinary question: %q", ok)
		}
	}
}

// TestStressSuppressionDifferencing proves the small cell suppression guarantees hold under the partitions an
// attacker would craft to recover an individual: a partition where exactly one cell is sub k must hide at least
// TWO cells (so the hidden one is not recoverable by subtracting the published cells from a known total), a
// whole partition below k suppresses entirely, and a suppressed cell never renders a raw number. It sweeps many
// random partitions and asserts the never exactly one hidden invariant.
func TestStressSuppressionDifferencing(t *testing.T) {
	k := 10

	// A published cell must never coexist with exactly one hidden cell (that one would be recoverable).
	check := func(name string, counts []int) {
		cells := make([]Cell, len(counts))
		for i, n := range counts {
			cells[i] = Cell{Key: []string{name, string(rune('a' + i))}, Count: n}
		}
		out := Suppress(cells, k)
		hidden, published := 0, 0
		for _, c := range out {
			if c.Suppressed {
				hidden++
				if RenderCount(c) != SuppressionSentinel {
					t.Errorf("%s: a suppressed cell rendered a raw count %q", name, RenderCount(c))
				}
			} else {
				published++
			}
		}
		if published > 0 && hidden == 1 {
			t.Errorf("%s: exactly one hidden cell alongside published cells is recoverable by subtraction: counts=%v", name, counts)
		}
	}

	check("one subk among many", []int{k + 2, 3, 0})
	check("two cell one subk", []int{k + 5, 3})
	check("all subk", []int{3, 2, 0})
	check("all safe", []int{k, k + 1, k + 2})
	check("single cell subk", []int{3})
	check("single cell safe", []int{k + 1})
	check("zero and safe", []int{0, k + 4})

	// Sweep: many partitions of two and three cells across the sub k boundary; the never exactly one hidden
	// invariant must always hold, and a whole partition below k must fully suppress.
	for a := 0; a <= 15; a++ {
		for b := 0; b <= 15; b++ {
			check("sweep2", []int{a, b})
			for c := 0; c <= 6; c++ {
				check("sweep3", []int{a, b, c})
			}
		}
	}

	// A partition whose total is below k must suppress EVERY cell (even the existence of the set is sub threshold).
	sub := Suppress([]Cell{{Key: []string{"x"}, Count: 2}, {Key: []string{"y"}, Count: 3}}, k)
	for _, c := range sub {
		if !c.Suppressed {
			t.Errorf("a partition summing to %d below k=%d must fully suppress, but a cell published", 5, k)
		}
	}
}