From bd755ba0ccd3354ca4574407c31b2ab9563f7bc2 Mon Sep 17 00:00:00 2001 From: David Ahmann Date: Thu, 19 Feb 2026 21:50:27 -0500 Subject: [PATCH] gate: fix context leakage and envelope parsing --- core/gate/approved_scripts.go | 12 +++++-- core/gate/approved_scripts_test.go | 12 +++++++ core/gate/context_wrkr.go | 16 +++++++-- core/gate/context_wrkr_test.go | 14 ++++++++ core/gate/policy.go | 13 +++++++ core/gate/policy_test.go | 57 ++++++++++++++++++++++++++++++ 6 files changed, 118 insertions(+), 6 deletions(-) diff --git a/core/gate/approved_scripts.go b/core/gate/approved_scripts.go index 64b4a1c..8121b7e 100644 --- a/core/gate/approved_scripts.go +++ b/core/gate/approved_scripts.go @@ -177,9 +177,15 @@ func ReadApprovedScriptRegistry(path string) ([]schemagate.ApprovedScriptEntry, type registryEnvelope struct { Entries []schemagate.ApprovedScriptEntry `json:"entries"` } - var envelope registryEnvelope - if err := json.Unmarshal(content, &envelope); err == nil && len(envelope.Entries) > 0 { - return normalizeApprovedScriptEntries(envelope.Entries) + var envelopeRaw map[string]json.RawMessage + if err := json.Unmarshal(content, &envelopeRaw); err == nil { + if rawEntries, ok := envelopeRaw["entries"]; ok { + var envelope registryEnvelope + if err := json.Unmarshal(rawEntries, &envelope.Entries); err != nil { + return nil, fmt.Errorf("parse approved script registry: %w", err) + } + return normalizeApprovedScriptEntries(envelope.Entries) + } } var entries []schemagate.ApprovedScriptEntry diff --git a/core/gate/approved_scripts_test.go b/core/gate/approved_scripts_test.go index ebacd7e..47b8106 100644 --- a/core/gate/approved_scripts_test.go +++ b/core/gate/approved_scripts_test.go @@ -141,6 +141,18 @@ func TestReadApprovedScriptRegistryVariants(t *testing.T) { t.Fatalf("expected empty entries for blank registry file, got %#v", entries) } + emptyEnvelopePath := filepath.Join(t.TempDir(), "empty_envelope.json") + if err := os.WriteFile(emptyEnvelopePath, []byte(`{"entries":[]}`), 0o600); err != nil { + t.Fatalf("write empty envelope registry fixture: %v", err) + } + entries, err = ReadApprovedScriptRegistry(emptyEnvelopePath) + if err != nil { + t.Fatalf("read empty envelope approved script registry: %v", err) + } + if len(entries) != 0 { + t.Fatalf("expected empty entries for empty envelope registry, got %#v", entries) + } + nowUTC := time.Date(2026, time.February, 5, 0, 0, 0, 0, time.UTC) legacyPath := filepath.Join(t.TempDir(), "legacy.json") legacyJSON := []byte(`[ diff --git a/core/gate/context_wrkr.go b/core/gate/context_wrkr.go index 01eb8d1..9e20788 100644 --- a/core/gate/context_wrkr.go +++ b/core/gate/context_wrkr.go @@ -94,9 +94,19 @@ func parseWrkrInventory(content []byte) (map[string]WrkrToolMetadata, error) { } entries := []item{} - var wrapped envelope - if err := json.Unmarshal(content, &wrapped); err == nil && len(wrapped.Items) > 0 { - entries = wrapped.Items + var wrappedRaw map[string]json.RawMessage + if err := json.Unmarshal(content, &wrappedRaw); err == nil { + if rawItems, ok := wrappedRaw["items"]; ok { + var wrapped envelope + if err := json.Unmarshal(rawItems, &wrapped.Items); err != nil { + return nil, fmt.Errorf("parse wrkr inventory: %w", err) + } + entries = wrapped.Items + } else { + if err := json.Unmarshal(content, &entries); err != nil { + return nil, fmt.Errorf("parse wrkr inventory: %w", err) + } + } } else { if err := json.Unmarshal(content, &entries); err != nil { return nil, fmt.Errorf("parse wrkr inventory: %w", err) diff --git a/core/gate/context_wrkr_test.go b/core/gate/context_wrkr_test.go index 9042f8c..45a7257 100644 --- a/core/gate/context_wrkr_test.go +++ b/core/gate/context_wrkr_test.go @@ -88,6 +88,20 @@ func TestLoadWrkrInventoryRejectsInvalidPayload(t *testing.T) { } } +func TestLoadWrkrInventoryAcceptsEmptyEnvelope(t *testing.T) { + workDir := t.TempDir() + path := filepath.Join(workDir, "wrkr_inventory.json") + mustWriteWrkrInventoryFile(t, path, `{"items":[]}`) + + inventory, err := LoadWrkrInventory(path) + if err != nil { + t.Fatalf("LoadWrkrInventory returned error for empty envelope: %v", err) + } + if len(inventory.Tools) != 0 { + t.Fatalf("expected no tools from empty envelope, got %#v", inventory.Tools) + } +} + func TestApplyWrkrContextAddsMetadata(t *testing.T) { intent := schemagate.IntentRequest{ Context: schemagate.IntentContext{}, diff --git a/core/gate/policy.go b/core/gate/policy.go index f9148b1..1165ea0 100644 --- a/core/gate/policy.go +++ b/core/gate/policy.go @@ -374,6 +374,8 @@ func evaluateScriptPolicyDetailed(policy Policy, intent schemagate.IntentRequest for index, step := range intent.Script.Steps { stepIntent := intent + stepIntent.Context = intent.Context + stepIntent.Context.AuthContext = cloneAuthContext(intent.Context.AuthContext) stepIntent.Script = nil stepIntent.ScriptHash = "" stepIntent.ToolName = step.ToolName @@ -501,6 +503,17 @@ func mergeRateLimitPolicy(current RateLimitPolicy, candidate RateLimitPolicy) Ra return current } +func cloneAuthContext(input map[string]any) map[string]any { + if input == nil { + return nil + } + output := make(map[string]any, len(input)) + for key, value := range input { + output[key] = value + } + return output +} + func normalizedWindowPriority(window string) int { switch strings.ToLower(strings.TrimSpace(window)) { case "minute": diff --git a/core/gate/policy_test.go b/core/gate/policy_test.go index 33fa1ff..599d73a 100644 --- a/core/gate/policy_test.go +++ b/core/gate/policy_test.go @@ -1361,6 +1361,63 @@ rules: } } +func TestEvaluateScriptIntentWrkrContextDoesNotLeakAcrossSteps(t *testing.T) { + policy, err := ParsePolicyYAML([]byte(` +default_verdict: allow +rules: + - name: block-write-with-pii-context + effect: block + match: + tool_names: [tool.write] + context_data_classes: [pii] +`)) + if err != nil { + t.Fatalf("parse policy: %v", err) + } + intent := baseIntent() + intent.ToolName = "script" + intent.Script = &schemagate.IntentScript{ + Steps: []schemagate.IntentScriptStep{ + { + ToolName: "tool.read", + Targets: []schemagate.IntentTarget{ + {Kind: "path", Value: "/tmp/in.txt", Operation: "read"}, + }, + }, + { + ToolName: "tool.write", + Targets: []schemagate.IntentTarget{ + {Kind: "path", Value: "/tmp/out.txt", Operation: "write"}, + }, + }, + }, + } + + outcome, err := EvaluatePolicyDetailed(policy, intent, EvalOptions{ + ProducerVersion: "test", + WrkrInventory: map[string]WrkrToolMetadata{ + "tool.read": { + ToolName: "tool.read", + DataClass: "pii", + EndpointClass: "fs.read", + AutonomyLevel: "assist", + }, + }, + }) + if err != nil { + t.Fatalf("evaluate script policy: %v", err) + } + if outcome.Result.Verdict != "allow" { + t.Fatalf("expected allow verdict without context leakage, got %#v", outcome.Result) + } + if len(outcome.StepVerdicts) != 2 { + t.Fatalf("expected two step verdicts, got %#v", outcome.StepVerdicts) + } + if outcome.StepVerdicts[1].Verdict != "allow" { + t.Fatalf("expected write step to remain allow without wrkr match, got %#v", outcome.StepVerdicts[1]) + } +} + func TestRuleMatchesWrkrContextFields(t *testing.T) { policy, err := ParsePolicyYAML([]byte(` default_verdict: allow