Skip to content

feat(middleware): inspect WebSocket text messages#2477

Open
pimlock wants to merge 2 commits into
mainfrom
2428-websocket-message-introspection/pimlock
Open

feat(middleware): inspect WebSocket text messages#2477
pimlock wants to merge 2 commits into
mainfrom
2428-websocket-message-introspection/pimlock

Conversation

@pimlock

@pimlock pimlock commented Jul 25, 2026

Copy link
Copy Markdown
Collaborator

Summary

Adds the PR1 supervisor middleware contract and relay path for forward-direction WebSocket text inspection. Operator-run middleware can now preflight a wss upgrade, inspect OpenAI WebSocket mode events before credential resolution, return a replacement payload, or deny the session.

Related Issue

Part of #2428

Changes

  • Add the bidirectional EvaluateWebSocket gRPC stream and exact WEBSOCKET_MESSAGE/PRE_CREDENTIALS binding validation.
  • Add ordered preflight/session/message handling with fail-open/fail-closed behavior, per-stage deadlines, a total chain budget, shared admission capacity, and bounded envelopes.
  • Integrate forward text inspection into the parsed WebSocket relay while leaving server-to-client traffic, binary messages, and no-middleware sessions on their existing paths.
  • Preserve built-in policy ordering: inspect unresolved content, apply external transformations, re-evaluate policy, then resolve credential placeholders.
  • Add safe OCSF invocation/failure/saturation events and protocol close behavior (1008, 1009, 1002, and 1012).
  • Add real tonic sample middleware coverage that redacts an OpenAI response.create event and verifies the upstream sees only [REDACTED]; also verify deny-after-allow closes both peers with 1008.
  • Update the RFC contract documentation and leave PRE_RETURN explicitly reserved for PR2.

Testing

  • mise run pre-commit passes
  • Unit tests added/updated
  • E2E tests added/updated (if applicable)
  • mise run ci
  • mise run e2e:websocket-conformance
  • mise run e2e:docker

Checklist

  • Follows Conventional Commits
  • Commits are signed off (DCO)
  • Architecture docs updated (if applicable)

Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>

@pimlock pimlock left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

PR Review Status

Validation: This PR is project-valid because it implements the maintainer-authored, review-ready plan in #2428 and stays within the planned PR1 forward-text contract.
Head SHA: 22866e79e76619221db4922193eecc5330870c1c

Review findings:

  • High — plaintext WebSocket middleware bypass (CWE-693): handle_forward_proxy runs only the unary HTTP chain, where WebSocket-only bindings are filtered out, and then upgrades ws:// without preflight or a middleware session. Add the same WebSocket preflight/session/subprotocol propagation used by the other relay paths, or reject matching plaintext upgrades until supported. Add a ws:// forward-proxy regression test.
  • High — admission does not bound buffered memory (CWE-400): WebSocket payloads are fully read/reassembled/decompressed before waiting for the shared permit, and HTTP bodies are likewise buffered before admission. Concurrent multi-megabyte requests can therefore accumulate while queued. Reserve bounded work or weighted bytes before buffering and bound or shed the waiting queue.
  • Medium — termination causes are stringly typed and misclassified: middleware deny and service failure collapse to the same error; close codes are inferred from error text; and session-end reasons conflate policy reload, protocol errors, disconnects, and middleware failures. Use a typed termination cause mapped directly to RFC 6455 close codes and WebSocketSessionEndReason (including 1007 for invalid UTF-8).
  • Medium — published operator docs and mapped skills remain HTTP-only: update the existing Fern middleware, policy-schema, gateway-config, sandbox-policy, and logging pages for the WebSocket binding, max_message_bytes, stream/failure behavior, OCSF events, and the valid zero max_body_bytes shape for WebSocket-only services. The existing pages are already in navigation. Also update the affected generate-sandbox-policy, openshell-cli, and debug-openshell-cluster guidance identified by the maintenance map.
  • Medium — PR1 fault behavior lacks regression coverage: add deterministic coverage for slow/hanging/closed/missing/duplicate/out-of-order/oversized responses, subsequent messages after fail-open, admission saturation, stream lifetime, plaintext upgrades, and transformed fragmented/compressed text.

Two additional line-specific findings are attached inline.

Docs: Missing for this direct operator-facing middleware contract change.

Next state: gator:in-review pending author changes. No local tests were run as part of this code-only review.

Comment thread crates/openshell-supervisor-middleware/src/websocket.rs
Comment thread crates/openshell-supervisor-middleware/src/websocket.rs Outdated
@pimlock pimlock added the gator:in-review Gator is reviewing or awaiting PR review feedback label Jul 25, 2026
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
@github-actions

Copy link
Copy Markdown

@pimlock

pimlock commented Jul 25, 2026

Copy link
Copy Markdown
Collaborator Author

Addressed the initial review in d818907.

  • Plaintext forward-proxy WebSocket upgrades now run the WebSocket middleware preflight, preserve the ws scheme and negotiated subprotocol, and carry the middleware session into the frame relay.
  • Shared middleware admission is reserved before HTTP body buffering and before WebSocket frame payload buffering, reassembly, or decompression. The queue is bounded at 32 active evaluations plus 32 unbuffered waiters and sheds overflow.
  • WebSocket termination is typed and maps directly to close codes and session-end reasons, including 1007 for invalid UTF-8, 1009 for oversized messages, 1008 for policy or middleware denial, and 1012 for policy reload.
  • The RPC-wide 61-minute stream deadline is removed; per-event deadlines remain.
  • Fail-open stream disablement is now explicit and observable, with regression coverage for later messages. This intentionally preserves the PR1 handoff behavior.
  • Operator docs and the generate-sandbox-policy, openshell-cli, and debug-openshell-cluster skills now cover WebSocket middleware configuration, limits, failure behavior, and telemetry.
  • Regression coverage was expanded for bounded admission, fail-open continuation, ws preflight metadata, typed termination, OpenAI event redaction and denial, and the existing fragmented, compressed, malformed, and capacity cases. The full-duplex return-path fault fixture remains scoped to PR2 as documented in the handoff.

Validation passed:

  • mise run pre-commit
  • cargo test -p openshell-supervisor-middleware: 62 passed
  • cargo test -p openshell-supervisor-network: 1033 passed, 1 PR2 fixture ignored

@pimlock pimlock left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

PR Review Status

Validation: This PR remains project-valid because it implements the maintainer-authored, review-ready PR1 plan in #2428.
Head SHA: d81890743a80e74962ab4fdf8a51b22e068f2477

Thanks @pimlock. I checked your update describing the plaintext-upgrade path, bounded admission, typed termination, stream lifetime, fail-open observability, docs, skills, and expanded coverage. The plaintext path, pre-buffer admission, typed close handling, stream-deadline removal, mapped skill updates, and much of the added coverage are present. The independent re-review found three high-severity and four medium-severity issues that remain; they are attached inline.

Review findings:

  • High: incomplete WebSocket messages can pin all shared middleware admission indefinitely.
  • High: middleware-only parsed WebSockets do not observe policy-generation reloads.
  • High: the operator docs claim the built-in regex covers WebSockets, but its manifest and implementation do not.
  • Medium: persistent middleware streams have no independent session bound; fully disabled fail-open sessions still require admission; one path performs token-grant work before preflight; and policy denial is reported as middleware denial.

Docs: Fern pages and all three mapped skills were updated, but the built-in WebSocket coverage claim must be corrected or implemented.

Next state: gator:in-review pending author changes. No local tests were run as part of this code-only review.

));
}
let admission = if let Some(session) = options.middleware_session.as_ref() {
Some(session.reserve_message().await.map_err(|error| {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

High (CWE-400): Admission is reserved before read_masked_payload, but neither that read nor continuation assembly has a deadline. A client can send 32 text-frame headers or first fragments and then stall, pinning every process-wide permit indefinitely so unrelated HTTP/WebSocket middleware work blocks or is shed. Put the complete message-assembly interval under an idle/total deadline and cover stalled initial and fragmented frames.

// Built-in transport/GraphQL inspection sees the original unresolved
// message. External transformations run next, then policy is re-evaluated
// before credential material is introduced.
if let Some(inspector) = options.inspector.as_ref() {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

High (CWE-693): Policy-generation checking is conditional on inspector. A REST upgrade whose parsed relay is enabled solely by middleware_session has no inspector, so policy or middleware hot reload is never observed and the old chain can remain active indefinitely. Pass the generation guard independently, check it before each client message, and close/end with 1012/POLICY_RELOAD. Please cover middleware-only REST upgrades over both ws:// and wss://.

| Operator-run service | Required in gateway TOML | Set by the operator, up to the service capability | Runs as a separate service reachable by the gateway and supervisors |

`openshell/regex` is an example built-in middleware. It replaces only simple, self-contained token patterns in UTF-8 request bodies; the initial pattern recognizes `sk-` tokens. It does not infer values from keyword assignments such as JSON `password` fields. This best-effort text transformation is not parser-aware and does not guarantee that it will detect or fully remove sensitive values. Its `config` accepts one field, `mode: redact`, which is also the default when the field is omitted. Unknown config fields and non-string values are rejected at policy validation. Custom expressions are not configurable yet.
`openshell/regex` is an example built-in middleware. It replaces only simple, self-contained token patterns in UTF-8 HTTP bodies and client WebSocket text messages; the initial pattern recognizes `sk-` tokens. It does not infer values from keyword assignments such as JSON `password` fields. This best-effort text transformation is not parser-aware and does not guarantee that it will detect or fully remove sensitive values. Its `config` accepts one field, `mode: redact`, which is also the default when the field is omitted. Unknown config fields and non-string values are rejected at policy validation. Custom expressions are not configurable yet.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

High (CWE-693): This promises that openshell/regex inspects WebSocket text, but its manifest advertises only HTTP_REQUEST, and evaluate_web_socket explicitly returns UNIMPLEMENTED. Because HTTP-only bindings are silently excluded from the WebSocket chain, an operator following this guidance receives no WebSocket redaction and no failure. Either implement and advertise the binding or correct every claim that the built-in covers WebSockets, with a regression proving the selected behavior.

Ok(WebSocketPreflightResult {
allowed: true,
reason: String::new(),
session: (!stages.is_empty()).then_some(WebSocketSession {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

Medium (CWE-400): The preflight admission permit is released when this function returns, while WebSocketSession retains up to ten long-lived gRPC streams. There is no separate session/stream limit, so a sandbox can create unbounded idle upgrades and accumulate streams, channel buffers, and relay tasks. Keep healthy streams deadline-free, but hold a distinct bounded session permit in WebSocketSession and add an idle-session fan-out regression.

};
}

let Some(permit) = admission else {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

Medium: Once every fail-open stage has been disabled, later messages should bypass middleware. Instead, admission is still mandatory before the loop discovers there are no active stages; the relay likewise reserves unconditionally. Under saturation, a connection documented as bypassing the broken stage can therefore be closed. Return the original payload before admission when no stages remain active, and test a fully disabled session while the semaphore is saturated.

return Ok(());
}
};
let mut middleware_session = if let Some(chain) = websocket_chain.as_deref() {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

Medium: Run WebSocket preflight before inject_if_needed. This branch obtains and injects a token grant before preflight, unlike the other TLS and plaintext paths. Although preflight does not expose the header, credential issuance, cache mutation, and token-endpoint traffic still occur before middleware can fail closed. Preflight the unresolved request, then obtain credentials only after it succeeds; add a resolver-call-order regression.

openshell_core::proto::WebSocketSessionEndReason::PeerDisconnect
}
Self::PolicyReload => openshell_core::proto::WebSocketSessionEndReason::PolicyReload,
Self::MiddlewareDenial | Self::PolicyDenial => {

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

gator-agent

Medium (CWE-778): PolicyDenial is reported to middleware as MiddlewareDenial, so the typed lifecycle still conflates two distinct enforcement sources. Add a POLICY_DENIAL session-end reason and map it separately, then extend the termination mapping test.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

gator:in-review Gator is reviewing or awaiting PR review feedback

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant