diff --git a/plans/egress-gate-message-ledger.md b/plans/egress-gate-message-ledger.md new file mode 100644 index 00000000..77888a11 --- /dev/null +++ b/plans/egress-gate-message-ledger.md @@ -0,0 +1,452 @@ +# Egress Gate message fingerprinting and denied-history ledger plan + +## Scope + +The first release is intentionally process-local and in-memory. It does not +provide cross-process, cross-replica, or restart persistence. Losing state may +cause an old denied message to be denied again, but must never cause content to +be silently allowed or an unrelated message to be removed. + +## Executive decision + +Build two separate features on one reusable fingerprint-state primitive: + +1. **Evaluation memoization** reuses immutable per-target gate results so a + gate does not repeat expensive work for unchanged historical content. +2. **Denied-history sanitation** records a denied current message and removes + that complete message when the harness later resubmits it as history. + +Never implement `seen fingerprint -> skip`. A cache hit must replay the earlier +evaluation semantics. Detect findings still count, replacement edits are still +applied to the current body, and denied current input is still denied. + +History sanitation is different from memoization. It requires a stable +conversation key, an explicit current-versus-history boundary, a stable message +position, and structural removal of a complete message object. Keep its state +and invalidation semantics separate even if both features use the same storage +implementation. + +## Goals + +1. Give built-in and custom gates a public, bounded, thread-safe fingerprint + store that retains no raw request content. +2. Avoid repeated deterministic evaluation of unchanged message text within + one prepared gate instance. +3. Deny disallowed current input the first time it appears. +4. If the harness later resubmits that denied input as history, remove the + complete historical message before evaluating and forwarding the request. +5. If the same content is submitted again as current input, evaluate and deny + it again. +6. Preserve source bytes outside structurally removed messages and explicit + text replacements. +7. Preserve the current gate, pipeline, timeout, mutation, finding, and + content-safe logging contracts. + +## Non-goals + +- Durable or distributed storage. +- Exactly-once evaluation across concurrent requests. Concurrent misses may + compute the same result more than once; correctness must remain identical. +- Guessing conversation identity from `sandbox_id`. +- Inferring current input from “the last message” without harness metadata. +- Modifying the harness's local session store. +- Returning body mutations with a terminal deny. Denied evaluations remain + mutation-free. +- Silently permitting content when state is missing, expired, or unavailable. +- Removing arbitrary JSON object members or general JSON tree surgery. The + first structural edit is removal of message objects from a configured array. +- Solving provider-specific tool-call dependency repair generically. + +## Required harness contract + +History sanitation must be explicitly enabled and requires two bounded request +headers supplied by the harness: + +- a stable conversation ID; +- the number of leading message-array entries that are history for this call. + +Make both header names configurable in the message-history policy. Do not +assign implicit global header names in the domain model. When the feature is +enabled, each configured header must occur exactly once. The conversation value +must be a bounded scalar string. The history length must be a canonical +non-negative decimal integer no greater than the parsed message count. + +Egress Gate must remove both coordination headers from every proceeding +request so they are not sent to the LLM provider. The regex gate must therefore +declare `MUTATE_HEADERS` when this feature is configured. A denied request is +not forwarded, so no header mutation is necessary in the terminal result. + +Do not use `request_id`: it changes on every call. Do not use `sandbox_id`: one +sandbox may host multiple conversations. If either required header is missing, +duplicated, malformed, or inconsistent with the body, fail closed as invalid +input rather than guessing. + +The harness remains the preferred owner of its own history. It should avoid +persisting denied input when it can observe the denial. The ledger is a +defensive recovery mechanism for harnesses that resubmit rejected messages. + +## Message identity and classification + +Add a normalized message envelope above `MessageBlock`: + +```python +class MessageOccurrence(StrEnum): + HISTORY = "history" + CURRENT = "current" + + +class MessageEnvelope(StrictDomainModel): + id: str + node_id: str + message_index: int + occurrence: MessageOccurrence + role: MessageRole + blocks: tuple[MessageBlock, ...] +``` + +`JsonMessageMapParser` should continue deriving ordered blocks, but return them +grouped by their owning message envelope. The history length supplied by the +harness classifies indices below the boundary as `HISTORY` and all remaining +indices as `CURRENT`. + +Each `TextTarget` exposed by `MessageBlocksParser` must carry the document-local +ID of its owning removable content unit: + +```python +@dataclass(frozen=True, slots=True) +class TextTarget: + id: str + text: str + unit_id: str | None = None +``` + +Whole-body and arbitrary JSON-field targets may leave `unit_id` as `None`. +Message-block targets use their `MessageEnvelope.id`. Do not place conversation +IDs, fingerprints, or persistent state keys on `TextTarget`. + +For the in-memory first release, identify a denied message by: + +```text +conversation fingerprint ++ message_index at first denial ++ semantic message fingerprint +``` + +The semantic message fingerprint must include the normalized role and the +ordered `(block kind, block text)` sequence. This avoids storing raw JSON and +matches equivalent JSON string escaping. Including `message_index` prevents a +different identical message elsewhere in the same conversation from being +removed. + +If a harness truncates leading history and changes message indices, the record +will not match and the message may be denied again. That conservative false +negative is acceptable in the in-memory release. Never fall back to removing +every historical message with the same text. + +## General fingerprint-state API + +Create a focused public module, tentatively `egress_gate.fingerprints`, with: + +- an opaque immutable `ContentFingerprint` value; +- a `ContentFingerprinter` that creates keyed BLAKE2b digests with explicit + domain separation; +- a generic `FingerprintStore[ValueT]` protocol; +- a bounded `InMemoryFingerprintStore[ValueT]` implementation. + +Use a process-random key generated when a fingerprinter is constructed. Never +log or serialize the key or resulting fingerprints. Callers provide a short, +developer-authored domain such as `regex-target-v1`, `conversation-v1`, or +`message-v1`; the domain must be included in the digest input. + +The store API should be small: + +```python +class FingerprintStore(Protocol, Generic[ValueT]): + def get( + self, + scope: ContentFingerprint, + key: ContentFingerprint, + *, + timeout: Timeout, + ) -> ValueT | None: ... + + def put( + self, + scope: ContentFingerprint, + key: ContentFingerprint, + value: ValueT, + *, + timeout: Timeout, + ) -> None: ... +``` + +The in-memory implementation must: + +- use monotonic time; +- provide bounded entry count and TTL; +- use LRU eviction among live entries; +- refresh recency, but not necessarily TTL, on reads; +- be safe for prepared gates shared by worker threads; +- honor the shared `Timeout` while waiting for its lock and before returning; +- perform hashing outside the lock; +- never store raw request text, headers, bodies, or conversation IDs; +- accept immutable caller-owned values and never mutate them; +- expose content-free hit, miss, insertion, expiration, and eviction counters + only if the project has an appropriate metrics surface by implementation + time. Do not add request fingerprints to logs or traces. + +Use separate store instances or namespaces for evaluation cache records and +denied-message records. Evaluation records are scoped to one prepared gate and +therefore one rule/configuration generation. Denied-message records are scoped +to a conversation within that prepared instance. A policy replacement may lose +both forms of in-memory state; the resulting behavior is conservative +re-evaluation. + +Make the protocol and implementation usable from custom `GateResources`. +Document composition rather than adding a hidden global state service or +changing the `Gate._evaluate()` signature. + +## Parsed-content mutation contract + +Generalize parsed-content rendering beyond text replacement without teaching +gates about JSON source spans: + +```python +@dataclass(frozen=True, slots=True) +class ContentRemoval: + unit_id: str + + +ContentMutation: TypeAlias = TextReplacement | ContentRemoval + + +class ParsedRequestContent(Protocol): + @property + def targets(self) -> tuple[TextTarget, ...]: ... + + def render( + self, + mutations: tuple[ContentMutation, ...], + *, + timeout: Timeout, + ) -> bytes: ... +``` + +Replace the current `replace_text()` protocol method rather than maintaining +parallel mutation APIs. `Utf8TextParser` supports its one text replacement and +rejects removals. `JsonFieldsParser` supports selected text replacements and +rejects removals. The parsed result owned by `MessageBlocksParser` supports +both selected text replacements and removal of complete message envelopes. + +Validation must reject: + +- duplicate mutations for one target or unit; +- replacement of an unselected target; +- removal of an unknown or non-removable unit; +- replacement of a target inside a unit removed by the same render call; +- overlapping structural edits; +- output beyond `MAX_BODY_BYTES`; +- unsupported mutation kinds for a parser. + +Keep `TextReplacement` and `ContentRemoval` named and immutable. Do not return +to raw `(id, value)` tuples. + +## Source-preserving JSON array removal + +Extend the private source-aware JSON engine and `JsonDocument` with one narrow +structural operation: removal of selected immediate items from a known array. + +The document must compute non-overlapping source spans that correctly handle: + +- the only array item; +- first, middle, and last items; +- adjacent and non-adjacent removals; +- whitespace before or after commas; +- multiple removals in one linear render pass; +- simultaneous text replacements outside removed items. + +Preserve every source byte outside the removed array-item spans and explicitly +replaced string tokens. Assemble the output once in source order, reusing the +existing linear-edit discipline. Check the shared deadline during span +planning, output sizing, and assembly. + +Do not expose raw token offsets publicly. `MessageBlocksParser` retains the +document-bound message-node handles needed to request removal. + +Provider message validity remains the harness parser's responsibility. The +configurable JSON message-map implementation removes exactly the selected +message object. A future harness-specific parser may expand a removal to a +dependency group such as an assistant tool call and its tool outputs. Do not +guess those relationships in the generic JSON document. + +## Regex-gate integration + +Add an optional history configuration only to `RegexMessageBlocksScan` in the +first release. Raw body, JSON-field, path, query, and header scans do not have a +message-history contract. + +During gate preparation: + +1. Construct one process-random `ContentFingerprinter`. +2. Construct bounded evaluation and denied-history stores. +3. Validate the conversation and history-boundary header names. +4. Retain the prepared regex rule identity implicitly through the lifetime of + the gate instance; no policy fingerprint needs to enter an evaluation-cache + key owned exclusively by that instance. + +During evaluation: + +1. Parse conversation metadata and the message document. +2. Fingerprint the conversation ID without retaining the raw value. +3. For each historical message, compute its message fingerprint and look up + `(conversation scope, message index, message fingerprint)` in the denied + ledger. +4. Remove matching historical envelopes and reparse the sanitized body. Favor + the simple two-parse implementation initially; optimize only with evidence. +5. Remove the coordination headers from a proceeding request. +6. For each remaining target, look up its immutable regex detections in the + evaluation cache. On miss, run `_match_text()` and insert the result. +7. Aggregate detections and findings per occurrence exactly as today. Cache + reuse must not collapse finding counts for repeated message occurrences. +8. Apply detect, replace, or deny behavior normally. +9. Before returning a deny, record every message envelope containing a detected + target in the denied ledger. This normally records current input. After + eviction, expiration, or restart, it may rediscover and deny a historical + message once; recording that historical occurrence ensures the next request + can remove it. Never record an envelope without a detection. + +Cache immutable detection spans and rule identities, not final +`GateEvaluation` objects. On a hit, reconstruct findings and replacements +against the current target occurrence. This preserves current mutation and +counting behavior. + +If the same semantic content appears as both history and current input, remove +only a historical occurrence whose full ledger key matches. Always evaluate +the current occurrence. + +## Failure and safety behavior + +- Missing or malformed required history metadata is invalid input. +- An expired shared deadline follows the existing runtime-limit behavior. +- Store construction or internal contract failures are configuration or gate + execution failures, not allow decisions. +- A cache miss, expiration, eviction, or new process performs normal + evaluation. +- A denied-ledger miss performs normal evaluation and may deny again. +- Hash collisions are treated as computationally infeasible through keyed + 256-bit BLAKE2b output. Do not add raw-content collision verification that + defeats data minimization. +- Do not include raw conversation IDs, message text, fingerprints, cache keys, + or stored detection spans in logs, public findings, traces, or errors. +- Bound hashing input by existing request and selected-text limits. Check the + shared deadline before and after encoding large text for hashing. + +## TDD implementation sequence + +Implement in contract-first slices. Keep each slice green before proceeding. + +### Slice 1: fingerprint primitives + +Write tests for deterministic domain-separated fingerprints within one +fingerprinter, different results across domains and keys, frozen opaque values, +no content in representations, bounded LRU/TTL behavior, lock timeout, and +concurrent access. Then implement the public protocol and in-memory store. + +### Slice 2: message envelopes and harness metadata + +Write parser tests for exact history/current classification, zero history, +multiple current messages, missing/duplicate headers, malformed and excessive +history lengths, and header stripping. Then add message envelopes and the +history configuration models. + +### Slice 3: structural message removal + +Write exact-byte tests for only/first/middle/last/adjacent removals with unusual +whitespace and escaped strings. Add overlap, unknown-unit, output-bound, and +forced-timeout tests. Then generalize `ParsedRequestContent.render()` and +implement source-preserving array-item removal. + +### Slice 4: denied-history ledger + +Write end-to-end regex tests proving: + +1. a disallowed current message is denied and recorded; +2. the same message at the same historical position is removed on the next + request; +3. the next request proceeds if nothing else is disallowed; +4. the same text submitted as current input is denied again; +5. an identical historical message at a different index is not removed; +6. eviction or expiration causes conservative re-evaluation; +7. a new prepared gate starts with empty state; +8. no denied result carries mutations. + +Then implement ledger lookup, sanitation, and recording. + +### Slice 5: evaluation memoization + +Use a recording matcher in tests to prove cache hits avoid repeated matching +while preserving detection findings, occurrence counts, deny behavior, and +source-preserving replacement. Test negative-result caching because clean +history is expected to dominate. Accept concurrent duplicate computation; do +not add single-flight coordination. + +### Slice 6: custom-gate adoption and documentation + +Add a focused custom-gate test or example resource that composes +`InMemoryFingerprintStore`. Document thread safety, data minimization, +expiration, process-local limitations, and conservative fallback behavior. +Update architecture, configuration, regex-gate, custom-gate, limits, and +failure documentation. Regenerate project documentation through the canonical +staging script. + +## Validation + +Run from `projects/egress-gate/`: + +```bash +make check +``` + +Because the plan requires documentation changes, also run from the repository +root: + +```bash +python3 tests/test_render_dev_notes.py +scripts/build-docs.sh +``` + +Run the offline evaluator against at least these multi-request fixtures: + +- deny current, then sanitize the same historical message; +- replace current, then replay cached detections for history; +- duplicate text at different message indices; +- multiple current messages after one history boundary; +- missing and malformed coordination headers; +- TTL/eviction fallback to normal evaluation. + +Before handoff, independently review concurrency, deadline enforcement, +source-span deletion, conversation isolation, duplicate-message behavior, +content-safe diagnostics, and custom-gate usability. + +## Acceptance criteria + +- Custom gates can use the public in-memory fingerprint store through ordinary + typed resources without changes to the gate evaluation signature. +- Repeated allowed or replaceable message text reuses cached gate evaluation + while preserving observable findings and mutations. +- A denied current message is denied on first submission. +- The same ledger-keyed message is removed when explicitly classified as + history in the same conversation. +- The same content submitted as current input is evaluated and denied again. +- Identical text at another historical index is not removed. +- Coordination headers never reach a proceeding upstream request. +- Missing state, eviction, expiration, restart, and concurrent misses preserve + fail-closed evaluation semantics. +- Memory, entry count, TTL, selected text, structural edits, and output bytes + are bounded and deadline-aware. +- No raw request content, conversation identity, or fingerprint appears in + logs, findings, traces, errors, or object representations. +- Existing whole-body, JSON-field, message-block, detect, replace, deny, and + custom-gate behavior remains available when history handling is disabled. +- Project checks, documentation checks, offline fixtures, and independent + review are clean. diff --git a/projects/egress-gate/AGENTS.md b/projects/egress-gate/AGENTS.md index 83c2a9fc..aaeb8db9 100644 --- a/projects/egress-gate/AGENTS.md +++ b/projects/egress-gate/AGENTS.md @@ -37,10 +37,13 @@ Run focused tests while working and `make check` before handoff. ## Project map -- `src/egress_gate/gates/`: `Gate`, helper bases, registry, and the regex gate +- `src/egress_gate/gates/`: `Gate`, helper bases, registry, regex policy models, + and regex matching - `src/egress_gate/config.py`: strict ordered `gates` and `default_decision` policy models - `src/egress_gate/request.py`: protobuf-free request and request-mutation models +- `src/egress_gate/request_content/`: reusable text parsers, strict JSON + selection, source-preserving edits, and normalized message blocks - `src/egress_gate/result.py`: gate evaluations, five-field findings, provenance, traces, metadata, and final results - `src/egress_gate/request_processor.py`: shared deadline, current-request @@ -93,10 +96,11 @@ registration state. ## Current built-ins and boundaries -This slice ships exactly one built-in. `regex` selects one typed body, path, -query, or header scan and preserves bounded catalog loading, matching, -overlap resolution, and detect/deny actions. Body scans also support strict -UTF-8 replacement. Deterministic network request policy belongs to OpenShell. +This slice ships exactly one built-in. `regex` selects one typed complete-body, +structured JSON, normalized message-block, path, query, or header scan and +preserves bounded catalog loading, matching, overlap resolution, and +detect/deny actions. Complete-body and structured scans also support bounded +replacement. Deterministic network request policy belongs to OpenShell. Do not add more built-ins speculatively. The OpenShell wire `Finding` remains the released five-field contract: diff --git a/projects/egress-gate/README.md b/projects/egress-gate/README.md index 6853d668..1471f75e 100644 --- a/projects/egress-gate/README.md +++ b/projects/egress-gate/README.md @@ -66,12 +66,20 @@ gates: default_decision: allow ``` -The shipped registry contains exactly `regex`. Its `scan` selects the body, -path, query, or selected request headers. Each scan contains its `action`. -Every scan supports `detect` and `deny`. A body scan also supports `replace`. -The typed configuration prevents unsupported combinations. A replace action -preserves an explicit body-replacement intent even when the resulting bytes -equal the input. Add custom trusted gates through `--registry`. +The shipped registry contains exactly `regex`. Its `scan` selects the complete +body, selected JSON string fields, normalized JSON message blocks, the path, +query, or selected request headers. Each scan contains its `action`. Every scan +supports `detect` and `deny`. Complete-body and structured JSON scans also +support source-preserving `replace`. The typed configuration prevents +unsupported combinations. A replace action preserves an explicit +body-replacement intent even when the resulting bytes equal the input. Add +custom trusted gates through `--registry`. + +Structured scans are explicit policy choices. `json-fields` uses bounded typed +selectors. `message-blocks` applies a configurable JSON message mapping and can +filter normalized roles and block kinds. Replacement re-encodes only selected +JSON string tokens; all unrelated request-body bytes remain unchanged. Existing +`body` scans continue to inspect and optionally replace the complete UTF-8 body. Small stateless gates can use the optional `registry.gate` helper. Gates that need initialization, helper bases, or typed resources use the full class-based @@ -120,6 +128,7 @@ timeout failures must deny. - [Overview](https://github.com/NVIDIA/OpenShell-Research/blob/main/projects/egress-gate/docs/index.md) - [Configuration](https://github.com/NVIDIA/OpenShell-Research/blob/main/projects/egress-gate/docs/configuration.md) +- [Request-content parsing](https://github.com/NVIDIA/OpenShell-Research/blob/main/projects/egress-gate/docs/request-content.md) - [Test policies offline](https://github.com/NVIDIA/OpenShell-Research/blob/main/projects/egress-gate/docs/evaluation.md) - [Operations](https://github.com/NVIDIA/OpenShell-Research/blob/main/projects/egress-gate/docs/operations.md) - [Gate authoring](https://github.com/NVIDIA/OpenShell-Research/blob/main/projects/egress-gate/docs/gates/custom.md) diff --git a/projects/egress-gate/docs/architecture/index.md b/projects/egress-gate/docs/architecture/index.md index 7b0f4f79..81ebd3ac 100644 --- a/projects/egress-gate/docs/architecture/index.md +++ b/projects/egress-gate/docs/architecture/index.md @@ -18,10 +18,12 @@ Egress Gate has one transport adapter and one protobuf-free pipeline processor. | Module | Responsibility | | --- | --- | | `request.py` | Immutable request, headers, and `RequestMutations` | +| `request_content/` | Reusable text parsers, strict JSON documents, typed selection, source-preserving edits, and normalized message blocks | | `result.py` | Gate evaluations, five-field findings, provenance, traces, and result invariants | | `gates/base.py` | Gate lifecycle, capabilities, output validation, and UTF-8 helper | | `gates/registry.py` | Trusted registration, exact pipeline schema, resources, discovery, and processor preparation | -| `gates/regex.py` | Typed scan and action selection, bounded matching, overlap handling, caching, and body replacement | +| `gates/regex_scans.py` | Typed regex scan and action policy configuration | +| `gates/regex.py` | Content-parser composition, non-body text adaptation, pattern catalogs, bounded matching, finding aggregation, and gate evaluation | | `config.py` | Strict ordered gates and required default decision | | `request_processor.py` | Shared deadline, immutable snapshot construction, control flow, aggregation, and provenance | | `service/` | Protobuf validation/conversion, worker slots, lifecycle, and wire serialization | @@ -33,6 +35,21 @@ does not add a second execution path or import the transport adapter. Only `service/` imports generated protobuf/gRPC bindings. The pipeline processor and gates receive domain values and can be tested offline. +The request body remains canonical immutable bytes. A configured gate can +interpret the current body snapshot as strict JSON and select text nodes, then +optionally adapt those nodes to normalized message blocks. These views remain +local to one gate evaluation. The pipeline processor does not parse bodies or +cache request state across reusable gate instances. + +Regex scan models remain declarative Pydantic configuration. During gate +preparation, body-based variants compose a reusable `RequestContentParser`: +`Utf8TextParser`, `JsonFieldsParser`, or `MessageBlocksParser`. Each parser owns +text extraction and how immutable `TextReplacement` values become bounded body +bytes. `MessageBlocksParser` composes a `MessageBlockExtractor` to normalize a +parsed JSON document without conflating that semantic step with body parsing. +The regex gate only adapts path, query, and header values itself, then applies +matching and actions uniformly to the text targets it receives. + ## Pipeline execution
diff --git a/projects/egress-gate/docs/architecture/request-lifecycle.md b/projects/egress-gate/docs/architecture/request-lifecycle.md index ee876949..f98fb6fd 100644 --- a/projects/egress-gate/docs/architecture/request-lifecycle.md +++ b/projects/egress-gate/docs/architecture/request-lifecycle.md @@ -31,12 +31,15 @@ For each configured gate, the Egress Gate pipeline processor: 1. Check the shared deadline. 2. Pass the current read-only `HttpRequest` snapshot to the gate. -3. Reconstruct and validate the returned `GateEvaluation`. -4. Add a content-safe `GateTrace` and `SourcedFinding` values owned by the +3. When configured by the gate, parse the current body as strict JSON, select + bounded string nodes, and optionally normalize those nodes as message + blocks. +4. Reconstruct and validate the returned `GateEvaluation`. +5. Add a content-safe `GateTrace` and `SourcedFinding` values owned by the pipeline processor. -5. On `proceed`, validate the request mutations and construct the next request +6. On `proceed`, validate the request mutations and construct the next request snapshot. -6. On terminal `allow` or `deny`, stop without invoking later gates. +7. On terminal `allow` or `deny`, stop without invoking later gates. The pipeline processor never changes a request object in place. It keeps the first snapshot private, constructs a new snapshot after each validated mutation @@ -45,6 +48,12 @@ combines these mutations in order. A denied result always has an empty mutation set. Body replacement `None` and `b""` remain distinct. Header mutation variants use the required `kind` values `write` and `remove`. +Structured JSON replacement works through the same complete-body mutation +contract. The JSON document renders selected string-token edits in one bounded +pass while preserving every byte outside those tokens. It returns one complete +replacement body. A later gate therefore parses the body snapshot produced by +earlier structured or raw replacements. + If every gate proceeds, `default_decision` controls the result. Default deny uses `egress_gate_default_deny`. Default allow has no reason code. diff --git a/projects/egress-gate/docs/configuration.md b/projects/egress-gate/docs/configuration.md index b0e48cdc..aa1ec106 100644 --- a/projects/egress-gate/docs/configuration.md +++ b/projects/egress-gate/docs/configuration.md @@ -49,10 +49,16 @@ fields, unknown gate types, missing defaults, and duplicate names. The shipped registry contains only `regex`. See [Regex gate](gates/regex.md) for scans, actions, catalogs, and replacement -templates. `scan.kind` selects the body, path, query, or named headers. -`scan.action.kind` selects `detect` or `deny`; a body scan can also select -`replace`. The schema does not permit `replace` for another scan kind. A -trusted application registry supplies other behavior. +templates. `scan.kind` selects the complete body, selected JSON string fields, +normalized message blocks, path, query, or named headers. Every scan supports +`detect` and `deny`. Complete-body, `json-fields`, and `message-blocks` scans +also support `replace`; path, query, and header schemas do not. A trusted +application registry supplies other behavior. + +Structured scans do not activate automatically based on headers or request +contents. Existing `kind: body` policies retain their complete-body UTF-8 +behavior. Choose `json-fields` or `message-blocks` explicitly when a policy +expects a strict JSON request body. ## Inspect the installed registry diff --git a/projects/egress-gate/docs/gates/custom.md b/projects/egress-gate/docs/gates/custom.md index 4fc9a701..f3fecdcd 100644 --- a/projects/egress-gate/docs/gates/custom.md +++ b/projects/egress-gate/docs/gates/custom.md @@ -12,6 +12,13 @@ protobuf, or `RequestProcessor` internals. Use the function helper for a small, stateless gate. Use the class-based API when a gate needs initialization, helper-base behavior, or operational resources. +Custom gates that inspect request-body text can compose the public +`egress_gate.request_content` surface. Prepared parsers are stateless and safe +to reuse; each parsed result and its text targets remain local to one +`evaluate` call. See [Parse request content](../request-content.md) for parser +selection, typed JSON paths, message mappings, replacement, and a custom-gate +example. + The repository includes runnable examples for both extension styles: - [Function-based custom gate](https://github.com/NVIDIA/OpenShell-Research/tree/main/projects/egress-gate/examples/custom-gate) diff --git a/projects/egress-gate/docs/gates/regex.md b/projects/egress-gate/docs/gates/regex.md index 4ae31031..3eacfdb2 100644 --- a/projects/egress-gate/docs/gates/regex.md +++ b/projects/egress-gate/docs/gates/regex.md @@ -7,8 +7,9 @@ agent_markdown: true # Regex gate The `regex` gate matches one configured part of the current request. It can -inspect the body, path, query, or selected header values. It returns audit-safe -findings with type `regex_match`. +inspect the complete body, selected JSON string fields, normalized JSON message +blocks, path, query, or selected header values. It returns audit-safe findings +with type `regex_match`. Choose what to scan with `scan.kind`, then choose what to do with `scan.action.kind`. This example replaces matches in the request body: @@ -53,6 +54,92 @@ general regex replacement cannot rewrite arbitrary selected headers. A custom gate can return supported header writes or removals when it declares the `GateCapability.MUTATE_HEADERS` capability. +## Structured JSON fields + +`json-fields` parses the current body as strict UTF-8 JSON and scans only string +values selected by typed paths. The selectors are part of the general +[`egress_gate.request_content` contract](../request-content.md), which is also +available to custom gates. A `json-fields` selector starts at the document +root: + +```yaml title="Selected JSON message content" +name: message-identifiers +kind: regex +scan: + kind: json-fields + selectors: + - segments: + - kind: key + value: messages + - kind: each + - kind: key + value: content + action: + kind: replace + template: '[{entity}]' +pattern_catalog: patterns.yaml +``` + +Missing paths and non-string terminal values produce no scan text. Overlapping +selectors select the same JSON string once. Matches cannot span two selected +string values. Structured replacement re-encodes each selected string token +and preserves every source byte outside selected tokens, including whitespace, +number spellings, key order, and escaping in unrelated strings. + +The JSON parser rejects invalid UTF-8, malformed JSON, duplicate object keys, +non-standard constants, invalid Unicode scalar values, and configured parsing +limits. Invalid UTF-8 produces `body_encoding_invalid`; invalid strict JSON +produces `body_format_invalid`. + +## Normalized message blocks + +`message-blocks` builds on the same JSON document. A `json-message-map` selects +one or more message arrays, reads each message role, and applies relative +`text_selectors`. Optional `tool_input_selectors` and `tool_output_selectors` +classify provider- or harness-specific fields explicitly. The mapping requires +at least one selector across those three groups. The scan can then filter +normalized roles and block kinds: + +```yaml title="Selected user and tool messages" +name: model-visible-identifiers +kind: regex +scan: + kind: message-blocks + message_mapping: + kind: json-message-map + messages: + segments: + - kind: key + value: request + - kind: key + value: messages + role_key: role + text_selectors: + - segments: + - kind: key + value: content + - segments: + - kind: key + value: content + - kind: each + - kind: key + value: text + roles: [system, developer, user, tool] + block_kinds: [text, tool_output] + action: + kind: deny +pattern_catalog: patterns.yaml +``` + +Known roles are `system`, `developer`, `user`, `assistant`, and `tool`; other or +missing roles normalize to `unknown`. Text selected from a `tool` message is a +`tool_output`; other selected text is `text`. Message mappings are policy +configuration, so harness-specific envelopes do not require changes to regex +matching. A custom gate can reuse the public `RequestContentParser`, +`MessageBlocksParser`, and `MessageBlockExtractor` surfaces when it needs the +same text-target and replacement contracts with a different semantic +extractor. + A catalog can be inline or in a relative `.yaml` or `.yml` file. Relative paths resolve from the Egress Gate process working directory, not from the policy file. Use an inline catalog when the process does not have a stable working @@ -74,8 +161,9 @@ detections. Replacement uses deterministic, non-overlapping matches. | `deny` | terminal `deny`, findings, `egress_gate_regex_denied` | | `replace` | `proceed`, findings, explicit body replacement | -`detect` and `deny` work with every scan kind. `replace` exists only in the -body scan schema. It cannot be configured for a path, query, or header scan. +`detect` and `deny` work with every scan kind. `replace` exists in the complete +body, `json-fields`, and `message-blocks` schemas. It cannot be configured for +a path, query, or header scan. This structure keeps unsupported combinations out of generated schemas and editor suggestions. OpenShell middleware results cannot rewrite a request path or query. Header replacement is not part of the built-in gate. @@ -83,7 +171,8 @@ or query. Header replacement is not part of the built-in gate. The replace action owns its template. It returns a body replacement even when there is no match. This preserves the operator's explicit intent to replace the current body. Invalid body UTF-8 is a stable `body_encoding_invalid` service -failure. +failure. A structured scan whose UTF-8 body is not strict JSON produces the +stable `body_format_invalid` failure. The regex gate does not edit the body in place. It returns the replacement in `RequestMutations`. The pipeline processor uses it to build the next immutable @@ -99,6 +188,8 @@ OpenShell body limit. | `scan.kind` | Additional fields | Supported `action.kind` values | | --- | --- | --- | | `body` | none | `detect`, `deny`, `replace` | +| `json-fields` | non-empty typed `selectors` | `detect`, `deny`, `replace` | +| `message-blocks` | `message_mapping`; optional `roles` and `block_kinds` | `detect`, `deny`, `replace` | | `path` | none | `detect`, `deny` | | `query` | none | `detect`, `deny` | | `header` | non-empty `names` list | `detect`, `deny` | diff --git a/projects/egress-gate/docs/index.md b/projects/egress-gate/docs/index.md index 1f969d05..fe98ed26 100644 --- a/projects/egress-gate/docs/index.md +++ b/projects/egress-gate/docs/index.md @@ -99,6 +99,7 @@ policy pipeline can run in offline tests. ## Further reading - [Configuration](configuration.md) +- [Parse request content](request-content.md) - [Test policies offline](evaluation.md) - [Operations](operations.md) - [Gate authoring](gates/custom.md) diff --git a/projects/egress-gate/docs/reference/limits-and-failures.md b/projects/egress-gate/docs/reference/limits-and-failures.md index ea5cdc0a..9bb6926f 100644 --- a/projects/egress-gate/docs/reference/limits-and-failures.md +++ b/projects/egress-gate/docs/reference/limits-and-failures.md @@ -14,6 +14,12 @@ applies its configured `on_error` behavior when that ceiling expires first. | Area | Limit | | --- | ---: | | Request body | 4 MiB | +| JSON nesting depth | 128 | +| JSON value nodes | 100,000 | +| JSON field selectors per scan | 32 | +| Message content selectors per mapping | 32, plus the required messages selector | +| JSON path segments per selector | 32 | +| Selected JSON nodes or normalized message blocks | 4,096 | | Pipeline gates | 10 | | Finding groups per gate/result | 32 | | Estimated finding wire size | 4 KiB | @@ -35,7 +41,7 @@ rejected value. | Condition | Outcome | | --- | --- | -| Invalid phase, envelope, policy, or input encoding | gRPC `INVALID_ARGUMENT` | +| Invalid phase, envelope, policy, input encoding, or configured JSON format | gRPC `INVALID_ARGUMENT` | | Gate contract or unexpected execution failure | gRPC `INTERNAL` | | Internal processing deadline or pipeline processor limit | deny, source `runtime_limit`, code `egress_gate_limit_exceeded` | | Gate terminal deny | deny, source `gate`, gate-owned reason code | diff --git a/projects/egress-gate/docs/request-content.md b/projects/egress-gate/docs/request-content.md new file mode 100644 index 00000000..53edb5e8 --- /dev/null +++ b/projects/egress-gate/docs/request-content.md @@ -0,0 +1,187 @@ +--- +title: Parse request content +description: Select and replace bounded text in complete bodies, JSON fields, and normalized message blocks. +agent_markdown: true +--- + +# Parse request content + +The public `egress_gate.request_content` package lets built-in and custom gates +interpret the current request body as independently inspectable text targets. +It supports a complete UTF-8 body, selected strings in a strict JSON document, +and normalized message blocks derived from harness-specific JSON envelopes. + +`Utf8TextParser`, `JsonFieldsParser`, and `MessageBlocksParser` implement the +same `RequestContentParser` contract. Each parser owns both text extraction and +rendering immutable `TextReplacement` values back into a complete replacement +body. The regex gate composes these parsers, but they are not regex-specific; +trusted custom gates can use the same public API. + +Here, *bounded* means that parsing and replacement enforce explicit limits on +body size, JSON depth, node and selector counts, selected text, and message +blocks. All work also uses the request's shared deadline. See +[Limits and failures](reference/limits-and-failures.md) for the failure +contract. + +## Use a parser from a custom gate + +Prepare a parser once with the gate, then parse the current body inside each +evaluation. The returned view contains request-local targets and owns rendering +their replacements back into complete body bytes: + +```python title="Parse and replace selected JSON text" +from egress_gate.request_content import ( + JsonFieldsParser, + JsonKeySegment, + JsonSelector, + TextReplacement, +) + +parser = JsonFieldsParser( + selectors=( + JsonSelector( + segments=(JsonKeySegment(kind="key", value="prompt"),), + ), + ), +) + +parsed = parser.parse(request.body, timeout=timeout) +replacements = tuple( + TextReplacement(target_id=target.id, text=transform(target.text)) + for target in parsed.targets +) +replacement_body = parsed.replace_text(replacements, timeout=timeout) +``` + +The custom gate can return `replacement_body` through `RequestMutations` after +declaring `GateCapability.REPLACE_BODY`. Keeping parsing and rendering on the +same request-local view prevents a replacement from targeting text that the +parser did not expose. + +## Select JSON values with typed paths + +A `JsonSelector` is an ordered sequence of typed path segments. It is not a +JSONPath string. Egress Gate starts with one JSON value and applies each segment +to the values selected by the preceding segment. + +| Segment | Required input | Selection | +| --- | --- | --- | +| `key` | JSON object | The value of the member whose key exactly equals `value` | +| `index` | JSON array | The item at position `value`; `0` is the first item, `1` is the second, and so on | +| `each` | JSON object or array | Every immediate object-member value or array item | + +For example, this selector chooses the `content` value from every item in the +top-level `messages` array: + +```yaml title="Select every message's content" +segments: + - kind: key + value: messages + - kind: each + - kind: key + value: content +``` + +Given this document: + +```json +{ + "messages": [ + {"content": "first"}, + {"content": "second"} + ] +} +``` + +the selector produces the two string values `first` and `second`. + +Use `index` when a policy needs one specific array position. The numeric +`value` is the position in the array: + +```yaml title="Select the first message's content" +segments: + - kind: key + value: messages + - kind: index + value: 0 + - kind: key + value: content +``` + +This selector produces only `first`. An index equal to or greater than the +array length selects nothing; it is not an input error. + +A segment applied to the wrong JSON kind also selects nothing. For example, +`key` does not select from an array, and `index` does not select from an object. +Missing paths therefore produce no selected values. When a parser requests +text, selected non-string terminal values also produce no text targets. + +Multiple selectors are evaluated in configuration order. When selectors reach +the same JSON node, Egress Gate returns that node once. Results produced by one +selector follow their order in the document. + +## Select JSON string fields + +`JsonFieldsParser` evaluates selectors from the JSON document root and exposes +each selected string as an independent `TextTarget`. A match or other custom +gate operation on one target cannot span another selected string. + +The regex gate exposes this parser through `scan.kind: json-fields`: + +```yaml title="Scan selected JSON strings" +scan: + kind: json-fields + selectors: + - segments: + - kind: key + value: messages + - kind: each + - kind: key + value: content + action: + kind: replace + template: '[{entity}]' +``` + +## Map normalized message blocks + +`MessageBlocksParser` parses the request body as the same strict JSON document, +then delegates message normalization to a `MessageBlockExtractor`. The built-in +`JsonMessageBlockExtractor` applies a `JsonMessageMapConfig` for ordinary JSON +harness envelopes. + +The mapping's `messages` selector starts at the document root. Its +`text_selectors`, `tool_input_selectors`, and `tool_output_selectors` start at +each selected message object. This distinction is important: relative text +selectors do not repeat the path to the message array. + +Message roles normalize to `system`, `developer`, `user`, `assistant`, `tool`, +or `unknown`. Blocks normalize to `text`, `tool_input`, or `tool_output`. +Provider- or harness-specific adapters can implement the public +`MessageBlockExtractor` protocol without changing request-content consumers +such as the regex gate. + +When selector groups reach the same JSON string, the extractor retains each +distinct block classification. `MessageBlocksParser` applies role and block +kind filters first, then exposes the shared string once. A broad +`text_selectors` entry therefore cannot hide the same node from an explicit +tool-input or tool-output filter. + +See the [regex gate](gates/regex.md#normalized-message-blocks) for a complete +policy example and [custom gates](gates/custom.md) for the trusted extension +contract. + +## Replacement and parsing behavior + +JSON parsing is strict and bounded. It rejects invalid UTF-8, malformed JSON, +duplicate object keys, non-standard constants, invalid Unicode scalar values, +and configured parsing limits. + +Structured replacement re-encodes only selected JSON string tokens. It +preserves every source byte outside those tokens, including whitespace, number +spellings, key order, and escaping in unrelated strings. A parsed-content view +rejects replacement of targets it did not expose. + +Request-content parsers are prepared stateless objects. Each parsed result and +its document-local target identities belong to one request evaluation and must +not be reused with another request body. diff --git a/projects/egress-gate/src/egress_gate/constants.py b/projects/egress-gate/src/egress_gate/constants.py index d9de72ce..791e75fa 100644 --- a/projects/egress-gate/src/egress_gate/constants.py +++ b/projects/egress-gate/src/egress_gate/constants.py @@ -33,6 +33,15 @@ MAX_EVALUATION_CASE_NAME_BYTES = 128 MAX_EVALUATION_TAGS = 16 +# Structured request-body parsing limits. +MAX_JSON_DEPTH = 128 +MAX_JSON_NODES = 100_000 +MAX_JSON_SELECTOR_SEGMENTS = 32 +MAX_JSON_SELECTORS = 32 +MAX_JSON_SELECTED_NODES = 4_096 +MAX_JSON_SELECTED_TEXT_BYTES = MAX_BODY_BYTES +MAX_MESSAGE_BLOCKS = 4_096 + # Gate and result limits. MAX_DETECTIONS_PER_GATE = 256 MAX_DIAGNOSTIC_TEXT_BYTES = 1024 diff --git a/projects/egress-gate/src/egress_gate/errors.py b/projects/egress-gate/src/egress_gate/errors.py index 547ce059..c7308cce 100644 --- a/projects/egress-gate/src/egress_gate/errors.py +++ b/projects/egress-gate/src/egress_gate/errors.py @@ -35,6 +35,7 @@ class ErrorCode(StrEnum): REQUEST_ENVELOPE_INVALID = "request_envelope_invalid" REQUEST_BODY_TOO_LARGE = "request_body_too_large" BODY_ENCODING_INVALID = "body_encoding_invalid" + BODY_FORMAT_INVALID = "body_format_invalid" GATE_OUTPUT_INVALID = "gate_output_invalid" GATE_EXECUTION_FAILED = "gate_execution_failed" SERVER_BIND_FAILED = "server_bind_failed" @@ -100,6 +101,10 @@ class GateInputError(GateError): """A gate could not interpret a bounded request input.""" +class BodyFormatError(GateInputError): + """A gate expected a structured request body that was malformed.""" + + class TimeoutExpiredError(Exception): """The shared request-processing timeout expired.""" @@ -183,6 +188,14 @@ class _ErrorSpec: "Request body encoding is invalid.", "Supply a valid UTF-8 request body.", ), + ErrorCode.BODY_FORMAT_INVALID: _ErrorSpec( + ErrorKind.INVALID_INPUT, + ErrorComponent.SERVICE, + "parse_body", + "Request body format is invalid.", + "Supply a strict JSON request body matching the configured structured " + "body scan.", + ), ErrorCode.GATE_OUTPUT_INVALID: _ErrorSpec( ErrorKind.INTERNAL, ErrorComponent.PROCESSOR, diff --git a/projects/egress-gate/src/egress_gate/gates/__init__.py b/projects/egress-gate/src/egress_gate/gates/__init__.py index 95ee1ae8..2243e2b4 100644 --- a/projects/egress-gate/src/egress_gate/gates/__init__.py +++ b/projects/egress-gate/src/egress_gate/gates/__init__.py @@ -9,20 +9,24 @@ ) from egress_gate.gates.regex import ( ConfidenceLevel, + RegexConfig, + RegexEntity, + RegexGate, + RegexPatternCatalog, + RegexRule, +) +from egress_gate.gates.regex_scans import ( RegexBodyAction, RegexBodyScan, - RegexConfig, RegexDenyAction, RegexDetectAction, - RegexEntity, - RegexGate, RegexHeaderScan, + RegexJsonFieldsScan, + RegexMessageBlocksScan, RegexPathScan, - RegexPatternCatalog, RegexQueryScan, RegexReadOnlyAction, RegexReplaceAction, - RegexRule, RegexScan, ) from egress_gate.gates.registry import ( @@ -49,6 +53,8 @@ "RegexEntity", "RegexGate", "RegexHeaderScan", + "RegexJsonFieldsScan", + "RegexMessageBlocksScan", "RegexPatternCatalog", "RegexPathScan", "RegexQueryScan", diff --git a/projects/egress-gate/src/egress_gate/gates/regex.py b/projects/egress-gate/src/egress_gate/gates/regex.py index 0dd10b9e..8e9422b8 100644 --- a/projects/egress-gate/src/egress_gate/gates/regex.py +++ b/projects/egress-gate/src/egress_gate/gates/regex.py @@ -1,4 +1,4 @@ -"""Typed, bounded regular-expression scans and actions for HTTP requests.""" +"""Bounded regex catalog compilation, matching, and gate evaluation.""" from __future__ import annotations @@ -10,7 +10,7 @@ from pathlib import Path from stat import S_ISREG from string import Formatter -from typing import Annotated, Literal, Protocol, Self, TypeAlias +from typing import Literal, Protocol, Self import regex import yaml @@ -24,9 +24,7 @@ from egress_gate.constants import ( MAX_BODY_BYTES, MAX_DETECTIONS_PER_GATE, - MAX_DIAGNOSTIC_TEXT_BYTES, MAX_PROTO_FINDING_GROUPS, - MAX_PROTO_HEADERS, MAX_REGEX_CATALOG_FILE_BYTES, MAX_REGEX_CATALOG_PATH_BYTES, MAX_REGEX_ENTITIES_PER_CATALOG, @@ -37,11 +35,31 @@ from egress_gate.errors import ( GateConfigurationError, GateContractError, - GateInputError, GateLimitExceededError, ) from egress_gate.gates.base import Gate, GateCapability, GateConfig -from egress_gate.request import HeaderName, HttpRequest, RequestMutations +from egress_gate.gates.regex_scans import ( + RegexBodyScan, + RegexDenyAction, + RegexHeaderScan, + RegexJsonFieldsScan, + RegexMessageBlocksScan, + RegexPathScan, + RegexQueryScan, + RegexReplaceAction, + RegexScan, +) +from egress_gate.request import HttpRequest, RequestMutations +from egress_gate.request_content import ( + JsonFieldsParser, + JsonMessageBlockExtractor, + MessageBlocksParser, + ParsedRequestContent, + RequestContentParser, + TextReplacement, + TextTarget, + Utf8TextParser, +) from egress_gate.result import Finding, FindingTypeDefinition, GateEvaluation from egress_gate.string_validators import ScalarString, validate_scalar_string from egress_gate.timeout import Timeout @@ -141,99 +159,6 @@ def _catalog_is_bounded_and_unambiguous(self) -> Self: return self -class RegexDetectAction(StrictDomainModel): - """Report matches and continue without changing the request.""" - - kind: Literal["detect"] - - -class RegexDenyAction(StrictDomainModel): - """Deny the request when the scan finds a match.""" - - kind: Literal["deny"] - - -class RegexReplaceAction(StrictDomainModel): - """Replace body matches with a constrained template.""" - - kind: Literal["replace"] - template: ScalarString = Field(default="[{entity}]", repr=False) - - @field_validator("template") - @classmethod - def _template_is_safe_and_bounded(cls, value: str) -> str: - if len(value.encode("utf-8")) > MAX_DIAGNOSTIC_TEXT_BYTES: - raise ValueError("replacement template exceeds the size limit") - try: - for _, field_name, format_spec, conversion in Formatter().parse(value): - if field_name is not None and field_name != "entity": - raise ValueError - if format_spec or conversion is not None: - raise ValueError - except ValueError: - raise ValueError("replacement template syntax is invalid") from None - return value - - -RegexReadOnlyAction: TypeAlias = Annotated[ - RegexDetectAction | RegexDenyAction, - Field(discriminator="kind"), -] -RegexBodyAction: TypeAlias = Annotated[ - RegexDetectAction | RegexDenyAction | RegexReplaceAction, - Field(discriminator="kind"), -] - - -class RegexBodyScan(StrictDomainModel): - """Scan the UTF-8 request body and apply a body-compatible action.""" - - kind: Literal["body"] - action: RegexBodyAction - - -class RegexPathScan(StrictDomainModel): - """Scan the request path and detect or deny matches.""" - - kind: Literal["path"] - action: RegexReadOnlyAction - - -class RegexQueryScan(StrictDomainModel): - """Scan the raw request query and detect or deny matches.""" - - kind: Literal["query"] - action: RegexReadOnlyAction - - -class RegexHeaderScan(StrictDomainModel): - """Scan values from named request headers and detect or deny matches.""" - - kind: Literal["header"] - names: tuple[HeaderName, ...] = Field(min_length=1, max_length=MAX_PROTO_HEADERS) - action: RegexReadOnlyAction - - @field_validator("names", mode="before") - @classmethod - def _names_are_a_tuple(cls, value: object) -> object: - if isinstance(value, list | tuple): - return tuple(value) - return value - - @model_validator(mode="after") - def _names_are_unique(self) -> Self: - normalized = tuple(name.casefold() for name in self.names) - if len(normalized) != len(set(normalized)): - raise ValueError("header scan names must be unique") - return self - - -RegexScan: TypeAlias = Annotated[ - RegexBodyScan | RegexPathScan | RegexQueryScan | RegexHeaderScan, - Field(discriminator="kind"), -] - - class RegexConfig(GateConfig): """Exact policy configuration owned by ``RegexGate``.""" @@ -272,7 +197,7 @@ def _patterns_are_valid(self) -> Self: class RegexGate(Gate[RegexConfig, None]): - """Scan the request body, path, query, or selected headers with regex rules.""" + """Scan one raw or structured request view with bounded regex rules.""" capabilities = frozenset( { @@ -287,6 +212,7 @@ class RegexGate(Gate[RegexConfig, None]): def _initialize(self, *, timeout: Timeout | None = None) -> None: try: + self._content_parser = _prepare_content_parser(self.config.scan) self._rules = _compile_pattern_catalog( self.config.pattern_catalog, timeout=timeout, @@ -302,10 +228,18 @@ def _evaluate( *, timeout: Timeout, ) -> GateEvaluation: - scan_texts = self._scan_texts(request) + text_view = _read_regex_text( + self.config.scan, + self._content_parser, + request, + timeout=timeout, + ) detections_with_identity: list[tuple[_RegexDetection, str]] = [] - for text in scan_texts: - detections_with_identity.extend(self._match_text(text, timeout=timeout)) + detections_by_target: dict[str, list[tuple[_RegexDetection, str]]] = {} + for target in text_view.targets: + target_detections = self._match_text(target.text, timeout=timeout) + detections_by_target[target.id] = target_detections + detections_with_identity.extend(target_detections) if len(detections_with_identity) > MAX_DETECTIONS_PER_GATE: raise GateLimitExceededError("regex detection count exceeds the limit") detections = tuple(item[0] for item in detections_with_identity) @@ -321,40 +255,26 @@ def _evaluate( if not isinstance(action, RegexReplaceAction): return GateEvaluation.proceed(findings=findings) - body_text = scan_texts[0] - output_text = body_text - if detections: - winners = _resolve_overlaps(detections_with_identity) - output_text = _render_bounded_replacement( - body_text, - winners, - action.template, + replacements = tuple( + TextReplacement( + target_id=target.id, + text=_render_bounded_replacement( + target.text, + _resolve_overlaps(detections_by_target[target.id]), + action.template, + ), ) + for target in text_view.targets + ) + replacement_body = text_view.replace_text( + replacements, + timeout=timeout, + ) return GateEvaluation.proceed( - request_mutations=RequestMutations( - replacement_body=output_text.encode("utf-8") - ), + request_mutations=RequestMutations(replacement_body=replacement_body), findings=findings, ) - def _scan_texts(self, request: HttpRequest) -> tuple[str, ...]: - scan = self.config.scan - if isinstance(scan, RegexBodyScan): - try: - return (request.body.decode("utf-8", errors="strict"),) - except UnicodeDecodeError: - raise GateInputError("regex body scan is not valid UTF-8") from None - if isinstance(scan, RegexPathScan): - return (request.target.path,) - if isinstance(scan, RegexQueryScan): - return (request.target.query,) - selected_names = frozenset(name.casefold() for name in scan.names) - return tuple( - header.value - for header in request.headers - if header.name.casefold() in selected_names - ) - def _match_text( self, text: str, @@ -426,6 +346,70 @@ class _RegexDetection: confidence: ConfidenceLevel +@dataclass(frozen=True) +class _RegexTextView: + targets: tuple[TextTarget, ...] + parsed_content: ParsedRequestContent | None = None + + def replace_text( + self, + replacements: tuple[TextReplacement, ...], + *, + timeout: Timeout, + ) -> bytes: + if self.parsed_content is None: + raise GateContractError("regex source does not support replacement") + return self.parsed_content.replace_text(replacements, timeout=timeout) + + +def _prepare_content_parser(scan: RegexScan) -> RequestContentParser | None: + match scan: + case RegexBodyScan(): + return Utf8TextParser() + case RegexJsonFieldsScan(): + return JsonFieldsParser(selectors=scan.selectors) + case RegexMessageBlocksScan(): + return MessageBlocksParser( + extractor=JsonMessageBlockExtractor(scan.message_mapping), + roles=scan.roles, + block_kinds=scan.block_kinds, + ) + case _: + return None + + +def _read_regex_text( + scan: RegexScan, + parser: RequestContentParser | None, + request: HttpRequest, + *, + timeout: Timeout, +) -> _RegexTextView: + if parser is not None: + parsed_content = parser.parse(request.body, timeout=timeout) + return _RegexTextView( + targets=parsed_content.targets, + parsed_content=parsed_content, + ) + + timeout.raise_if_expired() + match scan: + case RegexPathScan(): + targets = (TextTarget(id="path", text=request.target.path),) + case RegexQueryScan(): + targets = (TextTarget(id="query", text=request.target.query),) + case RegexHeaderScan(): + names = frozenset(name.casefold() for name in scan.names) + targets = tuple( + TextTarget(id=f"header-{index}", text=header.value) + for index, header in enumerate(request.headers) + if header.name.casefold() in names + ) + case _: + raise GateContractError("regex source preparation is invalid") + return _RegexTextView(targets=targets) + + def _aggregate_findings( detections: tuple[_RegexDetection, ...], ) -> tuple[Finding, ...]: @@ -798,19 +782,9 @@ def _rendered_template_size(template: str, entity: str) -> int: } __all__ = [ "ConfidenceLevel", - "RegexBodyAction", - "RegexBodyScan", "RegexConfig", - "RegexDenyAction", - "RegexDetectAction", "RegexEntity", "RegexGate", - "RegexHeaderScan", "RegexPatternCatalog", - "RegexPathScan", - "RegexQueryScan", - "RegexReadOnlyAction", - "RegexReplaceAction", "RegexRule", - "RegexScan", ] diff --git a/projects/egress-gate/src/egress_gate/gates/regex_scans.py b/projects/egress-gate/src/egress_gate/gates/regex_scans.py new file mode 100644 index 00000000..b1b3e611 --- /dev/null +++ b/projects/egress-gate/src/egress_gate/gates/regex_scans.py @@ -0,0 +1,200 @@ +"""Public action and source configuration for regex request scans.""" + +from __future__ import annotations + +from string import Formatter +from typing import Annotated, Literal, Self, TypeAlias + +from pydantic import Field, field_validator, model_validator + +from egress_gate.base import StrictDomainModel +from egress_gate.constants import ( + MAX_DIAGNOSTIC_TEXT_BYTES, + MAX_JSON_SELECTORS, + MAX_PROTO_HEADERS, +) +from egress_gate.request import HeaderName +from egress_gate.request_content import ( + JsonMessageMapConfig, + JsonSelector, + MessageBlockKind, + MessageRole, +) +from egress_gate.string_validators import ScalarString, validate_scalar_string + + +class RegexDetectAction(StrictDomainModel): + """Report matches and continue without changing the request.""" + + kind: Literal["detect"] + + +class RegexDenyAction(StrictDomainModel): + """Deny the request when the scan finds a match.""" + + kind: Literal["deny"] + + +class RegexReplaceAction(StrictDomainModel): + """Replace body matches with a constrained template.""" + + kind: Literal["replace"] + template: ScalarString = Field(default="[{entity}]", repr=False) + + @field_validator("template") + @classmethod + def _template_is_safe_and_bounded(cls, value: str) -> str: + if len(value.encode("utf-8")) > MAX_DIAGNOSTIC_TEXT_BYTES: + raise ValueError("replacement template exceeds the size limit") + try: + for _, field_name, format_spec, conversion in Formatter().parse(value): + if field_name is not None and field_name != "entity": + raise ValueError + if format_spec or conversion is not None: + raise ValueError + except ValueError: + raise ValueError("replacement template syntax is invalid") from None + return value + + +RegexReadOnlyAction: TypeAlias = Annotated[ + RegexDetectAction | RegexDenyAction, + Field(discriminator="kind"), +] +RegexBodyAction: TypeAlias = Annotated[ + RegexDetectAction | RegexDenyAction | RegexReplaceAction, + Field(discriminator="kind"), +] + + +class RegexBodyScan(StrictDomainModel): + """Scan the UTF-8 request body and apply a body-compatible action.""" + + kind: Literal["body"] + action: RegexBodyAction + + +class RegexPathScan(StrictDomainModel): + """Scan the request path and detect or deny matches.""" + + kind: Literal["path"] + action: RegexReadOnlyAction + + +class RegexQueryScan(StrictDomainModel): + """Scan the raw request query and detect or deny matches.""" + + kind: Literal["query"] + action: RegexReadOnlyAction + + +class RegexHeaderScan(StrictDomainModel): + """Scan values from named request headers and detect or deny matches.""" + + kind: Literal["header"] + names: tuple[HeaderName, ...] = Field(min_length=1, max_length=MAX_PROTO_HEADERS) + action: RegexReadOnlyAction + + @field_validator("names", mode="before") + @classmethod + def _names_are_a_tuple(cls, value: object) -> object: + if isinstance(value, list | tuple): + return tuple(value) + return value + + @model_validator(mode="after") + def _names_are_unique(self) -> Self: + normalized = tuple(name.casefold() for name in self.names) + if len(normalized) != len(set(normalized)): + raise ValueError("header scan names must be unique") + return self + + +class RegexJsonFieldsScan(StrictDomainModel): + """Scan selected JSON string values and optionally replace their matches.""" + + kind: Literal["json-fields"] + selectors: tuple[JsonSelector, ...] = Field( + min_length=1, + max_length=MAX_JSON_SELECTORS, + ) + action: RegexBodyAction + + @field_validator("selectors", mode="before") + @classmethod + def _selectors_are_a_tuple(cls, value: object) -> object: + if isinstance(value, list | tuple): + return tuple(value) + return value + + +class RegexMessageBlocksScan(StrictDomainModel): + """Scan normalized text-bearing JSON message blocks.""" + + kind: Literal["message-blocks"] + message_mapping: JsonMessageMapConfig + roles: tuple[MessageRole, ...] | None = None + block_kinds: tuple[MessageBlockKind, ...] | None = None + action: RegexBodyAction + + @field_validator("roles", mode="before") + @classmethod + def _parse_roles(cls, value: object) -> object: + if value is None: + return None + if not isinstance(value, list | tuple): + return value + return tuple( + item + if isinstance(item, MessageRole) + else MessageRole(validate_scalar_string(item)) + for item in value + ) + + @field_validator("block_kinds", mode="before") + @classmethod + def _parse_block_kinds(cls, value: object) -> object: + if value is None: + return None + if not isinstance(value, list | tuple): + return value + return tuple( + item + if isinstance(item, MessageBlockKind) + else MessageBlockKind(validate_scalar_string(item)) + for item in value + ) + + @model_validator(mode="after") + def _filters_are_unique(self) -> Self: + for values in (self.roles, self.block_kinds): + if values is not None and len(values) != len(set(values)): + raise ValueError("message block filters must be unique") + return self + + +RegexScan: TypeAlias = Annotated[ + RegexBodyScan + | RegexPathScan + | RegexQueryScan + | RegexHeaderScan + | RegexJsonFieldsScan + | RegexMessageBlocksScan, + Field(discriminator="kind"), +] + + +__all__ = [ + "RegexBodyAction", + "RegexBodyScan", + "RegexDenyAction", + "RegexDetectAction", + "RegexHeaderScan", + "RegexJsonFieldsScan", + "RegexMessageBlocksScan", + "RegexPathScan", + "RegexQueryScan", + "RegexReadOnlyAction", + "RegexReplaceAction", + "RegexScan", +] diff --git a/projects/egress-gate/src/egress_gate/request_content/__init__.py b/projects/egress-gate/src/egress_gate/request_content/__init__.py new file mode 100644 index 00000000..a39f9665 --- /dev/null +++ b/projects/egress-gate/src/egress_gate/request_content/__init__.py @@ -0,0 +1,56 @@ +"""Public structured request-content parsing surface.""" + +from egress_gate.request_content.json import ( + JsonDocument, + JsonEachSegment, + JsonIndexSegment, + JsonKeySegment, + JsonNode, + JsonNodeKind, + JsonPathSegment, + JsonSelector, + JsonTextNode, +) +from egress_gate.request_content.messages import ( + JsonMessageBlockExtractor, + JsonMessageMapConfig, + MessageBlock, + MessageBlockExtractor, + MessageBlockKind, + MessageDocument, + MessageRole, +) +from egress_gate.request_content.parsers import ( + JsonFieldsParser, + MessageBlocksParser, + ParsedRequestContent, + RequestContentParser, + Utf8TextParser, +) +from egress_gate.request_content.text import TextReplacement, TextTarget + +__all__ = [ + "JsonDocument", + "JsonEachSegment", + "JsonFieldsParser", + "JsonIndexSegment", + "JsonKeySegment", + "JsonMessageBlockExtractor", + "JsonMessageMapConfig", + "JsonNode", + "JsonNodeKind", + "JsonPathSegment", + "JsonSelector", + "JsonTextNode", + "MessageBlock", + "MessageBlockExtractor", + "MessageBlockKind", + "MessageBlocksParser", + "MessageDocument", + "MessageRole", + "ParsedRequestContent", + "RequestContentParser", + "TextReplacement", + "TextTarget", + "Utf8TextParser", +] diff --git a/projects/egress-gate/src/egress_gate/request_content/_json_parser.py b/projects/egress-gate/src/egress_gate/request_content/_json_parser.py new file mode 100644 index 00000000..72e51bd8 --- /dev/null +++ b/projects/egress-gate/src/egress_gate/request_content/_json_parser.py @@ -0,0 +1,227 @@ +"""Private bounded parser for source-aware strict JSON documents.""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass + +from egress_gate.constants import MAX_JSON_DEPTH, MAX_JSON_NODES +from egress_gate.errors import GateLimitExceededError +from egress_gate.request_content.json_values import JsonNodeKind +from egress_gate.timeout import Timeout + + +@dataclass(frozen=True) +class _JsonNode: + id: str + path: tuple[str | int, ...] + kind: JsonNodeKind + token_start: int + token_end: int + text: str | None = None + object_members: tuple[tuple[str, _JsonNode], ...] = () + array_items: tuple[_JsonNode, ...] = () + + +class _JsonParser: + def __init__(self, source: str, timeout: Timeout) -> None: + self.source = source + self.timeout = timeout + self.index = 0 + self.node_count = 0 + self.nodes: dict[str, _JsonNode] = {} + + def parse(self) -> _JsonNode: + self._skip_whitespace() + root = self._parse_value((), 0) + self._skip_whitespace() + if self.index != len(self.source): + raise ValueError("trailing JSON data") + return root + + def _parse_value(self, path: tuple[str | int, ...], depth: int) -> _JsonNode: + self.timeout.raise_if_expired() + if self.index >= len(self.source): + raise ValueError("JSON value is missing") + character = self.source[self.index] + if character == "{": + if depth >= MAX_JSON_DEPTH: + raise GateLimitExceededError("JSON nesting depth exceeds the limit") + return self._parse_object(path, depth) + if character == "[": + if depth >= MAX_JSON_DEPTH: + raise GateLimitExceededError("JSON nesting depth exceeds the limit") + return self._parse_array(path, depth) + if character == '"': + start = self.index + text = self._parse_string() + return self._new_node( + path, + JsonNodeKind.STRING, + start, + self.index, + text=text, + ) + if character in "-0123456789": + start = self.index + match = _NUMBER_PATTERN.match(self.source, self.index) + if match is None: + raise ValueError("JSON number is invalid") + self.index = match.end() + return self._new_node(path, JsonNodeKind.NUMBER, start, self.index) + for literal, kind in ( + ("true", JsonNodeKind.BOOLEAN), + ("false", JsonNodeKind.BOOLEAN), + ("null", JsonNodeKind.NULL), + ): + if self.source.startswith(literal, self.index): + start = self.index + self.index += len(literal) + return self._new_node(path, kind, start, self.index) + raise ValueError("JSON value is invalid") + + def _parse_object(self, path: tuple[str | int, ...], depth: int) -> _JsonNode: + start = self.index + self.index += 1 + self._skip_whitespace() + members: list[tuple[str, _JsonNode]] = [] + keys: set[str] = set() + if self._consume("}"): + return self._new_node( + path, JsonNodeKind.OBJECT, start, self.index, object_members=() + ) + while True: + if self.index >= len(self.source) or self.source[self.index] != '"': + raise ValueError("JSON object key is invalid") + key = self._parse_string() + if key in keys: + raise ValueError("JSON object keys must be unique") + keys.add(key) + self._skip_whitespace() + if not self._consume(":"): + raise ValueError("JSON object separator is missing") + self._skip_whitespace() + value = self._parse_value((*path, key), depth + 1) + members.append((key, value)) + self._skip_whitespace() + if self._consume("}"): + break + if not self._consume(","): + raise ValueError("JSON object delimiter is missing") + self._skip_whitespace() + return self._new_node( + path, + JsonNodeKind.OBJECT, + start, + self.index, + object_members=tuple(members), + ) + + def _parse_array(self, path: tuple[str | int, ...], depth: int) -> _JsonNode: + start = self.index + self.index += 1 + self._skip_whitespace() + items: list[_JsonNode] = [] + if self._consume("]"): + return self._new_node( + path, JsonNodeKind.ARRAY, start, self.index, array_items=() + ) + while True: + item = self._parse_value((*path, len(items)), depth + 1) + items.append(item) + self._skip_whitespace() + if self._consume("]"): + break + if not self._consume(","): + raise ValueError("JSON array delimiter is missing") + self._skip_whitespace() + return self._new_node( + path, + JsonNodeKind.ARRAY, + start, + self.index, + array_items=tuple(items), + ) + + def _parse_string(self) -> str: + start = self.index + self.index += 1 + next_timeout_check = self.index + _LEXICAL_TIMEOUT_CHECK_INTERVAL + while self.index < len(self.source): + if self.index >= next_timeout_check: + self.timeout.raise_if_expired() + next_timeout_check = self.index + _LEXICAL_TIMEOUT_CHECK_INTERVAL + character = self.source[self.index] + if character == '"': + self.index += 1 + token = self.source[start : self.index] + value = json.loads(token) + value.encode("utf-8", errors="strict") + return value + if character == "\\": + self.index += 1 + if self.index >= len(self.source): + raise ValueError("JSON string escape is incomplete") + escape = self.source[self.index] + if escape == "u": + digits = self.source[self.index + 1 : self.index + 5] + if len(digits) != 4 or _HEX_PATTERN.fullmatch(digits) is None: + raise ValueError("JSON Unicode escape is invalid") + self.index += 5 + continue + if escape not in '"\\/bfnrt': + raise ValueError("JSON string escape is invalid") + elif ord(character) < 0x20: + raise ValueError("JSON string contains a control character") + self.index += 1 + raise ValueError("JSON string is unterminated") + + def _new_node( + self, + path: tuple[str | int, ...], + kind: JsonNodeKind, + token_start: int, + token_end: int, + *, + text: str | None = None, + object_members: tuple[tuple[str, _JsonNode], ...] = (), + array_items: tuple[_JsonNode, ...] = (), + ) -> _JsonNode: + self.node_count += 1 + if self.node_count > MAX_JSON_NODES: + raise GateLimitExceededError("JSON node count exceeds the limit") + node = _JsonNode( + id=f"json-node-{self.node_count}", + path=path, + kind=kind, + token_start=token_start, + token_end=token_end, + text=text, + object_members=object_members, + array_items=array_items, + ) + self.nodes[node.id] = node + return node + + def _skip_whitespace(self) -> None: + next_timeout_check = self.index + _LEXICAL_TIMEOUT_CHECK_INTERVAL + while self.index < len(self.source) and self.source[self.index] in " \t\r\n": + self.index += 1 + if self.index >= next_timeout_check: + self.timeout.raise_if_expired() + next_timeout_check = self.index + _LEXICAL_TIMEOUT_CHECK_INTERVAL + + def _consume(self, token: str) -> bool: + if self.source.startswith(token, self.index): + self.index += len(token) + return True + return False + + +_NUMBER_PATTERN = re.compile(r"-?(?:0|[1-9][0-9]*)(?:\.[0-9]+)?(?:[eE][+-]?[0-9]+)?") +_HEX_PATTERN = re.compile(r"[0-9a-fA-F]{4}") +_LEXICAL_TIMEOUT_CHECK_INTERVAL = 256 + + +__all__ = ["_JsonNode", "_JsonParser"] diff --git a/projects/egress-gate/src/egress_gate/request_content/json.py b/projects/egress-gate/src/egress_gate/request_content/json.py new file mode 100644 index 00000000..7add8d7b --- /dev/null +++ b/projects/egress-gate/src/egress_gate/request_content/json.py @@ -0,0 +1,337 @@ +"""Strict source-aware JSON request-content parsing and selection.""" + +from __future__ import annotations + +import json +from typing import Annotated, Literal, TypeAlias + +from pydantic import Field, field_validator + +from egress_gate.base import StrictDomainModel +from egress_gate.constants import ( + MAX_BODY_BYTES, + MAX_JSON_SELECTED_NODES, + MAX_JSON_SELECTED_TEXT_BYTES, + MAX_JSON_SELECTOR_SEGMENTS, + MAX_JSON_SELECTORS, +) +from egress_gate.errors import BodyFormatError, GateInputError, GateLimitExceededError +from egress_gate.request_content._json_parser import _JsonNode, _JsonParser +from egress_gate.request_content.json_values import ( + JsonNode, + JsonNodeKind, + JsonTextNode, +) +from egress_gate.request_content.text import TextReplacement +from egress_gate.string_validators import ScalarString +from egress_gate.timeout import Timeout + + +class JsonKeySegment(StrictDomainModel): + """Select one object member by its exact decoded key.""" + + kind: Literal["key"] + value: ScalarString + + +class JsonIndexSegment(StrictDomainModel): + """Select one array item by its zero-based index.""" + + kind: Literal["index"] + value: int = Field(ge=0) + + +class JsonEachSegment(StrictDomainModel): + """Select every immediate array item or object member value.""" + + kind: Literal["each"] + + +JsonPathSegment: TypeAlias = Annotated[ + JsonKeySegment | JsonIndexSegment | JsonEachSegment, + Field(discriminator="kind"), +] + + +class JsonSelector(StrictDomainModel): + """One bounded path from a selected JSON node.""" + + segments: tuple[JsonPathSegment, ...] = Field( + min_length=1, + max_length=MAX_JSON_SELECTOR_SEGMENTS, + ) + + @field_validator("segments", mode="before") + @classmethod + def _segments_are_a_tuple(cls, value: object) -> object: + if isinstance(value, list | tuple): + return tuple(value) + return value + + +class JsonDocument: + """One strict JSON body with bounded traversal and source-preserving edits.""" + + __slots__ = ("_nodes", "_public_nodes", "_root", "_source") + + def __init__(self, source: str, root: _JsonNode, nodes: dict[str, _JsonNode]): + self._source = source + self._root = root + self._nodes = nodes + self._public_nodes: dict[str, JsonNode] = {} + + @classmethod + def parse(cls, body: bytes, *, timeout: Timeout) -> JsonDocument: + """Parse one complete strict UTF-8 JSON body under the shared deadline.""" + timeout.raise_if_expired() + if len(body) > MAX_BODY_BYTES: + raise GateLimitExceededError("JSON request body exceeds the size limit") + try: + source = body.decode("utf-8", errors="strict") + except UnicodeDecodeError: + raise GateInputError("request body is not valid UTF-8") from None + try: + parser = _JsonParser(source, timeout) + root = parser.parse() + except GateLimitExceededError: + raise + except (RecursionError, UnicodeError, ValueError): + raise BodyFormatError("request body is not strict JSON") from None + timeout.raise_if_expired() + return cls(source, root, parser.nodes) + + def select_nodes( + self, + selectors: tuple[JsonSelector, ...], + *, + timeout: Timeout, + ) -> tuple[JsonNode, ...]: + """Select unique nodes in selector order and then document order.""" + return self._select_from(self._root, selectors, timeout=timeout) + + def select_from( + self, + node: JsonNode, + selectors: tuple[JsonSelector, ...], + *, + timeout: Timeout, + ) -> tuple[JsonNode, ...]: + """Select unique nodes relative to one node from this document.""" + return self._select_from(self._resolve_node(node), selectors, timeout=timeout) + + def select_text( + self, + selectors: tuple[JsonSelector, ...], + *, + timeout: Timeout, + ) -> tuple[JsonTextNode, ...]: + """Select unique string nodes from the document root.""" + return self._text_nodes( + self.select_nodes(selectors, timeout=timeout), + timeout=timeout, + ) + + def select_text_from( + self, + node: JsonNode, + selectors: tuple[JsonSelector, ...], + *, + timeout: Timeout, + ) -> tuple[JsonTextNode, ...]: + """Select unique string nodes relative to one document node.""" + return self._text_nodes( + self.select_from(node, selectors, timeout=timeout), + timeout=timeout, + ) + + def array_items( + self, + node: JsonNode, + *, + timeout: Timeout, + ) -> tuple[JsonNode, ...]: + """Return the ordered immediate items of an array node.""" + internal = self._resolve_node(node) + if internal.kind is not JsonNodeKind.ARRAY: + return () + items: list[JsonNode] = [] + for index, item in enumerate(internal.array_items): + if index % _TIMEOUT_CHECK_INTERVAL == 0: + timeout.raise_if_expired() + items.append(self._public_node(item)) + timeout.raise_if_expired() + return tuple(items) + + def object_member(self, node: JsonNode, key: str) -> JsonNode | None: + """Return one exact object member without exposing mutable JSON values.""" + internal = self._resolve_node(node) + if internal.kind is not JsonNodeKind.OBJECT: + return None + for member_key, value in internal.object_members: + if member_key == key: + return self._public_node(value) + return None + + def text_value(self, node: JsonNode) -> str | None: + """Return the decoded value when a node is a JSON string.""" + internal = self._resolve_node(node) + return internal.text if internal.kind is JsonNodeKind.STRING else None + + def replace_text( + self, + replacements: tuple[TextReplacement, ...], + *, + timeout: Timeout, + ) -> bytes: + """Replace complete JSON string tokens and preserve all unrelated source.""" + timeout.raise_if_expired() + replacement_ids = tuple(item.target_id for item in replacements) + if len(replacement_ids) != len(set(replacement_ids)): + raise ValueError("JSON replacement node IDs must be unique") + edits: list[tuple[int, int, bytes]] = [] + for replacement in replacements: + timeout.raise_if_expired() + node = self._nodes.get(replacement.target_id) + if node is None or node.kind is not JsonNodeKind.STRING: + raise ValueError("JSON replacement node ID is unknown") + try: + rendered = json.dumps(replacement.text, ensure_ascii=False).encode( + "utf-8", errors="strict" + ) + except UnicodeEncodeError: + raise ValueError("JSON replacement text is invalid") from None + edits.append((node.token_start, node.token_end, rendered)) + + output_parts: list[bytes] = [] + output_size = 0 + source_cursor = 0 + for start, end, rendered in sorted(edits): + timeout.raise_if_expired() + unchanged = self._source[source_cursor:start].encode("utf-8") + output_size += len(unchanged) + len(rendered) + if output_size > MAX_BODY_BYTES: + raise GateLimitExceededError("JSON replacement body exceeds the limit") + output_parts.extend((unchanged, rendered)) + source_cursor = end + tail = self._source[source_cursor:].encode("utf-8") + output_size += len(tail) + if output_size > MAX_BODY_BYTES: + raise GateLimitExceededError("JSON replacement body exceeds the limit") + output_parts.append(tail) + timeout.raise_if_expired() + return b"".join(output_parts) + + def _select_from( + self, + start: _JsonNode, + selectors: tuple[JsonSelector, ...], + *, + timeout: Timeout, + ) -> tuple[JsonNode, ...]: + if len(selectors) > MAX_JSON_SELECTORS: + raise GateLimitExceededError("JSON selector count exceeds the limit") + selected: list[_JsonNode] = [] + seen: set[str] = set() + for selector in selectors: + current = (start,) + for segment in selector.segments: + timeout.raise_if_expired() + next_nodes: list[_JsonNode] = [] + for node in current: + next_nodes.extend(_select_segment(node, segment)) + if len(next_nodes) > MAX_JSON_SELECTED_NODES: + raise GateLimitExceededError( + "JSON selected node count exceeds the limit" + ) + current = tuple(next_nodes) + for node in current: + if node.id not in seen: + seen.add(node.id) + selected.append(node) + if len(selected) > MAX_JSON_SELECTED_NODES: + raise GateLimitExceededError( + "JSON selected node count exceeds the limit" + ) + timeout.raise_if_expired() + public_nodes: list[JsonNode] = [] + for index, node in enumerate(selected): + if index % _TIMEOUT_CHECK_INTERVAL == 0: + timeout.raise_if_expired() + public_nodes.append(self._public_node(node)) + timeout.raise_if_expired() + return tuple(public_nodes) + + def _text_nodes( + self, + nodes: tuple[JsonNode, ...], + *, + timeout: Timeout, + ) -> tuple[JsonTextNode, ...]: + selected: list[JsonTextNode] = [] + encoded_size = 0 + for index, node in enumerate(nodes): + if index % _TIMEOUT_CHECK_INTERVAL == 0: + timeout.raise_if_expired() + internal = self._resolve_node(node) + if internal.kind is not JsonNodeKind.STRING or internal.text is None: + continue + encoded_size += len(internal.text.encode("utf-8")) + if encoded_size > MAX_JSON_SELECTED_TEXT_BYTES: + raise GateLimitExceededError("JSON selected text exceeds the limit") + selected.append( + JsonTextNode(id=internal.id, path=internal.path, text=internal.text) + ) + timeout.raise_if_expired() + return tuple(selected) + + def _resolve_node(self, node: JsonNode) -> _JsonNode: + internal = self._nodes.get(node.id) + if internal is None or self._public_nodes.get(node.id) is not node: + raise ValueError("JSON node does not belong to this document") + return internal + + def _public_node(self, node: _JsonNode) -> JsonNode: + public = self._public_nodes.get(node.id) + if public is None: + public = JsonNode(id=node.id, path=node.path, kind=node.kind) + self._public_nodes[node.id] = public + return public + + +def _select_segment( + node: _JsonNode, + segment: JsonPathSegment, +) -> tuple[_JsonNode, ...]: + if isinstance(segment, JsonKeySegment): + if node.kind is not JsonNodeKind.OBJECT: + return () + return tuple( + value for key, value in node.object_members if key == segment.value + ) + if isinstance(segment, JsonIndexSegment): + if node.kind is not JsonNodeKind.ARRAY or segment.value >= len( + node.array_items + ): + return () + return (node.array_items[segment.value],) + if node.kind is JsonNodeKind.ARRAY: + return node.array_items + if node.kind is JsonNodeKind.OBJECT: + return tuple(value for _, value in node.object_members) + return () + + +_TIMEOUT_CHECK_INTERVAL = 256 + + +__all__ = [ + "JsonDocument", + "JsonEachSegment", + "JsonIndexSegment", + "JsonKeySegment", + "JsonNode", + "JsonNodeKind", + "JsonPathSegment", + "JsonSelector", + "JsonTextNode", +] diff --git a/projects/egress-gate/src/egress_gate/request_content/json_values.py b/projects/egress-gate/src/egress_gate/request_content/json_values.py new file mode 100644 index 00000000..a6c84fca --- /dev/null +++ b/projects/egress-gate/src/egress_gate/request_content/json_values.py @@ -0,0 +1,37 @@ +"""Shared JSON value kinds used by public documents and the private parser.""" + +from enum import StrEnum + +from pydantic import Field + +from egress_gate.base import StrictDomainModel + + +class JsonNodeKind(StrEnum): + """The JSON value kind at one immutable document node.""" + + OBJECT = "object" + ARRAY = "array" + STRING = "string" + NUMBER = "number" + BOOLEAN = "boolean" + NULL = "null" + + +class JsonNode(StrictDomainModel): + """Opaque immutable reference to one node in a ``JsonDocument``.""" + + id: str + path: tuple[str | int, ...] = Field(repr=False) + kind: JsonNodeKind + + +class JsonTextNode(StrictDomainModel): + """One selected JSON string with a stable document-local identity.""" + + id: str + path: tuple[str | int, ...] = Field(repr=False) + text: str = Field(repr=False) + + +__all__ = ["JsonNode", "JsonNodeKind", "JsonTextNode"] diff --git a/projects/egress-gate/src/egress_gate/request_content/messages.py b/projects/egress-gate/src/egress_gate/request_content/messages.py new file mode 100644 index 00000000..30373276 --- /dev/null +++ b/projects/egress-gate/src/egress_gate/request_content/messages.py @@ -0,0 +1,219 @@ +"""Normalized agent message blocks derived from structured JSON content.""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Literal, Protocol, Self + +from pydantic import Field, field_validator, model_validator + +from egress_gate.base import StrictDomainModel +from egress_gate.constants import MAX_JSON_SELECTORS, MAX_MESSAGE_BLOCKS +from egress_gate.errors import GateLimitExceededError +from egress_gate.request_content.json import ( + JsonDocument, + JsonNode, + JsonNodeKind, + JsonSelector, +) +from egress_gate.string_validators import ScalarString +from egress_gate.timeout import Timeout + + +class MessageRole(StrEnum): + """A normalized role for one model-visible message block.""" + + SYSTEM = "system" + DEVELOPER = "developer" + USER = "user" + ASSISTANT = "assistant" + TOOL = "tool" + UNKNOWN = "unknown" + + +class MessageBlockKind(StrEnum): + """The normalized purpose of one text-bearing message block.""" + + TEXT = "text" + TOOL_INPUT = "tool_input" + TOOL_OUTPUT = "tool_output" + + +class MessageBlock(StrictDomainModel): + """One normalized text block backed by a JSON string node.""" + + id: str + node_id: str + message_index: int | None = Field(ge=0) + role: MessageRole + kind: MessageBlockKind + text: str = Field(repr=False) + + +class MessageDocument(StrictDomainModel): + """The bounded normalized message view of one request body.""" + + blocks: tuple[MessageBlock, ...] = Field( + max_length=MAX_MESSAGE_BLOCKS, + repr=False, + ) + + +class JsonMessageMapConfig(StrictDomainModel): + """Map conventional JSON message objects to normalized text blocks.""" + + kind: Literal["json-message-map"] + messages: JsonSelector + role_key: ScalarString = "role" + text_selectors: tuple[JsonSelector, ...] = Field( + default=(), + max_length=MAX_JSON_SELECTORS, + ) + tool_input_selectors: tuple[JsonSelector, ...] = Field( + default=(), + max_length=MAX_JSON_SELECTORS, + ) + tool_output_selectors: tuple[JsonSelector, ...] = Field( + default=(), + max_length=MAX_JSON_SELECTORS, + ) + + @field_validator( + "text_selectors", + "tool_input_selectors", + "tool_output_selectors", + mode="before", + ) + @classmethod + def _text_selectors_are_a_tuple(cls, value: object) -> object: + if isinstance(value, list | tuple): + return tuple(value) + return value + + @model_validator(mode="after") + def _selectors_are_non_empty_and_bounded(self) -> Self: + selector_count = sum( + len(selectors) + for selectors in ( + self.text_selectors, + self.tool_input_selectors, + self.tool_output_selectors, + ) + ) + if selector_count == 0: + raise ValueError("message mapping requires at least one text selector") + if selector_count > MAX_JSON_SELECTORS: + raise ValueError("message mapping selector count exceeds the limit") + return self + + +class MessageBlockExtractor(Protocol): + """Extract normalized message blocks from one parsed JSON document.""" + + def extract( + self, + document: JsonDocument, + *, + timeout: Timeout, + ) -> MessageDocument: + """Return normalized blocks backed by nodes in ``document``.""" + ... + + +class JsonMessageBlockExtractor: + """Extract normalized blocks with one validated JSON message mapping.""" + + def __init__(self, config: JsonMessageMapConfig) -> None: + self._config = config + + def extract( + self, + document: JsonDocument, + *, + timeout: Timeout, + ) -> MessageDocument: + """Return normalized blocks in message and configured selector order.""" + containers = document.select_nodes( + (self._config.messages,), + timeout=timeout, + ) + blocks: list[MessageBlock] = [] + seen_classifications: set[tuple[str, MessageBlockKind]] = set() + message_index = 0 + for container in containers: + if container.kind is not JsonNodeKind.ARRAY: + continue + for message in document.array_items(container, timeout=timeout): + timeout.raise_if_expired() + role = _message_role(document, message, self._config.role_key) + default_kind = ( + MessageBlockKind.TOOL_OUTPUT + if role is MessageRole.TOOL + else MessageBlockKind.TEXT + ) + selector_groups = ( + (self._config.text_selectors, default_kind), + ( + self._config.tool_input_selectors, + MessageBlockKind.TOOL_INPUT, + ), + ( + self._config.tool_output_selectors, + MessageBlockKind.TOOL_OUTPUT, + ), + ) + for selectors, kind in selector_groups: + for text_node in document.select_text_from( + message, + selectors, + timeout=timeout, + ): + classification = (text_node.id, kind) + if classification in seen_classifications: + continue + seen_classifications.add(classification) + blocks.append( + MessageBlock( + id=f"message-block-{len(blocks) + 1}", + node_id=text_node.id, + message_index=message_index, + role=role, + kind=kind, + text=text_node.text, + ) + ) + if len(blocks) > MAX_MESSAGE_BLOCKS: + raise GateLimitExceededError( + "message block count exceeds the limit" + ) + message_index += 1 + timeout.raise_if_expired() + return MessageDocument(blocks=tuple(blocks)) + + +def _message_role( + document: JsonDocument, + message: JsonNode, + role_key: str, +) -> MessageRole: + role_node = document.object_member(message, role_key) + if role_node is None: + return MessageRole.UNKNOWN + value = document.text_value(role_node) + if value is None: + return MessageRole.UNKNOWN + try: + return MessageRole(value) + except ValueError: + return MessageRole.UNKNOWN + + +__all__ = [ + "JsonMessageBlockExtractor", + "JsonMessageMapConfig", + "MessageBlock", + "MessageBlockExtractor", + "MessageBlockKind", + "MessageDocument", + "MessageRole", +] diff --git a/projects/egress-gate/src/egress_gate/request_content/parsers.py b/projects/egress-gate/src/egress_gate/request_content/parsers.py new file mode 100644 index 00000000..7b3677da --- /dev/null +++ b/projects/egress-gate/src/egress_gate/request_content/parsers.py @@ -0,0 +1,183 @@ +"""Gate-agnostic parsers for text-bearing request content.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +from egress_gate.constants import MAX_BODY_BYTES +from egress_gate.errors import GateInputError, GateLimitExceededError +from egress_gate.request_content.json import JsonDocument, JsonSelector +from egress_gate.request_content.messages import ( + MessageBlockExtractor, + MessageBlockKind, + MessageRole, +) +from egress_gate.request_content.text import TextReplacement, TextTarget +from egress_gate.timeout import Timeout + + +class ParsedRequestContent(Protocol): + """A parsed text view that can render complete replacement body bytes.""" + + @property + def targets(self) -> tuple[TextTarget, ...]: + """Return independently inspected text values in parser-defined order.""" + ... + + def replace_text( + self, + replacements: tuple[TextReplacement, ...], + *, + timeout: Timeout, + ) -> bytes: + """Render complete body bytes after replacing complete target strings.""" + ... + + +class RequestContentParser(Protocol): + """A prepared stateless parser for one text-bearing request body format.""" + + def parse( + self, + body: bytes, + *, + timeout: Timeout, + ) -> ParsedRequestContent: + """Parse current body bytes into independently inspectable text targets.""" + ... + + +class Utf8TextParser: + """Expose one complete strict UTF-8 body as a text target.""" + + def parse( + self, + body: bytes, + *, + timeout: Timeout, + ) -> ParsedRequestContent: + timeout.raise_if_expired() + if len(body) > MAX_BODY_BYTES: + raise GateLimitExceededError("UTF-8 request body exceeds the size limit") + try: + text = body.decode("utf-8", errors="strict") + except UnicodeDecodeError: + raise GateInputError("request body is not valid UTF-8") from None + return _Utf8ParsedRequestContent(targets=(TextTarget(id="body", text=text),)) + + +@dataclass(frozen=True) +class JsonFieldsParser: + """Expose selected JSON string fields as independent text targets.""" + + selectors: tuple[JsonSelector, ...] + + def parse( + self, + body: bytes, + *, + timeout: Timeout, + ) -> ParsedRequestContent: + document = JsonDocument.parse(body, timeout=timeout) + nodes = document.select_text(self.selectors, timeout=timeout) + return _JsonParsedRequestContent( + targets=tuple(TextTarget(id=node.id, text=node.text) for node in nodes), + document=document, + ) + + +class MessageBlocksParser: + """Parse a JSON body into filtered normalized message-block targets.""" + + __slots__ = ("_block_kinds", "_extractor", "_roles") + + def __init__( + self, + *, + extractor: MessageBlockExtractor, + roles: tuple[MessageRole, ...] | None = None, + block_kinds: tuple[MessageBlockKind, ...] | None = None, + ) -> None: + self._extractor = extractor + self._roles = roles + self._block_kinds = block_kinds + + def parse( + self, + body: bytes, + *, + timeout: Timeout, + ) -> ParsedRequestContent: + document = JsonDocument.parse(body, timeout=timeout) + message_document = self._extractor.extract(document, timeout=timeout) + targets: list[TextTarget] = [] + seen_nodes: set[str] = set() + for block in message_document.blocks: + if self._roles is not None and block.role not in self._roles: + continue + if self._block_kinds is not None and block.kind not in self._block_kinds: + continue + if block.node_id in seen_nodes: + continue + seen_nodes.add(block.node_id) + targets.append(TextTarget(id=block.node_id, text=block.text)) + return _JsonParsedRequestContent( + targets=tuple(targets), + document=document, + ) + + +@dataclass(frozen=True) +class _Utf8ParsedRequestContent: + targets: tuple[TextTarget, ...] + + def replace_text( + self, + replacements: tuple[TextReplacement, ...], + *, + timeout: Timeout, + ) -> bytes: + timeout.raise_if_expired() + if len(replacements) != 1 or replacements[0].target_id != self.targets[0].id: + raise ValueError("UTF-8 body replacement target is invalid") + replacement_text = replacements[0].text + if len(replacement_text) > MAX_BODY_BYTES: + raise GateLimitExceededError("UTF-8 replacement body exceeds the limit") + try: + rendered = replacement_text.encode("utf-8", errors="strict") + except UnicodeEncodeError: + raise ValueError("UTF-8 body replacement text is invalid") from None + if len(rendered) > MAX_BODY_BYTES: + raise GateLimitExceededError("UTF-8 replacement body exceeds the limit") + timeout.raise_if_expired() + return rendered + + +@dataclass(frozen=True) +class _JsonParsedRequestContent: + targets: tuple[TextTarget, ...] + document: JsonDocument + + def replace_text( + self, + replacements: tuple[TextReplacement, ...], + *, + timeout: Timeout, + ) -> bytes: + replacement_ids = tuple(item.target_id for item in replacements) + if len(replacement_ids) != len(set(replacement_ids)): + raise ValueError("request-content replacement target IDs must be unique") + selected_ids = frozenset(target.id for target in self.targets) + if any(node_id not in selected_ids for node_id in replacement_ids): + raise ValueError("request-content replacement target was not selected") + return self.document.replace_text(replacements, timeout=timeout) + + +__all__ = [ + "JsonFieldsParser", + "MessageBlocksParser", + "ParsedRequestContent", + "RequestContentParser", + "Utf8TextParser", +] diff --git a/projects/egress-gate/src/egress_gate/request_content/text.py b/projects/egress-gate/src/egress_gate/request_content/text.py new file mode 100644 index 00000000..979240d3 --- /dev/null +++ b/projects/egress-gate/src/egress_gate/request_content/text.py @@ -0,0 +1,22 @@ +"""Shared immutable values for parsed request text and its replacements.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class TextTarget: + """One independently inspected text value and its content-local identity.""" + + id: str + text: str + + +@dataclass(frozen=True, slots=True) +class TextReplacement: + """The complete replacement text for one parsed target.""" + + target_id: str + text: str + + +__all__ = ["TextReplacement", "TextTarget"] diff --git a/projects/egress-gate/src/egress_gate/request_processor.py b/projects/egress-gate/src/egress_gate/request_processor.py index 06e39618..c900fa95 100644 --- a/projects/egress-gate/src/egress_gate/request_processor.py +++ b/projects/egress-gate/src/egress_gate/request_processor.py @@ -16,6 +16,7 @@ MAX_PROTO_FINDING_GROUPS, ) from egress_gate.errors import ( + BodyFormatError, EgressGateError, ErrorCode, GateConfigurationError, @@ -174,6 +175,8 @@ def process(self, request: HttpRequest, *, timeout: Timeout) -> EgressResult: except GateLimitExceededError: _LOGGER.info("egress_gate_processing_limit kind=resource") return _runtime_limit_result(self._policy_fingerprint) + except BodyFormatError: + raise EgressGateError(ErrorCode.BODY_FORMAT_INVALID) from None except GateInputError: raise EgressGateError(ErrorCode.BODY_ENCODING_INVALID) from None except GateConfigurationError: diff --git a/projects/egress-gate/tests/gates/test_regex.py b/projects/egress-gate/tests/gates/test_regex.py index d9a19bab..4bbfbbce 100644 --- a/projects/egress-gate/tests/gates/test_regex.py +++ b/projects/egress-gate/tests/gates/test_regex.py @@ -222,6 +222,195 @@ def test_non_body_scan_can_make_a_terminal_deny_decision() -> None: assert evaluation.reason_code == "egress_gate_regex_denied" +def test_json_fields_scan_only_matches_selected_string_nodes() -> None: + config = _config( + [{"pattern": "secret", "confidence": "high"}], + scan={ + "kind": "json-fields", + "selectors": [ + { + "segments": [ + {"kind": "key", "value": "messages"}, + {"kind": "each"}, + {"kind": "key", "value": "content"}, + ] + } + ], + }, + ) + + evaluation = RegexGate(config, None).evaluate( + _request(b'{"metadata":"secret","messages":[{"content":"secret"}]}'), + timeout=Timeout.from_seconds(1), + ) + + assert evaluation.findings[0].count == 1 + assert evaluation.request_mutations.is_empty + + +def test_json_fields_replace_preserves_unselected_json_source() -> None: + config = _config( + [{"pattern": "secret", "confidence": "high"}], + scan={ + "kind": "json-fields", + "selectors": [ + { + "segments": [ + {"kind": "key", "value": "messages"}, + {"kind": "each"}, + {"kind": "key", "value": "content"}, + ] + } + ], + }, + action_kind="replace", + template="[{entity}]", + ) + + evaluation = RegexGate(config, None).evaluate( + _request( + b'{ "messages" : [{"content":"a secret"}], "number":1.00, ' + b'"escaped":"\\u0078" }' + ), + timeout=Timeout.from_seconds(1), + ) + + assert evaluation.request_mutations.replacement_body == ( + b'{ "messages" : [{"content":"a [token]"}], "number":1.00, ' + b'"escaped":"\\u0078" }' + ) + + +def test_json_fields_replace_preserves_explicit_intent_without_matches() -> None: + body = b'{ "messages": [{"content":"safe"}], "number":1.00 }' + config = _config( + [{"pattern": "secret", "confidence": "high"}], + scan={ + "kind": "json-fields", + "selectors": [ + { + "segments": [ + {"kind": "key", "value": "messages"}, + {"kind": "each"}, + {"kind": "key", "value": "content"}, + ] + } + ], + }, + action_kind="replace", + ) + + evaluation = RegexGate(config, None).evaluate( + _request(body), + timeout=Timeout.from_seconds(1), + ) + + assert evaluation.request_mutations.replacement_body == body + + +def test_json_fields_matches_do_not_span_distinct_selected_nodes() -> None: + config = _config( + [{"pattern": "secretsecret", "confidence": "high"}], + scan={ + "kind": "json-fields", + "selectors": [ + { + "segments": [ + {"kind": "key", "value": "messages"}, + {"kind": "each"}, + ] + } + ], + }, + ) + + evaluation = RegexGate(config, None).evaluate( + _request(b'{"messages":["secret","secret"]}'), + timeout=Timeout.from_seconds(1), + ) + + assert evaluation.findings == () + + +def test_message_blocks_scan_filters_normalized_roles() -> None: + config = _config( + [{"pattern": "secret", "confidence": "high"}], + scan={ + "kind": "message-blocks", + "message_mapping": { + "kind": "json-message-map", + "messages": {"segments": [{"kind": "key", "value": "messages"}]}, + "role_key": "role", + "text_selectors": [{"segments": [{"kind": "key", "value": "content"}]}], + }, + "roles": ["user"], + }, + action_kind="replace", + ) + + evaluation = RegexGate(config, None).evaluate( + _request( + b'{"messages":[' + b'{"role":"assistant","content":"secret"},' + b'{"role":"user","content":"secret"}' + b"]}" + ), + timeout=Timeout.from_seconds(1), + ) + + assert evaluation.findings[0].count == 1 + assert evaluation.request_mutations.replacement_body == ( + b'{"messages":[' + b'{"role":"assistant","content":"secret"},' + b'{"role":"user","content":"[token]"}' + b"]}" + ) + + +def test_message_blocks_scan_filters_overlapping_explicit_tool_classification() -> None: + selector = {"segments": [{"kind": "key", "value": "content"}]} + config = _config( + [{"pattern": "secret", "confidence": "high"}], + scan={ + "kind": "message-blocks", + "message_mapping": { + "kind": "json-message-map", + "messages": {"segments": [{"kind": "key", "value": "messages"}]}, + "text_selectors": [selector], + "tool_input_selectors": [selector], + }, + "block_kinds": ["tool_input"], + }, + action_kind="deny", + ) + + evaluation = RegexGate(config, None).evaluate( + _request(b'{"messages":[{"role":"assistant","content":"secret"}]}'), + timeout=Timeout.from_seconds(1), + ) + + assert evaluation.control is GateControl.DENY + assert evaluation.findings[0].count == 1 + + +def test_message_block_filters_must_be_unique() -> None: + with pytest.raises(ValidationError, match="unique"): + _config( + [{"pattern": "secret", "confidence": "high"}], + scan={ + "kind": "message-blocks", + "message_mapping": { + "kind": "json-message-map", + "messages": {"segments": [{"kind": "key", "value": "messages"}]}, + "text_selectors": [ + {"segments": [{"kind": "key", "value": "content"}]} + ], + }, + "roles": ["user", "user"], + }, + ) + + @pytest.mark.parametrize("kind", ["path", "query", "header"]) def test_replace_action_is_structurally_unavailable_for_non_body_scans( kind: str, @@ -545,6 +734,24 @@ def test_patterns_compile_during_preparation_not_validation_or_each_run( assert recording_compile.call_count == prepared_count +def test_request_content_parser_is_prepared_once_and_reused( + monkeypatch: pytest.MonkeyPatch, +) -> None: + original_prepare = regex_module._prepare_content_parser + recording_prepare = Mock(wraps=original_prepare) + monkeypatch.setattr(regex_module, "_prepare_content_parser", recording_prepare) + + gate = RegexGate( + _config([{"pattern": "x", "confidence": "high"}]), + None, + timeout=Timeout.from_seconds(1), + ) + gate.evaluate(_request(b"x"), timeout=Timeout.from_seconds(1)) + gate.evaluate(_request(b"x"), timeout=Timeout.from_seconds(1)) + + recording_prepare.assert_called_once() + + def test_gate_preparation_honors_an_expired_timeout() -> None: config = _config([{"pattern": "x", "confidence": "high"}]) diff --git a/projects/egress-gate/tests/gates/test_registry.py b/projects/egress-gate/tests/gates/test_registry.py index 736219c4..6f55cf72 100644 --- a/projects/egress-gate/tests/gates/test_registry.py +++ b/projects/egress-gate/tests/gates/test_registry.py @@ -120,6 +120,16 @@ def test_builtin_registry_seals_on_first_use_and_contains_only_regex() -> None: ) assert "RegexReplaceAction" in str(body_scan_schema) assert "RegexReplaceAction" not in str(header_scan_schema) + json_scan_schema = next( + value for key, value in definitions.items() if key == "RegexJsonFieldsScan" + ) + message_scan_schema = next( + value for key, value in definitions.items() if key == "RegexMessageBlocksScan" + ) + assert "RegexReplaceAction" in str(json_scan_schema) + assert "RegexReplaceAction" in str(message_scan_schema) + assert "JsonSelector" in str(json_scan_schema) + assert "JsonMessageMapConfig" in str(message_scan_schema) with pytest.raises(EgressGateError): registry.validate_config( diff --git a/projects/egress-gate/tests/request_content/__init__.py b/projects/egress-gate/tests/request_content/__init__.py new file mode 100644 index 00000000..3370066b --- /dev/null +++ b/projects/egress-gate/tests/request_content/__init__.py @@ -0,0 +1 @@ +"""Structured request-content tests.""" diff --git a/projects/egress-gate/tests/request_content/test_json.py b/projects/egress-gate/tests/request_content/test_json.py new file mode 100644 index 00000000..2d017e22 --- /dev/null +++ b/projects/egress-gate/tests/request_content/test_json.py @@ -0,0 +1,356 @@ +"""Contracts for strict JSON selection and source-preserving replacement.""" + +from __future__ import annotations + +import pytest +from pydantic import ValidationError + +import egress_gate.request_content._json_parser as json_parser_module +import egress_gate.request_content.json as json_module +from egress_gate.errors import ( + BodyFormatError, + GateInputError, + GateLimitExceededError, + TimeoutExpiredError, +) +from egress_gate.request_content import ( + JsonDocument, + JsonEachSegment, + JsonIndexSegment, + JsonKeySegment, + JsonPathSegment, + JsonSelector, + TextReplacement, +) +from egress_gate.timeout import Timeout + + +def _selector(*segments: JsonPathSegment) -> JsonSelector: + return JsonSelector(segments=segments) + + +def _parse(body: bytes) -> JsonDocument: + return JsonDocument.parse(body, timeout=Timeout.from_seconds(1)) + + +def test_select_text_uses_typed_paths_and_deduplicates_overlapping_selectors() -> None: + document = _parse( + b'{"messages":[{"content":"first"},{"content":"second"}],"ignored":"x"}' + ) + all_content = _selector( + JsonKeySegment(kind="key", value="messages"), + JsonEachSegment(kind="each"), + JsonKeySegment(kind="key", value="content"), + ) + second_content = _selector( + JsonKeySegment(kind="key", value="messages"), + JsonIndexSegment(kind="index", value=1), + JsonKeySegment(kind="key", value="content"), + ) + + nodes = document.select_text( + (all_content, second_content), + timeout=Timeout.from_seconds(1), + ) + + assert tuple(node.text for node in nodes) == ("first", "second") + assert tuple(node.path for node in nodes) == ( + ("messages", 0, "content"), + ("messages", 1, "content"), + ) + assert len({node.id for node in nodes}) == 2 + + +def test_missing_paths_and_non_string_terminals_produce_no_text_nodes() -> None: + document = _parse(b'{"message":{"count":2}}') + + nodes = document.select_text( + ( + _selector(JsonKeySegment(kind="key", value="missing")), + _selector( + JsonKeySegment(kind="key", value="message"), + JsonKeySegment(kind="key", value="count"), + ), + ), + timeout=Timeout.from_seconds(1), + ) + + assert nodes == () + + +def test_replace_text_preserves_every_unselected_source_byte() -> None: + body = ( + b'{ "messages" : [ {"content":"secret\\nvalue", "count":1.00} ], ' + b'"unchanged":"\\u0078" }' + ) + document = _parse(body) + selector = _selector( + JsonKeySegment(kind="key", value="messages"), + JsonEachSegment(kind="each"), + JsonKeySegment(kind="key", value="content"), + ) + node = document.select_text((selector,), timeout=Timeout.from_seconds(1))[0] + + replacement = document.replace_text( + (TextReplacement(target_id=node.id, text='safe\n"value"'),), + timeout=Timeout.from_seconds(1), + ) + + assert replacement == ( + b'{ "messages" : [ {"content":"safe\\n\\"value\\"", "count":1.00} ], ' + b'"unchanged":"\\u0078" }' + ) + + +def test_replace_text_rejects_duplicate_or_unknown_node_ids() -> None: + document = _parse(b'{"value":"one"}') + node = document.select_text( + (_selector(JsonKeySegment(kind="key", value="value")),), + timeout=Timeout.from_seconds(1), + )[0] + + with pytest.raises(ValueError, match="unique"): + document.replace_text( + ( + TextReplacement(target_id=node.id, text="two"), + TextReplacement(target_id=node.id, text="three"), + ), + timeout=Timeout.from_seconds(1), + ) + with pytest.raises(ValueError, match="unknown"): + document.replace_text( + (TextReplacement(target_id="unknown", text="two"),), + timeout=Timeout.from_seconds(1), + ) + + +def test_replace_text_handles_the_maximum_selected_nodes_in_linear_time() -> None: + from egress_gate.constants import MAX_JSON_SELECTED_NODES + + values = ",".join( + f'"{index}":"{"x" * 900}"' for index in range(MAX_JSON_SELECTED_NODES) + ) + document = _parse(f"{{{values}}}".encode()) + nodes = document.select_text( + (_selector(JsonEachSegment(kind="each")),), + timeout=Timeout.from_seconds(1), + ) + + replaced = document.replace_text( + tuple(TextReplacement(target_id=node.id, text="safe") for node in nodes), + timeout=Timeout.from_seconds(1), + ) + + assert replaced.count(b'"safe"') == MAX_JSON_SELECTED_NODES + + +def test_node_references_are_bound_to_the_document_that_created_them() -> None: + first = _parse(b'{"value":"one"}') + second = _parse(b'{"value":"one"}') + selector = _selector(JsonKeySegment(kind="key", value="value")) + first_node = first.select_nodes( + (selector,), + timeout=Timeout.from_seconds(1), + )[0] + + with pytest.raises(ValueError, match="does not belong"): + second.select_from( + first_node, + (selector,), + timeout=Timeout.from_seconds(1), + ) + + +def test_long_strings_and_whitespace_check_the_shared_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + checks = 0 + + def record_check(_timeout: Timeout) -> None: + nonlocal checks + checks += 1 + + monkeypatch.setattr(Timeout, "raise_if_expired", record_check) + + _parse((" " * 10_000 + '"' + "x" * 10_000 + '"').encode()) + + assert checks >= 7 + + +def test_array_item_materialization_checks_the_shared_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + document = _parse( + ('{"items":[' + ",".join("0" for _ in range(1000)) + "]}").encode() + ) + root = document.select_nodes( + (_selector(JsonKeySegment(kind="key", value="items")),), + timeout=Timeout.from_seconds(1), + )[0] + checks = 0 + + def record_check(_timeout: Timeout) -> None: + nonlocal checks + checks += 1 + + monkeypatch.setattr(Timeout, "raise_if_expired", record_check) + + document.array_items(root, timeout=Timeout.from_seconds(1)) + + assert checks >= 4 + + +def test_text_node_projection_stops_when_the_shared_timeout_expires( + monkeypatch: pytest.MonkeyPatch, +) -> None: + document = _parse( + ('{"values":[' + ",".join('"value"' for _ in range(1000)) + "]}").encode() + ) + selector = _selector( + JsonKeySegment(kind="key", value="values"), + JsonEachSegment(kind="each"), + ) + checks = 0 + + def expire_during_projection(_timeout: Timeout) -> None: + nonlocal checks + checks += 1 + if checks == 5: + raise TimeoutExpiredError + + monkeypatch.setattr(Timeout, "raise_if_expired", expire_during_projection) + + with pytest.raises(TimeoutExpiredError): + document.select_text((selector,), timeout=Timeout.from_seconds(1)) + + assert checks == 5 + + +def test_public_node_projection_stops_when_the_shared_timeout_expires( + monkeypatch: pytest.MonkeyPatch, +) -> None: + document = _parse( + ('{"values":[' + ",".join("0" for _ in range(1000)) + "]}").encode() + ) + selector = _selector( + JsonKeySegment(kind="key", value="values"), + JsonEachSegment(kind="each"), + ) + checks = 0 + + def expire_during_projection(_timeout: Timeout) -> None: + nonlocal checks + checks += 1 + if checks == 5: + raise TimeoutExpiredError + + monkeypatch.setattr(Timeout, "raise_if_expired", expire_during_projection) + + with pytest.raises(TimeoutExpiredError): + document.select_nodes((selector,), timeout=Timeout.from_seconds(1)) + + assert checks == 5 + + +@pytest.mark.parametrize( + "body", + [ + b'{"duplicate":1,"duplicate":2}', + b'{"constant":NaN}', + b'{"trailing":true,}', + b'"\\ud800"', + ], +) +def test_parse_rejects_non_strict_json(body: bytes) -> None: + with pytest.raises(BodyFormatError): + _parse(body) + + +def test_parse_distinguishes_invalid_utf8_from_invalid_json() -> None: + with pytest.raises(GateInputError): + _parse(b'"\xff"') + + +def test_parse_enforces_the_request_body_size_limit( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(json_module, "MAX_BODY_BYTES", 4) + + _parse(b"null") + with pytest.raises(GateLimitExceededError, match="size limit"): + _parse(b" null") + + +def test_selector_models_are_strict_bounded_and_discriminated() -> None: + with pytest.raises(ValidationError): + JsonSelector.model_validate({"segments": []}) + with pytest.raises(ValidationError): + JsonSelector.model_validate({"segments": [{"kind": "unknown"}]}) + with pytest.raises(ValidationError): + JsonIndexSegment.model_validate({"kind": "index", "value": -1}) + with pytest.raises(ValidationError): + JsonKeySegment.model_validate({"kind": "key", "value": 1}) + + +def test_document_depth_is_bounded() -> None: + from egress_gate.constants import MAX_JSON_DEPTH + + accepted = ("[" * MAX_JSON_DEPTH + "0" + "]" * MAX_JSON_DEPTH).encode() + body = ("[" * (MAX_JSON_DEPTH + 1) + "0" + "]" * (MAX_JSON_DEPTH + 1)).encode() + + _parse(accepted) + with pytest.raises(GateLimitExceededError): + _parse(body) + + +def test_document_node_count_accepts_the_limit_and_rejects_the_next_node( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(json_parser_module, "MAX_JSON_NODES", 3) + + _parse(b'["one","two"]') + with pytest.raises(GateLimitExceededError, match="node count"): + _parse(b'["one","two","three"]') + + +def test_selector_segment_count_accepts_the_limit_and_rejects_the_next() -> None: + from egress_gate.constants import MAX_JSON_SELECTOR_SEGMENTS + + segment = JsonEachSegment(kind="each") + + JsonSelector(segments=(segment,) * MAX_JSON_SELECTOR_SEGMENTS) + with pytest.raises(ValidationError): + JsonSelector(segments=(segment,) * (MAX_JSON_SELECTOR_SEGMENTS + 1)) + + +def test_selector_count_accepts_the_limit_and_rejects_the_next() -> None: + from egress_gate.constants import MAX_JSON_SELECTORS + + document = _parse(b'{"value":"one"}') + selector = _selector(JsonKeySegment(kind="key", value="value")) + + document.select_nodes( + (selector,) * MAX_JSON_SELECTORS, + timeout=Timeout.from_seconds(1), + ) + with pytest.raises(GateLimitExceededError, match="selector count"): + document.select_nodes( + (selector,) * (MAX_JSON_SELECTORS + 1), + timeout=Timeout.from_seconds(1), + ) + + +def test_selected_node_count_accepts_the_limit_and_rejects_the_next( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(json_module, "MAX_JSON_SELECTED_NODES", 2) + selector = _selector( + JsonKeySegment(kind="key", value="values"), + JsonEachSegment(kind="each"), + ) + + accepted = _parse(b'{"values":["one","two"]}') + assert len(accepted.select_nodes((selector,), timeout=Timeout.from_seconds(1))) == 2 + rejected = _parse(b'{"values":["one","two","three"]}') + with pytest.raises(GateLimitExceededError, match="selected node count"): + rejected.select_nodes((selector,), timeout=Timeout.from_seconds(1)) diff --git a/projects/egress-gate/tests/request_content/test_messages.py b/projects/egress-gate/tests/request_content/test_messages.py new file mode 100644 index 00000000..2d83e75f --- /dev/null +++ b/projects/egress-gate/tests/request_content/test_messages.py @@ -0,0 +1,164 @@ +"""Contracts for normalized message blocks backed by JSON text nodes.""" + +from __future__ import annotations + +import pytest + +import egress_gate.request_content.messages as messages_module +from egress_gate.errors import GateLimitExceededError +from egress_gate.request_content import ( + JsonDocument, + JsonEachSegment, + JsonKeySegment, + JsonMessageBlockExtractor, + JsonMessageMapConfig, + JsonPathSegment, + JsonSelector, + MessageBlockKind, + MessageRole, +) +from egress_gate.timeout import Timeout + + +def _selector(*segments: JsonPathSegment) -> JsonSelector: + return JsonSelector(segments=segments) + + +def test_json_message_map_normalizes_roles_and_content_nodes() -> None: + document = JsonDocument.parse( + ( + b'{"request":{"messages":[' + b'{"role":"user","content":"hello"},' + b'{"role":"tool","content":[{"text":"result"}]},' + b'{"role":"future","content":"unknown"}' + b"]}}" + ), + timeout=Timeout.from_seconds(1), + ) + config = JsonMessageMapConfig( + kind="json-message-map", + messages=_selector( + JsonKeySegment(kind="key", value="request"), + JsonKeySegment(kind="key", value="messages"), + ), + role_key="role", + text_selectors=( + _selector(JsonKeySegment(kind="key", value="content")), + _selector( + JsonKeySegment(kind="key", value="content"), + JsonEachSegment(kind="each"), + JsonKeySegment(kind="key", value="text"), + ), + ), + ) + + messages = JsonMessageBlockExtractor(config).extract( + document, + timeout=Timeout.from_seconds(1), + ) + + assert tuple(block.text for block in messages.blocks) == ( + "hello", + "result", + "unknown", + ) + assert tuple(block.role for block in messages.blocks) == ( + MessageRole.USER, + MessageRole.TOOL, + MessageRole.UNKNOWN, + ) + assert tuple(block.kind for block in messages.blocks) == ( + MessageBlockKind.TEXT, + MessageBlockKind.TOOL_OUTPUT, + MessageBlockKind.TEXT, + ) + assert len({block.node_id for block in messages.blocks}) == 3 + + +def test_json_message_map_can_classify_explicit_tool_input_and_output_nodes() -> None: + document = JsonDocument.parse( + ( + b'{"messages":[' + b'{"role":"assistant","tool_call":{"arguments":"input"}},' + b'{"role":"tool","content":"output"}' + b"]}" + ), + timeout=Timeout.from_seconds(1), + ) + config = JsonMessageMapConfig( + kind="json-message-map", + messages=_selector(JsonKeySegment(kind="key", value="messages")), + tool_input_selectors=( + _selector( + JsonKeySegment(kind="key", value="tool_call"), + JsonKeySegment(kind="key", value="arguments"), + ), + ), + tool_output_selectors=(_selector(JsonKeySegment(kind="key", value="content")),), + ) + + messages = JsonMessageBlockExtractor(config).extract( + document, + timeout=Timeout.from_seconds(1), + ) + + assert tuple((block.kind, block.text) for block in messages.blocks) == ( + (MessageBlockKind.TOOL_INPUT, "input"), + (MessageBlockKind.TOOL_OUTPUT, "output"), + ) + + +def test_message_block_count_accepts_the_limit_and_rejects_the_next( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(messages_module, "MAX_MESSAGE_BLOCKS", 2) + extractor = JsonMessageBlockExtractor( + JsonMessageMapConfig( + kind="json-message-map", + messages=_selector(JsonKeySegment(kind="key", value="messages")), + text_selectors=(_selector(JsonKeySegment(kind="key", value="content")),), + ) + ) + accepted = JsonDocument.parse( + b'{"messages":[{"content":"one"},{"content":"two"}]}', + timeout=Timeout.from_seconds(1), + ) + rejected = JsonDocument.parse( + b'{"messages":[{"content":"one"},{"content":"two"},{"content":"three"}]}', + timeout=Timeout.from_seconds(1), + ) + + assert len(extractor.extract(accepted, timeout=Timeout.from_seconds(1)).blocks) == 2 + with pytest.raises(GateLimitExceededError, match="message block count"): + extractor.extract(rejected, timeout=Timeout.from_seconds(1)) + + +def test_message_block_limit_counts_unique_node_kind_classifications( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(messages_module, "MAX_MESSAGE_BLOCKS", 2) + selector = _selector(JsonKeySegment(kind="key", value="content")) + extractor = JsonMessageBlockExtractor( + JsonMessageMapConfig( + kind="json-message-map", + messages=_selector(JsonKeySegment(kind="key", value="messages")), + text_selectors=(selector,), + tool_output_selectors=(selector,), + ) + ) + document = JsonDocument.parse( + ( + b'{"messages":[' + b'{"role":"tool","content":"one"},' + b'{"role":"tool","content":"two"}' + b"]}" + ), + timeout=Timeout.from_seconds(1), + ) + + blocks = extractor.extract(document, timeout=Timeout.from_seconds(1)).blocks + + assert tuple((block.kind, block.text) for block in blocks) == ( + (MessageBlockKind.TOOL_OUTPUT, "one"), + (MessageBlockKind.TOOL_OUTPUT, "two"), + ) diff --git a/projects/egress-gate/tests/request_content/test_parsers.py b/projects/egress-gate/tests/request_content/test_parsers.py new file mode 100644 index 00000000..9c8e56d1 --- /dev/null +++ b/projects/egress-gate/tests/request_content/test_parsers.py @@ -0,0 +1,206 @@ +"""Contracts for reusable request-content text parsers.""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import pytest + +import egress_gate.request_content.parsers as parsers_module +from egress_gate.errors import GateInputError, GateLimitExceededError +from egress_gate.request_content import ( + JsonEachSegment, + JsonFieldsParser, + JsonKeySegment, + JsonMessageBlockExtractor, + JsonMessageMapConfig, + JsonSelector, + MessageBlockKind, + MessageBlocksParser, + MessageRole, + TextReplacement, + Utf8TextParser, +) +from egress_gate.timeout import Timeout + + +def _content_selector() -> JsonSelector: + return JsonSelector( + segments=( + JsonKeySegment(kind="key", value="messages"), + JsonEachSegment(kind="each"), + JsonKeySegment(kind="key", value="content"), + ) + ) + + +def test_text_replacement_is_named_and_immutable() -> None: + replacement = TextReplacement(target_id="target", text="safe") + + assert replacement.target_id == "target" + assert replacement.text == "safe" + with pytest.raises(FrozenInstanceError): + setattr(replacement, "text", "changed") + + +def test_utf8_parser_extracts_and_replaces_the_complete_body() -> None: + content = Utf8TextParser().parse( + b"secret", + timeout=Timeout.from_seconds(1), + ) + + assert tuple((target.id, target.text) for target in content.targets) == ( + ("body", "secret"), + ) + assert ( + content.replace_text( + (TextReplacement(target_id="body", text="safe"),), + timeout=Timeout.from_seconds(1), + ) + == b"safe" + ) + + +def test_utf8_parser_rejects_invalid_body_encoding() -> None: + with pytest.raises(GateInputError, match="not valid UTF-8"): + Utf8TextParser().parse( + b"\xff", + timeout=Timeout.from_seconds(1), + ) + + +def test_utf8_parser_enforces_input_and_replacement_body_limits( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(parsers_module, "MAX_BODY_BYTES", 4) + parser = Utf8TextParser() + + parser.parse(b"four", timeout=Timeout.from_seconds(1)) + with pytest.raises(GateLimitExceededError, match="size limit"): + parser.parse(b"three", timeout=Timeout.from_seconds(1)) + + content = parser.parse(b"x", timeout=Timeout.from_seconds(1)) + assert ( + content.replace_text( + (TextReplacement(target_id="body", text="éé"),), + timeout=Timeout.from_seconds(1), + ) + == "éé".encode() + ) + with pytest.raises(GateLimitExceededError, match="exceeds the limit"): + content.replace_text( + (TextReplacement(target_id="body", text="ééx"),), + timeout=Timeout.from_seconds(1), + ) + + +def test_json_fields_parser_owns_source_preserving_replacement() -> None: + parser = JsonFieldsParser(selectors=(_content_selector(),)) + body = b'{ "messages":[{"content":"secret"}], "number":1.00 }' + + content = parser.parse(body, timeout=Timeout.from_seconds(1)) + + assert tuple(target.text for target in content.targets) == ("secret",) + assert ( + content.replace_text( + (TextReplacement(target_id=content.targets[0].id, text="safe"),), + timeout=Timeout.from_seconds(1), + ) + == b'{ "messages":[{"content":"safe"}], "number":1.00 }' + ) + + +def test_json_fields_parser_rejects_replacement_of_an_unselected_node() -> None: + parser = JsonFieldsParser(selectors=(_content_selector(),)) + body = b'{"messages":[{"content":"selected"}],"ignored":"secret"}' + content = parser.parse(body, timeout=Timeout.from_seconds(1)) + ignored_parser = JsonFieldsParser( + selectors=( + JsonSelector(segments=(JsonKeySegment(kind="key", value="ignored"),)), + ) + ) + ignored = ignored_parser.parse(body, timeout=Timeout.from_seconds(1)).targets[0] + + with pytest.raises(ValueError, match="was not selected"): + content.replace_text( + (TextReplacement(target_id=ignored.id, text="exposed"),), + timeout=Timeout.from_seconds(1), + ) + + +def test_message_blocks_parser_applies_mapping_and_filters() -> None: + parser = MessageBlocksParser( + extractor=JsonMessageBlockExtractor( + JsonMessageMapConfig( + kind="json-message-map", + messages=JsonSelector( + segments=(JsonKeySegment(kind="key", value="messages"),) + ), + text_selectors=( + JsonSelector( + segments=(JsonKeySegment(kind="key", value="content"),) + ), + ), + ) + ), + roles=(MessageRole.USER,), + ) + + content = parser.parse( + ( + b'{"messages":[' + b'{"role":"assistant","content":"ignored"},' + b'{"role":"user","content":"selected"}' + b"]}" + ), + timeout=Timeout.from_seconds(1), + ) + + assert tuple(target.text for target in content.targets) == ("selected",) + + +def test_message_blocks_parser_deduplicates_shared_text_nodes() -> None: + selector = JsonSelector(segments=(JsonKeySegment(kind="key", value="content"),)) + parser = MessageBlocksParser( + extractor=JsonMessageBlockExtractor( + JsonMessageMapConfig( + kind="json-message-map", + messages=JsonSelector( + segments=(JsonKeySegment(kind="key", value="messages"),) + ), + text_selectors=(selector,), + tool_output_selectors=(selector,), + ) + ) + ) + + content = parser.parse( + b'{"messages":[{"role":"tool","content":"once"}]}', + timeout=Timeout.from_seconds(1), + ) + + assert tuple(target.text for target in content.targets) == ("once",) + + +def test_message_blocks_parser_filters_before_deduplicating_classifications() -> None: + selector = JsonSelector(segments=(JsonKeySegment(kind="key", value="content"),)) + parser = MessageBlocksParser( + extractor=JsonMessageBlockExtractor( + JsonMessageMapConfig( + kind="json-message-map", + messages=JsonSelector( + segments=(JsonKeySegment(kind="key", value="messages"),) + ), + text_selectors=(selector,), + tool_input_selectors=(selector,), + ) + ), + block_kinds=(MessageBlockKind.TOOL_INPUT,), + ) + + content = parser.parse( + b'{"messages":[{"role":"assistant","content":"selected"}]}', + timeout=Timeout.from_seconds(1), + ) + + assert tuple(target.text for target in content.targets) == ("selected",) diff --git a/projects/egress-gate/tests/service/test_servicer.py b/projects/egress-gate/tests/service/test_servicer.py index 44392d76..528bc219 100644 --- a/projects/egress-gate/tests/service/test_servicer.py +++ b/projects/egress-gate/tests/service/test_servicer.py @@ -56,14 +56,19 @@ def _values( - *, action_kind: str = "detect", default_decision: str = "allow" + *, + action_kind: str = "detect", + default_decision: str = "allow", + scan: dict[str, object] | None = None, ) -> dict[str, object]: action: dict[str, object] = {"kind": action_kind} if action_kind == "replace": action["template"] = "[{entity}]" + scan_values = {"kind": "body"} if scan is None else dict(scan) + scan_values["action"] = action config: dict[str, object] = { "kind": "regex", - "scan": {"kind": "body", "action": action}, + "scan": scan_values, "pattern_catalog": { "entities": [ { @@ -618,6 +623,66 @@ def test_invalid_utf8_is_an_input_failure_before_wire_evaluation() -> None: assert error.value.code is ErrorCode.BODY_ENCODING_INVALID +def test_structured_replacement_is_serialized_as_a_complete_body_mutation() -> None: + middleware = EgressGateMiddleware(create_builtin_registry()) + request = _request(body=b'{ "messages":[{"content":"secret"}], "n":1.00 }') + request.config.CopyFrom( + _proto_config( + _values( + action_kind="replace", + scan={ + "kind": "json-fields", + "selectors": [ + { + "segments": [ + {"kind": "key", "value": "messages"}, + {"kind": "each"}, + {"kind": "key", "value": "content"}, + ] + } + ], + }, + ) + ) + ) + try: + response, _ = asyncio.run( + middleware._evaluate_http_request(request, Timeout.from_seconds(1)) + ) + finally: + asyncio.run(middleware.close()) + + assert response.has_body is True + assert response.body == b'{ "messages":[{"content":"[token]"}], "n":1.00 }' + + +def test_invalid_json_is_an_input_failure_for_a_structured_scan() -> None: + middleware = EgressGateMiddleware(create_builtin_registry()) + request = _request(body=b"{") + request.config.CopyFrom( + _proto_config( + _values( + scan={ + "kind": "json-fields", + "selectors": [{"segments": [{"kind": "key", "value": "messages"}]}], + } + ) + ) + ) + try: + with pytest.raises(EgressGateError) as error: + asyncio.run( + middleware._evaluate_http_request( + request, + Timeout.from_seconds(1), + ) + ) + finally: + asyncio.run(middleware.close()) + + assert error.value.code is ErrorCode.BODY_FORMAT_INVALID + + def test_service_request_body_limit_is_checked_before_worker_execution() -> None: request = _request(body=b"x" * (MAX_BODY_BYTES + 1)) middleware = EgressGateMiddleware(create_builtin_registry()) diff --git a/projects/egress-gate/tests/test_request_processor.py b/projects/egress-gate/tests/test_request_processor.py index fe068535..5ea74ef9 100644 --- a/projects/egress-gate/tests/test_request_processor.py +++ b/projects/egress-gate/tests/test_request_processor.py @@ -471,6 +471,45 @@ def test_three_regex_gates_progressively_redact_the_current_body() -> None: ] +def test_structured_regex_scan_parses_the_body_replaced_by_an_earlier_gate() -> None: + processor = _processor( + ( + ( + "wrap", + { + "kind": "test-control", + "replacement": '{"messages":[{"content":"secret"}]}', + }, + ), + ( + "redact", + _regex_config( + "replace", + scan={ + "kind": "json-fields", + "selectors": [ + { + "segments": [ + {"kind": "key", "value": "messages"}, + {"kind": "each"}, + {"kind": "key", "value": "content"}, + ] + } + ], + }, + ), + ), + ), + include_regex=True, + ) + + result = processor.process(_request(), timeout=Timeout.from_seconds(1)) + + assert result.request_mutations.replacement_body == ( + b'{"messages":[{"content":"[token]"}]}' + ) + + def test_later_regex_gates_use_the_body_after_overlapping_text_is_redacted() -> None: original = _request(body=b"credential: alice@example.com") processor = _processor( @@ -981,6 +1020,30 @@ def test_invalid_utf8_is_translated_to_the_stable_input_error() -> None: assert error.value.code is ErrorCode.BODY_ENCODING_INVALID +def test_invalid_json_is_translated_to_the_stable_body_format_error() -> None: + processor = _processor( + ( + ( + "regex", + _regex_config( + scan={ + "kind": "json-fields", + "selectors": [ + {"segments": [{"kind": "key", "value": "messages"}]} + ], + } + ), + ), + ), + include_regex=True, + ) + + with pytest.raises(EgressGateError) as error: + processor.process(_request(body=b"{"), timeout=Timeout.from_seconds(1)) + + assert error.value.code is ErrorCode.BODY_FORMAT_INVALID + + def test_prepared_gate_type_is_part_of_the_processor_contract() -> None: registry = GateRegistry() registry.register(_ControlGate) diff --git a/tests/test_page_navigation.py b/tests/test_page_navigation.py index 2649c026..15aaae02 100644 --- a/tests/test_page_navigation.py +++ b/tests/test_page_navigation.py @@ -15,6 +15,9 @@ CONFIGURATION_GUIDE = ( ROOT / "site" / "documentation" / "egress-gate" / "configuration" / "index.html" ) +REQUEST_CONTENT_GUIDE = ( + ROOT / "site" / "documentation" / "egress-gate" / "request-content" / "index.html" +) class PageNavigationTests(unittest.TestCase): @@ -49,6 +52,7 @@ def test_rendered_links_follow_the_reading_path(self) -> None: documentation = DOCUMENTATION_LANDING.read_text(encoding="utf-8") egress_gate = EGRESS_GATE_LANDING.read_text(encoding="utf-8") configuration = CONFIGURATION_GUIDE.read_text(encoding="utf-8") + request_content = REQUEST_CONTENT_GUIDE.read_text(encoding="utf-8") self.assertIn("Back to OpenShell Research", documentation) self.assertIn("Next: Egress Gate", documentation) @@ -56,7 +60,9 @@ def test_rendered_links_follow_the_reading_path(self) -> None: self.assertIn("Previous: Documentation", egress_gate) self.assertIn("Next: Configure policies", egress_gate) self.assertIn("Previous: Egress Gate", configuration) - self.assertIn("Next: Test policies offline", configuration) + self.assertIn("Next: Parse request content", configuration) + self.assertIn("Previous: Configure policies", request_content) + self.assertIn("Next: Test policies offline", request_content) if __name__ == "__main__": diff --git a/zensical.toml b/zensical.toml index d9c7e499..ecc2d1bc 100644 --- a/zensical.toml +++ b/zensical.toml @@ -29,6 +29,7 @@ nav = [ "documentation/egress-gate/index.md", {"Guides" = [ {"Configure policies" = "documentation/egress-gate/configuration.md"}, + {"Parse request content" = "documentation/egress-gate/request-content.md"}, {"Test policies offline" = "documentation/egress-gate/evaluation.md"}, {"Run and operate Egress Gate" = "documentation/egress-gate/operations.md"} ]},