Skip to content

feat(middleware): inspect WebSocket text messages - #2477

Draft
pimlock wants to merge 12 commits into
mainfrom
2428-websocket-message-introspection/pimlock
Draft

feat(middleware): inspect WebSocket text messages#2477
pimlock wants to merge 12 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)

@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 Outdated
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
@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.

@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.

@pimlock pimlock Jul 27, 2026

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.

Fixed in b1a7c14a.

The network relay now owns an explicit TextMessageAssembly state containing the buffered text, compression/fragment state, middleware admission permit, and absolute deadline. Admission still happens before payload buffering. The same assembly-owned bounds cover the initial payload, continuation headers, continuation payloads, and sequentially relayed interleaved control frames. A 30-second progress-idle timeout resets only after client input, while a 2-minute total timeout never resets. Timeout and every terminal error drop the state, releasing the shared permit through RAII. Timeout telemetry is classified without payload contents; existing protocol close/session-end semantics remain unchanged.

We also investigated established defaults before selecting the values. NGINX uses a 60-second between-read body timeout, Envoy defaults to a 5-minute stream-idle timeout, and AWS API Gateway WebSockets uses a 10-minute connection-idle timeout, although those boundaries are broader than one logical message holding scarce middleware capacity. Since this hop is local to the sandbox and the message cap is 4 MiB, we chose deliberately permissive initial bounds of 30 seconds idle / 2 minutes total to minimize the chance of disrupting legitimate fragmented clients while still preventing indefinite permit pinning. The new architecture/sandbox-limits.md records this as an early platform limit. We can tighten it using operational evidence, or later allow operators to narrow it within a platform ceiling; this PR adds no operator configuration.

Deterministic Tokio paused-time coverage verifies stalled initial payloads, stalled continuation payloads, missing continuation headers, a non-resetting total deadline under continued progress, and recovery after all 32 shared permits are saturated. Existing fragmented, interleaved-control, compressed, and rewrite behavior remains covered.

// 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://.

@pimlock pimlock Jul 27, 2026

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.

Fixed in 47932ffa.

  • The admitted generation guard now travels independently through UpgradeRelayOptions and websocket::RelayOptions; it is no longer coupled to InspectionOptions.
  • TLS upgrades propagate the tunnel guard. The plaintext forward-proxy fallback propagates forward_generation_guard even when no TunnelPolicyEngine exists.
  • The parsed relay checks freshness after validating every client data-frame header and before middleware admission, payload processing, or upstream forwarding.
  • A stale guard returns typed PolicyReload. Existing terminal handling then sends best-effort 1012 closes to both peers, ends the middleware session with POLICY_RELOAD, and emits the existing safe stale-generation telemetry.
  • The raw fast path remains unchanged.

Coverage uses real OpaEngine generation changes without sleeps: middleware-only REST wss:// reload, plaintext ws:// reload during a fragmented text message, POLICY_RELOAD delivery, and assertions that stale payload reaches neither middleware nor upstream.

| 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.

@pimlock pimlock Jul 27, 2026

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.

Fixed in 9fd47141 and ef26419d.

  • 9fd47141 removes the in-process/remote capability split. Middleware registrations now hold a shared SupervisorMiddlewareEndpoint, and the WebSocket runner opens the same transport-neutral stream for built-ins and operator-run gRPC services. This removes MiddlewareServiceState.remote and keeps lifecycle, deadlines, backpressure, diagnostics, and failure handling in the common runner.
  • ef26419d advertises WEBSOCKET_MESSAGE/PRE_CREDENTIALS for openshell/regex with max_message_bytes = 256 KiB; the HTTP binding retains its separate max_body_bytes = 256 KiB limit.
  • The built-in now implements preflight, session start, ordered message/result sequencing, and session end through both the in-process endpoint and generated tonic method. A selected regex stage returns INSPECT and accepts only client-to-upstream text messages.
  • Redaction reuses the existing fixed-pattern replacement and finding/metadata logic. It operates on complete UTF-8 messages after relay fragmentation/reassembly, returns no replacement on a no-match, rejects invalid UTF-8 or oversized input, and does not inspect or log payloads for binary or upstream-to-client traffic.
  • The policy documentation and relevant policy/CLI skills now describe the actual built-in capability and its text-only direction and limits.

Regression coverage now verifies the exact manifest and operation-specific limits, config and UTF-8/capacity validation, ordered session behavior, findings/metadata, no-match behavior, fail-closed capacity handling, a fragmented OpenAI-shaped message through the real relay, binary pass-through, and unchanged HTTP 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.

@pimlock pimlock Jul 27, 2026

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.

Fixed in 7c05b7ea.

  • Added a process-wide MAX_CONCURRENT_MIDDLEWARE_SESSIONS = 32 platform constant and a registry-owned session semaphore independent of the short-lived work semaphore.
  • WebSocket preflight uses immediate/non-waiting session admission before opening any inspecting transport. One permit covers the whole session regardless of stage fan-out; the existing 10-stage chain limit bounds fan-out separately.
  • All-skip preflight releases the permit immediately. An admitted session retains it only while at least one stage remains active; session end/drop and final-stage disablement release it through ownership/RAII.
  • Capacity exhaustion opens no streams and is evaluated through each selected config’s on_error: any fail-closed stage denies the upgrade, while an all-fail-open chain admits the upgrade without a middleware session. Mixed-chain behavior is covered.
  • The primitives and names are protocol-neutral (MiddlewareWorkAdmission, registry session admission/permit), so future streaming HTTP middleware can reuse the same shared budget without putting protocol limits in the network proxy.

Deterministic coverage caps 32 idle inspecting sessions, verifies release and re-admission, verifies all-skip release, and covers all-fail-open and mixed fail-open/fail-closed exhaustion. Both the built-in regex and remote transport paths remain covered under the unified transport abstraction.

};
}

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.

@pimlock pimlock Jul 27, 2026

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.

Fixed in 7c05b7ea.

  • Added an explicit WebSocketMessageAdmission::{Bypass, Inspect(MiddlewareWorkAdmission)} lifecycle instead of overloading Option<...>.
  • The relay asks the session for this typed outcome before reserving per-message work or buffering the text message. When no stage remains active, it drops the empty middleware session and forwards the original payload without touching the work semaphore.
  • Each stage now owns its unified in-process/remote transport as an optional resource. Disabling a fail-open stage immediately drops its sender and response stream. When the final stage disables, lifecycle reconciliation also releases the persistent session permit.
  • An inspecting message retains its work permit through complete assembly and evaluation. Mixed sessions continue taking message admission while any stage remains active.

The network regression uses the remote transport path: the first message breaks a fail-open stream and disables the only stage, all 32 work permits are then occupied, and the next message still reaches upstream unchanged with no second middleware RPC. Unit coverage also verifies immediate bypass under saturation, session-capacity release on final disablement, and continued admission when a mixed chain retains one active built-in regex stage.

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.

@pimlock pimlock Jul 27, 2026

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.

Fixed in 47932ffa.

The TLS relay_rest order is now:

  1. L7 policy
  2. HTTP middleware
  3. WebSocket preflight on the unresolved request
  4. inject_if_needed / token grant
  5. Upstream relay

This keeps preflight ahead of credential issuance, cache mutation, and token-endpoint work while preserving the existing transport-specific failure response.

The new call-order tests hold preflight at an explicit notification boundary without sleeps:

  • The denied case verifies zero resolver requests and zero upstream bytes before and after denial.
  • The admitted case verifies zero resolver requests before preflight release, followed by one resolver request and the injected Authorization header upstream.

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.

@pimlock pimlock Jul 27, 2026

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.

Fixed in 47932ffa.

  • supervisor_middleware.proto now defines WEB_SOCKET_SESSION_END_REASON_POLICY_DENIAL with tag 9.
  • Typed terminal handling reports built-in PolicyDenial as POLICY_DENIAL.
  • External middleware rejection remains MIDDLEWARE_DENIAL.
  • Both remain workload-visible WebSocket close code 1008.

Coverage extends the direct typed-mapping assertions and adds a behavior test in which built-in GraphQL policy denies a message, both peers receive 1008, middleware receives POLICY_DENIAL, and the denied message reaches neither middleware nor upstream.

pimlock

This comment was marked as duplicate.

@pimlock
pimlock marked this pull request as draft July 27, 2026 21:00
@copy-pr-bot

copy-pr-bot Bot commented Jul 27, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@pimlock pimlock added gator:blocked Gator is blocked by process or repository gates and removed gator:in-review Gator is reviewing or awaiting PR review feedback labels Jul 27, 2026
pimlock added 6 commits July 27, 2026 15:14
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
@pimlock
pimlock force-pushed the 2428-websocket-message-introspection/pimlock branch from ef26419 to 83121b1 Compare July 27, 2026 22:16
pimlock added 2 commits July 27, 2026 16:11
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
@pimlock

pimlock commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator Author

/ok to test dea07e4

pimlock added 3 commits July 28, 2026 10:10
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
@pimlock

pimlock commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator Author

gator-agent

Blocked

Head SHA: 2fd92cd94c0445de52498d0e2e016d98ea061ef1

Gator is still blocked because this pull request remains a draft, so the independent code review must wait until it is ready for review.

The current head is new since the prior draft-blocker disposition. I confirmed that DCO is passing, but draft status remains the process blocker; pending branch checks do not change that disposition.

Next action: finish the draft updates and mark the pull request ready for review. Gator will then run one fresh independent review for the current head. I did not approve, merge, push, apply test:e2e, or post /ok to test.

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

pimlock commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator Author

gator-agent

Blocked

Head SHA: c6788ef57c939dc96491bd89424aa60fb3d54813

Gator is blocked because this pull request is still a draft, so the independent code review must wait until it is ready for review.

The head changed since the previous draft-blocker disposition. DCO is passing; the remaining process blocker is draft status.

Next action: mark the pull request ready for review. Gator will then run one fresh independent review for the current head. I did not approve, merge, push, apply test:e2e, or post /ok to test.

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

Labels

gator:blocked Gator is blocked by process or repository gates

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant