diff --git a/.agents/skills/debug-openshell-cluster/SKILL.md b/.agents/skills/debug-openshell-cluster/SKILL.md index 0de85fb0a4..ee4cad2f8e 100644 --- a/.agents/skills/debug-openshell-cluster/SKILL.md +++ b/.agents/skills/debug-openshell-cluster/SKILL.md @@ -96,15 +96,15 @@ The gateway calls each interceptor's `Describe` RPC and validates its manifest a For operator-run supervisor middleware, inspect `[[openshell.supervisor.middleware]]`, service reachability, and both gateway and supervisor logs: ```bash -rg -n 'supervisor|middleware|grpc_endpoint|max_body_bytes|timeout' /etc/openshell/gateway.toml +rg -n 'supervisor|middleware|grpc_endpoint|max_payload_bytes|timeout' /etc/openshell/gateway.toml journalctl -u --no-pager --lines=200 journalctl -u openshell-gateway --no-pager --lines=200 openshell logs --tail --source sandbox ``` -The middleware service must start before the gateway and be reachable from both the gateway and sandbox supervisors. Gateway startup fails if `Describe` is unavailable, a manifest exposes duplicate `HttpRequest/pre_credentials` bindings, the registration claims the reserved `openshell/` namespace, or body and timeout limits are invalid. Changing a registration requires a gateway restart. A policy update can also fail before persistence if the selected implementation rejects its `network_middlewares` config. +The middleware service must start before the gateway and be reachable from both the gateway and sandbox supervisors. Gateway startup fails if `Describe` is unavailable, a manifest exposes duplicate operation/phase bindings, the registration claims the reserved `openshell/` namespace, or body, message, and timeout limits are invalid. Supported V1 bindings are `HTTP_REQUEST/PRE_CREDENTIALS` and `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`. Changing a registration requires a gateway restart. A policy update can also fail before persistence if the selected implementation rejects its `network_middlewares` config. -At request time, distinguish an explicit `middleware_denied` result from `middleware_failed`. A denial is always enforced. A failure follows the policy-local `on_error`: `fail_closed` blocks the request, while `fail_open` bypasses only that stage and emits a detection finding. If a running supervisor cannot install a new registry, it preserves its last-known-good generation and emits a configuration failure event. +At request time, distinguish attachment, binding selection, coverage, denial, and failure. A host-matched HTTP-only attachment can inspect the upgrade GET but does not join the WebSocket chain; the connection proceeds under either `on_error` mode and emits `binding_not_selected` coverage. A selected WebSocket stage receives text messages only. Binary messages pass under both modes, emit `unsupported_message_type` coverage, and consume a session sequence without an RPC. An explicit `middleware_denied` result is always enforced. WebSocket preflight returns `INSPECT`, voluntary `SKIP`, or authoritative `DENY`; `DENY` rejects the upgrade before upstream contact under both `on_error` modes. A selected-stage failure follows the policy-local `on_error`: `fail_closed` blocks the HTTP request or closes the WebSocket, while `fail_open` bypasses only that stage and emits a detection finding. A fail-open per-message capacity failure bypasses that message without disabling the stage. A timeout, transport failure, stream closure, missing or invalid response, duplicate or regressed sequence, or other failure that makes an established WebSocket stream unreliable disables that stage for later messages on the connection and emits `openshell.middleware.websocket_stage_disabled`. Confirm preflight, session-start, and session-end in service logs. OpenShell best-effort sends at most one session-end to each still-writable opened stage, including a preflight that terminates before session start; distinguish `MIDDLEWARE_DENIAL` from `MIDDLEWARE_FAILURE`. WebSocket message sequences are allocated session-wide; each stage receives a strictly increasing subset, so gaps are valid when binary messages or other units are not delivered to that stage. Zero, duplicate, or regressed sequences are protocol errors. If a running supervisor cannot install a new registry, it preserves its last-known-good generation and emits a configuration failure event. ### Step 4: Check Docker-Backed Gateways @@ -403,7 +403,10 @@ openshell logs | Provider profiles disappear after enabling an interceptor catalog | `provider_profile_sources` selected only an authoritative interceptor or returned invalid/duplicate IDs | Inspect source list and interceptor `Describe`/catalog logs; include `builtin` and `user` when intended | | Gateway fails after registering supervisor middleware | Service unavailable, invalid manifest, duplicate binding, reserved name, or invalid body/timeout limit | Middleware service and gateway logs; `[[openshell.supervisor.middleware]]`; `Describe` response | | Policy update rejects `network_middlewares` | Unknown middleware name, implementation-owned config invalid, duplicate order, broad/invalid host selector, or fail-closed coverage of `tls: skip` | Policy error, gateway logs, middleware `ValidateConfig`, selector and order fields | -| HTTP request returns `middleware_failed` or `middleware_denied` | Selected stage failed or explicitly denied the admitted request | Sandbox OCSF logs; policy-local middleware config; service availability; `on_error` | +| HTTP request returns `middleware_failed`, or WebSocket closes with `1008` | Selected stage failed or explicitly denied admitted traffic | Sandbox OCSF logs; policy-local middleware config; service availability; binding operation; `on_error` | +| WebSocket upgrades but a host-matched middleware receives no preflight or message RPC | The implementation did not advertise `WEBSOCKET_MESSAGE/PRE_CREDENTIALS` | `WEBSOCKET_MIDDLEWARE_COVERAGE state=binding_not_selected`; service `Describe`; the upgrade GET may still have used its HTTP binding | +| Binary WebSocket message passes without a middleware RPC | Binary is unsupported by the V1 text-message binding under both `on_error` modes | `WEBSOCKET_MIDDLEWARE_COVERAGE state=unsupported_message_type`; the next text RPC may have a valid sequence gap | +| WebSocket messages stop reaching middleware after one failure | A fail-open stage stream was disabled for the rest of the connection | `openshell.middleware.websocket_stage_disabled`; middleware timeout/stream/protocol logs. A per-message capacity bypass alone leaves the stage active. Reconnect to create a fresh stream after a genuine stream failure | | Custom compute driver is unavailable | Driver process/socket missing, inaccessible, or configured with a reserved/mismatched name | Socket ownership/mode, driver service logs, gateway `GetCapabilities` logs | | Image pull failure | Gateway or sandbox image cannot be pulled | Runtime events and image pull credentials | | `K8s namespace not ready` with `envoy-gateway-openshell.yaml: the server could not find the requested resource` | Optional Gateway API manifest was applied without Envoy Gateway CRDs, or k3s Helm controller startup exceeded the namespace wait | Apply `deploy/kube/manifests/envoy-gateway-openshell.yaml` manually only after Envoy Gateway is installed and `grpcRoute` is enabled | diff --git a/.agents/skills/generate-sandbox-policy/SKILL.md b/.agents/skills/generate-sandbox-policy/SKILL.md index 3b261b5ded..c4010482b2 100644 --- a/.agents/skills/generate-sandbox-policy/SKILL.md +++ b/.agents/skills/generate-sandbox-policy/SKILL.md @@ -79,7 +79,7 @@ Regardless of tier, extract (or infer) these from the user's description: | **Paths** | Specific URL paths or patterns | Only for custom/fine-grained | | **Enforcement** | `enforce` or `audit`? Default to `enforce`. | No — has a default | | **Binary** | Which binary/process should have access | Yes — ask if not stated | -| **Middleware** | Whether admitted HTTP requests need an ordered built-in or operator-run processing stage | No | +| **Middleware** | Whether admitted HTTP requests or client WebSocket text messages need an ordered built-in or operator-run processing stage | No | If the host and access level are clear but binaries are not specified, ask the user which binary or process will be making the requests. Suggest common defaults like `/usr/bin/curl`, `/usr/local/bin/claude`, etc. @@ -216,10 +216,14 @@ Is L7 inspection needed? ### Middleware Decision -Add `network_middlewares` only when the user asks to inspect, transform, redact, or independently authorize admitted HTTP requests. Middleware runs after network and L7 policy admission and before provider credential injection. +Add `network_middlewares` only when the user asks to inspect, transform, redact, or independently authorize admitted HTTP requests or client WebSocket text messages. Middleware runs after network and L7 policy admission and before provider credential injection. -- Use a built-in name such as `openshell/regex` without gateway registration. +- Use `openshell/regex` without gateway registration for fixed-pattern redaction of UTF-8 HTTP request bodies or complete client-to-upstream WebSocket text messages. - Use an operator-owned middleware name only when it is already registered under `[[openshell.supervisor.middleware]]` and reachable from both the gateway and sandbox supervisors. +- Confirm that a requested WebSocket implementation exposes a `WEBSOCKET_MESSAGE/PRE_CREDENTIALS` binding. `openshell/regex` exposes this binding. A host-matched HTTP-only implementation may inspect the upgrade GET but does not join the post-upgrade chain; messages pass and OpenShell emits `binding_not_selected` coverage regardless of `on_error`. +- WebSocket middleware runs for both `ws://` and `wss://` and receives complete client text messages only. Binary messages pass under both error modes and emit `unsupported_message_type` coverage for active stages. Upstream-to-client messages remain uninspected. Do not claim that V1 provides all-message WebSocket inspection. +- Treat `fail_open` on WebSocket as a session-scoped bypass: if the stage stream fails, OpenShell disables it for later messages on that connection and emits a state-change finding. Prefer `fail_closed` for required redaction or authorization. +- `on_error` governs failures after an advertised operation binding is selected. It does not apply to an unadvertised WebSocket binding or binary-message pass-through. An explicit HTTP, WebSocket preflight, or WebSocket message denial is authoritative under both `fail_open` and `fail_closed`. - Default `on_error` to `fail_closed`. Use `fail_open` only when bypassing the stage preserves the user's stated security requirement. - Assign unique `order` values across the complete policy. Lower values run first, and at most 10 configs may be selected. - Match the narrowest destination hosts possible with `endpoints.include`; use `exclude` when a broad selector has trusted exceptions. @@ -376,6 +380,7 @@ Before presenting the policy to the user, verify correctness **and** flag breadt - [ ] Every middleware config has a non-empty `middleware` name and non-empty `endpoints.include` - [ ] Middleware `order` values are unique and no selected chain exceeds 10 stages - [ ] No fail-closed middleware selector can cover a `tls: skip` endpoint +- [ ] Any required WebSocket control advertises `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`, and the user understands that V1 does not inspect binary messages ### Schema Warnings (log-only, but should be fixed) @@ -406,7 +411,7 @@ Evaluate the generated policy for overly broad access and **include warnings in | **Hostless `allowed_ips`** (no `host` field) | "This endpoint has no `host` — any domain resolving to the allowed IP range on this port will be permitted. Consider adding a `host` field to restrict which domains can use this allowlist." | | **Broad CIDR** in `allowed_ips` (e.g., `10.0.0.0/8`) | "This `allowed_ips` entry covers a very broad range. Consider narrowing to a specific subnet (e.g., `10.0.5.0/24`) to minimize exposure." | | **`on_error: fail_open`** | "This middleware can be bypassed when it is unavailable, rejects configuration, returns an invalid result, or exceeds its body limit. Use `fail_closed` unless availability is more important than this control." | -| **Broad middleware host selector** | "This middleware applies independently of the admitting network rule to every matching HTTP destination. Narrow `endpoints.include` or add exclusions if the stage is not required for every matching host." | +| **Broad middleware host selector** | "This middleware attaches independently of the admitting network rule to every matching destination, then runs only for operation bindings its implementation advertises. Narrow `endpoints.include` or add exclusions if the attachment is not required for every matching host." | Format breadth warnings clearly in the output, e.g.: diff --git a/.agents/skills/openshell-cli/SKILL.md b/.agents/skills/openshell-cli/SKILL.md index 346672b912..8f0863641a 100644 --- a/.agents/skills/openshell-cli/SKILL.md +++ b/.agents/skills/openshell-cli/SKILL.md @@ -344,10 +344,12 @@ Edit `current-policy.yaml` to allow the blocked actions. **For policy content au - TLS termination configuration - Enforcement modes (`audit` vs `enforce`) - Binary matching patterns -- Ordered `network_middlewares`, host selection, and `fail_open` or `fail_closed` behavior +- Ordered `network_middlewares`, host selection, HTTP and WebSocket bindings, and `fail_open` or `fail_closed` behavior `network_policies` and `network_middlewares` can be modified at runtime. If `filesystem_policy`, `landlock`, or `process` need changes, the sandbox must be recreated. Built-in middleware such as `openshell/regex` needs no gateway registration. An operator-run middleware must already be registered under `[[openshell.supervisor.middleware]]`; changing that static registration requires a gateway restart. +Middleware can inspect parsed HTTP request bodies and complete client-to-upstream WebSocket text messages over both `ws://` and `wss://` when the implementation advertises the matching binding. The built-in `openshell/regex` advertises both bindings and applies its fixed patterns to UTF-8 text. A host-matched HTTP-only attachment can inspect the upgrade GET but does not join the WebSocket chain; look for `binding_not_selected` coverage. Binary messages pass under both `on_error` modes and active stages emit `unsupported_message_type` coverage; upstream-to-client messages remain uninspected. A broken fail-open WebSocket stage is disabled for the rest of that connection; inspect sandbox OCSF logs for `openshell.middleware.websocket_stage_disabled`. + ### Step 5: Push the updated policy ```bash diff --git a/Cargo.lock b/Cargo.lock index f028df5f42..9fb11207b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3900,6 +3900,7 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokio", + "tokio-stream", "tonic", "tonic-prost", "tonic-prost-build", @@ -4249,6 +4250,7 @@ dependencies = [ name = "openshell-supervisor-middleware" version = "0.0.0" dependencies = [ + "futures", "miette", "openshell-core", "openshell-supervisor-middleware-builtins", @@ -4270,6 +4272,7 @@ dependencies = [ "serde", "serde_json", "tokio", + "tokio-stream", "tonic", ] @@ -4314,6 +4317,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-rustls 0.26.4", + "tokio-stream", "tokio-tungstenite 0.26.2", "tonic", "tower-mcp-types", diff --git a/architecture/README.md b/architecture/README.md index 3fc72afd25..ecf3114a6a 100644 --- a/architecture/README.md +++ b/architecture/README.md @@ -163,6 +163,7 @@ that crate's `README.md`. |---|---| | [Gateway](gateway.md) | Gateway control plane, auth, APIs, persistence, settings, and relay coordination. | | [Sandbox](sandbox.md) | Sandbox supervisor, child process isolation, proxy, credentials, inference, connect, and logs. | +| [Sandbox Limits](sandbox-limits.md) | Sandbox supervisor and egress safety ceilings, ownership rules, current enforcement, and known gaps. | | [Security Policy](security-policy.md) | Policy model, enforcement layers, policy updates, policy advisor, and security logging. | | [Compute Runtimes](compute-runtimes.md) | Docker, Podman, Kubernetes, VM, sandbox images, and runtime-specific responsibilities. | | [Build](build.md) | Build artifacts, CI/E2E, docs site validation, and release packaging. | diff --git a/architecture/gateway.md b/architecture/gateway.md index 5a84177e1f..b89caf3e5f 100644 --- a/architecture/gateway.md +++ b/architecture/gateway.md @@ -397,6 +397,9 @@ each service and validates its described bindings and operator body limit. Policies attach a complete external middleware by its operator-owned registration name. Manifest bindings are identified by operation and phase, and each manifest may declare at most one binding for an operation and phase pair. +Attaching a registration does not require it to advertise every supported +operation. Supervisors select only the manifest bindings that match the current +operation; policy-local config identity remains internal audit metadata. Before persisting a policy, the gateway asks each selected implementation to validate its config. The effective sandbox config contains only the registered services required by that policy; supervisors invoke those services directly on diff --git a/architecture/sandbox-limits.md b/architecture/sandbox-limits.md new file mode 100644 index 0000000000..8390539201 --- /dev/null +++ b/architecture/sandbox-limits.md @@ -0,0 +1,172 @@ +# Sandbox Limits + +The sandbox supervisor processes untrusted agent traffic while sharing one +process across connections. Limits protect that process from unbounded memory, +work queues, parser effort, and waits. They are safety boundaries, not capacity +targets or compatibility guarantees. OpenShell is still early, so the values +and some ownership boundaries will change as production evidence improves. + +This document inventories durable sandbox supervisor and egress limits. It +intentionally omits retry cadence, test-only values, and ordinary buffer chunk +sizes that do not cap aggregate resource use. + +## Limit Model + +OpenShell uses three kinds of limits: + +| Kind | Purpose | Configuration | +|---|---|---| +| Platform ceiling | Protect the supervisor even when policy or an external service is hostile or mistaken. | Fixed in code by default. | +| Operator ceiling | Bound an operator-run integration below the platform maximum. | Configurable within platform validation bounds. | +| Policy inspection bound | State how much application data a policy needs the supervisor to buffer and inspect. | Configurable when the application protocol needs it, preferably below a platform ceiling. | + +The component that first allocates, queues, or waits on a resource owns its +limit. Network frame and message assembly belongs to the network supervisor; +middleware RPC and envelope limits belong to the middleware runner; application +inspection limits belong to the corresponding L7 parser. + +New limits should follow these rules: + +- Enforce the bound before allocation or admission whenever possible. +- Acquire shared capacity before buffering and retain it through the complete + buffered operation. +- Give partial-progress protocols both an idle bound and an absolute bound when + either one alone still permits resource pinning. +- Let operator or policy configuration narrow a platform ceiling, never raise + it silently. +- Define the terminal behavior: reject, close, truncate, shed, or backpressure. +- Do not let `fail_open` bypass a platform safety or protocol-integrity bound. +- Emit safe saturation or limit telemetry without request bodies, credentials, + query parameters, or external free-form diagnostics. +- Test time bounds with simulated time and test shared budgets under saturation. + +## Middleware + +Middleware limits are process-wide per sandbox. Registry replacement preserves +the shared work, waiter, and persistent-session admission state so activity +retained by an older generation still consumes the same process-lifetime +budgets as new activity. + +| Resource | Current bound | Scope and behavior | +|---|---:|---| +| Concurrent buffered work | 32 | Shared by HTTP requests, WebSocket messages, and WebSocket preflight. One permit covers one complete unit of work. | +| Admission waiters | 64 | Additional work is shed when both the active budget and waiter budget are full. HTTP receives a complete 503 response before its body is buffered. | +| Persistent middleware sessions | 32 | Shared process-wide session budget for streaming middleware protocols. WebSocket preflight uses immediate admission before opening streams and retains one permit while any stage remains active. | +| HTTP body or WebSocket text message | 4 MiB | Platform maximum for input and replacement payloads. Service, operator, and stage limits may narrow it. | +| Middleware configs and stages | 10 | At most 10 configs in policy and 10 selected stages in one chain. | +| Selector patterns | 32 | Combined include and exclude patterns per middleware config. | +| Per-stage RPC | 500 ms default, 10 ms–30 s | An operator timeout caps a binding timeout. | +| Complete message chain | 30 s | Starts after work admission; admission backpressure does not consume the chain budget. | +| WebSocket preflight | 1 s maximum | Caps handshake delay independently of the message RPC timeout. | +| Remote service connect | 5 s | Applies while establishing a middleware gRPC channel. | + +Middleware also validates every non-body envelope component. Important examples +include 64 KiB service config, 4 KiB request context, 32 KiB target data, 128 +request headers totaling 64 KiB, 64 header mutations, 32 findings per stage, +and 64 metadata entries. The detailed external contract lives in +[Supervisor Middleware](../docs/extensibility/supervisor-middleware.mdx). + +The work semaphore bounds aggregate buffered middleware input to approximately +`32 × 4 MiB`, plus bounded envelope and parser overhead. It is a concurrency +safety valve, not rate limiting or a promise that 32 simultaneous maximum-size +messages are inexpensive. + +The persistent session semaphore is independent from the work semaphore. One +WebSocket middleware session consumes one permit regardless of its active-stage +fan-out, which is separately capped at 10 stages. All-skip preflight releases +the permit immediately. A retained session releases it at connection end or as +soon as its last active stage is disabled. Session admission does not wait: +capacity exhaustion follows each selected config's `on_error` behavior before +any stream opens. The protocol-neutral registry ownership allows future +streaming HTTP middleware to use the same process-wide budget. + +## Egress Framing and Inspection + +| Path | Current bound | Terminal behavior | +|---|---:|---| +| Initial CONNECT request headers | 8 KiB | Reject the proxy request. | +| Inspected HTTP/1 request headers | 16 KiB | Reject the request. | +| Credential-rewritten HTTP body | 256 KiB | Reject when rewriting requires a larger buffered body. | +| SigV4 body signing | 10 MiB | Reject when signing requires a larger buffered body. | +| GraphQL request body | 64 KiB default | Policy can set a positive `graphql_max_body_bytes`; there is no shared platform ceiling yet. | +| MCP or JSON-RPC request body | 64 KiB default | Policy can set a positive `max_body_bytes`; there is no shared platform ceiling yet. | +| Parsed WebSocket client text message | 4 MiB | Close with `1009` when the complete or decompressed message is larger. | +| Concurrent parsed WebSocket text assemblies | 32 active, 64 waiters | Shared process-wide across all parsed relays. Additional messages close with `1013` before payload allocation or reading when both bounds are full. | +| Parsed WebSocket fragments per message | 4,096 | Close with `1002`. | +| Parsed WebSocket text assembly | 30 s input idle, 2 min total | Close with `1002`; the total includes initial and continuation payloads, continuation headers, and interleaved control frames. These deliberately permissive initial bounds can be tightened, or made operator-tunable within platform bounds, after production behavior is understood. | +| Parsed raw WebSocket binary frame | 16 MiB | Close with `1002`; binary messages are relayed rather than inspected. | +| HTTP relay waiting for EOF | 5 s input idle | End the relay with a timeout. | +| TLS certificate cache | 256 hosts | Clear the cache before inserting another host. | + +Ordinary allowed traffic is streamed rather than accumulated to a +connection-sized buffer. A parsing or transformation feature introduces a +buffer only when it owns an explicit bound. + +Every parsed WebSocket text message acquires network-owned assembly capacity before payload allocation or reading, including relays used only for native policy, credential rewriting, compression, or a disabled fail-open middleware session. The process-lifetime budget survives policy reloads, and the assembly retains its permit through decompression, policy and middleware evaluation, credential rewriting, and upstream forwarding. Active middleware sessions additionally acquire shared middleware work before buffering. Input progress resets only the idle deadline. Every timeout and terminal parser error releases both permits through ordinary ownership. Queue exhaustion emits a payload-free network denial event. + +The operator middleware `max_payload_bytes` ceiling applies to payloads exposed +through HTTP-body and WebSocket text-message bindings. It does not replace the +raw binary frame safety bound because binary messages are never delivered to V1 +middleware. A passed binary logical message still advances the active +middleware session sequence and emits coverage telemetry, so a later text RPC +can contain a valid sequence gap. + +## Inference and Upstream Proxying + +| Path | Current bound | Terminal behavior | +|---|---:|---| +| `inference.local` request parse buffer | 10 MiB | Return `413` for an oversized request. | +| Chunked inference request | 10 MiB and 4,096 chunks | Reject an invalid or over-limit request. | +| Streaming inference response | 32 MiB and 120 s chunk idle | Truncate the stream and attempt a safe SSE error. | +| Corporate proxy CONNECT response headers | 8 KiB | Fail the tunnel. | +| Corporate proxy CONNECT handshake | 30 s total | Fail the tunnel; validated-address attempts share the aggregate budget. | +| Token-grant HTTP request | 30 s request and connect | Fail credential resolution. | +| Response-derived token cache TTL | 5 min default; 1 h response cap; 30 s expiry margin | A positive profile `cache_ttl_seconds` override replaces the response-derived calculation. | + +Streaming response byte limits are integrity-relevant. Protocols whose clients +require one complete buffered object do not use the truncating SSE path. + +## Sandbox-Local Surfaces + +| Surface | Current bound | Scope and behavior | +|---|---:|---| +| `policy.local` request body | 64 KiB and 15 s read | Reject an oversized or stalled local request. | +| Policy proposal long-poll | 60 s default, 1–300 s | Clamp the requested hold time; clients can issue another poll. | +| `policy.local` denial read | 100 records, 4 KiB per surfaced line | Bound response and log parsing work. | +| Log push reconnect buffer | 200 records | Drop new records above the local batch ceiling while disconnected. | +| Log push reconnect backoff | 30 s maximum | Cap the delay between reconnect attempts. | +| Policy status outbox | No fixed capacity | Preserve FIFO revision status without blocking policy reconciliation. | +| Policy status retry backoff | 32 s maximum | Retain a retryable update and retry independently of enforcement. | + +The bounded log batch favors supervisor health over retaining an unbounded +diagnostic backlog. The policy status outbox makes the opposite tradeoff: +revision ordering and delivery survive an extended gateway outage at the cost +of potential queue growth. + +## Known Gaps and Review Triggers + +The current limits grew with individual features and are not yet a complete +resource model. Known gaps include: + +- GraphQL, MCP, and JSON-RPC policy body limits have defaults but no common + platform maximum. +- A positive token-cache TTL override replaces the response-derived one-hour + ceiling rather than narrowing it. +- Socket read deadlines are not expressed consistently as idle plus total + budgets across every parser. +- There is no documented aggregate connection budget or per-destination + fairness policy in the supervisor. +- The policy status outbox is intentionally unbounded and can grow if policy + revisions continue while its gateway endpoint remains unavailable. +- Limit telemetry is not yet uniform enough to derive saturation trends across + all paths. + +Revisit this document when adding a parser, body transformation, persistent +stream, shared queue, cache, or external call. A change should state: + +1. What untrusted resource can grow or wait. +2. Which component owns the bound. +3. Whether the scope is per message, connection, destination, or sandbox. +4. Whether configuration can narrow the limit. +5. How saturation or timeout terminates. +6. Which telemetry and deterministic tests prove the behavior. diff --git a/architecture/sandbox.md b/architecture/sandbox.md index 4f7dfe4272..9c702248a0 100644 --- a/architecture/sandbox.md +++ b/architecture/sandbox.md @@ -51,6 +51,9 @@ paths, such as proxy support files or GPU device paths when a GPU is present. ## Network and Inference +See [Sandbox Limits](sandbox-limits.md) for the current numeric safety ceilings, +their ownership, terminal behavior, and known gaps. + All ordinary agent egress is routed through the sandbox proxy. The proxy identifies the calling binary, checks trust-on-first-use binary identity, rejects unsafe internal destinations, and evaluates the active policy. On Linux, it @@ -60,12 +63,11 @@ socket inode. For inspected HTTP traffic, the proxy can enforce REST method/path rules, WebSocket upgrade and text-message rules, GraphQL operation rules, and MCP method, tool, and supported params rules or generic JSON-RPC method rules -on sandbox-to-server request bodies. MCP and JSON-RPC inspection buffers up to -the endpoint `mcp.max_body_bytes` or `json_rpc.max_body_bytes` limit. MCP -`tools/call` tool names are checked against the spec-recommended syntax by -default before policy evaluation, with a per-endpoint `mcp.strict_tool_names` -compatibility opt-out. Generic JSON-RPC policies do not support `params` -matchers; generic JSON-RPC rules match only the method. +on sandbox-to-server request bodies. MCP and JSON-RPC inspection buffers +bounded request bodies. MCP `tools/call` tool names are checked against the +spec-recommended syntax by default before policy evaluation, with a per-endpoint +`mcp.strict_tool_names` compatibility opt-out. Generic JSON-RPC policies do not +support `params` matchers; generic JSON-RPC rules match only the method. JSON-RPC responses and server-to-client MCP messages on response or SSE streams are relayed but are not currently parsed for policy enforcement. @@ -75,11 +77,20 @@ host selectors choose the chain independently of the network rule that admitted the request. Policy-local map keys identify configs, while built-in names or operator-owned registration names identify implementations. -Built-ins run in-process; operator services use the same bounded gRPC contract. -`openshell-policy` validates policy-owned structure, and the active middleware -registry validates implementation-owned config. The generic registry and chain -runner live in `openshell-supervisor-middleware`; first-party implementations -live in `openshell-supervisor-middleware-builtins`. +Built-ins run in-process; operator services use gRPC. Both implement the same +transport-neutral endpoint contract, including bidirectional WebSocket +sessions, so a manifest advertises capabilities independently of transport. +The runtime keeps three states distinct: host selection attaches policy configs, +manifest operation and phase bindings select the active chain, and the parsed +message type determines whether that chain can inspect an individual payload. +An attachment without a WebSocket binding is not a failed WebSocket stage. +Binary messages are outside the V1 text-message binding. Both cases pass through +with informational coverage telemetry rather than applying `on_error`. +The chain runner owns shared sequencing, deadlines, backpressure, and response +validation. `openshell-policy` validates policy-owned structure, and the active +middleware registry validates implementation-owned config. The generic +registry and chain runner live in `openshell-supervisor-middleware`; first-party +implementations live in `openshell-supervisor-middleware-builtins`. The supervisor installs policy and middleware registry changes as one runtime generation and preserves the last-known-good generation if preparation fails. @@ -91,8 +102,9 @@ credential, routing, or framing headers. Body transformations are re-evaluated against body-aware L7 policy before later stages or the upstream can observe them. Requests, results, chain length, execution time, and diagnostics are bounded; external free-form diagnostic text is not exposed in responses or -security logs. See [Supervisor Middleware](../docs/extensibility/supervisor-middleware.mdx) -for configuration and protocol details. +security logs. See +[Supervisor Middleware](../docs/extensibility/supervisor-middleware.mdx) for +configuration and protocol details. `https://inference.local` is special. It bypasses OPA network policy and is handled by the inference interception path: @@ -197,8 +209,7 @@ exchanges it for an OAuth2 access token, caches the token, and injects it as an `Authorization: Bearer` header before forwarding the request. Token grant endpoints are HTTPS-only except for loopback and Kubernetes service DNS hosts, and returned access tokens must be bearer-compatible before they are cached or -injected. Token response lifetimes are capped and cached with an expiry margin -unless a profile supplies an explicit cache TTL override. +injected. Token caching follows response-derived and profile override TTL rules. For AWS endpoints that require request-level signing, the proxy supports SigV4 re-signing. When `credential_signing: sigv4` is set on an L7 endpoint, the proxy @@ -206,9 +217,8 @@ strips the client's placeholder-based AWS auth headers, re-signs with real credentials from the provider, and forwards the request upstream. The signing mode is auto-detected from the client SDK's `x-amz-content-sha256` header: -- **Signed body** (hex hash): buffers the request body (up to 10 MiB), computes - its SHA-256, and includes the hash in the signature. Used by Bedrock and most - AWS services. +- **Signed body** (hex hash): buffers the request body, computes its SHA-256, + and includes the hash in the signature. Used by Bedrock and most AWS services. - **Streaming unsigned** (`STREAMING-UNSIGNED-PAYLOAD-TRAILER`): signs headers only and streams the body through without buffering. Used by S3 uploads with `aws-chunked` encoding. @@ -281,12 +291,10 @@ remains `Pending`. If the first poll returns a different revision, the superviso processes it through the normal reload path instead of treating it as already loaded. -Policy status delivery uses a FIFO background worker. Retryable delivery -failures retain the ordered update and retry with capped exponential backoff; -terminal errors are logged and discarded. The outbox is nonblocking and does -not discard updates because of a fixed queue capacity, so status endpoint -outages cannot block policy polling, enforcement, settings, or provider -refreshes and cannot permanently lose the initial acknowledgement. +Policy status delivery uses a FIFO background worker independently from policy +reconciliation. Retryable delivery failures retain the ordered update; terminal +errors are logged and discarded. This prevents status endpoint outages from +blocking policy polling, enforcement, settings, or provider refreshes. Only sandbox-scoped revisions (`PolicySource::Sandbox`, version greater than zero) are acknowledged. Global policies and local-file development policies do diff --git a/crates/openshell-core/Cargo.toml b/crates/openshell-core/Cargo.toml index 35a3732cf9..597dc72d65 100644 --- a/crates/openshell-core/Cargo.toml +++ b/crates/openshell-core/Cargo.toml @@ -17,6 +17,7 @@ prost-types = { workspace = true } tonic = { workspace = true, features = ["channel", "tls-native-roots"] } tonic-prost = { workspace = true } tokio = { workspace = true } +tokio-stream = { workspace = true } thiserror = { workspace = true } miette = { workspace = true } serde = { workspace = true } diff --git a/crates/openshell-core/src/middleware.rs b/crates/openshell-core/src/middleware.rs index 75ed66003b..85909dbe66 100644 --- a/crates/openshell-core/src/middleware.rs +++ b/crates/openshell-core/src/middleware.rs @@ -1,16 +1,70 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Platform-wide supervisor middleware limits. +//! Supervisor middleware endpoint contract and platform-wide limits. +use std::pin::Pin; use std::time::Duration; +use tokio::sync::mpsc; +use tonic::{Request, Response, Status}; + +use crate::proto::{ + HttpRequestEvaluation, HttpRequestResult, MiddlewareManifest, ValidateConfigRequest, + ValidateConfigResponse, WebSocketSessionEvent, WebSocketSessionEventResult, +}; + +/// Transport-neutral response stream for one WebSocket middleware stage. +pub type WebSocketResponseStream = Pin< + Box< + dyn tokio_stream::Stream> + + Send + + 'static, + >, +>; + +/// A middleware implementation reachable either in-process or over a transport. +/// +/// Capability is defined by the implementation's manifest. The endpoint hides +/// whether invocations are direct calls or serialized gRPC requests. +#[tonic::async_trait] +pub trait SupervisorMiddlewareEndpoint: Send + Sync { + async fn describe(&self, request: Request<()>) -> Result, Status>; + + async fn validate_config( + &self, + request: Request, + ) -> Result, Status>; + + async fn evaluate_http_request( + &self, + request: Request, + ) -> Result, Status>; + + async fn open_websocket_session( + &self, + requests: mpsc::Receiver, + ) -> Result; +} + /// Default timeout for one supervisor middleware RPC. pub const DEFAULT_MIDDLEWARE_TIMEOUT: Duration = Duration::from_millis(500); /// Smallest operator-configured supervisor middleware RPC timeout. pub const MIN_MIDDLEWARE_TIMEOUT: Duration = Duration::from_millis(10); /// Largest operator-configured supervisor middleware RPC timeout. pub const MAX_MIDDLEWARE_TIMEOUT: Duration = Duration::from_secs(30); +/// Maximum time a complete message may spend in all middleware stages. +pub const MAX_MIDDLEWARE_CHAIN_TIMEOUT: Duration = Duration::from_secs(30); +/// Maximum time WebSocket preflight may delay an upstream handshake. +pub const MAX_MIDDLEWARE_PREFLIGHT_TIMEOUT: Duration = Duration::from_secs(1); +/// Process-wide safety valve for concurrently buffered middleware work. +pub const MAX_CONCURRENT_MIDDLEWARE_WORK: usize = 32; +/// Process-wide safety valve for retained streaming middleware sessions. +/// +/// A session consumes one permit regardless of its protocol or stage fan-out. +/// Persistent middleware protocols must acquire from this shared budget before +/// opening streams and retain the permit while any stage remains active. +pub const MAX_CONCURRENT_MIDDLEWARE_SESSIONS: usize = 32; /// Largest number of middleware configurations accepted in one sandbox policy. pub const MAX_MIDDLEWARE_CONFIGS: usize = 10; diff --git a/crates/openshell-sandbox/src/lib.rs b/crates/openshell-sandbox/src/lib.rs index 5ba41b9d7e..86059f56d5 100644 --- a/crates/openshell-sandbox/src/lib.rs +++ b/crates/openshell-sandbox/src/lib.rs @@ -3473,7 +3473,7 @@ filesystem_policy: let invalid_external = openshell_core::proto::SupervisorMiddlewareService { name: "unavailable-guard".into(), grpc_endpoint: "http://127.0.0.1:1".into(), - max_body_bytes: 1024, + max_payload_bytes: 1024, ..Default::default() }; connect_middleware_registry(&[invalid_external]) diff --git a/crates/openshell-server/src/config_file.rs b/crates/openshell-server/src/config_file.rs index c4e0cbc959..117ba712a0 100644 --- a/crates/openshell-server/src/config_file.rs +++ b/crates/openshell-server/src/config_file.rs @@ -189,8 +189,10 @@ pub struct MiddlewareServiceFileConfig { pub name: String, /// Plaintext gRPC endpoint reachable by the gateway and supervisors. pub grpc_endpoint: String, - /// Operator-owned body limit for every binding exposed by this service. - pub max_body_bytes: u64, + /// Operator-owned logical payload limit for every binding exposed by this + /// service, including HTTP bodies and complete WebSocket messages. + #[serde(alias = "max_body_bytes")] + pub max_payload_bytes: u64, /// Default RPC timeout using an integer with an `ms` or `s` suffix. #[serde(default)] pub timeout: Option, @@ -201,7 +203,7 @@ impl From<&MiddlewareServiceFileConfig> for SupervisorMiddlewareService { Self { name: config.name.clone(), grpc_endpoint: config.grpc_endpoint.clone(), - max_body_bytes: config.max_body_bytes, + max_payload_bytes: config.max_payload_bytes, timeout: config.timeout.clone().unwrap_or_default(), } } @@ -454,7 +456,7 @@ allow_unauthenticated_users = true [[openshell.supervisor.middleware]] name = "local-guard" grpc_endpoint = "http://127.0.0.1:50051" -max_body_bytes = 262144 +max_payload_bytes = 262144 timeout = "2s" "#; let tmp = write_tmp(toml); @@ -464,7 +466,7 @@ timeout = "2s" vec![MiddlewareServiceFileConfig { name: "local-guard".into(), grpc_endpoint: "http://127.0.0.1:50051".into(), - max_body_bytes: 262_144, + max_payload_bytes: 262_144, timeout: Some("2s".into()), }] ); @@ -473,6 +475,22 @@ timeout = "2s" assert_eq!(registration.timeout, "2s"); } + #[test] + fn parses_legacy_supervisor_middleware_payload_limit() { + let toml = r#" +[[openshell.supervisor.middleware]] +name = "local-guard" +grpc_endpoint = "http://127.0.0.1:50051" +max_body_bytes = 262144 +"#; + let tmp = write_tmp(toml); + let file = load(tmp.path()).expect("legacy middleware registration parses"); + assert_eq!( + file.openshell.supervisor.middleware[0].max_payload_bytes, + 262_144 + ); + } + #[test] fn parses_provider_profile_source_composition() { let toml = r#" diff --git a/crates/openshell-server/src/grpc/policy.rs b/crates/openshell-server/src/grpc/policy.rs index f78bbf413f..1a9c642646 100644 --- a/crates/openshell-server/src/grpc/policy.rs +++ b/crates/openshell-server/src/grpc/policy.rs @@ -11750,7 +11750,7 @@ mod tests { let service = openshell_core::proto::SupervisorMiddlewareService { name: "local-guard".into(), grpc_endpoint: "http://127.0.0.1:50051".into(), - max_body_bytes: 1024, + max_payload_bytes: 1024, ..Default::default() }; diff --git a/crates/openshell-supervisor-middleware-builtins/Cargo.toml b/crates/openshell-supervisor-middleware-builtins/Cargo.toml index f892c718fd..eef4b4db82 100644 --- a/crates/openshell-supervisor-middleware-builtins/Cargo.toml +++ b/crates/openshell-supervisor-middleware-builtins/Cargo.toml @@ -18,10 +18,9 @@ prost-types = { workspace = true } regex = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } -tonic = { workspace = true, features = ["server"] } - -[dev-dependencies] tokio = { workspace = true } +tokio-stream = { workspace = true } +tonic = { workspace = true, features = ["server"] } [lints] workspace = true diff --git a/crates/openshell-supervisor-middleware-builtins/src/lib.rs b/crates/openshell-supervisor-middleware-builtins/src/lib.rs index e23a228f05..285ebc45b4 100644 --- a/crates/openshell-supervisor-middleware-builtins/src/lib.rs +++ b/crates/openshell-supervisor-middleware-builtins/src/lib.rs @@ -8,21 +8,26 @@ mod regex; use std::sync::Arc; use miette::{Result, miette}; +use openshell_core::middleware::{SupervisorMiddlewareEndpoint, WebSocketResponseStream}; use openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware; use openshell_core::proto::{ - HttpRequestEvaluation, HttpRequestResult, MiddlewareManifest, ValidateConfigRequest, - ValidateConfigResponse, + HttpRequestEvaluation, HttpRequestResult, MiddlewareManifest, SupervisorMiddlewarePhase, + ValidateConfigRequest, ValidateConfigResponse, WebSocketMessageType, WebSocketPreflightAction, + WebSocketPreflightDecision, WebSocketSessionEvent, WebSocketSessionEventResult, + web_socket_session_event, web_socket_session_event_result, }; +use tokio_stream::{Stream, StreamExt}; use tonic::{Request, Response, Status}; pub use regex::{NAME as BUILTIN_REGEX, RegexConfig, RegexMode}; /// Return the first-party services that the gateway and supervisor install. -pub fn services() -> Vec> { +pub fn services() -> Vec> { vec![Arc::new(BuiltinMiddlewareService)] } -/// Validate configuration for a first-party binding. +/// Resolve and validate a first-party config before the supervisor has selected +/// its service endpoint. pub fn validate_config(implementation: &str, config: &prost_types::Struct) -> Result<()> { match implementation { BUILTIN_REGEX => regex::validate_config(config), @@ -33,20 +38,147 @@ pub fn validate_config(implementation: &str, config: &prost_types::Struct) -> Re } fn evaluate_http_request(evaluation: &HttpRequestEvaluation) -> Result { - match evaluation.middleware_name.as_str() { - BUILTIN_REGEX => regex::evaluate_http_request(evaluation), - other => Err(miette!( - "middleware implementation '{other}' is not a registered OpenShell built-in" - )), - } + regex::evaluate_http_request(evaluation) } /// Built-in regex service exposed through the standard middleware contract. #[derive(Debug, Default)] pub struct BuiltinMiddlewareService; +impl BuiltinMiddlewareService { + fn websocket_stream(mut requests: S) -> WebSocketResponseStream + where + S: Stream> + + Send + + Unpin + + 'static, + { + let (responses_tx, responses_rx) = tokio::sync::mpsc::channel(4); + tokio::spawn(async move { + let mut config = None; + let mut started = false; + let mut sequence_lower_bound = Some(1u64); + + while let Some(request) = requests.next().await { + let request = match request { + Ok(request) => request, + Err(error) => { + let _ = responses_tx.send(Err(error)).await; + break; + } + }; + let response = match request.event { + Some(web_socket_session_event::Event::Preflight(preflight)) + if config.is_none() && !started => + { + if preflight.phase == SupervisorMiddlewarePhase::PreCredentials as i32 { + let selected_config = preflight.config.unwrap_or_default(); + match regex::validate_config(&selected_config) { + Ok(()) => { + config = Some(selected_config); + Ok(Some(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::PreflightDecision( + WebSocketPreflightDecision { + action: WebSocketPreflightAction::Inspect + as i32, + ..Default::default() + }, + ), + ), + })) + } + Err(error) => Err(Status::invalid_argument(error.to_string())), + } + } else { + Err(Status::invalid_argument( + "unsupported built-in WebSocket binding", + )) + } + } + Some(web_socket_session_event::Event::SessionStart(_)) + if config.is_some() && !started => + { + started = true; + Ok(None) + } + Some(web_socket_session_event::Event::Message(message)) if started => { + if let Err(error) = advance_sequence_lower_bound( + &mut sequence_lower_bound, + message.sequence, + ) { + Err(error) + } else if message.message_type != WebSocketMessageType::Text as i32 { + Err(Status::invalid_argument( + "openshell/regex supports only client-to-upstream WebSocket text messages", + )) + } else { + let selected_config = + config.as_ref().expect("started stream has config"); + match regex::evaluate_websocket_text( + message.sequence, + &message.payload, + selected_config, + ) { + Ok(result) => Ok(Some(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::MessageResult( + result, + ), + ), + })), + Err(error) => Err(Status::invalid_argument(error.to_string())), + } + } + } + Some(web_socket_session_event::Event::SessionEnd(_)) if config.is_some() => { + break; + } + _ => Err(Status::failed_precondition( + "invalid built-in WebSocket session lifecycle", + )), + }; + + match response { + Ok(Some(response)) => { + if responses_tx.send(Ok(response)).await.is_err() { + break; + } + } + Ok(None) => {} + Err(error) => { + let _ = responses_tx.send(Err(error)).await; + break; + } + } + } + }); + Box::pin(tokio_stream::wrappers::ReceiverStream::new(responses_rx)) + } +} + +fn advance_sequence_lower_bound( + lower_bound: &mut Option, + sequence: u64, +) -> std::result::Result<(), Status> { + let Some(current_lower_bound) = *lower_bound else { + return Err(Status::invalid_argument( + "WebSocket message sequence must be strictly increasing", + )); + }; + if sequence < current_lower_bound { + return Err(Status::invalid_argument( + "WebSocket message sequence must be strictly increasing", + )); + } + *lower_bound = sequence.checked_add(1); + Ok(()) +} + #[tonic::async_trait] impl SupervisorMiddleware for BuiltinMiddlewareService { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + async fn describe( &self, _request: Request<()>, @@ -54,7 +186,7 @@ impl SupervisorMiddleware for BuiltinMiddlewareService { Ok(Response::new(MiddlewareManifest { name: BUILTIN_REGEX.into(), service_version: env!("CARGO_PKG_VERSION").into(), - bindings: vec![regex::describe()], + bindings: regex::describe(), })) } @@ -64,18 +196,16 @@ impl SupervisorMiddleware for BuiltinMiddlewareService { ) -> Result, Status> { let request = request.into_inner(); let config = request.config.unwrap_or_default(); - Ok(Response::new( - match validate_config(&request.middleware_name, &config) { - Ok(()) => ValidateConfigResponse { - valid: true, - reason: String::new(), - }, - Err(error) => ValidateConfigResponse { - valid: false, - reason: error.to_string(), - }, + Ok(Response::new(match regex::validate_config(&config) { + Ok(()) => ValidateConfigResponse { + valid: true, + reason: String::new(), }, - )) + Err(error) => ValidateConfigResponse { + valid: false, + reason: error.to_string(), + }, + })) } async fn evaluate_http_request( @@ -86,13 +216,50 @@ impl SupervisorMiddleware for BuiltinMiddlewareService { .map(Response::new) .map_err(|error| Status::invalid_argument(error.to_string())) } + + async fn evaluate_web_socket_session( + &self, + request: Request>, + ) -> Result, Status> { + Ok(Response::new(Self::websocket_stream(request.into_inner()))) + } +} + +#[tonic::async_trait] +impl SupervisorMiddlewareEndpoint for BuiltinMiddlewareService { + async fn describe(&self, request: Request<()>) -> Result, Status> { + SupervisorMiddleware::describe(self, request).await + } + + async fn validate_config( + &self, + request: Request, + ) -> Result, Status> { + SupervisorMiddleware::validate_config(self, request).await + } + + async fn evaluate_http_request( + &self, + request: Request, + ) -> Result, Status> { + SupervisorMiddleware::evaluate_http_request(self, request).await + } + + async fn open_websocket_session( + &self, + receiver: tokio::sync::mpsc::Receiver, + ) -> Result { + Ok(Self::websocket_stream( + tokio_stream::wrappers::ReceiverStream::new(receiver).map(Ok), + )) + } } #[cfg(test)] mod tests { use super::*; use openshell_core::proto::{ - Decision, SupervisorMiddlewareOperation, SupervisorMiddlewarePhase, + Decision, SupervisorMiddlewareOperation, SupervisorMiddlewarePhase, WebSocketPreflight, }; fn string_config(key: &str, value: &str) -> prost_types::Struct { @@ -109,12 +276,11 @@ mod tests { #[tokio::test] async fn service_describes_regex_binding() { - let manifest = BuiltinMiddlewareService - .describe(Request::new(())) + let manifest = SupervisorMiddleware::describe(&BuiltinMiddlewareService, Request::new(())) .await .expect("describe") .into_inner(); - assert_eq!(manifest.bindings.len(), 1); + assert_eq!(manifest.bindings.len(), 2); assert_eq!( manifest.bindings[0].operation, SupervisorMiddlewareOperation::HttpRequest as i32 @@ -123,7 +289,16 @@ mod tests { manifest.bindings[0].phase, SupervisorMiddlewarePhase::PreCredentials as i32 ); - assert_eq!(manifest.bindings[0].max_body_bytes, 256 * 1024); + assert_eq!(manifest.bindings[0].max_payload_bytes, 256 * 1024); + assert_eq!( + manifest.bindings[1].operation, + SupervisorMiddlewareOperation::WebsocketMessage as i32 + ); + assert_eq!( + manifest.bindings[1].phase, + SupervisorMiddlewarePhase::PreCredentials as i32 + ); + assert_eq!(manifest.bindings[1].max_payload_bytes, 256 * 1024); } #[test] @@ -160,7 +335,7 @@ mod tests { #[test] fn regex_replacement_evaluates_through_binding() { let result = evaluate_http_request(&HttpRequestEvaluation { - middleware_name: BUILTIN_REGEX.into(), + middleware_name: "operator-assigned-name".into(), body: br#"{"password":"top-secret","token":"sk-ABCDEFGHIJKLMNOP"}"#.to_vec(), config: Some(prost_types::Struct::default()), ..Default::default() @@ -180,6 +355,49 @@ mod tests { ); } + #[tokio::test] + async fn service_config_validation_does_not_dispatch_on_registration_name() { + let response = SupervisorMiddleware::validate_config( + &BuiltinMiddlewareService, + Request::new(ValidateConfigRequest { + config: Some(prost_types::Struct::default()), + middleware_name: "operator-assigned-name".into(), + }), + ) + .await + .expect("validate config") + .into_inner(); + + assert!(response.valid); + } + + #[tokio::test] + async fn websocket_preflight_does_not_dispatch_on_registration_name() { + let requests = tokio_stream::iter([Ok::<_, Status>(WebSocketSessionEvent { + event: Some(web_socket_session_event::Event::Preflight( + WebSocketPreflight { + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + middleware_name: "operator-assigned-name".into(), + config: Some(prost_types::Struct::default()), + ..Default::default() + }, + )), + })]); + let mut responses = BuiltinMiddlewareService::websocket_stream(requests); + let response = responses + .next() + .await + .expect("preflight response") + .expect("valid preflight"); + + assert!(matches!( + response.result, + Some(web_socket_session_event_result::Result::PreflightDecision( + WebSocketPreflightDecision { action, .. } + )) if action == WebSocketPreflightAction::Inspect as i32 + )); + } + #[test] fn regex_replacement_does_not_parse_keyword_assignments() { let body = concat!( @@ -199,4 +417,73 @@ mod tests { assert_eq!(result.body, body.as_bytes()); assert!(result.findings.is_empty()); } + + #[test] + fn regex_websocket_text_reuses_findings_and_metadata_semantics() { + let payload = + br#"{"type":"response.create","input":"sk-ABCDEFGHIJKLMNOP sk-QRSTUVWXYZabcdef"}"#; + let result = regex::evaluate_websocket_text(7, payload, &prost_types::Struct::default()) + .expect("evaluate WebSocket text"); + + assert_eq!(result.sequence, 7); + assert_eq!(result.decision, Decision::Allow as i32); + assert!(result.has_replacement); + assert_eq!( + String::from_utf8(result.replacement).expect("replacement UTF-8"), + r#"{"type":"response.create","input":"[REDACTED] [REDACTED]"}"# + ); + assert_eq!(result.findings.len(), 1); + assert_eq!(result.findings[0].r#type, "regex.openai"); + assert_eq!(result.findings[0].count, 2); + assert_eq!( + result.metadata.get("regex_matches_replaced"), + Some(&"2".to_string()) + ); + } + + #[test] + fn regex_websocket_no_match_returns_no_replacement_or_findings() { + let result = regex::evaluate_websocket_text( + 1, + br#"{"type":"response.create","input":"public"}"#, + &prost_types::Struct::default(), + ) + .expect("evaluate WebSocket text"); + + assert!(!result.has_replacement); + assert!(result.replacement.is_empty()); + assert!(result.findings.is_empty()); + assert!(result.metadata.is_empty()); + } + + #[test] + fn regex_websocket_rejects_invalid_utf8_and_oversize_messages() { + assert!( + regex::evaluate_websocket_text(1, &[0xff], &prost_types::Struct::default()).is_err() + ); + assert!( + regex::evaluate_websocket_text( + 1, + &vec![b'a'; 256 * 1024 + 1], + &prost_types::Struct::default(), + ) + .is_err() + ); + } + + #[test] + fn websocket_sequence_lower_bound_accepts_gaps_and_rejects_reuse() { + let mut lower_bound = Some(1); + advance_sequence_lower_bound(&mut lower_bound, 2).expect("first delivered sequence"); + assert_eq!(lower_bound, Some(3)); + assert!(advance_sequence_lower_bound(&mut lower_bound, 2).is_err()); + assert!(advance_sequence_lower_bound(&mut lower_bound, 1).is_err()); + + advance_sequence_lower_bound(&mut lower_bound, 7).expect("forward gap"); + assert_eq!(lower_bound, Some(8)); + + advance_sequence_lower_bound(&mut lower_bound, u64::MAX).expect("last sequence"); + assert_eq!(lower_bound, None); + assert!(advance_sequence_lower_bound(&mut lower_bound, u64::MAX).is_err()); + } } diff --git a/crates/openshell-supervisor-middleware-builtins/src/regex.rs b/crates/openshell-supervisor-middleware-builtins/src/regex.rs index 34e727430a..f80d7b1087 100644 --- a/crates/openshell-supervisor-middleware-builtins/src/regex.rs +++ b/crates/openshell-supervisor-middleware-builtins/src/regex.rs @@ -2,7 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 //! Example built-in middleware that applies a fixed set of regular-expression -//! replacements to UTF-8 request bodies. +//! replacements to UTF-8 HTTP request bodies and WebSocket text messages. //! //! This is intentionally a best-effort text transformation, not a secret //! scanner or a parser-aware redactor. It provides no guarantee that sensitive @@ -14,13 +14,13 @@ use std::sync::LazyLock; use miette::{Result, miette}; use openshell_core::proto::{ Decision, Finding, HttpRequestEvaluation, HttpRequestResult, MiddlewareBinding, - SupervisorMiddlewareOperation, SupervisorMiddlewarePhase, + SupervisorMiddlewareOperation, SupervisorMiddlewarePhase, WebSocketMessageResult, }; use regex::Regex; use serde::Deserialize; pub const NAME: &str = "openshell/regex"; -const MAX_BODY_BYTES: u64 = 256 * 1024; +const MAX_PAYLOAD_BYTES: u64 = 256 * 1024; #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Deserialize)] #[serde(default, deny_unknown_fields)] @@ -46,13 +46,21 @@ impl RegexConfig { } } -pub fn describe() -> MiddlewareBinding { - MiddlewareBinding { - operation: SupervisorMiddlewareOperation::HttpRequest as i32, - phase: SupervisorMiddlewarePhase::PreCredentials as i32, - max_body_bytes: MAX_BODY_BYTES, - timeout: String::new(), - } +pub fn describe() -> Vec { + vec![ + MiddlewareBinding { + operation: SupervisorMiddlewareOperation::HttpRequest as i32, + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + max_payload_bytes: MAX_PAYLOAD_BYTES, + timeout: String::new(), + }, + MiddlewareBinding { + operation: SupervisorMiddlewareOperation::WebsocketMessage as i32, + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + max_payload_bytes: MAX_PAYLOAD_BYTES, + timeout: String::new(), + }, + ] } struct ReplacementPattern { @@ -85,34 +93,75 @@ pub fn evaluate_http_request(evaluation: &HttpRequestEvaluation) -> Result Result { + validate_config(config)?; + let payload_bytes = u64::try_from(payload.len()) + .map_err(|_| miette!("{NAME} WebSocket text message length is not representable"))?; + if payload_bytes > MAX_PAYLOAD_BYTES { + return Err(miette!( + "{NAME} WebSocket text message exceeds {MAX_PAYLOAD_BYTES} bytes" + )); + } + let text = std::str::from_utf8(payload) + .map_err(|_| miette!("{NAME} requires UTF-8 WebSocket text messages"))?; + let (replacement, matches) = apply_replacements(text); + let (findings, metadata) = findings_and_metadata(&matches); + let has_replacement = !matches.is_empty(); + Ok(WebSocketMessageResult { + sequence, + decision: Decision::Allow as i32, + replacement: if has_replacement { + replacement.into_bytes() + } else { + Vec::new() + }, + has_replacement, + reason: String::new(), + findings, + metadata, + reason_code: String::new(), + }) +} + +fn findings_and_metadata( + matches: &[(&'static str, u32)], +) -> (Vec, HashMap) { + let findings = matches + .iter() + .map(|(kind, count)| Finding { r#type: format!("regex.{kind}"), label: format!("{kind} regex match"), count: *count, confidence: "medium".into(), severity: "medium".into(), - }); - } + }) + .collect(); + let mut metadata = HashMap::new(); if !matches.is_empty() { - result - .metadata - .insert("regex_matches_replaced".into(), total.to_string()); + let total = matches + .iter() + .fold(0u32, |acc, (_, count)| acc.saturating_add(*count)); + metadata.insert("regex_matches_replaced".into(), total.to_string()); } - Ok(result) + (findings, metadata) } fn apply_replacements(input: &str) -> (String, Vec<(&'static str, u32)>) { diff --git a/crates/openshell-supervisor-middleware/Cargo.toml b/crates/openshell-supervisor-middleware/Cargo.toml index 9cdc53febb..8307883550 100644 --- a/crates/openshell-supervisor-middleware/Cargo.toml +++ b/crates/openshell-supervisor-middleware/Cargo.toml @@ -14,9 +14,11 @@ rust-version.workspace = true openshell-core = { path = "../openshell-core", default-features = false } miette = { workspace = true } +futures = { workspace = true } prost = { workspace = true } prost-types = { workspace = true } tokio = { workspace = true } +tokio-stream = { workspace = true } tonic = { workspace = true, features = ["channel", "server", "tls-native-roots"] } [dev-dependencies] diff --git a/crates/openshell-supervisor-middleware/src/lib.rs b/crates/openshell-supervisor-middleware/src/lib.rs index fe0f15f0a6..0d97e864b8 100644 --- a/crates/openshell-supervisor-middleware/src/lib.rs +++ b/crates/openshell-supervisor-middleware/src/lib.rs @@ -5,6 +5,13 @@ mod headers; mod remote; +mod websocket; + +pub use websocket::{ + WebSocketCoverage, WebSocketCoverageState, WebSocketInvocation, WebSocketInvocationOutcome, + WebSocketMessageAdmission, WebSocketMessageOutcome, WebSocketPreflightInput, + WebSocketPreflightResult, WebSocketSession, WebSocketSessionStartOutcome, +}; #[cfg(test)] use std::collections::HashMap; @@ -16,6 +23,7 @@ use std::time::Duration; use miette::{Result, miette}; use prost::Message; +use openshell_core::middleware::SupervisorMiddlewareEndpoint; use openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware; use openshell_core::proto::{ Decision, Finding, HeaderMutation, HttpHeader, HttpRequestEvaluation, HttpRequestTarget, @@ -23,18 +31,138 @@ use openshell_core::proto::{ SupervisorMiddlewareOperation, SupervisorMiddlewarePhase, SupervisorMiddlewareService, ValidateConfigRequest, }; -use tokio::sync::OnceCell; -use tonic::Request; +use tokio::sync::{OnceCell, OwnedSemaphorePermit, Semaphore}; +use tonic::{Request, Response as TonicResponse, Status as TonicStatus}; + +pub use openshell_core::middleware::WebSocketResponseStream; +pub type MiddlewareService = + dyn SupervisorMiddleware; +pub type MiddlewareServiceEndpoint = dyn SupervisorMiddlewareEndpoint; + +struct GeneratedMiddlewareEndpoint { + service: Arc, +} + +#[tonic::async_trait] +impl SupervisorMiddlewareEndpoint for GeneratedMiddlewareEndpoint { + async fn describe( + &self, + request: Request<()>, + ) -> std::result::Result, TonicStatus> { + self.service.describe(request).await + } + + async fn validate_config( + &self, + request: Request, + ) -> std::result::Result< + TonicResponse, + TonicStatus, + > { + self.service.validate_config(request).await + } + + async fn evaluate_http_request( + &self, + request: Request, + ) -> std::result::Result, TonicStatus> + { + self.service.evaluate_http_request(request).await + } + + async fn open_websocket_session( + &self, + _receiver: tokio::sync::mpsc::Receiver, + ) -> std::result::Result { + Err(TonicStatus::unimplemented( + "middleware service does not expose an in-process WebSocket stream", + )) + } +} + +/// Adapt a generated in-process service that only implements unary operations. +/// +/// Services that advertise WebSocket support must implement +/// [`SupervisorMiddlewareEndpoint`] directly. +pub fn http_only_endpoint(service: Arc) -> Arc { + Arc::new(GeneratedMiddlewareEndpoint { service }) +} + +/// Maximum short-lived middleware work items allowed to wait for active +/// capacity. +/// +/// Waiters do not buffer request or message bodies, so the queue can absorb a +/// larger burst without increasing the active payload-memory bound. +pub const MAX_QUEUED_MIDDLEWARE_WORK: usize = MAX_CONCURRENT_MIDDLEWARE_WORK * 2; + +/// One slot in the shared middleware work budget. +/// +/// Callers that buffer request or message bodies acquire this guard first and +/// retain it through evaluation, bounding aggregate buffered middleware input. +#[derive(Debug)] +pub struct MiddlewareWorkAdmission { + _work: OwnedSemaphorePermit, + saturated: bool, +} + +impl MiddlewareWorkAdmission { + pub fn saturated(&self) -> bool { + self.saturated + } +} + +/// Result of attempting to enter the bounded middleware work queue. +/// +/// Active-capacity saturation is ordinary backpressure: callers that obtain a +/// waiter slot eventually receive [`Self::Admitted`]. [`Self::QueueExhausted`] +/// is immediate load shedding after both active capacity and the waiter queue +/// are full. +#[derive(Debug)] +pub enum MiddlewareWorkAdmissionOutcome { + Admitted(MiddlewareWorkAdmission), + QueueExhausted, +} + +impl MiddlewareWorkAdmissionOutcome { + /// Preserve the existing failure behavior for protocols whose outer layer + /// already translates middleware admission errors into a stable response + /// or typed termination. + pub fn into_admission(self) -> Result { + match self { + Self::Admitted(admission) => Ok(admission), + Self::QueueExhausted => Err(miette!( + "middleware admission queue is full; refusing additional buffered work" + )), + } + } +} + +/// One slot in the shared persistent middleware session budget. +/// +/// Protocol-specific session runners retain this guard while at least one +/// streaming stage remains active. Registry replacement preserves the shared +/// admission state so future streaming HTTP middleware can use the same +/// process-wide bound. +#[derive(Debug)] +struct MiddlewareSessionPermit { + _session: OwnedSemaphorePermit, +} + +enum MiddlewareSessionAdmission { + Admitted(MiddlewareSessionPermit), + AtCapacity, +} pub use openshell_core::middleware::{ - DEFAULT_MIDDLEWARE_TIMEOUT, MAX_MIDDLEWARE_CHAIN_FINDINGS, MAX_MIDDLEWARE_CHAIN_STAGES, - MAX_MIDDLEWARE_CONFIGS, MAX_MIDDLEWARE_FINDINGS_PER_STAGE, MAX_MIDDLEWARE_SELECTOR_PATTERNS, - MAX_MIDDLEWARE_TIMEOUT, MIN_MIDDLEWARE_TIMEOUT, middleware_timeout_or_default, - parse_middleware_timeout, + DEFAULT_MIDDLEWARE_TIMEOUT, MAX_CONCURRENT_MIDDLEWARE_SESSIONS, MAX_CONCURRENT_MIDDLEWARE_WORK, + MAX_MIDDLEWARE_CHAIN_FINDINGS, MAX_MIDDLEWARE_CHAIN_STAGES, MAX_MIDDLEWARE_CHAIN_TIMEOUT, + MAX_MIDDLEWARE_CONFIGS, MAX_MIDDLEWARE_FINDINGS_PER_STAGE, MAX_MIDDLEWARE_PREFLIGHT_TIMEOUT, + MAX_MIDDLEWARE_SELECTOR_PATTERNS, MAX_MIDDLEWARE_TIMEOUT, MIN_MIDDLEWARE_TIMEOUT, + middleware_timeout_or_default, parse_middleware_timeout, }; -/// Largest request or replacement body accepted by the middleware platform. -pub const MAX_MIDDLEWARE_BODY_BYTES: usize = 4 * 1024 * 1024; +/// Largest logical payload or replacement accepted by the middleware platform. +pub const MAX_MIDDLEWARE_PAYLOAD_BYTES: usize = 4 * 1024 * 1024; /// Largest encoded service-specific configuration attached to one evaluation. pub const MAX_MIDDLEWARE_CONFIG_BYTES: usize = 64 * 1024; /// Largest encoded request identity context attached to one evaluation. @@ -69,22 +197,24 @@ const MAX_MIDDLEWARE_RESPONSE_ENVELOPE_BYTES: usize = MAX_MIDDLEWARE_REASON_BYTE + MAX_MIDDLEWARE_FINDINGS_PER_STAGE * MAX_MIDDLEWARE_FINDING_BYTES + MAX_MIDDLEWARE_METADATA_BYTES + MAX_MIDDLEWARE_PROTOBUF_OVERHEAD_BYTES; -/// gRPC envelope headroom derived from every bounded non-body component. +/// gRPC envelope headroom derived from every bounded non-payload component. pub const MIDDLEWARE_GRPC_ENVELOPE_BYTES: usize = if MAX_MIDDLEWARE_REQUEST_ENVELOPE_BYTES > MAX_MIDDLEWARE_RESPONSE_ENVELOPE_BYTES { MAX_MIDDLEWARE_REQUEST_ENVELOPE_BYTES } else { MAX_MIDDLEWARE_RESPONSE_ENVELOPE_BYTES }; -/// gRPC message limit derived from the body and bounded protobuf components. +/// gRPC message limit derived from the payload and bounded protobuf components. pub const MIDDLEWARE_GRPC_MESSAGE_BYTES: usize = - MAX_MIDDLEWARE_BODY_BYTES + MIDDLEWARE_GRPC_ENVELOPE_BYTES; + MAX_MIDDLEWARE_PAYLOAD_BYTES + MIDDLEWARE_GRPC_ENVELOPE_BYTES; +const MAX_STABLE_IDENTIFIER_BYTES: usize = 128; +const EXTERNAL_FINDING_LABEL: &str = "External middleware finding"; +#[cfg(test)] const HTTP_REQUEST_OPERATION: SupervisorMiddlewareOperation = SupervisorMiddlewareOperation::HttpRequest; +#[cfg(test)] const PRE_CREDENTIALS_PHASE: SupervisorMiddlewarePhase = SupervisorMiddlewarePhase::PreCredentials; -const MAX_STABLE_IDENTIFIER_BYTES: usize = 128; -const EXTERNAL_FINDING_LABEL: &str = "External middleware finding"; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum OnError { FailClosed, @@ -136,20 +266,28 @@ impl TryFrom<(&str, &NetworkMiddlewareConfig)> for ChainEntry { } /// A policy-selected middleware config joined with metadata reported by its -/// service's `Describe` call. A missing binding is retained so `on_error` can -/// decide whether the request fails open or closed. +/// service's `Describe` call. +/// +/// An unregistered implementation is retained so `on_error` can decide whether +/// the request fails open or closed. A registered implementation without the +/// requested binding is not part of this chain. #[derive(Clone)] pub struct DescribedChainEntry { entry: ChainEntry, service: Option>, binding: Option, - max_body_bytes: usize, + max_payload_bytes: usize, timeout: Duration, } +struct DescribedChain { + entries: Vec, + unbound: Vec, +} + impl DescribedChainEntry { - pub fn max_body_bytes(&self) -> usize { - self.max_body_bytes + pub fn max_payload_bytes(&self) -> usize { + self.max_payload_bytes } pub fn on_error(&self) -> OnError { @@ -163,7 +301,7 @@ impl DescribedChainEntry { /// True when this entry resolved to a registered binding and will be /// evaluated. When false, the binding is absent from the current registry /// and the entry is handled entirely by its `on_error` policy, so it - /// imposes no body-buffering limit on the chain. + /// imposes no payload-buffering limit on the chain. pub fn is_resolved(&self) -> bool { self.binding.is_some() } @@ -303,10 +441,10 @@ struct MiddlewareServiceState { /// single-service test constructor leaves this empty and uses the manifest /// name after Describe. attachment_name: Option, - service: Arc, + endpoint: Arc, manifest: OnceCell, diagnostic_policy: MiddlewareDiagnosticPolicy, - operator_max_body_bytes: Option, + operator_max_payload_bytes: Option, operator_timeout: Duration, } @@ -375,6 +513,9 @@ pub struct MiddlewareRegistry { services: Arc>>, registered_services: Arc>, middleware_names: Arc>, + work_admission: Arc, + work_admission_waiters: Arc, + session_admission: Arc, } impl std::fmt::Debug for MiddlewareRegistry { @@ -384,6 +525,14 @@ impl std::fmt::Debug for MiddlewareRegistry { .field("service_count", &self.services.len()) .field("registered_service_count", &self.registered_services.len()) .field("middleware_count", &self.middleware_names.len()) + .field( + "available_work_permits", + &self.work_admission.available_permits(), + ) + .field( + "available_session_permits", + &self.session_admission.available_permits(), + ) .finish() } } @@ -399,6 +548,9 @@ impl Default for MiddlewareRegistry { services: Arc::new(Vec::new()), registered_services: Arc::new(Vec::new()), middleware_names: Arc::new(HashSet::new()), + work_admission: Arc::new(Semaphore::new(MAX_CONCURRENT_MIDDLEWARE_WORK)), + work_admission_waiters: Arc::new(Semaphore::new(MAX_QUEUED_MIDDLEWARE_WORK)), + session_admission: Arc::new(Semaphore::new(MAX_CONCURRENT_MIDDLEWARE_SESSIONS)), } } } @@ -423,15 +575,9 @@ fn validate_registration(registration: &SupervisorMiddlewareService) -> Result MAX_MIDDLEWARE_BODY_BYTES as u64 { + if registration.max_payload_bytes > MAX_MIDDLEWARE_PAYLOAD_BYTES as u64 { return Err(miette!( - "middleware registration '{}' max_body_bytes exceeds the platform maximum of {MAX_MIDDLEWARE_BODY_BYTES}", + "middleware registration '{}' max_payload_bytes exceeds the platform maximum of {MAX_MIDDLEWARE_PAYLOAD_BYTES}", registration.name )); } @@ -477,23 +623,54 @@ fn middleware_denial_reason(config_name: &str, reason_code: Option<&str>) -> Str ) } -fn validate_body_limit(source: &str, binding: &MiddlewareBinding) -> Result { - if binding.max_body_bytes == 0 { - return Err(miette!("{source} must advertise a non-zero body limit")); +fn validate_payload_limit(source: &str, binding: &MiddlewareBinding) -> Result { + if binding.max_payload_bytes == 0 { + return Err(miette!("{source} must advertise a non-zero payload limit")); } - if binding.max_body_bytes > MAX_MIDDLEWARE_BODY_BYTES as u64 { + if binding.max_payload_bytes > MAX_MIDDLEWARE_PAYLOAD_BYTES as u64 { return Err(miette!( - "{source} body limit exceeds the platform maximum of {MAX_MIDDLEWARE_BODY_BYTES}" + "{source} payload limit exceeds the platform maximum of {MAX_MIDDLEWARE_PAYLOAD_BYTES}" )); } - usize::try_from(binding.max_body_bytes) - .map_err(|_| miette!("{source} reports a body limit too large for this platform")) + usize::try_from(binding.max_payload_bytes) + .map_err(|_| miette!("{source} reports a payload limit too large for this platform")) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SupportedBinding { + HttpPreCredentials, + WebSocketPreCredentials, +} + +fn supported_binding(source: &str, binding: &MiddlewareBinding) -> Result { + match ( + SupervisorMiddlewareOperation::try_from(binding.operation).ok(), + SupervisorMiddlewarePhase::try_from(binding.phase).ok(), + ) { + ( + Some(SupervisorMiddlewareOperation::HttpRequest), + Some(SupervisorMiddlewarePhase::PreCredentials), + ) => Ok(SupportedBinding::HttpPreCredentials), + ( + Some(SupervisorMiddlewareOperation::WebsocketMessage), + Some(SupervisorMiddlewarePhase::PreCredentials), + ) => Ok(SupportedBinding::WebSocketPreCredentials), + ( + Some(SupervisorMiddlewareOperation::WebsocketMessage), + Some(SupervisorMiddlewarePhase::PreReturn), + ) => Err(miette!( + "{source} advertises WEBSOCKET_MESSAGE/PRE_RETURN, which is reserved for PR 2" + )), + _ => Err(miette!( + "{source} advertises an unsupported middleware operation/phase pair" + )), + } } fn validate_manifest_bindings( source: &str, manifest: &MiddlewareManifest, - operator_max_body_bytes: Option, + operator_max_payload_bytes: Option, ) -> Result<()> { if manifest.bindings.is_empty() { return Err(miette!("{source} describes no bindings")); @@ -501,27 +678,26 @@ fn validate_manifest_bindings( let mut described_pairs = HashSet::with_capacity(manifest.bindings.len()); for binding in &manifest.bindings { - if binding.operation != HTTP_REQUEST_OPERATION as i32 - || binding.phase != PRE_CREDENTIALS_PHASE as i32 - { - return Err(miette!( - "{source} must support HTTP_REQUEST/PRE_CREDENTIALS" - )); - } + supported_binding(source, binding)?; if !described_pairs.insert((binding.operation, binding.phase)) { return Err(miette!( - "{source} describes more than one binding for HTTP_REQUEST/PRE_CREDENTIALS" + "{source} describes a duplicate middleware operation/phase pair" )); } - let advertised = validate_body_limit(source, binding)?; + let advertised = validate_payload_limit(source, binding)?; if !binding.timeout.trim().is_empty() { parse_middleware_timeout(&binding.timeout) .map_err(|reason| miette!("{source} has invalid timeout for binding: {reason}"))?; } - if operator_max_body_bytes.is_some_and(|limit| limit > advertised) { + if operator_max_payload_bytes.is_some_and(|limit| limit > advertised) { + return Err(miette!( + "{source} max_payload_bytes ({}) exceeds the binding capability ({advertised})", + operator_max_payload_bytes.expect("operator limit checked above") + )); + } + if operator_max_payload_bytes == Some(0) { return Err(miette!( - "{source} max_body_bytes ({}) exceeds the binding capability ({advertised})", - operator_max_body_bytes.expect("operator limit checked above") + "{source} must configure max_payload_bytes for every payload-bearing binding" )); } } @@ -531,12 +707,12 @@ fn validate_manifest_bindings( fn validate_external_manifest( registration: &SupervisorMiddlewareService, manifest: &MiddlewareManifest, - operator_max_body_bytes: usize, + operator_max_payload_bytes: Option, ) -> Result<()> { validate_manifest_bindings( &format!("external middleware registration '{}'", registration.name), manifest, - Some(operator_max_body_bytes), + operator_max_payload_bytes, ) } @@ -576,7 +752,7 @@ fn normalize_untrusted_diagnostics( fn validate_request_envelope( evaluation: &HttpRequestEvaluation, ) -> std::result::Result<(), &'static str> { - if evaluation.body.len() > MAX_MIDDLEWARE_BODY_BYTES { + if evaluation.body.len() > MAX_MIDDLEWARE_PAYLOAD_BYTES { return Err("request_body_over_capacity"); } if evaluation @@ -618,7 +794,7 @@ fn validate_request_envelope( fn validate_response_envelope( result: &openshell_core::proto::HttpRequestResult, ) -> std::result::Result<(), &'static str> { - if result.body.len() > MAX_MIDDLEWARE_BODY_BYTES { + if result.body.len() > MAX_MIDDLEWARE_PAYLOAD_BYTES { return Err("response_body_over_capacity"); } if result.reason.len() > MAX_MIDDLEWARE_REASON_BYTES { @@ -668,7 +844,7 @@ impl MiddlewareRegistry { /// Describe in-process services, then connect and validate every /// operator-provided service registration. pub async fn connect_services( - in_process_services: Vec>, + in_process_services: Vec>, registrations: Vec, ) -> Result { let mut services = Vec::with_capacity(in_process_services.len() + registrations.len()); @@ -713,10 +889,10 @@ impl MiddlewareRegistry { .map_err(|_| miette!("middleware manifest cache initialized twice"))?; services.push(Arc::new(MiddlewareServiceState { attachment_name: Some(attachment_name), - service, + endpoint: service, manifest: manifest_cell, diagnostic_policy: MiddlewareDiagnosticPolicy::Preserve, - operator_max_body_bytes: None, + operator_max_payload_bytes: None, operator_timeout: DEFAULT_MIDDLEWARE_TIMEOUT, })); } @@ -730,10 +906,10 @@ impl MiddlewareRegistry { )); } - let operator_max_body_bytes = - usize::try_from(registration.max_body_bytes).map_err(|_| { + let operator_max_payload_bytes = usize::try_from(registration.max_payload_bytes) + .map_err(|_| { miette!( - "middleware registration '{}' body limit is too large for this platform", + "middleware registration '{}' payload limit is too large for this platform", registration.name ) })?; @@ -758,17 +934,17 @@ impl MiddlewareRegistry { safe_reason(&error.to_string()) ) })?; - validate_external_manifest(®istration, &manifest, operator_max_body_bytes)?; + validate_external_manifest(®istration, &manifest, Some(operator_max_payload_bytes))?; let manifest_cell = OnceCell::new(); manifest_cell .set(manifest) .map_err(|_| miette!("middleware manifest cache initialized twice"))?; services.push(Arc::new(MiddlewareServiceState { attachment_name: Some(registration.name.clone()), - service, + endpoint: service, manifest: manifest_cell, diagnostic_policy: MiddlewareDiagnosticPolicy::Normalize, - operator_max_body_bytes: Some(operator_max_body_bytes), + operator_max_payload_bytes: Some(operator_max_payload_bytes), operator_timeout, })); registered_services.push(RegisteredMiddlewareService { registration }); @@ -778,6 +954,9 @@ impl MiddlewareRegistry { services: Arc::new(services), registered_services: Arc::new(registered_services), middleware_names: Arc::new(middleware_names), + work_admission: Arc::new(Semaphore::new(MAX_CONCURRENT_MIDDLEWARE_WORK)), + work_admission_waiters: Arc::new(Semaphore::new(MAX_QUEUED_MIDDLEWARE_WORK)), + session_admission: Arc::new(Semaphore::new(MAX_CONCURRENT_MIDDLEWARE_SESSIONS)), }) } @@ -846,19 +1025,26 @@ impl Default for ChainRunner { } impl ChainRunner { - pub fn new(service: Arc) -> Self { + pub fn new(service: Arc) -> Self { + Self::from_endpoint(http_only_endpoint(service)) + } + + pub fn from_endpoint(endpoint: Arc) -> Self { Self { registry: Arc::new(MiddlewareRegistry { services: Arc::new(vec![Arc::new(MiddlewareServiceState { attachment_name: None, - service, + endpoint, manifest: OnceCell::new(), diagnostic_policy: MiddlewareDiagnosticPolicy::Preserve, - operator_max_body_bytes: None, + operator_max_payload_bytes: None, operator_timeout: DEFAULT_MIDDLEWARE_TIMEOUT, })]), registered_services: Arc::new(Vec::new()), middleware_names: Arc::new(HashSet::new()), + work_admission: Arc::new(Semaphore::new(MAX_CONCURRENT_MIDDLEWARE_WORK)), + work_admission_waiters: Arc::new(Semaphore::new(MAX_QUEUED_MIDDLEWARE_WORK)), + session_admission: Arc::new(Semaphore::new(MAX_CONCURRENT_MIDDLEWARE_SESSIONS)), }), } } @@ -869,6 +1055,67 @@ impl ChainRunner { } } + /// Build a runner for a replacement registry while preserving process-wide + /// admission budgets across registry generations. + #[must_use] + pub fn with_replacement_registry(&self, mut registry: MiddlewareRegistry) -> Self { + registry.work_admission = Arc::clone(&self.registry.work_admission); + registry.work_admission_waiters = Arc::clone(&self.registry.work_admission_waiters); + registry.session_admission = Arc::clone(&self.registry.session_admission); + Self::from_registry(registry) + } + + /// Reserve one unit of short-lived middleware work. + /// + /// The bounded waiter queue provides backpressure for work expected to + /// complete promptly, such as HTTP evaluations, WebSocket messages, and + /// streaming-session preflight. + pub async fn reserve_middleware_work(&self) -> Result { + if let Ok(permit) = Arc::clone(&self.registry.work_admission).try_acquire_owned() { + Ok(MiddlewareWorkAdmissionOutcome::Admitted( + MiddlewareWorkAdmission { + _work: permit, + saturated: false, + }, + )) + } else { + let Ok(waiter) = Arc::clone(&self.registry.work_admission_waiters).try_acquire_owned() + else { + return Ok(MiddlewareWorkAdmissionOutcome::QueueExhausted); + }; + let permit = Arc::clone(&self.registry.work_admission) + .acquire_owned() + .await + .map_err(|_| miette!("middleware admission semaphore closed"))?; + drop(waiter); + Ok(MiddlewareWorkAdmissionOutcome::Admitted( + MiddlewareWorkAdmission { + _work: permit, + saturated: true, + }, + )) + } + } + + /// Reserve middleware work for a caller whose established external + /// behavior treats queue exhaustion as a middleware processing failure. + pub async fn reserve_middleware_work_admission(&self) -> Result { + self.reserve_middleware_work().await?.into_admission() + } + + /// Attempt to reserve one persistent middleware session without waiting. + /// + /// Long-lived sessions have no useful queueing bound because their release + /// time is unrelated to middleware latency. Protocol-specific runners apply + /// their own `on_error` semantics when the shared session budget is full. + fn try_reserve_middleware_session(&self) -> MiddlewareSessionAdmission { + Arc::clone(&self.registry.session_admission) + .try_acquire_owned() + .map_or(MiddlewareSessionAdmission::AtCapacity, |permit| { + MiddlewareSessionAdmission::Admitted(MiddlewareSessionPermit { _session: permit }) + }) + } + async fn manifests(&self) -> Result, MiddlewareManifest)>> { let mut manifests = Vec::with_capacity(self.registry.services.len()); for state in self.registry.services.iter() { @@ -878,7 +1125,7 @@ impl ChainRunner { call_with_timeout( state.operator_timeout, "Describe", - state.service.describe(Request::new(())), + state.endpoint.describe(Request::new(())), ) .await .map(tonic::Response::into_inner) @@ -905,62 +1152,89 @@ impl ChainRunner { .unwrap_or(manifest.name.as_str()) } - fn http_pre_credentials_binding(manifest: &MiddlewareManifest) -> Option<&MiddlewareBinding> { - manifest.bindings.iter().find(|binding| { - binding.operation == HTTP_REQUEST_OPERATION as i32 - && binding.phase == PRE_CREDENTIALS_PHASE as i32 - }) + fn binding( + manifest: &MiddlewareManifest, + operation: SupervisorMiddlewareOperation, + phase: SupervisorMiddlewarePhase, + ) -> Option<&MiddlewareBinding> { + manifest + .bindings + .iter() + .find(|binding| binding.operation == operation as i32 && binding.phase == phase as i32) } pub async fn describe_chain(&self, entries: &[ChainEntry]) -> Result> { + Ok(self + .describe_chain_for( + entries, + SupervisorMiddlewareOperation::HttpRequest, + SupervisorMiddlewarePhase::PreCredentials, + ) + .await? + .entries) + } + + pub async fn describe_websocket_chain( + &self, + entries: &[ChainEntry], + ) -> Result> { + Ok(self + .describe_chain_for( + entries, + SupervisorMiddlewareOperation::WebsocketMessage, + SupervisorMiddlewarePhase::PreCredentials, + ) + .await? + .entries) + } + + async fn describe_chain_for( + &self, + entries: &[ChainEntry], + operation: SupervisorMiddlewareOperation, + phase: SupervisorMiddlewarePhase, + ) -> Result { ensure_chain_capacity(entries.len())?; let manifests = self.manifests().await?; let mut entries = entries.to_vec(); sort_chain_entries(&mut entries); - entries - .iter() - .map(|entry| { - let described = manifests - .iter() - .find(|(state, manifest)| { - Self::attachment_name(state, manifest) == entry.implementation - }) - .and_then(|(state, manifest)| { - Self::http_pre_credentials_binding(manifest) - .cloned() - .map(|binding| (Arc::clone(state), binding)) - }); - let (service, binding) = described.map_or((None, None), |(service, binding)| { - (Some(service), Some(binding)) + let mut described_entries = Vec::with_capacity(entries.len()); + let mut unbound = Vec::new(); + for entry in entries { + let Some((state, manifest)) = manifests.iter().find(|(state, manifest)| { + Self::attachment_name(state, manifest) == entry.implementation + }) else { + described_entries.push(DescribedChainEntry { + entry, + service: None, + binding: None, + max_payload_bytes: 0, + timeout: DEFAULT_MIDDLEWARE_TIMEOUT, }); - let max_body_bytes = binding - .as_ref() - .map(|binding| { - let advertised = validate_body_limit("middleware manifest", binding)?; - Ok::<_, miette::Report>( - service - .as_ref() - .and_then(|state| state.operator_max_body_bytes) - .unwrap_or(advertised), - ) - }) - .transpose()? - .unwrap_or(0); - let timeout = service - .as_ref() - .zip(binding.as_ref()) - .map(|(state, binding)| state.timeout_for_binding(binding)) - .transpose()? - .unwrap_or(DEFAULT_MIDDLEWARE_TIMEOUT); - Ok(DescribedChainEntry { - entry: entry.clone(), - service, - binding, - max_body_bytes, - timeout, - }) - }) - .collect() + continue; + }; + let Some(binding) = Self::binding(manifest, operation, phase).cloned() else { + // The config remains globally ordered, but it does not + // participate in this exact operation/phase chain. + unbound.push(entry); + continue; + }; + let timeout = state.timeout_for_binding(&binding)?; + let advertised = validate_payload_limit("middleware manifest", &binding)?; + let max_payload_bytes = state.operator_max_payload_bytes.unwrap_or(advertised); + described_entries.push(DescribedChainEntry { + entry, + service: Some(Arc::clone(state)), + binding: Some(binding), + max_payload_bytes, + timeout, + }); + } + ensure_chain_capacity(described_entries.len())?; + Ok(DescribedChain { + entries: described_entries, + unbound, + }) } pub async fn validate_config( @@ -974,19 +1248,17 @@ impl ChainRunner { )); } let manifests = self.manifests().await?; - let Some((state, binding)) = manifests.iter().find_map(|(state, manifest)| { - (Self::attachment_name(state, manifest) == middleware_name) - .then(|| Self::http_pre_credentials_binding(manifest)) - .flatten() - .map(|binding| (state, binding)) - }) else { + let Some((state, _manifest)) = manifests + .iter() + .find(|(state, manifest)| Self::attachment_name(state, manifest) == middleware_name) + else { return Err(miette!("middleware '{middleware_name}' is not registered")); }; let response = call_with_timeout( - state.timeout_for_binding(binding)?, + state.operator_timeout, "ValidateConfig", state - .service + .endpoint .validate_config(Request::new(ValidateConfigRequest { config: Some(config), middleware_name: middleware_name.into(), @@ -1041,6 +1313,29 @@ impl ChainRunner { entries: &[DescribedChainEntry], input: HttpRequestInput, transformed_body_policy: TransformedBodyPolicy<'_>, + ) -> Result { + let admission = if entries.is_empty() { + None + } else { + Some(self.reserve_middleware_work_admission().await?) + }; + self.evaluate_described_with_policy_admitted( + entries, + input, + transformed_body_policy, + admission, + ) + .await + } + + /// Evaluate a chain using capacity reserved before its request body was + /// buffered. The guard is retained until the ordered chain completes. + pub async fn evaluate_described_with_policy_admitted( + &self, + entries: &[DescribedChainEntry], + input: HttpRequestInput, + transformed_body_policy: TransformedBodyPolicy<'_>, + admission: Option, ) -> Result { ensure_chain_capacity(entries.len())?; let mut headers = input.headers.clone(); @@ -1049,6 +1344,8 @@ impl ChainRunner { let mut findings = Vec::new(); let mut metadata = BTreeMap::new(); let mut applied = Vec::new(); + let _admission = admission; + let chain_deadline = tokio::time::Instant::now() + MAX_MIDDLEWARE_CHAIN_TIMEOUT; for entry in entries { let Some(binding) = entry.binding.as_ref() else { @@ -1068,7 +1365,7 @@ impl ChainRunner { } } }; - if body.len() > entry.max_body_bytes { + if body.len() > entry.max_payload_bytes { match apply_on_error(entry, "request_body_over_capacity", &mut applied) { OnErrorAction::FailOpen => continue, OnErrorAction::FailClosed(reason) => { @@ -1106,11 +1403,29 @@ impl ChainRunner { let Some(service) = entry.service.as_ref() else { unreachable!("described binding always has a service") }; + let remaining = chain_deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + match apply_on_error(entry, "middleware_chain_timeout", &mut applied) { + OnErrorAction::FailOpen => continue, + OnErrorAction::FailClosed(reason) => { + return Ok(ChainOutcome { + allowed: false, + reason, + body, + header_mutations, + findings, + metadata, + applied, + denial: None, + }); + } + } + } let mut result = match call_with_timeout( - entry.timeout, + entry.timeout.min(remaining), "EvaluateHttpRequest", service - .service + .endpoint .evaluate_http_request(Request::new(evaluation)), ) .await @@ -1224,7 +1539,7 @@ impl ChainRunner { }); } - if result.has_body && result.body.len() > entry.max_body_bytes { + if result.has_body && result.body.len() > entry.max_payload_bytes { match apply_on_error(entry, "response_body_over_capacity", &mut applied) { OnErrorAction::FailOpen => continue, OnErrorAction::FailClosed(reason) => { @@ -1416,15 +1731,17 @@ pub(crate) fn safe_reason(reason: &str) -> String { #[cfg(test)] mod tests { use super::*; + use futures::{FutureExt, Stream, StreamExt}; use openshell_core::proto::middleware::v1::supervisor_middleware_server::{ SupervisorMiddleware, SupervisorMiddlewareServer, }; use openshell_core::proto::{ExistingHeaderAction, header_mutation}; use openshell_supervisor_middleware_builtins::{BUILTIN_REGEX, services}; + use tokio_stream::wrappers::TcpListenerStream; fn builtin_runner() -> ChainRunner { - ChainRunner::new( + ChainRunner::from_endpoint( services() .into_iter() .next() @@ -1644,20 +1961,23 @@ mod tests { #[tokio::test] async fn injected_services_cannot_duplicate_middleware_names() { - let first: Arc = Arc::new(ScriptedService { + let first: Arc = Arc::new(ScriptedService { manifest_name: "openshell/test".into(), max_body_bytes: 1024, result: allow_result(), }); - let second: Arc = Arc::new(ScriptedService { + let second: Arc = Arc::new(ScriptedService { manifest_name: "openshell/test".into(), max_body_bytes: 1024, result: allow_result(), }); - let error = MiddlewareRegistry::connect_services(vec![first, second], Vec::new()) - .await - .expect_err("duplicate injected middleware name must fail registry construction"); + let error = MiddlewareRegistry::connect_services( + vec![http_only_endpoint(first), http_only_endpoint(second)], + Vec::new(), + ) + .await + .expect_err("duplicate injected middleware name must fail registry construction"); assert!( error .to_string() @@ -1676,6 +1996,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for ScriptedService { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + + async fn evaluate_web_socket_session( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -1686,7 +2016,7 @@ mod tests { bindings: vec![MiddlewareBinding { operation: SupervisorMiddlewareOperation::HttpRequest as i32, phase: SupervisorMiddlewarePhase::PreCredentials as i32, - max_body_bytes: self.max_body_bytes, + max_payload_bytes: self.max_body_bytes, timeout: String::new(), }], })) @@ -1725,6 +2055,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for SlowService { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + + async fn evaluate_web_socket_session( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -1735,7 +2075,7 @@ mod tests { bindings: vec![MiddlewareBinding { operation: SupervisorMiddlewareOperation::HttpRequest as i32, phase: SupervisorMiddlewarePhase::PreCredentials as i32, - max_body_bytes: 4096, + max_payload_bytes: 4096, timeout: self.binding_timeout.clone(), }], })) @@ -1778,6 +2118,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for TwoStageService { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + + async fn evaluate_web_socket_session( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -1788,7 +2138,7 @@ mod tests { bindings: vec![MiddlewareBinding { operation: SupervisorMiddlewareOperation::HttpRequest as i32, phase: SupervisorMiddlewarePhase::PreCredentials as i32, - max_body_bytes: 256 * 1024, + max_payload_bytes: 256 * 1024, timeout: String::new(), }], })) @@ -1842,7 +2192,7 @@ mod tests { // must stop there: the second stage never runs, so it never sees a // payload the policy would reject. let second_ran = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let service: Arc = Arc::new(TwoStageService { + let service: Arc = Arc::new(TwoStageService { second_ran: Arc::clone(&second_ran), }); let runner = ChainRunner::new(service); @@ -1904,7 +2254,7 @@ mod tests { // A validator that accepts every body lets both stages run; the second // stage sees the first stage's output. let second_ran = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let service: Arc = Arc::new(TwoStageService { + let service: Arc = Arc::new(TwoStageService { second_ran: Arc::clone(&second_ran), }); let runner = ChainRunner::new(service); @@ -1956,7 +2306,7 @@ mod tests { #[tokio::test] async fn per_stage_validator_error_becomes_structured_denial() { let second_ran = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let service: Arc = Arc::new(TwoStageService { + let service: Arc = Arc::new(TwoStageService { second_ran: Arc::clone(&second_ran), }); let runner = ChainRunner::new(service); @@ -2042,6 +2392,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for RecordingService { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + + async fn evaluate_web_socket_session( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -2052,7 +2412,7 @@ mod tests { bindings: vec![MiddlewareBinding { operation: SupervisorMiddlewareOperation::HttpRequest as i32, phase: SupervisorMiddlewarePhase::PreCredentials as i32, - max_body_bytes: 4096, + max_payload_bytes: 4096, timeout: String::new(), }], })) @@ -2101,6 +2461,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for HeaderChainService { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + + async fn evaluate_web_socket_session( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -2111,7 +2481,7 @@ mod tests { bindings: vec![MiddlewareBinding { operation: SupervisorMiddlewareOperation::HttpRequest as i32, phase: SupervisorMiddlewarePhase::PreCredentials as i32, - max_body_bytes: 4096, + max_payload_bytes: 4096, timeout: String::new(), }], })) @@ -2226,7 +2596,7 @@ mod tests { validated: std::sync::Mutex::new(Vec::new()), received: std::sync::Mutex::new(Vec::new()), }); - let recorder: Arc = service.clone(); + let recorder: Arc = service.clone(); let runner = ChainRunner::new(recorder); runner .validate_config("test/recorder", prost_types::Struct::default()) @@ -2275,17 +2645,17 @@ mod tests { ); } - fn external_registration(max_body_bytes: u64) -> SupervisorMiddlewareService { + fn external_registration(max_payload_bytes: u64) -> SupervisorMiddlewareService { SupervisorMiddlewareService { name: "local-guard-service".into(), grpc_endpoint: "http://127.0.0.1:50051".into(), - max_body_bytes, + max_payload_bytes, ..Default::default() } } async fn registry_with_external( - service: Arc, + service: Arc, registration: SupervisorMiddlewareService, ) -> MiddlewareRegistry { let builtin_service = services() @@ -2310,9 +2680,9 @@ mod tests { .await .expect("describe test service") .into_inner(); - let operator_max_body_bytes = usize::try_from(registration.max_body_bytes).unwrap(); + let operator_max_payload_bytes = usize::try_from(registration.max_payload_bytes).unwrap(); let operator_timeout = validate_registration(®istration).expect("valid registration"); - validate_external_manifest(®istration, &manifest, operator_max_body_bytes) + validate_external_manifest(®istration, &manifest, Some(operator_max_payload_bytes)) .expect("valid external manifest"); let manifest_cell = OnceCell::new(); manifest_cell.set(manifest).expect("manifest cache"); @@ -2321,23 +2691,26 @@ mod tests { services: Arc::new(vec![ Arc::new(MiddlewareServiceState { attachment_name: Some(builtin_name.clone()), - service: builtin_service, + endpoint: builtin_service, manifest: builtin_manifest_cell, diagnostic_policy: MiddlewareDiagnosticPolicy::Preserve, - operator_max_body_bytes: None, + operator_max_payload_bytes: None, operator_timeout: DEFAULT_MIDDLEWARE_TIMEOUT, }), Arc::new(MiddlewareServiceState { attachment_name: Some(registration_name.clone()), - service, + endpoint: http_only_endpoint(service), manifest: manifest_cell, diagnostic_policy: MiddlewareDiagnosticPolicy::Normalize, - operator_max_body_bytes: Some(operator_max_body_bytes), + operator_max_payload_bytes: Some(operator_max_payload_bytes), operator_timeout, }), ]), registered_services: Arc::new(vec![RegisteredMiddlewareService { registration }]), middleware_names: Arc::new(HashSet::from([builtin_name, registration_name])), + work_admission: Arc::new(Semaphore::new(MAX_CONCURRENT_MIDDLEWARE_WORK)), + work_admission_waiters: Arc::new(Semaphore::new(MAX_QUEUED_MIDDLEWARE_WORK)), + session_admission: Arc::new(Semaphore::new(MAX_CONCURRENT_MIDDLEWARE_SESSIONS)), } } @@ -2357,7 +2730,7 @@ mod tests { // The built-in resolves and reports its real limit; the missing binding // does not resolve and must not contribute a body limit. assert!(described[0].is_resolved()); - assert_eq!(described[0].max_body_bytes(), 256 * 1024); + assert_eq!(described[0].max_payload_bytes(), 256 * 1024); assert!(!described[1].is_resolved()); } @@ -2380,7 +2753,7 @@ mod tests { .describe_chain(std::slice::from_ref(&entry)) .await .expect("describe external middleware"); - assert_eq!(described[0].max_body_bytes(), 4096); + assert_eq!(described[0].max_payload_bytes(), 4096); assert_eq!( described[0] .binding @@ -2419,8 +2792,8 @@ mod tests { .describe_chain(&entries) .await .expect("describe chain"); - assert_eq!(described[0].max_body_bytes(), 256 * 1024); - assert_eq!(described[1].max_body_bytes(), 1024); + assert_eq!(described[0].max_payload_bytes(), 256 * 1024); + assert_eq!(described[1].max_payload_bytes(), 1024); let outcome = runner .evaluate_described(&described, input(r#"token="sk-ABCDEFGHIJKLMNOP""#)) @@ -2526,25 +2899,25 @@ mod tests { bindings: vec![MiddlewareBinding { operation: HTTP_REQUEST_OPERATION as i32, phase: PRE_CREDENTIALS_PHASE as i32, - max_body_bytes: 4096, + max_payload_bytes: 4096, timeout: String::new(), }], }; - let error = validate_external_manifest(®istration, &manifest, 4097) + let error = validate_external_manifest(®istration, &manifest, Some(4097)) .expect_err("operator limit must fit capability"); assert!(error.to_string().contains("exceeds")); } #[test] - fn external_registration_rejects_body_limit_above_platform_maximum() { + fn external_registration_rejects_payload_limit_above_platform_maximum() { let registration = external_registration(u64::MAX); let error = validate_registration(®istration) - .expect_err("extreme body limit must be rejected before allocation"); + .expect_err("extreme payload limit must be rejected before allocation"); assert!(error.to_string().contains("platform maximum")); } #[test] - fn manifest_rejects_body_limit_above_platform_maximum() { + fn manifest_rejects_payload_limit_above_platform_maximum() { let registration = external_registration(4096); let manifest = MiddlewareManifest { name: "example/service".into(), @@ -2552,12 +2925,12 @@ mod tests { bindings: vec![MiddlewareBinding { operation: HTTP_REQUEST_OPERATION as i32, phase: PRE_CREDENTIALS_PHASE as i32, - max_body_bytes: u64::MAX, + max_payload_bytes: u64::MAX, timeout: String::new(), }], }; - let error = validate_external_manifest(®istration, &manifest, 4096) - .expect_err("extreme advertised body limit must be rejected"); + let error = validate_external_manifest(®istration, &manifest, Some(4096)) + .expect_err("extreme advertised payload limit must be rejected"); assert!(error.to_string().contains("platform maximum")); } @@ -2567,7 +2940,7 @@ mod tests { let binding = || MiddlewareBinding { operation: HTTP_REQUEST_OPERATION as i32, phase: PRE_CREDENTIALS_PHASE as i32, - max_body_bytes: 4096, + max_payload_bytes: 4096, timeout: String::new(), }; let manifest = MiddlewareManifest { @@ -2576,15 +2949,79 @@ mod tests { bindings: vec![binding(), binding()], }; - let error = validate_external_manifest(®istration, &manifest, 4096) + let error = validate_external_manifest(®istration, &manifest, Some(4096)) .expect_err("one service cannot advertise two bindings for the same pair"); assert!( error .to_string() - .contains("more than one binding for HTTP_REQUEST/PRE_CREDENTIALS") + .contains("duplicate middleware operation/phase pair") + ); + } + + #[test] + fn manifest_accepts_forward_websocket_binding_and_reserves_return_phase() { + let binding = |phase| MiddlewareBinding { + operation: SupervisorMiddlewareOperation::WebsocketMessage as i32, + phase: phase as i32, + max_payload_bytes: MAX_MIDDLEWARE_PAYLOAD_BYTES as u64, + timeout: "500ms".into(), + }; + let mut manifest = MiddlewareManifest { + name: "example/websocket".into(), + service_version: "test".into(), + bindings: vec![binding(SupervisorMiddlewarePhase::PreCredentials)], + }; + validate_manifest_bindings("test WebSocket service", &manifest, None) + .expect("forward WebSocket binding is supported"); + + manifest.bindings = vec![binding(SupervisorMiddlewarePhase::PreReturn)]; + let error = validate_manifest_bindings("test WebSocket service", &manifest, None) + .expect_err("return-path binding stays reserved for PR 2"); + assert!(error.to_string().contains("reserved for PR 2")); + } + + #[test] + fn external_websocket_binding_requires_operator_payload_limit() { + let registration = external_registration(0); + let manifest = MiddlewareManifest { + name: "example/websocket".into(), + service_version: "test".into(), + bindings: vec![MiddlewareBinding { + operation: SupervisorMiddlewareOperation::WebsocketMessage as i32, + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + max_payload_bytes: 4096, + timeout: String::new(), + }], + }; + + let error = validate_external_manifest(®istration, &manifest, Some(0)) + .expect_err("WebSocket bindings require an operator payload ceiling"); + assert!( + error + .to_string() + .contains("must configure max_payload_bytes") ); } + #[test] + fn external_websocket_binding_rejects_operator_limit_above_capability() { + let registration = external_registration(4097); + let manifest = MiddlewareManifest { + name: "example/websocket".into(), + service_version: "test".into(), + bindings: vec![MiddlewareBinding { + operation: SupervisorMiddlewareOperation::WebsocketMessage as i32, + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + max_payload_bytes: 4096, + timeout: String::new(), + }], + }; + + let error = validate_external_manifest(®istration, &manifest, Some(4097)) + .expect_err("operator payload limit must fit WebSocket capability"); + assert!(error.to_string().contains("exceeds")); + } + #[test] fn external_registration_accepts_http_and_https_grpc_endpoints() { for grpc_endpoint in [ @@ -2648,11 +3085,11 @@ mod tests { bindings: vec![MiddlewareBinding { operation: HTTP_REQUEST_OPERATION as i32, phase: PRE_CREDENTIALS_PHASE as i32, - max_body_bytes: 4096, + max_payload_bytes: 4096, timeout: timeout.into(), }], }; - let error = validate_external_manifest(®istration, &manifest, 4096) + let error = validate_external_manifest(®istration, &manifest, Some(4096)) .expect_err("out-of-bounds binding timeout must be rejected"); assert!(error.to_string().contains("invalid timeout")); } @@ -2887,11 +3324,11 @@ mod tests { .add_service( SupervisorMiddlewareServer::new(ScriptedService { manifest_name: "test/middleware".into(), - max_body_bytes: MAX_MIDDLEWARE_BODY_BYTES as u64, + max_body_bytes: MAX_MIDDLEWARE_PAYLOAD_BYTES as u64, result: openshell_core::proto::HttpRequestResult { reason: "r".repeat(MAX_MIDDLEWARE_REASON_BYTES - 128), reason_code: "r".repeat(MAX_MIDDLEWARE_REASON_CODE_BYTES), - body: vec![b'x'; MAX_MIDDLEWARE_BODY_BYTES], + body: vec![b'x'; MAX_MIDDLEWARE_PAYLOAD_BYTES], has_body: true, header_mutations: vec![write_header( "x-openshell-middleware-envelope", @@ -2915,7 +3352,7 @@ mod tests { }); let server_task = tokio::spawn(server); - let mut registration = external_registration(MAX_MIDDLEWARE_BODY_BYTES as u64); + let mut registration = external_registration(MAX_MIDDLEWARE_PAYLOAD_BYTES as u64); registration.grpc_endpoint = format!("http://{address}"); let registry = MiddlewareRegistry::connect_services(Vec::new(), vec![registration]) .await @@ -2939,7 +3376,7 @@ mod tests { "x-large-envelope".into(), "v".repeat(MAX_MIDDLEWARE_HEADER_BYTES - 256), )]; - request.body = vec![b'b'; MAX_MIDDLEWARE_BODY_BYTES]; + request.body = vec![b'b'; MAX_MIDDLEWARE_PAYLOAD_BYTES]; let outcome = ChainRunner::from_registry(registry) .evaluate( &[ChainEntry { @@ -2955,7 +3392,7 @@ mod tests { .expect("maximum bounded envelopes should fit configured transport limit"); assert!(outcome.allowed); - assert_eq!(outcome.body.len(), MAX_MIDDLEWARE_BODY_BYTES); + assert_eq!(outcome.body.len(), MAX_MIDDLEWARE_PAYLOAD_BYTES); assert_eq!(outcome.header_mutations.len(), 1); assert_eq!(outcome.findings.len(), MAX_MIDDLEWARE_FINDINGS_PER_STAGE); let _ = shutdown_tx.send(()); @@ -2970,7 +3407,7 @@ mod tests { assert_eq!(MIDDLEWARE_GRPC_ENVELOPE_BYTES, 292 * 1024 + 64); assert_eq!( MIDDLEWARE_GRPC_MESSAGE_BYTES, - MAX_MIDDLEWARE_BODY_BYTES + 292 * 1024 + 64 + MAX_MIDDLEWARE_PAYLOAD_BYTES + 292 * 1024 + 64 ); } @@ -3511,4 +3948,1316 @@ mod tests { ); assert!(outcome.applied[0].failed); } + + #[derive(Clone, Default)] + struct OpenAiRedactionService { + preflight: Arc>>, + manifest_name: String, + describe_calls: Arc, + skip: bool, + deny: bool, + preflight_reason: String, + preflight_reason_code: String, + preflight_findings: Vec, + preflight_metadata: HashMap, + close_on_first_message: bool, + messages: Arc, + session_ends: Option< + tokio::sync::mpsc::UnboundedSender, + >, + } + + impl OpenAiRedactionService { + fn websocket_stream(&self, mut requests: S) -> WebSocketResponseStream + where + S: Stream< + Item = std::result::Result< + openshell_core::proto::WebSocketSessionEvent, + tonic::Status, + >, + > + Send + + Unpin + + 'static, + { + use openshell_core::proto::{ + WebSocketMessageResult, WebSocketPreflightAction, WebSocketPreflightDecision, + WebSocketSessionEventResult, web_socket_session_event, + web_socket_session_event_result, + }; + + let preflight = Arc::clone(&self.preflight); + let skip = self.skip; + let deny = self.deny; + let preflight_reason = self.preflight_reason.clone(); + let preflight_reason_code = self.preflight_reason_code.clone(); + let preflight_findings = self.preflight_findings.clone(); + let preflight_metadata = self.preflight_metadata.clone(); + let close_on_first_message = self.close_on_first_message; + let messages = Arc::clone(&self.messages); + let session_ends = self.session_ends.clone(); + let (responses_tx, responses_rx) = tokio::sync::mpsc::channel(4); + tokio::spawn(async move { + while let Some(Ok(request)) = requests.next().await { + let response = match request.event { + Some(web_socket_session_event::Event::Preflight(value)) => { + *preflight.lock().expect("preflight lock") = Some(value); + Some(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::PreflightDecision( + WebSocketPreflightDecision { + action: if deny { + WebSocketPreflightAction::Deny as i32 + } else if skip { + WebSocketPreflightAction::Skip as i32 + } else { + WebSocketPreflightAction::Inspect as i32 + }, + reason: preflight_reason.clone(), + findings: preflight_findings.clone(), + metadata: preflight_metadata.clone(), + reason_code: preflight_reason_code.clone(), + }, + ), + ), + }) + } + Some(web_socket_session_event::Event::Message(value)) => { + messages.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + if close_on_first_message { + break; + } + let payload = String::from_utf8(value.payload) + .expect("test OpenAI event must be UTF-8") + .replace("customer-secret", "[REDACTED]"); + Some(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::MessageResult( + WebSocketMessageResult { + sequence: value.sequence, + decision: Decision::Allow as i32, + replacement: payload.into_bytes(), + has_replacement: true, + reason_code: "redacted".into(), + ..Default::default() + }, + ), + ), + }) + } + Some(web_socket_session_event::Event::SessionStart(_)) | None => None, + Some(web_socket_session_event::Event::SessionEnd(end)) => { + if let Some(session_ends) = &session_ends + && let Ok(reason) = + openshell_core::proto::WebSocketSessionEndReason::try_from( + end.reason, + ) + { + let _ = session_ends.send(reason); + } + None + } + }; + if let Some(response) = response + && responses_tx.send(Ok(response)).await.is_err() + { + break; + } + } + }); + Box::pin(tokio_stream::wrappers::ReceiverStream::new(responses_rx)) + } + } + + fn websocket_preflight_input(session_id: impl Into) -> WebSocketPreflightInput { + WebSocketPreflightInput { + session_id: session_id.into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: Vec::new(), + } + } + + #[tonic::async_trait] + impl SupervisorMiddleware for OpenAiRedactionService { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + + async fn describe( + &self, + _request: Request<()>, + ) -> std::result::Result, tonic::Status> { + self.describe_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(tonic::Response::new(MiddlewareManifest { + name: if self.manifest_name.is_empty() { + "test/openai-websocket-redactor".into() + } else { + self.manifest_name.clone() + }, + service_version: "test".into(), + bindings: vec![MiddlewareBinding { + operation: SupervisorMiddlewareOperation::WebsocketMessage as i32, + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + max_payload_bytes: MAX_MIDDLEWARE_PAYLOAD_BYTES as u64, + timeout: "1s".into(), + }], + })) + } + + async fn validate_config( + &self, + _request: Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + Ok(tonic::Response::new( + openshell_core::proto::ValidateConfigResponse { + valid: true, + reason: String::new(), + }, + )) + } + + async fn evaluate_http_request( + &self, + _request: Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + Err(tonic::Status::unimplemented( + "WebSocket-only test middleware", + )) + } + + async fn evaluate_web_socket_session( + &self, + request: Request>, + ) -> std::result::Result, tonic::Status> + { + Ok(tonic::Response::new( + self.websocket_stream(request.into_inner()), + )) + } + } + + #[tonic::async_trait] + impl SupervisorMiddlewareEndpoint for OpenAiRedactionService { + async fn describe( + &self, + request: Request<()>, + ) -> std::result::Result, tonic::Status> { + SupervisorMiddleware::describe(self, request).await + } + + async fn validate_config( + &self, + request: Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + SupervisorMiddleware::validate_config(self, request).await + } + + async fn evaluate_http_request( + &self, + request: Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + SupervisorMiddleware::evaluate_http_request(self, request).await + } + + async fn open_websocket_session( + &self, + receiver: tokio::sync::mpsc::Receiver, + ) -> std::result::Result { + Ok( + self.websocket_stream( + tokio_stream::wrappers::ReceiverStream::new(receiver).map(Ok), + ), + ) + } + } + + fn websocket_entry( + name: &str, + implementation: &str, + order: i32, + on_error: OnError, + ) -> ChainEntry { + ChainEntry { + name: name.into(), + implementation: implementation.into(), + order, + config: prost_types::Struct::default(), + on_error, + } + } + + #[tokio::test] + async fn empty_websocket_chain_skips_describe_and_preflight() { + let service = OpenAiRedactionService::default(); + let describe_calls = Arc::clone(&service.describe_calls); + let observed_preflight = Arc::clone(&service.preflight); + let runner = ChainRunner::from_endpoint(Arc::new(service)); + + let outcome = runner + .preflight_websocket(&[], websocket_preflight_input("empty-chain")) + .await + .expect("empty preflight"); + + assert!(outcome.allowed); + assert_eq!(outcome.terminal_reason, None); + assert!(outcome.session.is_none()); + assert!(outcome.invocations.is_empty()); + assert_eq!(describe_calls.load(std::sync::atomic::Ordering::SeqCst), 0); + assert!(observed_preflight.lock().expect("preflight lock").is_none()); + } + + #[tokio::test] + async fn http_only_attachments_are_reported_but_do_not_select_websocket_stages() { + for on_error in [OnError::FailClosed, OnError::FailOpen] { + let runner = ChainRunner::new(Arc::new(ScriptedService { + manifest_name: "test/http-only".into(), + max_body_bytes: 4096, + result: allow_result(), + })); + let chain = [ChainEntry { + name: "http-guard".into(), + implementation: "test/http-only".into(), + order: 0, + config: prost_types::Struct::default(), + on_error, + }]; + + let outcome = runner + .preflight_websocket( + &chain, + websocket_preflight_input(format!("http-only-{on_error:?}")), + ) + .await + .expect("HTTP-only attachment coverage"); + + assert!(outcome.allowed); + assert_eq!(outcome.terminal_reason, None); + assert!(outcome.session.is_none()); + assert!(outcome.invocations.is_empty()); + assert_eq!( + outcome.coverage, + [WebSocketCoverage { + config_name: "http-guard".into(), + implementation: "test/http-only".into(), + state: WebSocketCoverageState::BindingNotSelected, + sequence: None, + message_type: None, + original_size: 0, + }] + ); + } + } + + #[tokio::test] + async fn http_only_attachment_allows_33_requested_subprotocols() { + let runner = ChainRunner::new(Arc::new(ScriptedService { + manifest_name: "test/http-only".into(), + max_body_bytes: 4096, + result: allow_result(), + })); + let chain = [ChainEntry { + name: "http-guard".into(), + implementation: "test/http-only".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailClosed, + }]; + let mut input = websocket_preflight_input("http-only-many-subprotocols"); + input.requested_subprotocols = (0..33) + .map(|index| format!("subprotocol-{index}")) + .collect(); + + let outcome = runner + .preflight_websocket(&chain, input) + .await + .expect("HTTP-only attachment must not validate a WebSocket middleware envelope"); + + assert!(outcome.allowed); + assert_eq!(outcome.terminal_reason, None); + assert!(outcome.session.is_none()); + assert!(outcome.invocations.is_empty()); + assert_eq!( + outcome.coverage, + [WebSocketCoverage { + config_name: "http-guard".into(), + implementation: "test/http-only".into(), + state: WebSocketCoverageState::BindingNotSelected, + sequence: None, + message_type: None, + original_size: 0, + }] + ); + } + + #[tokio::test] + async fn unsupported_binary_messages_advance_sequence_without_applying_on_error() { + for on_error in [OnError::FailClosed, OnError::FailOpen] { + let runner = builtin_runner(); + let chain = [entry("regex-redactor", on_error)]; + let preflight = runner + .preflight_websocket( + &chain, + websocket_preflight_input(format!("binary-{on_error:?}")), + ) + .await + .expect("preflight"); + let mut session = preflight.session.expect("built-in inspects text"); + assert!(session.start("").await.allowed); + + let coverage = session.observe_unsupported_message( + openshell_core::proto::WebSocketMessageType::Binary, + 23, + ); + assert_eq!( + coverage, + [WebSocketCoverage { + config_name: "regex-redactor".into(), + implementation: BUILTIN_REGEX.into(), + state: WebSocketCoverageState::UnsupportedMessageType, + sequence: Some(1), + message_type: Some(openshell_core::proto::WebSocketMessageType::Binary), + original_size: 23, + }] + ); + + let text = session.evaluate_text(br#"{"input":"safe"}"#.to_vec()).await; + assert!(text.allowed); + assert_eq!(text.invocations[0].sequence, Some(2)); + assert!(!text.invocations[0].failed); + + session + .end(openshell_core::proto::WebSocketSessionEndReason::NormalClose) + .await; + } + } + + #[tokio::test] + async fn explicit_websocket_preflight_denial_is_authoritative_for_both_error_modes() { + use openshell_core::proto::WebSocketSessionEndReason; + + for on_error in [OnError::FailOpen, OnError::FailClosed] { + let (session_ends_tx, mut session_ends_rx) = tokio::sync::mpsc::unbounded_channel(); + let runner = ChainRunner::from_endpoint(Arc::new(OpenAiRedactionService { + deny: true, + preflight_reason: "contains sensitive request data".into(), + preflight_reason_code: "upgrade_blocked".into(), + preflight_findings: vec![Finding { + r#type: "content.sensitive".into(), + label: "Sensitive content".into(), + count: 1, + confidence: "high".into(), + severity: "high".into(), + }], + preflight_metadata: HashMap::from([("policy_version".into(), "1".into())]), + session_ends: Some(session_ends_tx), + ..Default::default() + })); + + let outcome = runner + .preflight_websocket( + &[websocket_entry( + "deny-upgrade", + "test/openai-websocket-redactor", + 0, + on_error, + )], + websocket_preflight_input("explicit-denial"), + ) + .await + .expect("denied preflight"); + + assert!(!outcome.allowed); + assert_eq!( + outcome.terminal_reason, + Some(WebSocketSessionEndReason::MiddlewareDenial) + ); + assert_eq!( + outcome.reason, + "middleware_denied:deny-upgrade:upgrade_blocked" + ); + assert_eq!( + outcome + .denial + .as_ref() + .map(|denial| denial.config_name.as_str()), + Some("deny-upgrade") + ); + assert_eq!( + outcome + .denial + .as_ref() + .and_then(|denial| denial.reason_code.as_deref()), + Some("upgrade_blocked") + ); + assert_eq!( + outcome.invocations[0].reason_code.as_deref(), + Some("upgrade_blocked") + ); + assert_eq!(outcome.findings.len(), 1); + assert_eq!(outcome.findings[0].middleware, "deny-upgrade"); + assert_eq!(outcome.findings[0].finding.r#type, "content.sensitive"); + assert_eq!(outcome.metadata["deny-upgrade"]["policy_version"], "1"); + assert!(!outcome.reason.contains("sensitive request data")); + assert!(outcome.session.is_none()); + assert_eq!(outcome.invocations.len(), 1); + assert_eq!( + outcome.invocations[0].outcome, + WebSocketInvocationOutcome::Deny + ); + assert!(!outcome.invocations[0].failed); + assert_eq!( + session_ends_rx.recv().await, + Some(WebSocketSessionEndReason::MiddlewareDenial) + ); + assert!( + session_ends_rx.try_recv().is_err(), + "each opened stream receives at most one session_end" + ); + } + } + + #[tokio::test] + async fn mixed_websocket_preflight_denial_ends_every_opened_stage() { + use openshell_core::proto::WebSocketSessionEndReason; + + let (first_end_tx, mut first_end_rx) = tokio::sync::mpsc::unbounded_channel(); + let (denier_end_tx, mut denier_end_rx) = tokio::sync::mpsc::unbounded_channel(); + let (last_end_tx, mut last_end_rx) = tokio::sync::mpsc::unbounded_channel(); + let endpoints: Vec> = vec![ + Arc::new(OpenAiRedactionService { + manifest_name: "test/first-inspector".into(), + session_ends: Some(first_end_tx), + ..Default::default() + }), + Arc::new(OpenAiRedactionService { + manifest_name: "test/denier".into(), + deny: true, + session_ends: Some(denier_end_tx), + ..Default::default() + }), + Arc::new(OpenAiRedactionService { + manifest_name: "test/last-inspector".into(), + session_ends: Some(last_end_tx), + ..Default::default() + }), + ]; + let runner = ChainRunner::from_registry( + MiddlewareRegistry::connect_services(endpoints, Vec::new()) + .await + .expect("connect mixed middleware services"), + ); + let chain = [ + websocket_entry("first", "test/first-inspector", 0, OnError::FailClosed), + websocket_entry("deny", "test/denier", 1, OnError::FailOpen), + websocket_entry("last", "test/last-inspector", 2, OnError::FailClosed), + ]; + + let outcome = runner + .preflight_websocket(&chain, websocket_preflight_input("mixed-denial")) + .await + .expect("mixed preflight"); + + assert!(!outcome.allowed); + assert_eq!( + outcome.terminal_reason, + Some(WebSocketSessionEndReason::MiddlewareDenial) + ); + assert_eq!( + outcome + .invocations + .iter() + .map(|invocation| invocation.outcome) + .collect::>(), + vec![ + WebSocketInvocationOutcome::Inspect, + WebSocketInvocationOutcome::Deny, + WebSocketInvocationOutcome::Inspect, + ] + ); + for receiver in [&mut first_end_rx, &mut denier_end_rx, &mut last_end_rx] { + assert_eq!( + receiver.recv().await, + Some(WebSocketSessionEndReason::MiddlewareDenial) + ); + assert!( + receiver.try_recv().is_err(), + "each opened stream receives at most one session_end" + ); + } + assert_eq!( + runner.registry.session_admission.available_permits(), + MAX_CONCURRENT_MIDDLEWARE_SESSIONS + ); + } + + #[tokio::test] + async fn builtin_regex_redacts_ordered_websocket_text_messages() { + let runner = builtin_runner(); + let chain = [entry("regex-redactor", OnError::FailClosed)]; + let preflight = runner + .preflight_websocket( + &chain, + WebSocketPreflightInput { + session_id: "builtin-regex-session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: vec!["realtime".into()], + }, + ) + .await + .expect("preflight"); + assert!(preflight.allowed); + assert_eq!( + preflight.invocations[0].outcome, + WebSocketInvocationOutcome::Inspect + ); + let mut session = preflight.session.expect("built-in chose to inspect"); + assert!(session.start("realtime").await.allowed); + + let original = br#"{"type":"response.create","response":{"input":"sk-ABCDEFGHIJKLMNOP"}}"#; + let redacted = session.evaluate_text(original.to_vec()).await; + assert!(redacted.allowed); + assert_eq!( + String::from_utf8(redacted.payload).expect("transformed UTF-8"), + r#"{"type":"response.create","response":{"input":"[REDACTED]"}}"# + ); + assert_eq!(redacted.invocations[0].sequence, Some(1)); + assert!(redacted.invocations[0].transformed); + assert_eq!(redacted.findings.len(), 1); + assert_eq!(redacted.findings[0].middleware, "regex-redactor"); + assert_eq!(redacted.findings[0].finding.r#type, "regex.openai"); + assert_eq!( + redacted.metadata["regex-redactor"]["regex_matches_replaced"], + "1" + ); + + let unchanged = session + .evaluate_text(br#"{"type":"response.cancel"}"#.to_vec()) + .await; + assert!(unchanged.allowed); + assert_eq!(unchanged.payload, br#"{"type":"response.cancel"}"#); + assert_eq!(unchanged.invocations[0].sequence, Some(2)); + assert!(!unchanged.invocations[0].transformed); + assert!(unchanged.findings.is_empty()); + assert!(unchanged.metadata.is_empty()); + + let oversized = session.evaluate_text(vec![b'a'; 256 * 1024 + 1]).await; + assert!(!oversized.allowed); + assert_eq!( + oversized.reason, + "middleware_failed: request_message_over_capacity" + ); + session + .end(openshell_core::proto::WebSocketSessionEndReason::NormalClose) + .await; + } + + #[tokio::test] + async fn builtin_regex_redacts_after_fail_open_message_capacity_gap() { + let runner = builtin_runner(); + let chain = [entry("regex-redactor", OnError::FailOpen)]; + let preflight = runner + .preflight_websocket( + &chain, + WebSocketPreflightInput { + session_id: "builtin-regex-gap-session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: vec!["realtime".into()], + }, + ) + .await + .expect("preflight"); + assert!(preflight.allowed); + let mut session = preflight.session.expect("built-in chose to inspect"); + assert!(session.start("realtime").await.allowed); + + let oversized_payload = vec![b'a'; 256 * 1024 + 1]; + let oversized = session.evaluate_text(oversized_payload.clone()).await; + assert!(oversized.allowed); + assert_eq!(oversized.payload, oversized_payload); + assert_eq!(oversized.invocations[0].sequence, Some(1)); + assert_eq!( + oversized.invocations[0].outcome, + WebSocketInvocationOutcome::FailOpen + ); + assert!(!oversized.invocations[0].stage_disabled); + + let original = br#"{"type":"response.create","response":{"input":"sk-ABCDEFGHIJKLMNOP"}}"#; + let redacted = session.evaluate_text(original.to_vec()).await; + assert!(redacted.allowed); + assert_eq!( + String::from_utf8(redacted.payload).expect("transformed UTF-8"), + r#"{"type":"response.create","response":{"input":"[REDACTED]"}}"# + ); + assert_eq!(redacted.invocations[0].sequence, Some(2)); + assert_eq!( + redacted.invocations[0].outcome, + WebSocketInvocationOutcome::Allow + ); + assert!(redacted.invocations[0].transformed); + assert!(!redacted.invocations[0].stage_disabled); + + session + .end(openshell_core::proto::WebSocketSessionEndReason::NormalClose) + .await; + } + + #[tokio::test] + async fn in_process_websocket_endpoint_redacts_openai_event() { + let service = OpenAiRedactionService::default(); + let runner = ChainRunner::from_endpoint(Arc::new(service)); + let chain = [ChainEntry { + name: "openai-redactor".into(), + implementation: "test/openai-websocket-redactor".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailClosed, + }]; + let preflight = runner + .preflight_websocket( + &chain, + WebSocketPreflightInput { + session_id: "in-process-session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: vec!["realtime".into()], + }, + ) + .await + .expect("preflight"); + assert!(preflight.allowed); + let mut session = preflight.session.expect("middleware chose to inspect"); + assert!(session.start("realtime").await.allowed); + + let original = br#"{"type":"response.create","response":{"input":"customer-secret"}}"#; + let outcome = session.evaluate_text(original.to_vec()).await; + assert!(outcome.allowed); + assert_eq!( + String::from_utf8(outcome.payload).expect("transformed UTF-8"), + r#"{"type":"response.create","response":{"input":"[REDACTED]"}}"# + ); + assert!(outcome.invocations[0].transformed); + session + .end(openshell_core::proto::WebSocketSessionEndReason::NormalClose) + .await; + } + + #[tokio::test] + async fn openai_websocket_event_is_introspected_and_redacted() { + let service = OpenAiRedactionService::default(); + let observed_preflight = Arc::clone(&service.preflight); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket middleware"); + let address = listener.local_addr().expect("middleware address"); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tonic::transport::Server::builder() + .add_service(SupervisorMiddlewareServer::new(service)) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }); + let server_task = tokio::spawn(server); + + let mut registration = external_registration(1024); + registration.grpc_endpoint = format!("http://{address}"); + let registry = MiddlewareRegistry::connect_services(Vec::new(), vec![registration]) + .await + .expect("connect WebSocket middleware"); + let runner = ChainRunner::from_registry(registry); + let chain = [ChainEntry { + name: "openai-redactor".into(), + implementation: "local-guard-service".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailClosed, + }]; + let described = runner + .describe_websocket_chain(&chain) + .await + .expect("describe WebSocket chain"); + assert_eq!( + described[0].max_payload_bytes(), + 1024, + "operator max_payload_bytes must cap WebSocket messages" + ); + let preflight = runner + .preflight_websocket( + &chain, + WebSocketPreflightInput { + session_id: "ws-session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: vec!["realtime".into()], + }, + ) + .await + .expect("preflight"); + assert!(preflight.allowed); + let mut session = preflight.session.expect("middleware chose to inspect"); + assert!(session.start("realtime").await.allowed); + + let original = br#"{"type":"response.create","response":{"input":"customer-secret"}}"#; + let outcome = session.evaluate_text(original.to_vec()).await; + assert!(outcome.allowed); + let transformed = String::from_utf8(outcome.payload).expect("transformed UTF-8"); + assert!(transformed.contains("[REDACTED]")); + assert!(!transformed.contains("customer-secret")); + assert!(outcome.invocations[0].transformed); + + let observed = observed_preflight + .lock() + .expect("preflight lock") + .clone() + .expect("preflight observed"); + let target = observed.target.expect("preflight target"); + assert_eq!(target.scheme, "wss"); + assert_eq!(target.host, "api.openai.com"); + assert_eq!(target.port, 443); + assert_eq!(target.method, "GET"); + assert_eq!(target.path, "/v1/responses"); + assert!(target.query.is_empty()); + assert_eq!(observed.requested_subprotocols, ["realtime"]); + session + .end(openshell_core::proto::WebSocketSessionEndReason::NormalClose) + .await; + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join test middleware") + .expect("serve middleware"); + } + + #[tokio::test] + async fn fail_open_disables_broken_websocket_stage_for_later_messages() { + let service = OpenAiRedactionService { + close_on_first_message: true, + ..Default::default() + }; + let observed_preflight = Arc::clone(&service.preflight); + let message_count = Arc::clone(&service.messages); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket middleware"); + let address = listener.local_addr().expect("middleware address"); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tonic::transport::Server::builder() + .add_service(SupervisorMiddlewareServer::new(service)) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }); + let server_task = tokio::spawn(server); + + let mut registration = external_registration(1024); + registration.grpc_endpoint = format!("http://{address}"); + let registry = MiddlewareRegistry::connect_services(Vec::new(), vec![registration]) + .await + .expect("connect WebSocket middleware"); + let runner = ChainRunner::from_registry(registry); + let chain = [ChainEntry { + name: "openai-redactor".into(), + implementation: "local-guard-service".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailOpen, + }]; + let preflight = runner + .preflight_websocket( + &chain, + WebSocketPreflightInput { + session_id: "ws-session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "ws".into(), + host: "api.openai.com".into(), + port: 80, + path: "/v1/responses".into(), + requested_subprotocols: Vec::new(), + }, + ) + .await + .expect("preflight"); + let mut session = preflight.session.expect("middleware chose to inspect"); + assert!(session.start("").await.allowed); + assert_eq!( + observed_preflight + .lock() + .expect("preflight lock") + .as_ref() + .expect("preflight observed") + .target + .as_ref() + .expect("preflight target") + .scheme, + "ws" + ); + + let first = session + .evaluate_text(br#"{"type":"response.create"}"#.to_vec()) + .await; + assert!(first.allowed, "fail-open should bypass the broken stage"); + assert_eq!(first.invocations.len(), 1); + assert!(first.invocations[0].failed); + assert!(first.invocations[0].stage_disabled); + assert_eq!( + runner.registry.session_admission.available_permits(), + MAX_CONCURRENT_MIDDLEWARE_SESSIONS, + "the final disabled stage must release persistent session capacity" + ); + + let mut work = Vec::new(); + for _ in 0..MAX_CONCURRENT_MIDDLEWARE_WORK { + work.push( + runner + .reserve_middleware_work_admission() + .await + .expect("fill middleware work budget"), + ); + } + + let second = session + .evaluate_text(br#"{"type":"response.cancel"}"#.to_vec()) + .now_or_never() + .expect("fully disabled session must bypass without waiting for work admission"); + assert!(second.allowed); + assert!( + second.invocations.is_empty(), + "disabled stage must not be called again in this session" + ); + assert_eq!(message_count.load(std::sync::atomic::Ordering::SeqCst), 1); + drop(work); + + session + .end(openshell_core::proto::WebSocketSessionEndReason::NormalClose) + .await; + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join test middleware") + .expect("serve middleware"); + } + + #[tokio::test] + async fn mixed_websocket_session_keeps_message_admission_while_one_stage_is_active() { + let broken = Arc::new(OpenAiRedactionService { + close_on_first_message: true, + ..Default::default() + }); + let mut endpoints = services(); + endpoints.push(broken); + let runner = ChainRunner::from_registry( + MiddlewareRegistry::connect_services(endpoints, Vec::new()) + .await + .expect("connect mixed middleware services"), + ); + let broken_entry = ChainEntry { + name: "best-effort-remote".into(), + implementation: "test/openai-websocket-redactor".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailOpen, + }; + let mut regex_entry = entry("required-regex", OnError::FailClosed); + regex_entry.order = 1; + let preflight = runner + .preflight_websocket( + &[broken_entry, regex_entry], + websocket_preflight_input("mixed-active"), + ) + .await + .expect("mixed preflight"); + let mut session = preflight.session.expect("both stages inspect"); + assert!(session.start("").await.allowed); + + let first = session + .evaluate_text(br#"{"input":"sk-ABCDEFGHIJKLMNOP"}"#.to_vec()) + .await; + assert!(first.allowed); + assert!(first.invocations[0].stage_disabled); + assert_eq!( + first.invocations[1].outcome, + WebSocketInvocationOutcome::Allow + ); + + let mut work = Vec::new(); + for _ in 0..MAX_CONCURRENT_MIDDLEWARE_WORK { + work.push( + runner + .reserve_middleware_work_admission() + .await + .expect("fill middleware work budget"), + ); + } + assert!( + session.admit_message().now_or_never().is_none(), + "an active remaining stage must still wait for message work admission" + ); + drop(work); + + let second = session + .evaluate_text(br#"{"input":"sk-QRSTUVWXYZabcdef"}"#.to_vec()) + .await; + assert!(second.allowed); + assert_eq!(second.invocations.len(), 1); + assert_eq!( + second.invocations[0].config_name, "required-regex", + "disabled stage must stay bypassed while the active stage continues" + ); + } + + #[tokio::test] + async fn websocket_admission_wait_queue_is_bounded() { + let runner = ChainRunner::default(); + let mut active = Vec::new(); + for _ in 0..MAX_CONCURRENT_MIDDLEWARE_WORK { + active.push( + runner + .reserve_middleware_work_admission() + .await + .expect("active admission"), + ); + } + + let mut waiters = Vec::new(); + for _ in 0..MAX_QUEUED_MIDDLEWARE_WORK { + let runner = runner.clone(); + waiters.push(tokio::spawn(async move { + runner.reserve_middleware_work().await + })); + } + while runner.registry.work_admission_waiters.available_permits() != 0 { + tokio::task::yield_now().await; + } + + let overflow = runner + .reserve_middleware_work() + .await + .expect("admission outcome"); + assert!(matches!( + overflow, + MiddlewareWorkAdmissionOutcome::QueueExhausted + )); + runner + .reserve_middleware_work_admission() + .await + .expect_err("WebSocket callers retain their existing failure path"); + + drop(active); + for waiter in waiters { + let admission = waiter + .await + .expect("waiter task") + .expect("queued admission after capacity is released"); + assert!(matches!( + admission, + MiddlewareWorkAdmissionOutcome::Admitted(_) + )); + } + assert_eq!( + runner.registry.work_admission.available_permits(), + MAX_CONCURRENT_MIDDLEWARE_WORK + ); + assert_eq!( + runner.registry.work_admission_waiters.available_permits(), + MAX_QUEUED_MIDDLEWARE_WORK + ); + + let mut recovered = Vec::new(); + for _ in 0..MAX_CONCURRENT_MIDDLEWARE_WORK { + recovered.push( + runner + .reserve_middleware_work_admission() + .await + .expect("recovered active admission"), + ); + } + assert_eq!(recovered.len(), MAX_CONCURRENT_MIDDLEWARE_WORK); + } + + #[tokio::test] + async fn websocket_queue_exhaustion_remains_a_protocol_failure() { + let runner = builtin_runner(); + let chain = [entry("regex-redactor", OnError::FailClosed)]; + let preflight = runner + .preflight_websocket(&chain, websocket_preflight_input("established")) + .await + .expect("initial preflight"); + let mut session = preflight.session.expect("built-in inspects session"); + assert!(session.start("").await.allowed); + + let mut active = Vec::new(); + for _ in 0..MAX_CONCURRENT_MIDDLEWARE_WORK { + active.push( + runner + .reserve_middleware_work_admission() + .await + .expect("fill active work"), + ); + } + let mut waiters = Vec::new(); + for _ in 0..MAX_QUEUED_MIDDLEWARE_WORK { + let runner = runner.clone(); + waiters.push(tokio::spawn(async move { + runner.reserve_middleware_work().await + })); + } + while runner.registry.work_admission_waiters.available_permits() != 0 { + tokio::task::yield_now().await; + } + + let preflight_overflow = runner + .preflight_websocket(&chain, websocket_preflight_input("preflight-overflow")) + .await; + assert!( + preflight_overflow.is_err(), + "preflight exhaustion remains an outer HTTP failure" + ); + session + .admit_message() + .await + .expect_err("established message exhaustion remains a typed termination input"); + + for waiter in waiters { + waiter.abort(); + } + drop(active); + } + + #[tokio::test] + async fn websocket_preflight_skip_removes_stage_without_message_calls() { + let service = OpenAiRedactionService { + skip: true, + ..Default::default() + }; + let message_count = Arc::clone(&service.messages); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket middleware"); + let address = listener.local_addr().expect("middleware address"); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tonic::transport::Server::builder() + .add_service(SupervisorMiddlewareServer::new(service)) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }); + let server_task = tokio::spawn(server); + let mut registration = external_registration(1024); + registration.grpc_endpoint = format!("http://{address}"); + let runner = ChainRunner::from_registry( + MiddlewareRegistry::connect_services(Vec::new(), vec![registration]) + .await + .expect("connect middleware"), + ); + let result = runner + .preflight_websocket( + &[ChainEntry { + name: "scope".into(), + implementation: "local-guard-service".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailClosed, + }], + WebSocketPreflightInput { + session_id: "session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: Vec::new(), + }, + ) + .await + .expect("preflight"); + assert!(result.allowed); + assert!(result.session.is_none()); + assert_eq!( + result.invocations[0].outcome, + WebSocketInvocationOutcome::Skip + ); + assert_eq!(message_count.load(std::sync::atomic::Ordering::SeqCst), 0); + assert_eq!( + runner.registry.session_admission.available_permits(), + MAX_CONCURRENT_MIDDLEWARE_SESSIONS, + "all-skip preflight must not retain session capacity" + ); + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware") + .expect("serve middleware"); + } + + #[tokio::test] + async fn websocket_session_budget_caps_idle_inspecting_sessions_and_releases_on_end() { + let runner = builtin_runner(); + let chain = [entry("regex-redactor", OnError::FailClosed)]; + let mut sessions = Vec::new(); + for index in 0..MAX_CONCURRENT_MIDDLEWARE_SESSIONS { + let preflight = runner + .preflight_websocket( + &chain, + websocket_preflight_input(format!("session-{index}")), + ) + .await + .expect("admit inspecting session"); + assert!(preflight.allowed); + sessions.push(preflight.session.expect("built-in inspects session")); + } + assert_eq!(runner.registry.session_admission.available_permits(), 0); + + let overflow = runner + .preflight_websocket(&chain, websocket_preflight_input("overflow")) + .await + .expect("capacity exhaustion is a typed preflight outcome"); + assert!(!overflow.allowed); + assert!(overflow.session.is_none()); + assert!(overflow.session_capacity_exhausted); + assert_eq!( + overflow.invocations[0].outcome, + WebSocketInvocationOutcome::FailClosed + ); + + sessions + .pop() + .expect("retained session") + .end(openshell_core::proto::WebSocketSessionEndReason::NormalClose) + .await; + assert_eq!(runner.registry.session_admission.available_permits(), 1); + + let replacement = runner + .preflight_websocket(&chain, websocket_preflight_input("replacement")) + .await + .expect("released session capacity is reusable"); + assert!(replacement.allowed); + assert!(replacement.session.is_some()); + } + + #[tokio::test] + async fn websocket_session_budget_survives_registry_replacement() { + let runner = builtin_runner(); + let chain = [entry("regex-redactor", OnError::FailClosed)]; + let mut sessions = Vec::new(); + for index in 0..MAX_CONCURRENT_MIDDLEWARE_SESSIONS { + let preflight = runner + .preflight_websocket( + &chain, + websocket_preflight_input(format!("old-generation-{index}")), + ) + .await + .expect("admit old-generation session"); + sessions.push(preflight.session.expect("built-in inspects session")); + } + + let replacement_registry = MiddlewareRegistry::connect_services(services(), Vec::new()) + .await + .expect("connect replacement registry"); + let replacement = runner.with_replacement_registry(replacement_registry); + let overflow = replacement + .preflight_websocket(&chain, websocket_preflight_input("new-generation-overflow")) + .await + .expect("capacity exhaustion is a typed preflight outcome"); + assert!(!overflow.allowed); + assert!(overflow.session_capacity_exhausted); + + sessions + .pop() + .expect("retained old-generation session") + .end(openshell_core::proto::WebSocketSessionEndReason::PolicyReload) + .await; + let admitted = replacement + .preflight_websocket(&chain, websocket_preflight_input("new-generation-admitted")) + .await + .expect("released capacity is reusable after replacement"); + assert!(admitted.allowed); + assert!(admitted.session.is_some()); + } + + #[tokio::test] + async fn websocket_session_capacity_exhaustion_honors_mixed_on_error() { + let runner = builtin_runner(); + let mut held = Vec::new(); + for _ in 0..MAX_CONCURRENT_MIDDLEWARE_SESSIONS { + match runner.try_reserve_middleware_session() { + MiddlewareSessionAdmission::Admitted(admission) => held.push(admission), + MiddlewareSessionAdmission::AtCapacity => { + panic!("session budget exhausted before platform limit") + } + } + } + + let mut first = entry("best-effort-a", OnError::FailOpen); + first.order = 1; + let mut second = entry("best-effort-b", OnError::FailOpen); + second.order = 2; + let all_fail_open = runner + .preflight_websocket( + &[first.clone(), second.clone()], + websocket_preflight_input("all-fail-open"), + ) + .await + .expect("all-fail-open capacity outcome"); + assert!(all_fail_open.allowed); + assert!(all_fail_open.session.is_none()); + assert!(all_fail_open.session_capacity_exhausted); + assert!( + all_fail_open + .invocations + .iter() + .all(|invocation| invocation.outcome == WebSocketInvocationOutcome::FailOpen) + ); + + second.on_error = OnError::FailClosed; + let mixed = runner + .preflight_websocket(&[first, second], websocket_preflight_input("mixed")) + .await + .expect("mixed capacity outcome"); + assert!(!mixed.allowed); + assert!(mixed.session.is_none()); + assert!(mixed.session_capacity_exhausted); + assert_eq!( + mixed + .invocations + .iter() + .map(|invocation| invocation.outcome) + .collect::>(), + [ + WebSocketInvocationOutcome::FailOpen, + WebSocketInvocationOutcome::FailClosed, + ] + ); + drop(held); + } } diff --git a/crates/openshell-supervisor-middleware/src/remote.rs b/crates/openshell-supervisor-middleware/src/remote.rs index 378dac3ec4..cadd6dd3f6 100644 --- a/crates/openshell-supervisor-middleware/src/remote.rs +++ b/crates/openshell-supervisor-middleware/src/remote.rs @@ -4,11 +4,11 @@ use std::time::Duration; use miette::{IntoDiagnostic, Result, WrapErr}; +use openshell_core::middleware::{SupervisorMiddlewareEndpoint, WebSocketResponseStream}; use openshell_core::proto::middleware::v1::supervisor_middleware_client::SupervisorMiddlewareClient; -use openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware; use openshell_core::proto::{ HttpRequestEvaluation, HttpRequestResult, MiddlewareManifest, ValidateConfigRequest, - ValidateConfigResponse, + ValidateConfigResponse, WebSocketSessionEvent, }; use tonic::transport::{Channel, ClientTlsConfig, Endpoint}; use tonic::{Request, Response, Status}; @@ -58,7 +58,7 @@ impl RemoteMiddlewareService { } #[tonic::async_trait] -impl SupervisorMiddleware for RemoteMiddlewareService { +impl SupervisorMiddlewareEndpoint for RemoteMiddlewareService { async fn describe( &self, request: Request<()>, @@ -82,4 +82,18 @@ impl SupervisorMiddleware for RemoteMiddlewareService { let mut client = self.client.clone(); client.evaluate_http_request(request).await } + + async fn open_websocket_session( + &self, + receiver: tokio::sync::mpsc::Receiver, + ) -> std::result::Result { + let mut client = self.client.clone(); + let responses = client + .evaluate_web_socket_session(Request::new(tokio_stream::wrappers::ReceiverStream::new( + receiver, + ))) + .await? + .into_inner(); + Ok(Box::pin(responses)) + } } diff --git a/crates/openshell-supervisor-middleware/src/websocket.rs b/crates/openshell-supervisor-middleware/src/websocket.rs new file mode 100644 index 0000000000..76627cb8e7 --- /dev/null +++ b/crates/openshell-supervisor-middleware/src/websocket.rs @@ -0,0 +1,1339 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Forward-direction WebSocket middleware session runner. + +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; + +use futures::{StreamExt, future::join_all}; +use prost::Message as _; +use tokio::sync::mpsc; +use tokio::time::Instant; + +use openshell_core::proto::{ + Decision, HttpRequestTarget, RequestContext, SupervisorMiddlewarePhase, WebSocketMessage, + WebSocketMessageResult, WebSocketMessageType, WebSocketPreflight, WebSocketPreflightAction, + WebSocketPreflightDecision, WebSocketSessionEnd, WebSocketSessionEndReason, + WebSocketSessionEvent, WebSocketSessionStart, web_socket_session_event, + web_socket_session_event_result, +}; + +use super::{ + ChainEntry, ChainRunner, DescribedChainEntry, MAX_MIDDLEWARE_CHAIN_TIMEOUT, + MAX_MIDDLEWARE_CONFIG_BYTES, MAX_MIDDLEWARE_CONTEXT_BYTES, MAX_MIDDLEWARE_FINDING_BYTES, + MAX_MIDDLEWARE_FINDINGS_PER_STAGE, MAX_MIDDLEWARE_METADATA_BYTES, + MAX_MIDDLEWARE_METADATA_ENTRIES, MAX_MIDDLEWARE_PAYLOAD_BYTES, + MAX_MIDDLEWARE_PREFLIGHT_TIMEOUT, MAX_MIDDLEWARE_REASON_BYTES, MIDDLEWARE_GRPC_MESSAGE_BYTES, + MiddlewareDenial, MiddlewareSessionAdmission, MiddlewareSessionPermit, MiddlewareWorkAdmission, + NamespacedFinding, OnError, is_stable_reason_code, middleware_denial_reason, +}; + +const STREAM_CHANNEL_CAPACITY: usize = 4; +const MAX_REQUESTED_SUBPROTOCOLS: usize = 32; +const MAX_SUBPROTOCOL_BYTES: usize = 4 * 1024; +const MAX_SELECTED_SUBPROTOCOL_BYTES: usize = 256; +#[derive(Debug, Clone)] +pub struct WebSocketPreflightInput { + pub session_id: String, + pub request_id: String, + pub sandbox_id: String, + pub scheme: String, + pub host: String, + pub port: u16, + /// Raw request path without a query string. + pub path: String, + pub requested_subprotocols: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WebSocketInvocationOutcome { + Inspect, + Skip, + Allow, + Deny, + FailOpen, + FailClosed, +} + +#[derive(Debug, Clone)] +pub struct WebSocketInvocation { + pub config_name: String, + pub implementation: String, + pub outcome: WebSocketInvocationOutcome, + pub sequence: Option, + pub original_size: usize, + pub replacement_size: Option, + pub transformed: bool, + pub failed: bool, + /// The stage stream became unusable and will be bypassed for the rest of + /// this session when its policy is `fail_open`. + pub stage_disabled: bool, + pub reason_code: Option, +} + +/// Coverage of a selected policy attachment that did not result in a +/// WebSocket middleware invocation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WebSocketCoverageState { + /// The attached implementation did not advertise the exact + /// `WEBSOCKET_MESSAGE/PRE_CREDENTIALS` binding. + BindingNotSelected, + /// The binding was selected, but the V1 relay does not expose this message + /// class to middleware. + UnsupportedMessageType, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WebSocketCoverage { + pub config_name: String, + pub implementation: String, + pub state: WebSocketCoverageState, + pub sequence: Option, + pub message_type: Option, + pub original_size: usize, +} + +pub struct WebSocketPreflightResult { + pub allowed: bool, + /// Typed terminal reason when preflight denied the upgrade. `None` means + /// the request may continue, including voluntary skip and fail-open. + pub terminal_reason: Option, + pub reason: String, + pub denial: Option, + pub session: Option, + pub findings: Vec, + pub metadata: BTreeMap>, + pub invocations: Vec, + pub coverage: Vec, + pub saturated: bool, + pub session_capacity_exhausted: bool, +} + +#[derive(Debug)] +pub struct WebSocketSessionStartOutcome { + pub allowed: bool, + /// Typed terminal reason when session start cannot continue. + pub terminal_reason: Option, + pub reason: String, + pub invocations: Vec, +} + +#[derive(Debug)] +pub struct WebSocketMessageOutcome { + pub allowed: bool, + pub reason: String, + pub payload: Vec, + pub findings: Vec, + pub metadata: BTreeMap>, + pub invocations: Vec, + pub denial: Option, + pub saturated: bool, + pub platform_oversize: bool, +} + +struct WebSocketStageTransport { + sender: mpsc::Sender, + responses: super::WebSocketResponseStream, +} + +struct WebSocketStage { + entry: DescribedChainEntry, + transport: Option, +} + +impl WebSocketStage { + fn is_active(&self) -> bool { + self.transport.is_some() + } + + async fn disable(&mut self) { + self.end(WebSocketSessionEndReason::MiddlewareFailure).await; + } + + async fn end(&mut self, reason: WebSocketSessionEndReason) { + if let Some(transport) = self.transport.take() { + let _ = tokio::time::timeout( + Duration::from_millis(10), + transport.sender.send(session_end_request(reason)), + ) + .await; + } + } +} + +pub struct WebSocketSession { + runner: ChainRunner, + stages: Vec, + next_sequence: u64, + session_admission: Option, +} + +/// Whether a WebSocket text message needs the shared short-lived work budget. +/// +/// A fully disabled fail-open session returns `Bypass` without touching the +/// work semaphore. Other parsed WebSocket features may continue processing the +/// original payload independently. +#[derive(Debug)] +pub enum WebSocketMessageAdmission { + Bypass, + Inspect(MiddlewareWorkAdmission), +} + +enum OpenStage { + Inspect(Box, PreflightStageOutcome), + Deny(Box, PreflightStageOutcome), + Skip(PreflightStageOutcome), + Failed(DescribedChainEntry, String), +} + +struct PreflightStageOutcome { + invocation: WebSocketInvocation, + findings: Vec, + metadata: Option<(String, BTreeMap)>, +} + +impl ChainRunner { + pub async fn preflight_websocket( + &self, + entries: &[ChainEntry], + input: WebSocketPreflightInput, + ) -> miette::Result { + if entries.is_empty() { + return Ok(empty_preflight_result()); + } + let description = self + .describe_chain_for( + entries, + openshell_core::proto::SupervisorMiddlewareOperation::WebsocketMessage, + SupervisorMiddlewarePhase::PreCredentials, + ) + .await?; + let coverage = description + .unbound + .iter() + .map(binding_not_selected_coverage) + .collect::>(); + let described = description.entries; + if described.is_empty() { + let mut result = empty_preflight_result(); + result.coverage = coverage; + return Ok(result); + } + validate_preflight_input(&input)?; + + // One permit covers the complete concurrent preflight fan-out. Permit + // wait is deliberate backpressure and is excluded from every deadline. + let preflight_work = self.reserve_middleware_work_admission().await?; + let saturated = preflight_work.saturated(); + let session_admission = match self.try_reserve_middleware_session() { + MiddlewareSessionAdmission::Admitted(admission) => admission, + MiddlewareSessionAdmission::AtCapacity => { + return Ok(session_capacity_exhausted(described, coverage, saturated)); + } + }; + let opened = join_all( + described + .into_iter() + .map(|entry| open_stage(entry, input.clone())), + ) + .await; + + let mut stages = Vec::new(); + let mut invocations = Vec::new(); + let mut findings = Vec::new(); + let mut metadata = BTreeMap::new(); + let mut denial = None; + let mut fail_closed_reason = None; + for result in opened { + match result { + OpenStage::Inspect(stage, outcome) => { + stages.push(*stage); + collect_preflight_outcome( + outcome, + &mut invocations, + &mut findings, + &mut metadata, + ); + } + OpenStage::Deny(stage, outcome) => { + denial.get_or_insert_with(|| MiddlewareDenial { + config_name: stage.entry.entry.name.clone(), + reason_code: outcome.invocation.reason_code.clone(), + }); + stages.push(*stage); + collect_preflight_outcome( + outcome, + &mut invocations, + &mut findings, + &mut metadata, + ); + } + OpenStage::Skip(outcome) => collect_preflight_outcome( + outcome, + &mut invocations, + &mut findings, + &mut metadata, + ), + OpenStage::Failed(entry, reason) => { + let invocation = failure_invocation(&entry, None, 0, &reason); + if entry.entry.on_error == OnError::FailClosed { + fail_closed_reason + .get_or_insert_with(|| format!("middleware_failed: {reason}")); + } + invocations.push(invocation); + } + } + } + + if let Some(denial) = denial { + end_stages(&mut stages, WebSocketSessionEndReason::MiddlewareDenial).await; + return Ok(WebSocketPreflightResult { + allowed: false, + terminal_reason: Some(WebSocketSessionEndReason::MiddlewareDenial), + reason: middleware_denial_reason( + &denial.config_name, + denial.reason_code.as_deref(), + ), + denial: Some(denial), + session: None, + findings, + metadata, + invocations, + coverage, + saturated, + session_capacity_exhausted: false, + }); + } + + if let Some(reason) = fail_closed_reason { + end_stages(&mut stages, WebSocketSessionEndReason::MiddlewareFailure).await; + return Ok(WebSocketPreflightResult { + allowed: false, + terminal_reason: Some(WebSocketSessionEndReason::MiddlewareFailure), + reason, + denial: None, + session: None, + findings, + metadata, + invocations, + coverage, + saturated, + session_capacity_exhausted: false, + }); + } + + if stages.is_empty() { + drop(session_admission); + return Ok(WebSocketPreflightResult { + allowed: true, + terminal_reason: None, + reason: String::new(), + denial: None, + session: None, + findings, + metadata, + invocations, + coverage, + saturated, + session_capacity_exhausted: false, + }); + } + + Ok(WebSocketPreflightResult { + allowed: true, + terminal_reason: None, + reason: String::new(), + denial: None, + session: Some(WebSocketSession { + runner: self.clone(), + stages, + next_sequence: 1, + session_admission: Some(session_admission), + }), + findings, + metadata, + invocations, + coverage, + saturated, + session_capacity_exhausted: false, + }) + } +} + +fn empty_preflight_result() -> WebSocketPreflightResult { + WebSocketPreflightResult { + allowed: true, + terminal_reason: None, + reason: String::new(), + denial: None, + session: None, + findings: Vec::new(), + metadata: BTreeMap::new(), + invocations: Vec::new(), + coverage: Vec::new(), + saturated: false, + session_capacity_exhausted: false, + } +} + +fn session_capacity_exhausted( + described: Vec, + coverage: Vec, + saturated: bool, +) -> WebSocketPreflightResult { + let reason = "middleware_session_capacity_exhausted"; + let mut fail_closed = false; + let invocations = described + .iter() + .map(|entry| { + fail_closed |= entry.entry.on_error == OnError::FailClosed; + failure_invocation(entry, None, 0, reason) + }) + .collect(); + WebSocketPreflightResult { + allowed: !fail_closed, + terminal_reason: fail_closed.then_some(WebSocketSessionEndReason::MiddlewareFailure), + reason: if fail_closed { + format!("middleware_failed: {reason}") + } else { + String::new() + }, + denial: None, + session: None, + findings: Vec::new(), + metadata: BTreeMap::new(), + invocations, + coverage, + saturated, + session_capacity_exhausted: true, + } +} + +impl WebSocketSession { + fn has_active_stages(&self) -> bool { + self.stages.iter().any(WebSocketStage::is_active) + } + + fn reconcile_lifecycle(&mut self) { + if !self.has_active_stages() { + self.session_admission.take(); + } + } + + pub async fn admit_message(&mut self) -> miette::Result { + self.reconcile_lifecycle(); + if !self.has_active_stages() { + return Ok(WebSocketMessageAdmission::Bypass); + } + self.runner + .reserve_middleware_work_admission() + .await + .map(WebSocketMessageAdmission::Inspect) + } + + pub async fn start(&mut self, selected_subprotocol: &str) -> WebSocketSessionStartOutcome { + if selected_subprotocol.len() > MAX_SELECTED_SUBPROTOCOL_BYTES { + return WebSocketSessionStartOutcome { + allowed: false, + terminal_reason: Some(WebSocketSessionEndReason::MiddlewareFailure), + reason: "middleware_failed: selected_subprotocol_over_capacity".to_string(), + invocations: Vec::new(), + }; + } + let mut invocations = Vec::new(); + let mut fail_closed = None; + for stage in &mut self.stages { + let Some(transport) = stage.transport.as_mut() else { + continue; + }; + let request = WebSocketSessionEvent { + event: Some(web_socket_session_event::Event::SessionStart( + WebSocketSessionStart { + selected_subprotocol: selected_subprotocol.to_string(), + }, + )), + }; + let sent = + tokio::time::timeout(stage.entry.timeout, transport.sender.send(request)).await; + if !matches!(sent, Ok(Ok(()))) { + stage.disable().await; + let reason = "session_start_send_failed"; + let mut invocation = failure_invocation(&stage.entry, None, 0, reason); + invocation.stage_disabled = true; + if stage.entry.entry.on_error == OnError::FailClosed { + fail_closed.get_or_insert_with(|| format!("middleware_failed: {reason}")); + } + invocations.push(invocation); + } + } + self.reconcile_lifecycle(); + WebSocketSessionStartOutcome { + allowed: fail_closed.is_none(), + terminal_reason: fail_closed + .as_ref() + .map(|_| WebSocketSessionEndReason::MiddlewareFailure), + reason: fail_closed.unwrap_or_default(), + invocations, + } + } + + /// Record one logical message whose type the selected V1 WebSocket binding + /// cannot inspect. This is a coverage state, not a middleware failure: + /// `on_error` is not applied and active stages remain available for later + /// text messages. + pub fn observe_unsupported_message( + &mut self, + message_type: WebSocketMessageType, + original_size: usize, + ) -> Vec { + self.reconcile_lifecycle(); + if !self.has_active_stages() { + return Vec::new(); + } + let sequence = self.next_sequence; + self.next_sequence = self.next_sequence.saturating_add(1); + self.stages + .iter() + .filter(|stage| stage.is_active()) + .map(|stage| WebSocketCoverage { + config_name: stage.entry.entry.name.clone(), + implementation: stage.entry.entry.implementation.clone(), + state: WebSocketCoverageState::UnsupportedMessageType, + sequence: Some(sequence), + message_type: Some(message_type), + original_size, + }) + .collect() + } + + pub async fn evaluate_text(&mut self, payload: Vec) -> WebSocketMessageOutcome { + if payload.len() > MAX_MIDDLEWARE_PAYLOAD_BYTES { + return platform_oversize_outcome(payload); + } + match self.admit_message().await { + Ok(WebSocketMessageAdmission::Bypass) => bypassed_message_outcome(payload), + Ok(WebSocketMessageAdmission::Inspect(admission)) => { + self.evaluate_text_admitted(payload, admission).await + } + Err(_) => admission_failure_outcome(payload), + } + } + + pub async fn evaluate_text_admitted( + &mut self, + payload: Vec, + admission: MiddlewareWorkAdmission, + ) -> WebSocketMessageOutcome { + let outcome = self.evaluate_text_admitted_inner(payload, admission).await; + self.reconcile_lifecycle(); + outcome + } + + async fn evaluate_text_admitted_inner( + &mut self, + payload: Vec, + admission: MiddlewareWorkAdmission, + ) -> WebSocketMessageOutcome { + if payload.len() > MAX_MIDDLEWARE_PAYLOAD_BYTES { + return platform_oversize_outcome(payload); + } + + let saturated = admission.saturated(); + let sequence = self.next_sequence; + self.next_sequence = self.next_sequence.saturating_add(1); + let chain_deadline = Instant::now() + MAX_MIDDLEWARE_CHAIN_TIMEOUT; + let mut current = payload; + let mut findings = Vec::new(); + let mut metadata = BTreeMap::new(); + let mut invocations = Vec::new(); + + for stage in &mut self.stages { + if !stage.is_active() { + continue; + } + let original_size = current.len(); + if original_size > stage.entry.max_payload_bytes { + let reason = "request_message_over_capacity"; + let invocation = + failure_invocation(&stage.entry, Some(sequence), original_size, reason); + let fail_closed = stage.entry.entry.on_error == OnError::FailClosed; + invocations.push(invocation); + if fail_closed { + return denied_message_outcome( + current, + findings, + metadata, + invocations, + format!("middleware_failed: {reason}"), + None, + saturated, + ); + } + continue; + } + + let remaining = chain_deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + let reason = "middleware_chain_timeout"; + let mut invocation = + failure_invocation(&stage.entry, Some(sequence), original_size, reason); + let fail_closed = stage.entry.entry.on_error == OnError::FailClosed; + stage.disable().await; + invocation.stage_disabled = true; + invocations.push(invocation); + if fail_closed { + return denied_message_outcome( + current, + findings, + metadata, + invocations, + format!("middleware_failed: {reason}"), + None, + saturated, + ); + } + continue; + } + let stage_timeout = stage.entry.timeout.min(remaining); + let request = WebSocketSessionEvent { + event: Some(web_socket_session_event::Event::Message(WebSocketMessage { + sequence, + message_type: WebSocketMessageType::Text as i32, + payload: current.clone(), + })), + }; + let response = { + let transport = stage + .transport + .as_mut() + .expect("active WebSocket stage has transport"); + tokio::time::timeout(stage_timeout, async { + transport + .sender + .send(request) + .await + .map_err(|_| tonic::Status::unavailable("request stream closed"))?; + transport.responses.next().await.transpose() + }) + .await + }; + let result = match response { + Ok(Ok(Some(response))) => match response.result { + Some(web_socket_session_event_result::Result::MessageResult(result)) => result, + Some(web_socket_session_event_result::Result::PreflightDecision(_)) | None => { + if let Some(outcome) = handle_stage_failure( + stage, + sequence, + original_size, + "unexpected_websocket_response", + ¤t, + &findings, + &metadata, + &mut invocations, + saturated, + ) + .await + { + return outcome; + } + continue; + } + }, + Ok(Ok(None)) => { + if let Some(outcome) = handle_stage_failure( + stage, + sequence, + original_size, + "missing_message_result", + ¤t, + &findings, + &metadata, + &mut invocations, + saturated, + ) + .await + { + return outcome; + } + continue; + } + Ok(Err(error)) => { + let reason = stage.entry.service.as_ref().map_or_else( + || "binding_not_described".to_string(), + |service| service.diagnostic_policy.error_reason(&error), + ); + if let Some(outcome) = handle_stage_failure( + stage, + sequence, + original_size, + &reason, + ¤t, + &findings, + &metadata, + &mut invocations, + saturated, + ) + .await + { + return outcome; + } + continue; + } + Err(_) => { + if let Some(outcome) = handle_stage_failure( + stage, + sequence, + original_size, + "middleware_timeout", + ¤t, + &findings, + &metadata, + &mut invocations, + saturated, + ) + .await + { + return outcome; + } + continue; + } + }; + + let result = + match validate_message_result(result, sequence, stage.entry.max_payload_bytes) { + Ok(result) => result, + Err(reason) => { + if let Some(outcome) = handle_stage_failure( + stage, + sequence, + original_size, + reason, + ¤t, + &findings, + &metadata, + &mut invocations, + saturated, + ) + .await + { + return outcome; + } + continue; + } + }; + + let decision = Decision::try_from(result.decision).expect("validated decision"); + let reason_code = (!result.reason_code.is_empty()).then(|| result.reason_code.clone()); + for finding in result.findings { + findings.push(NamespacedFinding { + middleware: stage.entry.entry.name.clone(), + finding, + }); + } + if !result.metadata.is_empty() { + metadata.insert( + stage.entry.entry.name.clone(), + result.metadata.into_iter().collect(), + ); + } + if decision == Decision::Deny { + let denial = MiddlewareDenial { + config_name: stage.entry.entry.name.clone(), + reason_code, + }; + invocations.push(success_invocation( + &stage.entry, + WebSocketInvocationOutcome::Deny, + sequence, + original_size, + None, + false, + denial.reason_code.clone(), + )); + return denied_message_outcome( + current, + findings, + metadata, + invocations, + middleware_denial_reason(&denial.config_name, denial.reason_code.as_deref()), + Some(denial), + saturated, + ); + } + + let replacement_size = result.has_replacement.then_some(result.replacement.len()); + if result.has_replacement { + current = result.replacement; + } + invocations.push(success_invocation( + &stage.entry, + WebSocketInvocationOutcome::Allow, + sequence, + original_size, + replacement_size, + result.has_replacement, + reason_code, + )); + } + + WebSocketMessageOutcome { + allowed: true, + reason: String::new(), + payload: current, + findings, + metadata, + invocations, + denial: None, + saturated, + platform_oversize: false, + } + } + + pub async fn end(mut self, reason: WebSocketSessionEndReason) { + end_stages(&mut self.stages, reason).await; + self.reconcile_lifecycle(); + } +} + +impl Drop for WebSocketSession { + fn drop(&mut self) { + end_stages_now(&mut self.stages, WebSocketSessionEndReason::Cancellation); + self.reconcile_lifecycle(); + } +} + +fn platform_oversize_outcome(payload: Vec) -> WebSocketMessageOutcome { + WebSocketMessageOutcome { + allowed: false, + reason: "websocket_message_over_platform_capacity".to_string(), + payload, + findings: Vec::new(), + metadata: BTreeMap::new(), + invocations: Vec::new(), + denial: None, + saturated: false, + platform_oversize: true, + } +} + +fn admission_failure_outcome(payload: Vec) -> WebSocketMessageOutcome { + WebSocketMessageOutcome { + allowed: false, + reason: "middleware_admission_over_capacity".to_string(), + payload, + findings: Vec::new(), + metadata: BTreeMap::new(), + invocations: Vec::new(), + denial: None, + saturated: true, + platform_oversize: false, + } +} + +fn bypassed_message_outcome(payload: Vec) -> WebSocketMessageOutcome { + WebSocketMessageOutcome { + allowed: true, + reason: String::new(), + payload, + findings: Vec::new(), + metadata: BTreeMap::new(), + invocations: Vec::new(), + denial: None, + saturated: false, + platform_oversize: false, + } +} + +fn binding_not_selected_coverage(entry: &ChainEntry) -> WebSocketCoverage { + WebSocketCoverage { + config_name: entry.name.clone(), + implementation: entry.implementation.clone(), + state: WebSocketCoverageState::BindingNotSelected, + sequence: None, + message_type: None, + original_size: 0, + } +} + +fn preflight_stage_outcome( + entry: &DescribedChainEntry, + outcome: WebSocketInvocationOutcome, + decision: WebSocketPreflightDecision, +) -> PreflightStageOutcome { + let reason_code = (!decision.reason_code.is_empty()).then(|| decision.reason_code.clone()); + let findings = decision + .findings + .into_iter() + .map(|finding| NamespacedFinding { + middleware: entry.entry.name.clone(), + finding, + }) + .collect(); + let metadata = (!decision.metadata.is_empty()).then(|| { + ( + entry.entry.name.clone(), + decision.metadata.into_iter().collect(), + ) + }); + PreflightStageOutcome { + invocation: success_invocation(entry, outcome, 0, 0, None, false, reason_code), + findings, + metadata, + } +} + +fn collect_preflight_outcome( + outcome: PreflightStageOutcome, + invocations: &mut Vec, + findings: &mut Vec, + metadata: &mut BTreeMap>, +) { + invocations.push(outcome.invocation); + findings.extend(outcome.findings); + if let Some((middleware, values)) = outcome.metadata { + metadata.insert(middleware, values); + } +} + +async fn open_stage(entry: DescribedChainEntry, input: WebSocketPreflightInput) -> OpenStage { + let Some(service) = entry.service.as_ref() else { + return OpenStage::Failed(entry, "binding_not_described".into()); + }; + let endpoint = Arc::clone(&service.endpoint); + let diagnostic_policy = service.diagnostic_policy; + let preflight = WebSocketPreflight { + session_id: input.session_id, + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + context: Some(RequestContext { + request_id: input.request_id, + sandbox_id: input.sandbox_id, + originating_process: None, + }), + target: Some(HttpRequestTarget { + scheme: input.scheme, + host: input.host, + port: u32::from(input.port), + method: "GET".into(), + path: input.path, + query: String::new(), + }), + requested_subprotocols: input.requested_subprotocols, + middleware_name: entry.entry.implementation.clone(), + config: Some(entry.entry.config.clone()), + }; + if validate_preflight_envelope(&preflight).is_err() { + return OpenStage::Failed(entry, "preflight_envelope_over_capacity".into()); + } + + let timeout = entry.timeout.min(MAX_MIDDLEWARE_PREFLIGHT_TIMEOUT); + let opened = tokio::time::timeout(timeout, async { + let (sender, receiver) = mpsc::channel(STREAM_CHANNEL_CAPACITY); + sender + .send(WebSocketSessionEvent { + event: Some(web_socket_session_event::Event::Preflight(preflight)), + }) + .await + .map_err(|_| tonic::Status::unavailable("request stream closed"))?; + let mut responses = endpoint.open_websocket_session(receiver).await?; + let response = responses.next().await.transpose()?; + Ok::<_, tonic::Status>((sender, responses, response)) + }) + .await; + + let (sender, responses, response) = match opened { + Ok(Ok(opened)) => opened, + Ok(Err(error)) => { + return OpenStage::Failed(entry, diagnostic_policy.error_reason(&error)); + } + Err(_) => return OpenStage::Failed(entry, "middleware_timeout".into()), + }; + let Some(response) = response else { + let _ = sender.try_send(session_end_request( + WebSocketSessionEndReason::MiddlewareFailure, + )); + return OpenStage::Failed(entry, "missing_preflight_decision".into()); + }; + let Some(web_socket_session_event_result::Result::PreflightDecision(decision)) = + response.result + else { + let _ = sender.try_send(session_end_request( + WebSocketSessionEndReason::MiddlewareFailure, + )); + return OpenStage::Failed(entry, "invalid_preflight_decision".into()); + }; + let decision = match validate_preflight_decision(decision) { + Ok(decision) => decision, + Err(reason) => { + let _ = sender.try_send(session_end_request( + WebSocketSessionEndReason::MiddlewareFailure, + )); + return OpenStage::Failed(entry, reason.into()); + } + }; + let action = + WebSocketPreflightAction::try_from(decision.action).expect("validated preflight action"); + match action { + WebSocketPreflightAction::Inspect => { + let outcome = + preflight_stage_outcome(&entry, WebSocketInvocationOutcome::Inspect, decision); + OpenStage::Inspect( + Box::new(WebSocketStage { + entry, + transport: Some(WebSocketStageTransport { sender, responses }), + }), + outcome, + ) + } + WebSocketPreflightAction::Deny => { + let outcome = + preflight_stage_outcome(&entry, WebSocketInvocationOutcome::Deny, decision); + OpenStage::Deny( + Box::new(WebSocketStage { + entry, + transport: Some(WebSocketStageTransport { sender, responses }), + }), + outcome, + ) + } + WebSocketPreflightAction::Skip => { + let outcome = + preflight_stage_outcome(&entry, WebSocketInvocationOutcome::Skip, decision); + let _ = sender.try_send(session_end_request(WebSocketSessionEndReason::Cancellation)); + OpenStage::Skip(outcome) + } + WebSocketPreflightAction::Unspecified => { + unreachable!("validated preflight action cannot be unspecified") + } + } +} + +fn validate_preflight_input(input: &WebSocketPreflightInput) -> miette::Result<()> { + if input.session_id.is_empty() || input.session_id.len() > 128 { + return Err(miette::miette!("invalid WebSocket middleware session id")); + } + if input.path.len() > super::MAX_MIDDLEWARE_TARGET_BYTES { + return Err(miette::miette!( + "WebSocket middleware preflight path exceeds platform capacity" + )); + } + if input.path.contains('?') { + return Err(miette::miette!( + "WebSocket middleware preflight path must not contain a query string" + )); + } + if input.requested_subprotocols.len() > MAX_REQUESTED_SUBPROTOCOLS + || input + .requested_subprotocols + .iter() + .map(String::len) + .sum::() + > MAX_SUBPROTOCOL_BYTES + { + return Err(miette::miette!( + "WebSocket middleware requested subprotocols exceed platform capacity" + )); + } + Ok(()) +} + +fn validate_preflight_envelope(preflight: &WebSocketPreflight) -> Result<(), &'static str> { + let Some(target) = preflight.target.as_ref() else { + return Err("preflight_target_missing"); + }; + if target.method != "GET" { + return Err("preflight_target_method_invalid"); + } + if !target.query.is_empty() || target.path.contains('?') { + return Err("preflight_target_query_present"); + } + if target.encoded_len() > super::MAX_MIDDLEWARE_TARGET_BYTES { + return Err("preflight_target_over_capacity"); + } + if preflight + .config + .as_ref() + .is_some_and(|config| config.encoded_len() > MAX_MIDDLEWARE_CONFIG_BYTES) + { + return Err("preflight_config_over_capacity"); + } + if preflight + .context + .as_ref() + .is_some_and(|context| context.encoded_len() > MAX_MIDDLEWARE_CONTEXT_BYTES) + { + return Err("preflight_context_over_capacity"); + } + if preflight.encoded_len() > MIDDLEWARE_GRPC_MESSAGE_BYTES { + return Err("preflight_envelope_over_capacity"); + } + Ok(()) +} + +fn validate_preflight_decision( + decision: WebSocketPreflightDecision, +) -> Result { + if !matches!( + WebSocketPreflightAction::try_from(decision.action), + Ok(WebSocketPreflightAction::Inspect + | WebSocketPreflightAction::Skip + | WebSocketPreflightAction::Deny) + ) { + return Err("invalid_preflight_decision"); + } + if decision.reason.len() > MAX_MIDDLEWARE_REASON_BYTES { + return Err("response_reason_over_capacity"); + } + if !decision.reason_code.is_empty() && !is_stable_reason_code(&decision.reason_code) { + return Err("response_reason_code_invalid"); + } + if decision.findings.len() > MAX_MIDDLEWARE_FINDINGS_PER_STAGE + || decision + .findings + .iter() + .any(|finding| finding.encoded_len() > MAX_MIDDLEWARE_FINDING_BYTES) + { + return Err("response_findings_over_capacity"); + } + if decision.metadata.len() > MAX_MIDDLEWARE_METADATA_ENTRIES { + return Err("response_metadata_count_over_capacity"); + } + let metadata_bytes = decision + .metadata + .iter() + .fold(0usize, |total, (key, value)| { + total.saturating_add(key.len()).saturating_add(value.len()) + }); + if metadata_bytes > MAX_MIDDLEWARE_METADATA_BYTES { + return Err("response_metadata_bytes_over_capacity"); + } + if decision.encoded_len() > MIDDLEWARE_GRPC_MESSAGE_BYTES { + return Err("response_envelope_over_capacity"); + } + Ok(decision) +} + +fn validate_message_result( + result: WebSocketMessageResult, + sequence: u64, + stage_limit: usize, +) -> Result { + if result.sequence != sequence { + return Err("message_result_sequence_mismatch"); + } + if !matches!( + Decision::try_from(result.decision), + Ok(Decision::Allow | Decision::Deny) + ) { + return Err("invalid_response_decision"); + } + if result.reason.len() > MAX_MIDDLEWARE_REASON_BYTES { + return Err("response_reason_over_capacity"); + } + if !result.reason_code.is_empty() && !is_stable_reason_code(&result.reason_code) { + return Err("response_reason_code_invalid"); + } + if !result.has_replacement && !result.replacement.is_empty() { + return Err("unsolicited_replacement"); + } + if result.has_replacement { + if result.replacement.len() > MAX_MIDDLEWARE_PAYLOAD_BYTES { + return Err("response_message_over_platform_capacity"); + } + if result.replacement.len() > stage_limit { + return Err("response_message_over_capacity"); + } + if std::str::from_utf8(&result.replacement).is_err() { + return Err("text_replacement_invalid_utf8"); + } + } + if result.findings.len() > MAX_MIDDLEWARE_FINDINGS_PER_STAGE + || result + .findings + .iter() + .any(|finding| finding.encoded_len() > MAX_MIDDLEWARE_FINDING_BYTES) + { + return Err("response_findings_over_capacity"); + } + if result.metadata.len() > MAX_MIDDLEWARE_METADATA_ENTRIES { + return Err("response_metadata_count_over_capacity"); + } + let metadata_bytes = result.metadata.iter().fold(0usize, |total, (key, value)| { + total.saturating_add(key.len()).saturating_add(value.len()) + }); + if metadata_bytes > MAX_MIDDLEWARE_METADATA_BYTES { + return Err("response_metadata_bytes_over_capacity"); + } + if result.encoded_len() > MIDDLEWARE_GRPC_MESSAGE_BYTES { + return Err("response_envelope_over_capacity"); + } + Ok(result) +} + +#[allow(clippy::too_many_arguments)] +async fn handle_stage_failure( + stage: &mut WebSocketStage, + sequence: u64, + original_size: usize, + reason: &str, + current: &[u8], + findings: &[NamespacedFinding], + metadata: &BTreeMap>, + invocations: &mut Vec, + saturated: bool, +) -> Option { + stage.disable().await; + let mut invocation = failure_invocation(&stage.entry, Some(sequence), original_size, reason); + invocation.stage_disabled = true; + invocations.push(invocation); + (stage.entry.entry.on_error == OnError::FailClosed).then(|| { + denied_message_outcome( + current.to_vec(), + findings.to_vec(), + metadata.clone(), + invocations.clone(), + format!("middleware_failed: {reason}"), + None, + saturated, + ) + }) +} + +fn denied_message_outcome( + payload: Vec, + findings: Vec, + metadata: BTreeMap>, + invocations: Vec, + reason: String, + denial: Option, + saturated: bool, +) -> WebSocketMessageOutcome { + WebSocketMessageOutcome { + allowed: false, + reason, + payload, + findings, + metadata, + invocations, + denial, + saturated, + platform_oversize: false, + } +} + +fn success_invocation( + entry: &DescribedChainEntry, + outcome: WebSocketInvocationOutcome, + sequence: u64, + original_size: usize, + replacement_size: Option, + transformed: bool, + reason_code: Option, +) -> WebSocketInvocation { + WebSocketInvocation { + config_name: entry.entry.name.clone(), + implementation: entry.entry.implementation.clone(), + outcome, + sequence: (sequence != 0).then_some(sequence), + original_size, + replacement_size, + transformed, + failed: false, + stage_disabled: false, + reason_code, + } +} + +fn failure_invocation( + entry: &DescribedChainEntry, + sequence: Option, + original_size: usize, + _reason: &str, +) -> WebSocketInvocation { + let outcome = match entry.entry.on_error { + OnError::FailOpen => WebSocketInvocationOutcome::FailOpen, + OnError::FailClosed => WebSocketInvocationOutcome::FailClosed, + }; + WebSocketInvocation { + config_name: entry.entry.name.clone(), + implementation: entry.entry.implementation.clone(), + outcome, + sequence, + original_size, + replacement_size: None, + transformed: false, + failed: true, + stage_disabled: false, + reason_code: None, + } +} + +async fn end_stages(stages: &mut [WebSocketStage], reason: WebSocketSessionEndReason) { + for stage in stages { + stage.end(reason).await; + } +} + +fn end_stages_now(stages: &mut [WebSocketStage], reason: WebSocketSessionEndReason) { + for stage in stages { + if let Some(transport) = stage.transport.take() { + let _ = transport.sender.try_send(session_end_request(reason)); + } + } +} + +fn session_end_request(reason: WebSocketSessionEndReason) -> WebSocketSessionEvent { + WebSocketSessionEvent { + event: Some(web_socket_session_event::Event::SessionEnd( + WebSocketSessionEnd { + reason: reason as i32, + }, + )), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn disabling_stage_sends_middleware_failure_once() { + let (sender, mut requests) = mpsc::channel(1); + let mut stage = WebSocketStage { + entry: DescribedChainEntry { + entry: ChainEntry { + name: "best-effort".into(), + implementation: "test/middleware".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailOpen, + }, + service: None, + binding: None, + max_payload_bytes: 1024, + timeout: Duration::from_secs(1), + }, + transport: Some(WebSocketStageTransport { + sender, + responses: Box::pin(tokio_stream::empty()), + }), + }; + + stage.disable().await; + stage.disable().await; + + let event = requests.recv().await.expect("session end event"); + let Some(web_socket_session_event::Event::SessionEnd(end)) = event.event else { + panic!("disabled stage must receive session end"); + }; + assert_eq!( + WebSocketSessionEndReason::try_from(end.reason), + Ok(WebSocketSessionEndReason::MiddlewareFailure) + ); + assert!( + requests.try_recv().is_err(), + "disabled stage must receive at most one session end" + ); + } +} diff --git a/crates/openshell-supervisor-network/Cargo.toml b/crates/openshell-supervisor-network/Cargo.toml index 3dead2b8be..6608738094 100644 --- a/crates/openshell-supervisor-network/Cargo.toml +++ b/crates/openshell-supervisor-network/Cargo.toml @@ -54,9 +54,11 @@ openshell-supervisor-middleware-builtins = { path = "../openshell-supervisor-mid tempfile = "3" tonic = { workspace = true } temp-env = "0.3" +tokio = { workspace = true, features = ["test-util"] } tokio-tungstenite = { workspace = true } futures = { workspace = true } tracing-subscriber = { workspace = true } +tokio-stream = { workspace = true, features = ["net"] } [target.'cfg(unix)'.dependencies] libc = "0.2" diff --git a/crates/openshell-supervisor-network/src/l7/middleware.rs b/crates/openshell-supervisor-network/src/l7/middleware.rs index 28f91c3bb7..5b4f9a202f 100644 --- a/crates/openshell-supervisor-network/src/l7/middleware.rs +++ b/crates/openshell-supervisor-network/src/l7/middleware.rs @@ -8,7 +8,8 @@ use crate::opa::PolicyGenerationGuard; use miette::{Result, miette}; use openshell_ocsf::{ ActionId, ActivityId, DetectionFindingBuilder, DispositionId, Endpoint, FindingInfo, - HttpActivityBuilder, HttpRequest, SeverityId, StatusId, Url as OcsfUrl, ocsf_emit, + HttpActivityBuilder, HttpRequest, NetworkActivityBuilder, SeverityId, StatusId, Url as OcsfUrl, + ocsf_emit, }; use std::path::PathBuf; use tokio::io::{AsyncRead, AsyncWrite}; @@ -18,6 +19,11 @@ pub enum MiddlewareApplyResult { Denied { denial: Option, }, + /// The platform's shared active-work and waiter capacities were both full. + /// + /// This is platform load shedding, not a selected middleware-stage failure, + /// so callers must not apply a stage's `on_error` policy. + AdmissionExhausted, } /// How traffic a middleware chain can never inspect (h2c, non-HTTP TCP, @@ -83,6 +89,295 @@ pub fn emit_middleware_uninspectable(ctx: &L7EvalContext, detail: &str, denied: ocsf_emit!(event); } +pub fn emit_websocket_preflight_events( + ctx: &L7EvalContext, + outcome: &openshell_supervisor_middleware::WebSocketPreflightResult, +) { + emit_websocket_invocations(ctx, &outcome.invocations); + emit_websocket_coverage(ctx, &outcome.coverage); + emit_websocket_findings(ctx, &outcome.findings); + if outcome.saturated { + emit_websocket_saturation(ctx); + } + if outcome.session_capacity_exhausted { + emit_middleware_session_capacity_exhausted(ctx); + } +} + +pub(super) fn emit_websocket_coverage( + ctx: &L7EvalContext, + coverage: &[openshell_supervisor_middleware::WebSocketCoverage], +) { + for event in websocket_coverage_events(ctx, coverage) { + ocsf_emit!(event); + } +} + +/// Build coverage events separately from the tracing pipeline so tests can +/// assert exact coverage telemetry without process-global callsite cache races. +pub(super) fn websocket_coverage_events( + ctx: &L7EvalContext, + coverage: &[openshell_supervisor_middleware::WebSocketCoverage], +) -> Vec { + coverage + .iter() + .map(|coverage| { + use openshell_supervisor_middleware::WebSocketCoverageState as State; + + let state = match coverage.state { + State::BindingNotSelected => "binding_not_selected", + State::UnsupportedMessageType => "unsupported_message_type", + }; + let sequence = coverage + .sequence + .map_or_else(|| "-".to_string(), |sequence| sequence.to_string()); + let message_type = match coverage.message_type { + Some(openshell_core::proto::WebSocketMessageType::Text) => "text", + Some(openshell_core::proto::WebSocketMessageType::Binary) => "binary", + Some(openshell_core::proto::WebSocketMessageType::Unspecified) | None => "-", + }; + NetworkActivityBuilder::new(openshell_ocsf::ctx::ctx()) + .activity(ActivityId::Other) + .activity_name("WebSocket middleware coverage") + .action(ActionId::Allowed) + .disposition(DispositionId::Allowed) + .severity(SeverityId::Informational) + .status(StatusId::Success) + .dst_endpoint(Endpoint::from_domain(&ctx.host, ctx.port)) + .firewall_rule(&ctx.policy_name, "supervisor-middleware") + .unmapped("middleware_config", coverage.config_name.as_str()) + .unmapped( + "middleware_implementation", + coverage.implementation.as_str(), + ) + .unmapped("coverage_state", state) + .unmapped("websocket_message_type", message_type) + .unmapped("websocket_sequence", sequence.as_str()) + .unmapped("input_bytes", coverage.original_size) + .message(format!( + "WEBSOCKET_MIDDLEWARE_COVERAGE state={state} config={} implementation={} sequence={sequence} message_type={message_type} input_bytes={}", + coverage.config_name, coverage.implementation, coverage.original_size, + )) + .build() + }) + .collect() +} + +pub(super) fn emit_websocket_session_start_events( + ctx: &L7EvalContext, + outcome: &openshell_supervisor_middleware::WebSocketSessionStartOutcome, +) { + emit_websocket_invocations(ctx, &outcome.invocations); +} + +pub(super) fn emit_websocket_message_events( + ctx: &L7EvalContext, + outcome: &openshell_supervisor_middleware::WebSocketMessageOutcome, +) { + emit_websocket_invocations(ctx, &outcome.invocations); + if outcome.saturated { + emit_websocket_saturation(ctx); + } + emit_websocket_findings(ctx, &outcome.findings); +} + +fn emit_websocket_findings( + ctx: &L7EvalContext, + findings: &[openshell_supervisor_middleware::NamespacedFinding], +) { + for finding in findings { + let event = DetectionFindingBuilder::new(openshell_ocsf::ctx::ctx()) + .severity(SeverityId::Medium) + .finding_info(FindingInfo::new( + "openshell.middleware.websocket_finding", + "WebSocket middleware finding", + )) + .evidence_pairs(&[ + ("policy", ctx.policy_name.as_str()), + ("host", ctx.host.as_str()), + ("middleware", finding.middleware.as_str()), + ]) + .message("WebSocket middleware reported a finding") + .build(); + ocsf_emit!(event); + } +} + +fn emit_websocket_invocations( + ctx: &L7EvalContext, + invocations: &[openshell_supervisor_middleware::WebSocketInvocation], +) { + for invocation in invocations { + use openshell_supervisor_middleware::WebSocketInvocationOutcome as Outcome; + let (action, disposition, severity, status, outcome_name) = match invocation.outcome { + Outcome::Inspect => ( + ActionId::Other, + DispositionId::Allowed, + SeverityId::Informational, + StatusId::Success, + "inspect", + ), + Outcome::Skip => ( + ActionId::Other, + DispositionId::Other, + SeverityId::Informational, + StatusId::Success, + "voluntary_skip", + ), + Outcome::Allow => ( + ActionId::Allowed, + DispositionId::Allowed, + SeverityId::Informational, + StatusId::Success, + "allow", + ), + Outcome::Deny => ( + ActionId::Denied, + DispositionId::Blocked, + SeverityId::Medium, + StatusId::Failure, + "deny", + ), + Outcome::FailOpen => ( + ActionId::Other, + DispositionId::Allowed, + SeverityId::Medium, + StatusId::Failure, + "fail_open", + ), + Outcome::FailClosed => ( + ActionId::Denied, + DispositionId::Blocked, + SeverityId::High, + StatusId::Failure, + "fail_closed", + ), + }; + let sequence = invocation + .sequence + .map_or_else(|| "-".to_string(), |sequence| sequence.to_string()); + let replacement_size = invocation + .replacement_size + .map_or_else(|| "-".to_string(), |size| size.to_string()); + let reason_code = invocation.reason_code.as_deref().unwrap_or("-"); + let event = NetworkActivityBuilder::new(openshell_ocsf::ctx::ctx()) + .activity(ActivityId::Other) + .activity_name("WebSocket middleware") + .action(action) + .disposition(disposition) + .severity(severity) + .status(status) + .dst_endpoint(Endpoint::from_domain(&ctx.host, ctx.port)) + .firewall_rule(&ctx.policy_name, "supervisor-middleware") + .message(format!( + "WEBSOCKET_MIDDLEWARE {outcome_name} config={} implementation={} sequence={sequence} input_bytes={} replacement_bytes={replacement_size} transformed={} reason_code={reason_code}", + invocation.config_name, + invocation.implementation, + invocation.original_size, + invocation.transformed, + )) + .build(); + ocsf_emit!(event); + if invocation.failed { + let event = DetectionFindingBuilder::new(openshell_ocsf::ctx::ctx()) + .severity(if invocation.outcome == Outcome::FailClosed { + SeverityId::High + } else { + SeverityId::Medium + }) + .finding_info(FindingInfo::new( + "openshell.middleware.websocket_failure", + "WebSocket middleware processing failure", + )) + .evidence_pairs(&[ + ("policy", ctx.policy_name.as_str()), + ("host", ctx.host.as_str()), + ("config", invocation.config_name.as_str()), + ("implementation", invocation.implementation.as_str()), + ("disposition", outcome_name), + ]) + .message("WebSocket middleware stage failed") + .build(); + ocsf_emit!(event); + } + if invocation.stage_disabled { + let event = DetectionFindingBuilder::new(openshell_ocsf::ctx::ctx()) + .severity(SeverityId::Medium) + .finding_info(FindingInfo::new( + "openshell.middleware.websocket_stage_disabled", + "WebSocket middleware stage disabled", + )) + .evidence_pairs(&[ + ("policy", ctx.policy_name.as_str()), + ("host", ctx.host.as_str()), + ("config", invocation.config_name.as_str()), + ("implementation", invocation.implementation.as_str()), + ("disposition", outcome_name), + ]) + .message( + "WebSocket middleware stage stream became unusable and was disabled for this session", + ) + .build(); + ocsf_emit!(event); + } + } +} + +fn emit_websocket_saturation(ctx: &L7EvalContext) { + let event = DetectionFindingBuilder::new(openshell_ocsf::ctx::ctx()) + .severity(SeverityId::Medium) + .finding_info(FindingInfo::new( + "openshell.middleware.admission_saturated", + "Supervisor middleware admission saturated", + )) + .evidence_pairs(&[ + ("policy", ctx.policy_name.as_str()), + ("host", ctx.host.as_str()), + ("operation", "websocket_message"), + ]) + .message("WebSocket middleware work waited for admission capacity") + .build(); + ocsf_emit!(event); +} + +fn emit_middleware_session_capacity_exhausted(ctx: &L7EvalContext) { + let event = DetectionFindingBuilder::new(openshell_ocsf::ctx::ctx()) + .severity(SeverityId::Medium) + .finding_info(FindingInfo::new( + "openshell.middleware.session_capacity_exhausted", + "Supervisor middleware session capacity exhausted", + )) + .evidence_pairs(&[ + ("policy", ctx.policy_name.as_str()), + ("host", ctx.host.as_str()), + ("operation", "websocket_message"), + ]) + .message("Persistent middleware session admission was refused at process capacity") + .build(); + ocsf_emit!(event); +} + +fn middleware_admission_exhausted_event(ctx: &L7EvalContext) -> openshell_ocsf::OcsfEvent { + DetectionFindingBuilder::new(openshell_ocsf::ctx::ctx()) + .severity(SeverityId::Medium) + .finding_info(FindingInfo::new( + "openshell.middleware.admission_exhausted", + "Supervisor middleware admission exhausted", + )) + .evidence_pairs(&[ + ("policy", ctx.policy_name.as_str()), + ("host", ctx.host.as_str()), + ("operation", "http_request"), + ("disposition", "shed"), + ]) + .message("HTTP request shed because middleware work admission was exhausted") + .build() +} + +fn emit_middleware_admission_exhausted(ctx: &L7EvalContext) { + ocsf_emit!(middleware_admission_exhausted_event(ctx)); +} + /// Largest body-buffering limit across the entries that actually resolved to a /// registered binding. Buffering for the most capable stage lets every stage /// that can handle the body run; stages whose own limit is smaller are failed @@ -98,7 +393,7 @@ pub(super) fn middleware_chain_body_limit( chain .iter() .filter(|entry| entry.is_resolved()) - .map(openshell_supervisor_middleware::DescribedChainEntry::max_body_bytes) + .map(openshell_supervisor_middleware::DescribedChainEntry::max_payload_bytes) .max() } @@ -139,6 +434,20 @@ pub async fn apply_middleware_chain_for_scheme Some(admission), + openshell_supervisor_middleware::MiddlewareWorkAdmissionOutcome::QueueExhausted => { + emit_middleware_admission_exhausted(ctx); + return Ok(MiddlewareApplyResult::AdmissionExhausted); + } + } + }; let Some(max_body_bytes) = middleware_chain_body_limit(&chain) else { // No entry resolved to a registered binding, so nothing inspects the // body. Apply each entry's `on_error` policy without buffering (an @@ -153,7 +462,14 @@ pub async fn apply_middleware_chain_for_scheme( + req: &crate::l7::provider::L7Request, + client: &mut C, + ctx: &L7EvalContext, + redacted_target: &str, +) -> Result<()> { + crate::l7::rest::send_middleware_unavailable_response( + req, + &ctx.policy_name, + client, + Some(redacted_target), + Some(crate::l7::rest::DenyResponseContext::from_l7_context(ctx)), + ) + .await +} + pub(super) fn middleware_request_input( scheme: &str, req: &crate::l7::provider::L7Request, @@ -544,10 +887,87 @@ fn emit_middleware_events( #[cfg(test)] mod tests { - use super::{safe_middleware_headers, send_middleware_rejection_response}; + use super::{ + middleware_admission_exhausted_event, safe_middleware_headers, + send_middleware_rejection_response, websocket_coverage_events, + }; use crate::l7::relay::L7EvalContext; use tokio::io::AsyncReadExt; + #[test] + fn admission_exhaustion_event_contains_only_platform_context() { + let ctx = L7EvalContext { + host: "api.example.test".into(), + port: 443, + policy_name: "api-policy".into(), + binary_path: "DO_NOT_LOG_BINARY".into(), + ..Default::default() + }; + + let serialized = serde_json::to_string(&middleware_admission_exhausted_event(&ctx)) + .expect("serialize admission event"); + assert!(serialized.contains("openshell.middleware.admission_exhausted")); + assert!(serialized.contains("api-policy")); + assert!(serialized.contains("api.example.test")); + assert!(serialized.contains("http_request")); + assert!(serialized.contains("shed")); + assert!(!serialized.contains("DO_NOT_LOG_BINARY")); + assert!(!serialized.contains("middleware_config")); + assert!(!serialized.contains("request_body")); + assert!(!serialized.contains("query")); + } + + #[test] + fn websocket_coverage_events_distinguish_selection_from_message_support() { + use openshell_core::proto::WebSocketMessageType; + use openshell_ocsf::SeverityId; + use openshell_supervisor_middleware::{WebSocketCoverage, WebSocketCoverageState as State}; + + let ctx = L7EvalContext { + host: "api.example.test".into(), + port: 443, + policy_name: "api-policy".into(), + ..Default::default() + }; + let events = websocket_coverage_events( + &ctx, + &[ + WebSocketCoverage { + config_name: "http-guard".into(), + implementation: "example/http-only".into(), + state: State::BindingNotSelected, + sequence: None, + message_type: None, + original_size: 0, + }, + WebSocketCoverage { + config_name: "text-guard".into(), + implementation: "example/websocket-text".into(), + state: State::UnsupportedMessageType, + sequence: Some(7), + message_type: Some(WebSocketMessageType::Binary), + original_size: 23, + }, + ], + ); + + assert_eq!(events.len(), 2); + assert!(events.iter().all(|event| event.class_uid() == 4001)); + assert!( + events + .iter() + .all(|event| event.base().severity == SeverityId::Informational) + ); + let serialized = serde_json::to_string(&events).expect("serialize coverage events"); + assert!(serialized.contains("binding_not_selected")); + assert!(serialized.contains("unsupported_message_type")); + assert!(serialized.contains("\"websocket_message_type\":\"binary\"")); + assert!(serialized.contains("\"websocket_sequence\":\"7\"")); + assert!(serialized.contains("\"input_bytes\":23")); + assert!(!serialized.contains("fail_open")); + assert!(!serialized.contains("fail_closed")); + } + #[tokio::test] async fn direct_denial_uses_middleware_response_without_service_text() { let ctx = L7EvalContext { diff --git a/crates/openshell-supervisor-network/src/l7/provider.rs b/crates/openshell-supervisor-network/src/l7/provider.rs index 864d94ad26..3a51bd1d8c 100644 --- a/crates/openshell-supervisor-network/src/l7/provider.rs +++ b/crates/openshell-supervisor-network/src/l7/provider.rs @@ -30,6 +30,7 @@ pub enum RelayOutcome { Upgraded { overflow: Vec, websocket_permessage_deflate: bool, + websocket_subprotocol: Option, }, } diff --git a/crates/openshell-supervisor-network/src/l7/relay.rs b/crates/openshell-supervisor-network/src/l7/relay.rs index fa2eab4ad7..00a2ca3f2a 100644 --- a/crates/openshell-supervisor-network/src/l7/relay.rs +++ b/crates/openshell-supervisor-network/src/l7/relay.rs @@ -71,13 +71,17 @@ pub struct L7EvalContext { pub(crate) struct UpgradeRelayOptions<'a> { pub(crate) websocket_request: bool, pub(crate) websocket: WebSocketUpgradeBehavior, + pub(crate) assembly_budget: Option, pub(crate) secret_resolver: Option>, + pub(crate) generation_guard: Option<&'a PolicyGenerationGuard>, pub(crate) engine: Option<&'a TunnelPolicyEngine>, pub(crate) ctx: Option<&'a L7EvalContext>, pub(crate) enforcement: EnforcementMode, pub(crate) target: String, pub(crate) query_params: std::collections::HashMap>, pub(crate) policy_name: String, + pub(crate) middleware_session: Option, + pub(crate) selected_subprotocol: Option, } #[derive(Default)] @@ -470,6 +474,7 @@ where if allowed || (config.enforcement == EnforcementMode::Audit && !force_deny) { let chain = engine.query_middleware_chain(&middleware_network_input(ctx))?; + let websocket_chain = websocket_request.then(|| chain.clone()); // Route selection resolved `config` per request, so re-check the // body against that protocol's policy after every transforming // stage (a no-op for REST and websocket, whose policy inputs the @@ -505,8 +510,58 @@ where .await?; return Ok(()); } + MiddlewareApplyResult::AdmissionExhausted => { + let unavailable_request = crate::l7::provider::L7Request { + action: request_info.action.clone(), + target: redacted_target.clone(), + query_params: request_info.query_params.clone(), + raw_header: Vec::new(), + body_length: crate::l7::provider::BodyLength::None, + }; + crate::l7::middleware::send_middleware_admission_exhausted_response( + &unavailable_request, + client, + ctx, + &redacted_target, + ) + .await?; + return Ok(()); + } }; - let outcome = crate::l7::rest::relay_http_request_with_options_guarded( + let mut middleware_session = if let Some(chain) = websocket_chain.as_deref() { + let preflight = websocket_middleware_preflight( + &req, + chain, + engine.middleware_runner(), + ctx, + "wss", + ) + .await; + let preflight = match preflight { + Ok(preflight) => preflight, + Err(error) => { + warn!(error = %error, "WebSocket middleware preflight failed"); + write_bad_gateway_response(client).await?; + return Ok(()); + } + }; + crate::l7::middleware::emit_websocket_preflight_events(ctx, &preflight); + if preflight.terminal_reason.is_some() { + crate::l7::middleware::send_middleware_rejection_response( + &req, + client, + ctx, + preflight.denial.as_ref(), + &redacted_target, + ) + .await?; + return Ok(()); + } + preflight.session + } else { + None + }; + let outcome_result = crate::l7::rest::relay_http_request_with_options_guarded( &req, client, upstream, @@ -523,13 +578,36 @@ where port: ctx.port, }, ) + .await; + let outcome = finalize_websocket_pre_upgrade( + &mut middleware_session, + engine.generation_guard(), + &ctx.host, + ctx.port, + &ctx.policy_name, + outcome_result, + ) .await?; match outcome { - RelayOutcome::Reusable => {} - RelayOutcome::Consumed => return Ok(()), + RelayOutcome::Reusable => { + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::UpstreamRejected) + .await; + } + } + RelayOutcome::Consumed => { + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::UpstreamRejected) + .await; + } + return Ok(()); + } RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + websocket_subprotocol, } => { let mut options = upgrade_options( config, @@ -540,6 +618,8 @@ where Some(&engine), ); options.websocket.permessage_deflate = websocket_permessage_deflate; + options.middleware_session = middleware_session.take(); + options.selected_subprotocol = websocket_subprotocol; return handle_upgrade( client, upstream, overflow, &ctx.host, ctx.port, options, ) @@ -639,6 +719,37 @@ fn emit_activity(ctx: &L7EvalContext, denied: bool, deny_group: &'static str) { } } +pub(crate) async fn websocket_middleware_preflight( + req: &crate::l7::provider::L7Request, + chain: &[openshell_supervisor_middleware::ChainEntry], + runner: &openshell_supervisor_middleware::ChainRunner, + ctx: &L7EvalContext, + scheme: &str, +) -> Result { + let header_end = req + .raw_header + .windows(4) + .position(|window| window == b"\r\n\r\n") + .map_or(req.raw_header.len(), |position| position + 4); + let requested_subprotocols = + crate::l7::rest::websocket_requested_subprotocols(&req.raw_header[..header_end])?; + runner + .preflight_websocket( + chain, + openshell_supervisor_middleware::WebSocketPreflightInput { + session_id: uuid::Uuid::new_v4().to_string(), + request_id: uuid::Uuid::new_v4().to_string(), + sandbox_id: openshell_ocsf::ctx::ctx().sandbox_id.clone(), + scheme: scheme.to_string(), + host: ctx.host.clone(), + port: ctx.port, + path: req.target.clone(), + requested_subprotocols, + }, + ) + .await +} + /// Handle an upgraded connection (101 Switching Protocols). /// /// Forwards any overflow bytes from the upgrade response to the client, then @@ -650,16 +761,37 @@ pub(crate) async fn handle_upgrade( overflow: Vec, host: &str, port: u16, - options: UpgradeRelayOptions<'_>, + mut options: UpgradeRelayOptions<'_>, ) -> Result<()> where C: AsyncRead + AsyncWrite + Unpin + Send, U: AsyncRead + AsyncWrite + Unpin + Send, { + ensure_upgrade_generation_current(client, upstream, host, port, &mut options).await?; + let start_terminal_reason = if let Some(session) = options.middleware_session.as_mut() { + let start = session + .start(options.selected_subprotocol.as_deref().unwrap_or_default()) + .await; + if let Some(ctx) = options.ctx { + crate::l7::middleware::emit_websocket_session_start_events(ctx, &start); + } + start.terminal_reason + } else { + None + }; + ensure_upgrade_generation_current(client, upstream, host, port, &mut options).await?; + if let Some(reason) = start_terminal_reason { + if let Some(session) = options.middleware_session.take() { + session.end(reason).await; + } + send_websocket_close(client, upstream, 1008).await; + return Ok(()); + } let use_websocket_relay = options.websocket_request && (options.websocket.message_policy.inspects_messages() || options.websocket.permessage_deflate - || (options.websocket.credential_rewrite && options.secret_resolver.is_some())); + || (options.websocket.credential_rewrite && options.secret_resolver.is_some()) + || options.middleware_session.is_some()); let relay_mode = if use_websocket_relay { "websocket parsed relay" } else { @@ -678,6 +810,10 @@ where .build() ); if use_websocket_relay { + let assembly_budget = options + .assembly_budget + .take() + .ok_or_else(|| miette!("websocket parsed relay missing assembly budget"))?; let resolver = if options.websocket.credential_rewrite { options.secret_resolver.as_deref() } else { @@ -715,9 +851,13 @@ where port, crate::l7::websocket::RelayOptions { policy_name: &options.policy_name, + assembly_budget, resolver, + generation_guard: options.generation_guard, inspector, compression, + middleware_session: options.middleware_session.take(), + middleware_context: options.ctx, }, ) .await; @@ -732,6 +872,45 @@ where Ok(()) } +async fn ensure_upgrade_generation_current( + client: &mut C, + upstream: &mut U, + host: &str, + port: u16, + options: &mut UpgradeRelayOptions<'_>, +) -> Result<()> +where + C: AsyncWrite + Unpin, + U: AsyncWrite + Unpin, +{ + let Some(guard) = options.generation_guard else { + return Ok(()); + }; + if let Err(error) = guard.ensure_current() { + emit_policy_reload(guard, host, port, &options.policy_name); + if let Some(session) = options.middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::PolicyReload) + .await; + } + send_websocket_close(client, upstream, 1012).await; + return Err(error); + } + Ok(()) +} + +async fn send_websocket_close(client: &mut C, upstream: &mut U, code: u16) +where + C: AsyncWrite + Unpin, + U: AsyncWrite + Unpin, +{ + let payload = code.to_be_bytes(); + let _ = crate::l7::websocket::write_unmasked_close(client, &payload).await; + let _ = crate::l7::websocket::write_masked_close(upstream, &payload).await; + let _ = client.shutdown().await; + let _ = upstream.shutdown().await; +} + pub(crate) fn upgrade_options<'a>( config: &L7EndpointConfig, ctx: &'a L7EvalContext, @@ -759,17 +938,21 @@ pub(crate) fn upgrade_options<'a>( message_policy: websocket_message_policy, permessage_deflate: false, }, + assembly_budget: engine.map(TunnelPolicyEngine::websocket_assembly_budget), secret_resolver: if websocket_credential_rewrite { ctx.secret_resolver.clone() } else { None }, + generation_guard: engine.map(TunnelPolicyEngine::generation_guard), engine, ctx: engine.map(|_| ctx), enforcement: config.enforcement, target: target.to_string(), query_params: query_params.clone(), policy_name: ctx.policy_name.clone(), + middleware_session: None, + selected_subprotocol: None, } } @@ -956,6 +1139,7 @@ where if allowed || config.enforcement == EnforcementMode::Audit { let chain = engine.query_middleware_chain(&middleware_network_input(ctx))?; + let websocket_chain = websocket_request.then(|| chain.clone()); // REST and websocket-upgrade policy evaluates only the method, // path, and query, which a middleware result cannot mutate, so no // per-stage body re-check is needed. @@ -989,6 +1173,56 @@ where .await?; return Ok(()); } + MiddlewareApplyResult::AdmissionExhausted => { + let unavailable_request = crate::l7::provider::L7Request { + action: request_info.action.clone(), + target: redacted_target.clone(), + query_params: request_info.query_params.clone(), + raw_header: Vec::new(), + body_length: crate::l7::provider::BodyLength::None, + }; + crate::l7::middleware::send_middleware_admission_exhausted_response( + &unavailable_request, + client, + ctx, + &redacted_target, + ) + .await?; + return Ok(()); + } + }; + let mut middleware_session = if let Some(chain) = websocket_chain.as_deref() { + let preflight = websocket_middleware_preflight( + &req, + chain, + engine.middleware_runner(), + ctx, + "wss", + ) + .await; + let preflight = match preflight { + Ok(preflight) => preflight, + Err(error) => { + warn!(error = %error, "WebSocket middleware preflight failed"); + write_bad_gateway_response(client).await?; + return Ok(()); + } + }; + crate::l7::middleware::emit_websocket_preflight_events(ctx, &preflight); + if preflight.terminal_reason.is_some() { + crate::l7::middleware::send_middleware_rejection_response( + &req, + client, + ctx, + preflight.denial.as_ref(), + &redacted_target, + ) + .await?; + return Ok(()); + } + preflight.session + } else { + None }; let req_with_auth = match crate::l7::token_grant_injection::inject_if_needed(req, ctx).await { @@ -1006,7 +1240,7 @@ where }; // Forward request to upstream and relay response - let outcome = crate::l7::rest::relay_http_request_with_options_guarded( + let outcome_result = crate::l7::rest::relay_http_request_with_options_guarded( &req_with_auth, client, upstream, @@ -1023,10 +1257,30 @@ where port: ctx.port, }, ) + .await; + let outcome = finalize_websocket_pre_upgrade( + &mut middleware_session, + engine.generation_guard(), + &ctx.host, + ctx.port, + &ctx.policy_name, + outcome_result, + ) .await?; match outcome { - RelayOutcome::Reusable => {} // continue loop + RelayOutcome::Reusable => { + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::UpstreamRejected) + .await; + } + } RelayOutcome::Consumed => { + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::UpstreamRejected) + .await; + } debug!( host = %ctx.host, port = ctx.port, @@ -1037,6 +1291,7 @@ where RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + websocket_subprotocol, } => { let mut options = upgrade_options( config, @@ -1047,6 +1302,8 @@ where Some(engine), ); options.websocket.permessage_deflate = websocket_permessage_deflate; + options.middleware_session = middleware_session.take(); + options.selected_subprotocol = websocket_subprotocol; return handle_upgrade( client, upstream, overflow, &ctx.host, ctx.port, options, ) @@ -1075,6 +1332,16 @@ fn close_if_stale(guard: &PolicyGenerationGuard, ctx: &L7EvalContext) -> bool { return false; } + emit_policy_reload(guard, &ctx.host, ctx.port, &ctx.policy_name); + true +} + +pub(crate) fn emit_policy_reload( + guard: &PolicyGenerationGuard, + host: &str, + port: u16, + policy_name: &str, +) { ocsf_emit!( NetworkActivityBuilder::new(openshell_ocsf::ctx::ctx()) .activity(ActivityId::Open) @@ -1082,18 +1349,54 @@ fn close_if_stale(guard: &PolicyGenerationGuard, ctx: &L7EvalContext) -> bool { .disposition(DispositionId::Blocked) .severity(SeverityId::Medium) .status(StatusId::Failure) - .dst_endpoint(Endpoint::from_domain(&ctx.host, ctx.port)) - .firewall_rule(&ctx.policy_name, "l7") + .dst_endpoint(Endpoint::from_domain(host, port)) + .firewall_rule(policy_name, "l7") .message(format!( "L7 tunnel closed after policy reload [host:{} port:{} captured_generation:{} current_generation:{}]", - ctx.host, - ctx.port, + host, + port, guard.captured_generation(), guard.current_generation(), )) .build() ); - true +} + +pub(crate) async fn finalize_websocket_pre_upgrade( + session: &mut Option, + guard: &PolicyGenerationGuard, + host: &str, + port: u16, + policy_name: &str, + result: Result, +) -> Result { + match result { + Ok(value) => { + if let Err(error) = guard.ensure_current() { + emit_policy_reload(guard, host, port, policy_name); + if let Some(session) = session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::PolicyReload) + .await; + } + Err(error) + } else { + Ok(value) + } + } + Err(error) => { + let reason = if guard.is_stale() { + emit_policy_reload(guard, host, port, policy_name); + openshell_core::proto::WebSocketSessionEndReason::PolicyReload + } else { + openshell_core::proto::WebSocketSessionEndReason::UpstreamRejected + }; + if let Some(session) = session.take() { + session.end(reason).await; + } + Err(error) + } + } } async fn relay_jsonrpc( @@ -1254,6 +1557,23 @@ where .await?; return Ok(()); } + MiddlewareApplyResult::AdmissionExhausted => { + let unavailable_request = crate::l7::provider::L7Request { + action: request_info.action.clone(), + target: redacted_target.clone(), + query_params: request_info.query_params.clone(), + raw_header: Vec::new(), + body_length: crate::l7::provider::BodyLength::None, + }; + crate::l7::middleware::send_middleware_admission_exhausted_response( + &unavailable_request, + client, + ctx, + &redacted_target, + ) + .await?; + return Ok(()); + } }; // Future MCP response/SSE introspection or rewrite would hook here // before returning upstream bytes. The current policy schema has no @@ -1478,6 +1798,23 @@ where .await?; return Ok(()); } + MiddlewareApplyResult::AdmissionExhausted => { + let unavailable_request = crate::l7::provider::L7Request { + action: request_info.action.clone(), + target: redacted_target.clone(), + query_params: request_info.query_params.clone(), + raw_header: Vec::new(), + body_length: crate::l7::provider::BodyLength::None, + }; + crate::l7::middleware::send_middleware_admission_exhausted_response( + &unavailable_request, + client, + ctx, + &redacted_target, + ) + .await?; + return Ok(()); + } }; let outcome = crate::l7::rest::relay_http_request_with_resolver_guarded( &req, @@ -1500,8 +1837,12 @@ where RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + .. } => { let options = UpgradeRelayOptions { + assembly_budget: Some( + crate::l7::websocket::WebSocketAssemblyBudget::default(), + ), websocket: WebSocketUpgradeBehavior { permessage_deflate: websocket_permessage_deflate, ..Default::default() @@ -2110,6 +2451,23 @@ where .await?; return Ok(()); } + MiddlewareApplyResult::AdmissionExhausted => { + let unavailable_request = crate::l7::provider::L7Request { + action: "HTTP".into(), + target: redacted_target.clone(), + query_params: std::collections::HashMap::new(), + raw_header: Vec::new(), + body_length: crate::l7::provider::BodyLength::None, + }; + crate::l7::middleware::send_middleware_admission_exhausted_response( + &unavailable_request, + client, + ctx, + &redacted_target, + ) + .await?; + return Ok(()); + } } } else { req @@ -2192,12 +2550,14 @@ mod tests { const TEST_POLICY: &str = include_str!("../../data/sandbox-policy.rego"); fn install_builtin_middleware(engine: &OpaEngine) { - engine.set_middleware_runner_for_tests(openshell_supervisor_middleware::ChainRunner::new( - openshell_supervisor_middleware_builtins::services() - .into_iter() - .next() - .expect("built-in middleware service"), - )); + engine.set_middleware_runner_for_tests( + openshell_supervisor_middleware::ChainRunner::from_endpoint( + openshell_supervisor_middleware_builtins::services() + .into_iter() + .next() + .expect("built-in middleware service"), + ), + ); } fn assert_middleware_failure_response(response: &str, policy_name: &str) { @@ -2216,6 +2576,35 @@ mod tests { assert!(body.get("agent_guidance").is_none()); } + fn assert_middleware_unavailable_response(response: &str, policy_name: &str) { + assert!( + response.starts_with("HTTP/1.1 503 Service Unavailable\r\n"), + "{response}" + ); + assert!(!response.contains("100 Continue"), "{response}"); + assert!(!response.to_ascii_lowercase().contains("retry-after")); + let (headers, body) = response.split_once("\r\n\r\n").expect("HTTP response"); + let content_length = headers + .lines() + .find_map(|line| { + line.strip_prefix("Content-Length: ") + .and_then(|value| value.parse::().ok()) + }) + .expect("Content-Length"); + assert_eq!(content_length, body.len()); + let body: serde_json::Value = serde_json::from_str(body).expect("JSON response"); + assert_eq!(body["error"], "middleware_failed"); + assert_eq!( + body["detail"], + "Request could not be processed by configured middleware" + ); + assert_eq!(body["policy"], policy_name); + assert!(body.get("middleware").is_none()); + assert!(body.get("reason_code").is_none()); + assert!(body.get("rule_missing").is_none()); + assert!(body.get("next_steps").is_none()); + } + fn rest_token_grant_relay_context( resolver_response: std::result::Result<&str, &str>, ) -> ( @@ -2393,6 +2782,299 @@ network_policies: (generation_guard, ctx, fixture) } + #[derive(Clone)] + struct ControllableWebSocketPreflight { + seen: tokio::sync::mpsc::UnboundedSender<()>, + release: Arc, + action: openshell_core::proto::WebSocketPreflightAction, + } + + #[tonic::async_trait] + impl openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware + for ControllableWebSocketPreflight + { + type EvaluateWebSocketSessionStream = + openshell_supervisor_middleware::WebSocketResponseStream; + + async fn describe( + &self, + _request: tonic::Request<()>, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + Ok(tonic::Response::new( + openshell_core::proto::MiddlewareManifest { + name: "test/controllable-websocket-preflight".into(), + service_version: "test".into(), + bindings: vec![openshell_core::proto::MiddlewareBinding { + operation: + openshell_core::proto::SupervisorMiddlewareOperation::WebsocketMessage + as i32, + phase: openshell_core::proto::SupervisorMiddlewarePhase::PreCredentials + as i32, + max_payload_bytes: + openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES as u64, + timeout: "2s".into(), + }], + }, + )) + } + + async fn validate_config( + &self, + _request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + Ok(tonic::Response::new( + openshell_core::proto::ValidateConfigResponse { + valid: true, + reason: String::new(), + }, + )) + } + + async fn evaluate_http_request( + &self, + _request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + Err(tonic::Status::unimplemented("WebSocket-only middleware")) + } + + async fn evaluate_web_socket_session( + &self, + request: tonic::Request>, + ) -> std::result::Result, tonic::Status> + { + use openshell_core::proto::{ + WebSocketPreflightDecision, WebSocketSessionEventResult, web_socket_session_event, + web_socket_session_event_result, + }; + use tokio_stream::wrappers::ReceiverStream; + + let mut requests = request.into_inner(); + let seen = self.seen.clone(); + let release = Arc::clone(&self.release); + let action = self.action; + let (responses_tx, responses_rx) = tokio::sync::mpsc::channel(4); + tokio::spawn(async move { + while let Ok(Some(request)) = requests.message().await { + if matches!( + request.event, + Some(web_socket_session_event::Event::Preflight(_)) + ) { + let _ = seen.send(()); + release.notified().await; + let _ = responses_tx + .send(Ok(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::PreflightDecision( + WebSocketPreflightDecision { + action: action as i32, + ..Default::default() + }, + ), + ), + })) + .await; + } + } + }); + Ok(tonic::Response::new(Box::pin(ReceiverStream::new( + responses_rx, + )))) + } + } + + async fn assert_websocket_preflight_precedes_token_grant( + action: openshell_core::proto::WebSocketPreflightAction, + admitted: bool, + ) { + use openshell_core::proto::SupervisorMiddlewareService; + use openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddlewareServer; + use openshell_supervisor_middleware::{ChainRunner, MiddlewareRegistry}; + use tokio_stream::wrappers::TcpListenerStream; + + let (seen_tx, mut seen_rx) = tokio::sync::mpsc::unbounded_channel(); + let release = Arc::new(tokio::sync::Notify::new()); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind middleware"); + let address = listener.local_addr().expect("middleware address"); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tonic::transport::Server::builder() + .add_service(SupervisorMiddlewareServer::new( + ControllableWebSocketPreflight { + seen: seen_tx, + release: Arc::clone(&release), + action, + }, + )) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }); + let server_task = tokio::spawn(server); + let registry = MiddlewareRegistry::connect_services( + Vec::new(), + vec![SupervisorMiddlewareService { + name: "preflight-service".into(), + grpc_endpoint: format!("http://{address}"), + max_payload_bytes: openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES + as u64, + timeout: "2s".into(), + }], + ) + .await + .expect("connect middleware"); + + let data = r#" +network_middlewares: + websocket-preflight: + middleware: preflight-service + on_error: fail_closed + endpoints: + include: ["api.example.test"] +network_policies: + rest_api: + name: rest_api + endpoints: + - host: api.example.test + port: 8080 + protocol: rest + enforcement: enforce + rules: + - allow: + method: GET + path: "/v1/**" + binaries: + - { path: /usr/bin/curl } +"#; + let engine = OpaEngine::from_strings(TEST_POLICY, data).expect("test policy"); + engine.set_middleware_runner_for_tests(ChainRunner::from_registry(registry)); + let input = NetworkInput { + host: "api.example.test".into(), + port: 8080, + binary_path: PathBuf::from("/usr/bin/curl"), + binary_sha256: "unused".into(), + ancestors: vec![], + cmdline_paths: vec![], + }; + let (endpoint_config, generation) = engine + .query_endpoint_config_with_generation(&input) + .expect("endpoint config"); + let config = crate::l7::parse_l7_config(&endpoint_config.expect("configured endpoint")) + .expect("REST config"); + let tunnel_engine = engine + .clone_engine_for_tunnel(generation) + .expect("tunnel engine"); + let provider_key = "api.example.test\t8080\t/v1/**\tprovider:access_token"; + let fixture = + crate::l7::token_grant_injection::test_support::TokenGrantTestFixture::success( + provider_key, + "grant-token", + ); + let ctx = L7EvalContext { + host: "api.example.test".into(), + port: 8080, + policy_name: "rest_api".into(), + binary_path: "/usr/bin/curl".into(), + dynamic_credentials: Some(fixture.dynamic_credentials()), + token_grant_resolver: Some(fixture.resolver()), + ..Default::default() + }; + let (mut app, mut relay_client) = tokio::io::duplex(8192); + let (mut relay_upstream, mut upstream) = tokio::io::duplex(8192); + let relay = tokio::spawn(async move { + relay_rest( + &config, + &tunnel_engine, + &mut relay_client, + &mut relay_upstream, + &ctx, + ) + .await + }); + + app.write_all( + b"GET /v1/ws HTTP/1.1\r\nHost: api.example.test\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n\r\n", + ) + .await + .expect("send upgrade request"); + seen_rx.recv().await.expect("preflight reached middleware"); + fixture.assert_no_requests(); + release.notify_one(); + + if admitted { + let mut forwarded = Vec::new(); + let mut buffer = [0u8; 512]; + while !forwarded.windows(4).any(|window| window == b"\r\n\r\n") { + let count = upstream.read(&mut buffer).await.expect("read upstream"); + assert!(count > 0, "upstream closed before request headers"); + forwarded.extend_from_slice(&buffer[..count]); + } + let forwarded = String::from_utf8_lossy(&forwarded); + assert!(forwarded.contains("Authorization: Bearer grant-token\r\n")); + upstream + .write_all( + b"HTTP/1.1 400 Bad Request\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await + .expect("reject upgrade"); + let mut response = [0u8; 512]; + let count = app.read(&mut response).await.expect("read client response"); + assert!( + String::from_utf8_lossy(&response[..count]).contains("400 Bad Request"), + "unexpected client response" + ); + fixture.assert_one_request(provider_key); + } else { + let mut response = [0u8; 1024]; + let count = app.read(&mut response).await.expect("read client response"); + let response = String::from_utf8_lossy(&response[..count]); + assert!(response.contains("403 Forbidden"), "{response}"); + assert!(response.contains(r#""error":"middleware_denied""#)); + assert!(response.contains(r#""middleware":"websocket-preflight""#)); + let mut byte = [0u8; 1]; + let count = upstream.read(&mut byte).await.expect("upstream close"); + assert_eq!(count, 0, "denied preflight must not reach upstream"); + fixture.assert_no_requests(); + } + + drop(app); + relay + .await + .expect("join REST relay") + .expect("REST relay result"); + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + #[tokio::test] + async fn denied_websocket_preflight_has_no_token_grant_side_effects() { + assert_websocket_preflight_precedes_token_grant( + openshell_core::proto::WebSocketPreflightAction::Deny, + false, + ) + .await; + } + + #[tokio::test] + async fn admitted_websocket_preflight_precedes_token_grant() { + assert_websocket_preflight_precedes_token_grant( + openshell_core::proto::WebSocketPreflightAction::Inspect, + true, + ) + .await; + } + fn jsonrpc_test_relay_context() -> (L7EndpointConfig, TunnelPolicyEngine, L7EvalContext) { let data = r" network_policies: @@ -2766,6 +3448,94 @@ network_policies: .unwrap(); } + #[tokio::test] + async fn l7_rest_exhausted_middleware_admission_returns_503_before_body_or_upstream() { + let (config, tunnel_engine, ctx) = middleware_relay_context("openshell/regex", "fail_open"); + let runner = tunnel_engine.middleware_runner().clone(); + + let mut active = Vec::new(); + for _ in 0..openshell_supervisor_middleware::MAX_CONCURRENT_MIDDLEWARE_WORK { + active.push( + runner + .reserve_middleware_work_admission() + .await + .expect("fill active middleware work"), + ); + } + let mut waiters = Vec::new(); + for _ in 0..openshell_supervisor_middleware::MAX_QUEUED_MIDDLEWARE_WORK { + let runner = runner.clone(); + waiters.push(Box::pin( + async move { runner.reserve_middleware_work().await }, + )); + } + for waiter in &mut waiters { + assert!( + futures::poll!(waiter.as_mut()).is_pending(), + "every bounded waiter slot must be occupied" + ); + } + + let (mut app, mut relay_client) = tokio::io::duplex(8192); + let (mut relay_upstream, mut upstream) = tokio::io::duplex(8192); + let relay = tokio::spawn(async move { + relay_with_inspection( + &config, + tunnel_engine, + &mut relay_client, + &mut relay_upstream, + &ctx, + ) + .await + }); + + // The declared body is within the built-in regex HTTP capability, but + // the client intentionally withholds it behind Expect: 100-continue. + // Queue exhaustion must be answered before buffering begins. + app.write_all( + b"POST /v1/messages HTTP/1.1\r\nHost: api.example.test\r\nContent-Length: 32\r\nExpect: 100-continue\r\nConnection: close\r\n\r\n", + ) + .await + .expect("send headers without body"); + + tokio::time::timeout(std::time::Duration::from_secs(1), relay) + .await + .expect("relay should shed immediately") + .expect("join relay") + .expect("relay returns a complete HTTP response"); + + let mut response = Vec::new(); + tokio::time::timeout( + std::time::Duration::from_secs(1), + app.read_to_end(&mut response), + ) + .await + .expect("client should receive 503 without sending its body") + .expect("read client response"); + let response = String::from_utf8(response).expect("UTF-8 response"); + assert_middleware_unavailable_response(&response, "rest_api"); + + let mut upstream_bytes = Vec::new(); + tokio::time::timeout( + std::time::Duration::from_secs(1), + upstream.read_to_end(&mut upstream_bytes), + ) + .await + .expect("upstream side should close") + .expect("read upstream"); + assert!( + upstream_bytes.is_empty(), + "admission-exhausted request must not reach upstream" + ); + + drop(waiters); + drop(active); + runner + .reserve_middleware_work_admission() + .await + .expect("work capacity recovers after saturation fixture"); + } + #[tokio::test] async fn l7_rest_middleware_fail_closed_does_not_reach_upstream() { let (config, tunnel_engine, ctx) = @@ -3175,7 +3945,7 @@ network_policies: // A single unresolved (0-limit) entry must not drag the chain limit to // zero: the buffer limit reflects only the resolved built-in. - let mixed = ChainRunner::new( + let mixed = ChainRunner::from_endpoint( openshell_supervisor_middleware_builtins::services() .into_iter() .next() @@ -3206,6 +3976,21 @@ network_policies: impl openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware for BodyReplacingService { + type EvaluateWebSocketSessionStream = + openshell_supervisor_middleware::WebSocketResponseStream; + + async fn evaluate_web_socket_session( + &self, + _request: tonic::Request< + tonic::Streaming, + >, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented( + "test service does not inspect WebSocket messages", + )) + } + async fn describe( &self, _request: tonic::Request<()>, @@ -3222,7 +4007,7 @@ network_policies: as i32, phase: openshell_core::proto::SupervisorMiddlewarePhase::PreCredentials as i32, - max_body_bytes: 8192, + max_payload_bytes: 8192, timeout: String::new(), }], }, @@ -3693,6 +4478,21 @@ network_policies: impl openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware for LimitService { + type EvaluateWebSocketSessionStream = + openshell_supervisor_middleware::WebSocketResponseStream; + + async fn evaluate_web_socket_session( + &self, + _request: tonic::Request< + tonic::Streaming, + >, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented( + "test service does not inspect WebSocket messages", + )) + } + async fn describe( &self, _request: tonic::Request<()>, @@ -3710,7 +4510,7 @@ network_policies: bindings: vec![MiddlewareBinding { operation: SupervisorMiddlewareOperation::HttpRequest as i32, phase: SupervisorMiddlewarePhase::PreCredentials as i32, - max_body_bytes: self.max_body_bytes, + max_payload_bytes: self.max_body_bytes, timeout: String::new(), }], })) @@ -3764,16 +4564,16 @@ network_policies: middleware_relay_context("openshell/regex", "fail_closed"); let registry = openshell_supervisor_middleware::MiddlewareRegistry::connect_services( vec![ - Arc::new(LimitService { + openshell_supervisor_middleware::http_only_endpoint(Arc::new(LimitService { name: "test/redactor", max_body_bytes: 8192, replacement: Some(b"[SCRUBBED BY TEST REDACTOR]"), - }), - Arc::new(LimitService { + })), + openshell_supervisor_middleware::http_only_endpoint(Arc::new(LimitService { name: "test/guard", max_body_bytes: 16, replacement: None, - }), + })), ], Vec::new(), ) @@ -3838,6 +4638,9 @@ network_policies: MiddlewareApplyResult::Denied { .. } => { panic!("body within the largest stage limit must not fail the chain") } + MiddlewareApplyResult::AdmissionExhausted => { + panic!("test middleware work admission must be available") + } } } diff --git a/crates/openshell-supervisor-network/src/l7/rest.rs b/crates/openshell-supervisor-network/src/l7/rest.rs index fd9d373f81..b763124949 100644 --- a/crates/openshell-supervisor-network/src/l7/rest.rs +++ b/crates/openshell-supervisor-network/src/l7/rest.rs @@ -30,7 +30,7 @@ const MAX_REWRITE_BODY_BYTES: usize = 256 * 1024; const MAX_SIGV4_BODY_BYTES: usize = 10 * 1024 * 1024; #[cfg(test)] async fn max_middleware_body_bytes() -> usize { - let chain = openshell_supervisor_middleware::ChainRunner::new( + let chain = openshell_supervisor_middleware::ChainRunner::from_endpoint( openshell_supervisor_middleware_builtins::services() .into_iter() .next() @@ -45,7 +45,7 @@ async fn max_middleware_body_bytes() -> usize { }]) .await .expect("describe built-in middleware"); - chain[0].max_body_bytes() + chain[0].max_payload_bytes() } const RELAY_BUF_SIZE: usize = 8192; const HTTP_METHOD_PREFIXES: &[&[u8]] = &[ @@ -2020,6 +2020,12 @@ fn validate_websocket_upgrade_request(raw_header: &[u8]) -> Result { parse_websocket_upgrade_request(raw_header).map(|request| request.is_some()) } +pub(crate) fn websocket_requested_subprotocols(raw_header: &[u8]) -> Result> { + Ok(parse_websocket_upgrade_request(raw_header)? + .map(|request| request.subprotocols) + .unwrap_or_default()) +} + fn parse_websocket_upgrade_request(raw_header: &[u8]) -> Result> { let header_str = std::str::from_utf8(raw_header) .map_err(|_| miette!("HTTP headers contain invalid UTF-8"))?; @@ -2219,14 +2225,36 @@ pub(crate) async fn send_middleware_failure_response( send_forbidden_json(policy_name, body, client).await } +/// Send a platform-owned availability response when shared middleware work +/// admission is exhausted before the request body is buffered. +pub(crate) async fn send_middleware_unavailable_response( + req: &L7Request, + policy_name: &str, + client: &mut C, + redacted_target: Option<&str>, + context: Option>, +) -> Result<()> { + let body = middleware_failure_response_body(req, policy_name, redacted_target, context); + send_json_response(policy_name, body, client, "503 Service Unavailable").await +} + async fn send_forbidden_json( policy_name: &str, body: serde_json::Value, client: &mut C, +) -> Result<()> { + send_json_response(policy_name, body, client, "403 Forbidden").await +} + +async fn send_json_response( + policy_name: &str, + body: serde_json::Value, + client: &mut C, + status: &str, ) -> Result<()> { let body_bytes = body.to_string(); let response = format!( - "HTTP/1.1 403 Forbidden\r\n\ + "HTTP/1.1 {status}\r\n\ Content-Type: application/json\r\n\ Content-Length: {}\r\n\ X-OpenShell-Policy: {}\r\n\ @@ -2760,7 +2788,7 @@ where if !options.client_requested_upgrade { return Ok(RelayOutcome::Consumed); } - let websocket_permessage_deflate = validate_websocket_response( + let (websocket_permessage_deflate, websocket_subprotocol) = validate_websocket_response( &header_str, options.websocket_extensions, options.websocket.as_ref(), @@ -2779,6 +2807,7 @@ where return Ok(RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + websocket_subprotocol, }); } @@ -2908,9 +2937,10 @@ fn validate_websocket_response( headers: &str, mode: WebSocketExtensionMode, websocket: Option<&WebSocketResponseValidation>, -) -> Result { +) -> Result<(bool, Option)> { let Some(validation) = websocket else { - return validate_websocket_response_extensions_preserved(headers, mode); + return validate_websocket_response_extensions_preserved(headers, mode) + .map(|compressed| (compressed, None)); }; let mut upgrade_websocket = false; @@ -2970,11 +3000,11 @@ fn validate_websocket_response( "websocket upgrade response has multiple Sec-WebSocket-Protocol headers" )); } - if let Some(protocol) = selected_subprotocol + if let Some(ref protocol) = selected_subprotocol && !validation .offered_subprotocols .iter() - .any(|offered| offered == &protocol) + .any(|offered| offered == protocol) { return Err(miette!( "upstream selected WebSocket subprotocol that was not offered" @@ -2986,8 +3016,10 @@ fn validate_websocket_response( (None, Some(_)) => Err(miette!( "upstream negotiated WebSocket extension that was not offered" )), - (None | Some(_), None) => Ok(false), - (Some(expected), Some(actual)) if expected.eq_ignore_ascii_case(actual) => Ok(true), + (None | Some(_), None) => Ok((false, selected_subprotocol)), + (Some(expected), Some(actual)) if expected.eq_ignore_ascii_case(actual) => { + Ok((true, selected_subprotocol)) + } (Some(_), Some(_)) => Err(miette!( "upstream negotiated WebSocket extension that does not match the safe offer" )), @@ -3562,6 +3594,7 @@ mod tests { let RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + .. } = outcome else { panic!("expected upgraded relay outcome"); @@ -3575,6 +3608,7 @@ mod tests { 443, crate::l7::relay::UpgradeRelayOptions { websocket_request: true, + assembly_budget: Some(crate::l7::websocket::WebSocketAssemblyBudget::default()), websocket: crate::l7::relay::WebSocketUpgradeBehavior { credential_rewrite, permessage_deflate: websocket_permessage_deflate, @@ -4380,7 +4414,7 @@ mod tests { &mut client, &[], None, - Some(openshell_supervisor_middleware::MAX_MIDDLEWARE_BODY_BYTES), + Some(openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES), ) .await .expect("chunked body should decode"); @@ -4408,7 +4442,7 @@ mod tests { &req, &mut client, None, - openshell_supervisor_middleware::MAX_MIDDLEWARE_BODY_BYTES, + openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES, ) .await .expect("oversized body should produce a capacity result"); @@ -6075,7 +6109,7 @@ mod tests { offered_subprotocols: Vec::new(), }; - let negotiated = validate_websocket_response( + let (negotiated, subprotocol) = validate_websocket_response( "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\nSec-WebSocket-Extensions: permessage-deflate; server_no_context_takeover; client_no_context_takeover\r\n\r\n", WebSocketExtensionMode::PermessageDeflate, Some(&validation), @@ -6083,6 +6117,7 @@ mod tests { .expect("reordered safe extension params should canonicalize"); assert!(negotiated); + assert_eq!(subprotocol, None); } #[test] @@ -6105,7 +6140,7 @@ mod tests { #[test] fn preserve_mode_leaves_malformed_extension_response_raw() { - let negotiated = validate_websocket_response( + let (negotiated, subprotocol) = validate_websocket_response( "HTTP/1.1 101 Switching Protocols\r\nSec-WebSocket-Extensions: permessage-deflate; client_no_context_takeover=\"true\"\r\n\r\n", WebSocketExtensionMode::Preserve, None, @@ -6113,6 +6148,7 @@ mod tests { .expect("preserve mode should not parse or reject raw extension negotiation"); assert!(!negotiated); + assert_eq!(subprotocol, None); } #[test] @@ -6136,7 +6172,7 @@ mod tests { offered_subprotocols: vec!["chat".to_string(), "superchat".to_string()], }; - let negotiated = validate_websocket_response( + let (negotiated, subprotocol) = validate_websocket_response( "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\nSec-WebSocket-Protocol: superchat\r\n\r\n", WebSocketExtensionMode::PermessageDeflate, Some(&validation), @@ -6144,6 +6180,7 @@ mod tests { .expect("offered subprotocol should validate"); assert!(!negotiated); + assert_eq!(subprotocol.as_deref(), Some("superchat")); } #[test] diff --git a/crates/openshell-supervisor-network/src/l7/token_grant_injection.rs b/crates/openshell-supervisor-network/src/l7/token_grant_injection.rs index 2ca25e6e81..ce6c581bc8 100644 --- a/crates/openshell-supervisor-network/src/l7/token_grant_injection.rs +++ b/crates/openshell-supervisor-network/src/l7/token_grant_injection.rs @@ -419,6 +419,14 @@ pub mod test_support { self.resolver.clone() } + pub fn assert_no_requests(&self) { + let requests = self + .requests + .lock() + .expect("fake token grant requests lock poisoned"); + assert!(requests.is_empty(), "unexpected token grant requests"); + } + pub fn assert_one_request(&self, expected_provider_key: &str) { let requests = self .requests diff --git a/crates/openshell-supervisor-network/src/l7/websocket.rs b/crates/openshell-supervisor-network/src/l7/websocket.rs index bb59c5227b..6989d479e9 100644 --- a/crates/openshell-supervisor-network/src/l7/websocket.rs +++ b/crates/openshell-supervisor-network/src/l7/websocket.rs @@ -8,7 +8,7 @@ use crate::l7::relay::{L7EvalContext, evaluate_l7_request}; use crate::l7::{EnforcementMode, L7RequestInfo}; -use crate::opa::TunnelPolicyEngine; +use crate::opa::{PolicyGenerationGuard, TunnelPolicyEngine}; use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress, Status}; use miette::{IntoDiagnostic, Result, miette}; use openshell_core::secrets::SecretResolver; @@ -17,10 +17,19 @@ use openshell_ocsf::{ ocsf_emit, }; use std::collections::HashMap; +use std::future::Future; +use std::sync::Arc; +use std::time::Duration as StdDuration; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; -const MAX_TEXT_MESSAGE_BYTES: usize = 1024 * 1024; +const MAX_TEXT_MESSAGE_BYTES: usize = openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES; +pub const MAX_CONCURRENT_WEBSOCKET_ASSEMBLIES: usize = 32; +pub const MAX_QUEUED_WEBSOCKET_ASSEMBLIES: usize = MAX_CONCURRENT_WEBSOCKET_ASSEMBLIES * 2; const MAX_RAW_FRAME_PAYLOAD_BYTES: u64 = 16 * 1024 * 1024; +const MAX_MESSAGE_FRAGMENTS: usize = 4096; +const TEXT_MESSAGE_ASSEMBLY_IDLE_TIMEOUT: StdDuration = StdDuration::from_secs(30); +const TEXT_MESSAGE_ASSEMBLY_TOTAL_TIMEOUT: StdDuration = StdDuration::from_secs(120); const COPY_BUF_SIZE: usize = 8192; const OPCODE_CONTINUATION: u8 = 0x0; const OPCODE_TEXT: u8 = 0x1; @@ -29,6 +38,274 @@ const OPCODE_CLOSE: u8 = 0x8; const OPCODE_PING: u8 = 0x9; const OPCODE_PONG: u8 = 0xA; +#[derive(Clone, Debug)] +pub struct WebSocketAssemblyBudget { + active: Arc, + waiters: Arc, +} + +impl Default for WebSocketAssemblyBudget { + fn default() -> Self { + Self::new( + MAX_CONCURRENT_WEBSOCKET_ASSEMBLIES, + MAX_QUEUED_WEBSOCKET_ASSEMBLIES, + ) + } +} + +impl WebSocketAssemblyBudget { + pub fn new(active: usize, waiters: usize) -> Self { + Self { + active: Arc::new(Semaphore::new(active)), + waiters: Arc::new(Semaphore::new(waiters)), + } + } + + pub async fn reserve(&self) -> Result { + if let Ok(active) = Arc::clone(&self.active).try_acquire_owned() { + return Ok(WebSocketAssemblyAdmissionOutcome::Admitted( + WebSocketAssemblyAdmission { _active: active }, + )); + } + let Ok(waiter) = Arc::clone(&self.waiters).try_acquire_owned() else { + return Ok(WebSocketAssemblyAdmissionOutcome::QueueExhausted); + }; + let active = Arc::clone(&self.active) + .acquire_owned() + .await + .map_err(|_| miette!("websocket assembly admission semaphore closed"))?; + drop(waiter); + Ok(WebSocketAssemblyAdmissionOutcome::Admitted( + WebSocketAssemblyAdmission { _active: active }, + )) + } +} + +#[derive(Debug)] +pub struct WebSocketAssemblyAdmission { + _active: OwnedSemaphorePermit, +} + +#[derive(Debug)] +pub enum WebSocketAssemblyAdmissionOutcome { + Admitted(WebSocketAssemblyAdmission), + QueueExhausted, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum WebSocketTerminationCause { + PeerDisconnect, + PolicyReload, + CapacityExhausted, + MiddlewareDenial, + MiddlewareFailure, + PolicyDenial, + InvalidUtf8, + ProtocolError, + MessageTooBig, +} + +impl WebSocketTerminationCause { + fn close_code(self) -> Option { + match self { + Self::PeerDisconnect => None, + Self::PolicyReload => Some(1012), + Self::CapacityExhausted => Some(1013), + Self::MiddlewareDenial | Self::MiddlewareFailure | Self::PolicyDenial => Some(1008), + Self::InvalidUtf8 => Some(1007), + Self::ProtocolError => Some(1002), + Self::MessageTooBig => Some(1009), + } + } + + fn session_end_reason(self) -> openshell_core::proto::WebSocketSessionEndReason { + match self { + Self::PeerDisconnect => { + openshell_core::proto::WebSocketSessionEndReason::PeerDisconnect + } + Self::PolicyReload => openshell_core::proto::WebSocketSessionEndReason::PolicyReload, + Self::MiddlewareDenial => { + openshell_core::proto::WebSocketSessionEndReason::MiddlewareDenial + } + Self::PolicyDenial => openshell_core::proto::WebSocketSessionEndReason::PolicyDenial, + Self::CapacityExhausted | Self::MiddlewareFailure => { + openshell_core::proto::WebSocketSessionEndReason::MiddlewareFailure + } + Self::InvalidUtf8 | Self::ProtocolError | Self::MessageTooBig => { + openshell_core::proto::WebSocketSessionEndReason::ProtocolError + } + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum FrameFailureClass { + MessageAssemblyIdleTimeout, + MessageAssemblyTotalTimeout, + InvalidUtf8, + InvalidFragmentation, + InvalidCloseFrame, + InvalidControlFrame, + InvalidLength, + ReservedOpcode, + UnmaskedClientFrame, + RsvBits, + ProtocolError, +} + +impl FrameFailureClass { + fn as_str(self) -> &'static str { + match self { + Self::MessageAssemblyIdleTimeout => "message_assembly_idle_timeout", + Self::MessageAssemblyTotalTimeout => "message_assembly_total_timeout", + Self::InvalidUtf8 => "invalid_utf8", + Self::InvalidFragmentation => "invalid_fragmentation", + Self::InvalidCloseFrame => "invalid_close_frame", + Self::InvalidControlFrame => "invalid_control_frame", + Self::InvalidLength => "invalid_length", + Self::ReservedOpcode => "reserved_opcode", + Self::UnmaskedClientFrame => "unmasked_client_frame", + Self::RsvBits => "rsv_bits", + Self::ProtocolError => "protocol_error", + } + } +} + +#[derive(Debug)] +struct WebSocketTermination { + cause: WebSocketTerminationCause, + failure_class: Option, + error: miette::Report, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AssemblyTimeoutKind { + Idle, + Total, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum FrameErrorKind { + PeerDisconnect, + Protocol(FrameFailureClass), + InvalidUtf8, + MessageTooBig, + AssemblyTimeout(AssemblyTimeoutKind), +} + +#[derive(Debug)] +struct FrameError { + kind: FrameErrorKind, + error: miette::Report, +} + +impl std::fmt::Display for FrameError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + self.error.fmt(formatter) + } +} + +impl FrameError { + fn peer_io(context: &str, error: std::io::Error) -> Self { + Self { + kind: FrameErrorKind::PeerDisconnect, + error: miette!("{context}: {error}"), + } + } + + fn peer_disconnect(error: miette::Report) -> Self { + Self { + kind: FrameErrorKind::PeerDisconnect, + error, + } + } + + fn protocol(failure_class: FrameFailureClass, error: miette::Report) -> Self { + Self { + kind: FrameErrorKind::Protocol(failure_class), + error, + } + } + + fn invalid_utf8(error: miette::Report) -> Self { + Self { + kind: FrameErrorKind::InvalidUtf8, + error, + } + } + + fn message_too_big(error: miette::Report) -> Self { + Self { + kind: FrameErrorKind::MessageTooBig, + error, + } + } + + fn assembly_timeout(kind: AssemblyTimeoutKind) -> Self { + let error = match kind { + AssemblyTimeoutKind::Idle => { + miette!("websocket text message assembly idle timeout") + } + AssemblyTimeoutKind::Total => { + miette!("websocket text message assembly total timeout") + } + }; + Self { + kind: FrameErrorKind::AssemblyTimeout(kind), + error, + } + } +} + +impl From for WebSocketTermination { + fn from(frame_error: FrameError) -> Self { + let (cause, failure_class) = match frame_error.kind { + FrameErrorKind::PeerDisconnect => (WebSocketTerminationCause::PeerDisconnect, None), + FrameErrorKind::Protocol(failure_class) => ( + WebSocketTerminationCause::ProtocolError, + Some(failure_class), + ), + FrameErrorKind::InvalidUtf8 => ( + WebSocketTerminationCause::InvalidUtf8, + Some(FrameFailureClass::InvalidUtf8), + ), + FrameErrorKind::MessageTooBig => ( + WebSocketTerminationCause::MessageTooBig, + Some(FrameFailureClass::InvalidLength), + ), + FrameErrorKind::AssemblyTimeout(kind) => ( + WebSocketTerminationCause::ProtocolError, + Some(match kind { + AssemblyTimeoutKind::Idle => FrameFailureClass::MessageAssemblyIdleTimeout, + AssemblyTimeoutKind::Total => FrameFailureClass::MessageAssemblyTotalTimeout, + }), + ), + }; + Self { + cause, + failure_class, + error: frame_error.error, + } + } +} + +#[derive(Debug)] +enum WebSocketDecompressionError { + MessageTooBig(miette::Report), + Protocol(miette::Report), +} + +type WebSocketRelayResult = std::result::Result; +type FrameResult = std::result::Result; + +fn terminate(cause: WebSocketTerminationCause, error: miette::Report) -> WebSocketTermination { + WebSocketTermination { + cause, + failure_class: None, + error, + } +} + #[derive(Debug)] struct FrameHeader { fin: bool, @@ -43,8 +320,146 @@ struct FrameHeader { #[derive(Debug)] enum FragmentState { None, - Text { payload: Vec, compressed: bool }, - Binary, + Text(TextMessageAssembly), + Binary { + fragment_count: usize, + coverage: Vec, + }, +} + +/// One admitted client text message while its complete logical payload is +/// being assembled. +/// +/// Keeping the admission permit with the buffer and deadlines makes every +/// timeout and terminal parser error release shared middleware capacity through +/// ordinary ownership. The total deadline includes interleaved control frames; +/// the idle deadline resets only when a client read makes progress. +#[derive(Debug)] +struct TextMessageAssembly { + payload: Vec, + compressed: bool, + fragment_count: usize, + assembly_admission: WebSocketAssemblyAdmission, + admission: Option, + total_deadline: tokio::time::Instant, +} + +impl TextMessageAssembly { + fn new( + compressed: bool, + assembly_admission: WebSocketAssemblyAdmission, + admission: Option, + ) -> Self { + Self { + payload: Vec::new(), + compressed, + fragment_count: 1, + assembly_admission, + admission, + total_deadline: tokio::time::Instant::now() + TEXT_MESSAGE_ASSEMBLY_TOTAL_TIMEOUT, + } + } + + async fn read_some( + &self, + reader: &mut R, + buffer: &mut [u8], + ) -> FrameResult { + if buffer.is_empty() { + return Ok(0); + } + let idle_deadline = tokio::time::Instant::now() + TEXT_MESSAGE_ASSEMBLY_IDLE_TIMEOUT; + tokio::select! { + biased; + () = tokio::time::sleep_until(self.total_deadline) => { + Err(FrameError::assembly_timeout(AssemblyTimeoutKind::Total)) + } + () = tokio::time::sleep_until(idle_deadline) => { + Err(FrameError::assembly_timeout(AssemblyTimeoutKind::Idle)) + } + result = reader.read(buffer) => { + result.map_err(|error| FrameError::peer_io("websocket client read failed", error)) + }, + } + } + + async fn read_exact( + &self, + reader: &mut R, + buffer: &mut [u8], + ) -> FrameResult<()> { + let mut filled = 0; + while filled < buffer.len() { + let read = self.read_some(reader, &mut buffer[filled..]).await?; + if read == 0 { + return Err(FrameError::peer_disconnect(miette!( + "websocket payload ended before declared length" + ))); + } + filled += read; + } + Ok(()) + } + + async fn read_payload( + &mut self, + reader: &mut R, + frame: &FrameHeader, + ) -> FrameResult<()> { + let next = read_masked_payload(reader, frame, Some(self)).await?; + append_text_fragment(&mut self.payload, next) + } + + fn ensure_payload_fits(&self, frame: &FrameHeader) -> FrameResult<()> { + let frame_len = usize::try_from(frame.payload_len).map_err(|_| { + FrameError::message_too_big(miette!("websocket text frame is too large to buffer")) + })?; + let complete_len = self.payload.len().checked_add(frame_len).ok_or_else(|| { + FrameError::message_too_big(miette!("websocket text message length overflow")) + })?; + if complete_len > MAX_TEXT_MESSAGE_BYTES { + return Err(FrameError::message_too_big(miette!( + "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" + ))); + } + Ok(()) + } + + fn add_fragment(&mut self) -> FrameResult<()> { + self.fragment_count = self.fragment_count.saturating_add(1); + if self.fragment_count > MAX_MESSAGE_FRAGMENTS { + return Err(FrameError::protocol( + FrameFailureClass::InvalidFragmentation, + miette!("websocket message exceeds {MAX_MESSAGE_FRAGMENTS} fragment limit"), + )); + } + Ok(()) + } + + async fn within_total( + &self, + future: impl Future>, + ) -> FrameResult { + tokio::time::timeout_at(self.total_deadline, future) + .await + .map_err(|_| FrameError::assembly_timeout(AssemblyTimeoutKind::Total))? + } + + fn into_parts( + self, + ) -> ( + Vec, + bool, + WebSocketAssemblyAdmission, + Option, + ) { + ( + self.payload, + self.compressed, + self.assembly_admission, + self.admission, + ) + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -64,9 +479,13 @@ pub(super) struct InspectionOptions<'a> { pub(super) struct RelayOptions<'a> { pub(super) policy_name: &'a str, + pub(super) assembly_budget: WebSocketAssemblyBudget, pub(super) resolver: Option<&'a SecretResolver>, + pub(super) generation_guard: Option<&'a PolicyGenerationGuard>, pub(super) inspector: Option>, pub(super) compression: WebSocketCompression, + pub(super) middleware_session: Option, + pub(super) middleware_context: Option<&'a L7EvalContext>, } /// Relay an upgraded WebSocket connection with optional client text inspection, @@ -77,7 +496,7 @@ pub(super) async fn relay_with_options( overflow: Vec, host: &str, port: u16, - options: RelayOptions<'_>, + mut options: RelayOptions<'_>, ) -> Result<()> where C: AsyncRead + AsyncWrite + Unpin + Send, @@ -91,23 +510,57 @@ where client_write.flush().await.into_diagnostic()?; } - let client_to_server = - relay_client_to_server(&mut client_read, &mut upstream_write, host, port, &options); + let client_to_server = relay_client_to_server( + &mut client_read, + &mut upstream_write, + host, + port, + &mut options, + ); let server_to_client = async { tokio::io::copy(&mut upstream_read, &mut client_write) .await - .into_diagnostic()?; - client_write.flush().await.into_diagnostic()?; - Ok::<(), miette::Report>(()) + .map_err(|error| { + terminate( + WebSocketTerminationCause::PeerDisconnect, + miette!("websocket upstream relay ended: {error}"), + ) + })?; + client_write.flush().await.map_err(|error| { + terminate( + WebSocketTerminationCause::PeerDisconnect, + miette!("websocket client relay ended: {error}"), + ) + })?; + Ok::<_, WebSocketTermination>( + openshell_core::proto::WebSocketSessionEndReason::PeerDisconnect, + ) }; let result = tokio::select! { result = client_to_server => result, result = server_to_client => result, }; + if let Err(termination) = &result { + observe_termination(host, port, options.policy_name, termination); + } + if let Err(termination) = &result + && let Some(code) = termination.cause.close_code() + { + let payload = code.to_be_bytes(); + let _ = write_masked_close(&mut upstream_write, &payload).await; + let _ = write_unmasked_close(&mut client_write, &payload).await; + } + if let Some(session) = options.middleware_session.take() { + let reason = match &result { + Ok(reason) => *reason, + Err(termination) => termination.cause.session_end_reason(), + }; + session.end(reason).await; + } let _ = upstream_write.shutdown().await; let _ = client_write.shutdown().await; - result + result.map(|_| ()).map_err(|termination| termination.error) } async fn relay_client_to_server( @@ -115,166 +568,242 @@ async fn relay_client_to_server( writer: &mut W, host: &str, port: u16, - options: &RelayOptions<'_>, -) -> Result<()> + options: &mut RelayOptions<'_>, +) -> WebSocketRelayResult where R: AsyncRead + Unpin, W: AsyncWrite + Unpin, { let mut fragments = FragmentState::None; let mut close_seen = false; - loop { - let Some(frame) = read_frame_header(reader).await.inspect_err(|e| { - emit_protocol_failure(host, port, options.policy_name, protocol_failure_class(e)); - })? + let assembly = match &fragments { + FragmentState::Text(assembly) => Some(assembly), + FragmentState::None | FragmentState::Binary { .. } => None, + }; + let Some(frame) = read_frame_header(reader, assembly) + .await + .map_err(WebSocketTermination::from)? else { - writer.shutdown().await.into_diagnostic()?; - return Ok(()); + let _ = writer.shutdown().await; + return Ok(if close_seen { + openshell_core::proto::WebSocketSessionEndReason::NormalClose + } else { + openshell_core::proto::WebSocketSessionEndReason::PeerDisconnect + }); }; if close_seen { - let e = miette!("websocket frame received after close frame"); - emit_protocol_failure(host, port, options.policy_name, protocol_failure_class(&e)); - return Err(e); + return Err(FrameError::protocol( + FrameFailureClass::InvalidCloseFrame, + miette!("websocket frame received after close frame"), + ) + .into()); } - if let Err(e) = validate_frame_header(&frame, &fragments, options.compression) { - emit_protocol_failure(host, port, options.policy_name, protocol_failure_class(&e)); - return Err(e); + validate_frame_header(&frame, &fragments, options.compression) + .map_err(WebSocketTermination::from)?; + + if matches!( + frame.opcode, + OPCODE_TEXT | OPCODE_BINARY | OPCODE_CONTINUATION + ) { + ensure_generation_current(host, port, options)?; } match frame.opcode { OPCODE_TEXT => { - let payload = read_masked_payload(reader, &frame).await.inspect_err(|e| { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(e), - ); - })?; + if frame.payload_len > MAX_TEXT_MESSAGE_BYTES as u64 { + return Err(FrameError::message_too_big(miette!( + "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" + )) + .into()); + } + let assembly_outcome = + options.assembly_budget.reserve().await.map_err(|error| { + terminate(WebSocketTerminationCause::CapacityExhausted, error) + })?; + let assembly_admission = match assembly_outcome { + WebSocketAssemblyAdmissionOutcome::Admitted(admission) => admission, + WebSocketAssemblyAdmissionOutcome::QueueExhausted => { + return Err(terminate( + WebSocketTerminationCause::CapacityExhausted, + miette!( + "websocket assembly admission queue is full; refusing additional buffered work" + ), + )); + } + }; + ensure_generation_current(host, port, options)?; + let admission = if let Some(session) = options.middleware_session.as_mut() { + let message_admission = session.admit_message().await.map_err(|error| { + terminate(WebSocketTerminationCause::MiddlewareFailure, error) + })?; + match message_admission { + openshell_supervisor_middleware::WebSocketMessageAdmission::Bypass => { + options.middleware_session.take(); + None + } + openshell_supervisor_middleware::WebSocketMessageAdmission::Inspect( + admission, + ) => Some(admission), + } + } else { + None + }; + ensure_generation_current(host, port, options)?; let compressed = frame.rsv == 0x40; + let mut assembly = + TextMessageAssembly::new(compressed, assembly_admission, admission); + assembly + .read_payload(reader, &frame) + .await + .map_err(WebSocketTermination::from)?; + ensure_generation_current(host, port, options)?; if frame.fin { + let (payload, compressed, assembly_admission, admission) = + assembly.into_parts(); relay_text_payload( - writer, &frame, payload, false, compressed, host, port, options, - ) - .await - .inspect_err(|e| { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(e), - ); - })?; - } else { - fragments = FragmentState::Text { + writer, + &frame, payload, + assembly_admission, + admission, + false, compressed, - }; + host, + port, + options, + ) + .await?; + } else { + fragments = FragmentState::Text(assembly); } } - OPCODE_CONTINUATION => match &mut fragments { - FragmentState::Text { - payload, - compressed, - } => { - let next = read_masked_payload(reader, &frame).await.inspect_err(|e| { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(e), - ); - })?; - if let Err(e) = append_text_fragment(payload, next) { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(&e), - ); - return Err(e); + OPCODE_CONTINUATION => { + if let FragmentState::Text(assembly) = &mut fragments { + assembly + .add_fragment() + .map_err(WebSocketTermination::from)?; + if frame.payload_len > MAX_TEXT_MESSAGE_BYTES as u64 { + return Err(FrameError::message_too_big(miette!( + "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" + )) + .into()); } + assembly + .ensure_payload_fits(&frame) + .map_err(WebSocketTermination::from)?; + assembly + .read_payload(reader, &frame) + .await + .map_err(WebSocketTermination::from)?; + ensure_generation_current(host, port, options)?; if frame.fin { - let complete = std::mem::take(payload); - let was_compressed = *compressed; - fragments = FragmentState::None; + let FragmentState::Text(assembly) = + std::mem::replace(&mut fragments, FragmentState::None) + else { + unreachable!("validated text continuation state") + }; + let (complete, was_compressed, assembly_admission, admission) = + assembly.into_parts(); relay_text_payload( writer, &frame, complete, + assembly_admission, + admission, true, was_compressed, host, port, options, ) - .await - .inspect_err(|e| { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(e), - ); - })?; + .await?; + } + } else if let FragmentState::Binary { + fragment_count, + coverage, + } = &mut fragments + { + *fragment_count = fragment_count.saturating_add(1); + if *fragment_count > MAX_MESSAGE_FRAGMENTS { + return Err(FrameError::protocol( + FrameFailureClass::InvalidFragmentation, + miette!( + "websocket message exceeds {MAX_MESSAGE_FRAGMENTS} fragment limit" + ), + ) + .into()); } - } - FragmentState::Binary => { copy_raw_frame_payload(reader, writer, &frame) .await - .inspect_err(|e| { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(e), - ); - })?; + .map_err(WebSocketTermination::from)?; + let fragment_size = usize::try_from(frame.payload_len).unwrap_or(usize::MAX); + for record in coverage.iter_mut() { + record.original_size = record.original_size.saturating_add(fragment_size); + } if frame.fin { - fragments = FragmentState::None; + let FragmentState::Binary { coverage, .. } = + std::mem::replace(&mut fragments, FragmentState::None) + else { + unreachable!("validated binary continuation state") + }; + if let Some(ctx) = options.middleware_context { + crate::l7::middleware::emit_websocket_coverage(ctx, &coverage); + } } + } else { + return Err(FrameError::protocol( + FrameFailureClass::InvalidFragmentation, + miette!("websocket continuation frame without active fragmented message"), + ) + .into()); } - FragmentState::None => { - let e = - miette!("websocket continuation frame without active fragmented message"); - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(&e), - ); - return Err(e); - } - }, + } OPCODE_BINARY => { - if !frame.fin { - fragments = FragmentState::Binary; - } + let initial_size = usize::try_from(frame.payload_len).unwrap_or(usize::MAX); + let coverage = options + .middleware_session + .as_mut() + .map(|session| { + session.observe_unsupported_message( + openshell_core::proto::WebSocketMessageType::Binary, + initial_size, + ) + }) + .unwrap_or_default(); copy_raw_frame_payload(reader, writer, &frame) .await - .inspect_err(|e| { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(e), - ); - })?; + .map_err(WebSocketTermination::from)?; + if frame.fin + && let Some(ctx) = options.middleware_context + { + crate::l7::middleware::emit_websocket_coverage(ctx, &coverage); + } else if !frame.fin { + fragments = FragmentState::Binary { + fragment_count: 1, + coverage, + }; + } } OPCODE_CLOSE | OPCODE_PING | OPCODE_PONG => { - relay_control_frame(reader, writer, &frame) - .await - .inspect_err(|e| { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(e), - ); - })?; + let control_result = match &fragments { + FragmentState::Text(assembly) => { + assembly + .within_total(relay_control_frame( + reader, + writer, + &frame, + Some(assembly), + )) + .await + } + FragmentState::None | FragmentState::Binary { .. } => { + relay_control_frame(reader, writer, &frame, None).await + } + }; + control_result.map_err(WebSocketTermination::from)?; if frame.opcode == OPCODE_CLOSE { close_seen = true; } @@ -284,25 +813,41 @@ where } } -async fn read_frame_header(reader: &mut R) -> Result> { - let first = match reader.read_u8().await { - Ok(byte) => byte, - Err(e) - if matches!( - e.kind(), - std::io::ErrorKind::UnexpectedEof - | std::io::ErrorKind::ConnectionReset - | std::io::ErrorKind::BrokenPipe - ) => - { - return Ok(None); - } - Err(e) => return Err(miette!("{e}")), +async fn read_exact_for_assembly( + reader: &mut R, + buffer: &mut [u8], + assembly: Option<&TextMessageAssembly>, +) -> FrameResult<()> { + match assembly { + Some(assembly) => assembly.read_exact(reader, buffer).await, + None => reader + .read_exact(buffer) + .await + .map(|_| ()) + .map_err(|error| FrameError::peer_io("websocket client read failed", error)), + } +} + +async fn read_frame_header( + reader: &mut R, + assembly: Option<&TextMessageAssembly>, +) -> FrameResult> { + let mut first = [0u8; 1]; + let first_read = match assembly { + Some(assembly) => assembly.read_some(reader, &mut first).await, + None => reader + .read(&mut first) + .await + .map_err(|error| FrameError::peer_io("websocket client read failed", error)), }; - let second = reader - .read_u8() - .await - .map_err(|e| miette!("malformed websocket frame header: {e}"))?; + let first = match first_read { + Ok(0) => return Ok(None), + Ok(_) => first[0], + Err(error) => return Err(error), + }; + let mut second = [0u8; 1]; + read_exact_for_assembly(reader, &mut second, assembly).await?; + let second = second[0]; let mut raw_header = vec![first, second]; let len_code = second & 0x7F; @@ -310,33 +855,32 @@ async fn read_frame_header(reader: &mut R) -> Result u64::from(len_code), 126 => { let mut bytes = [0u8; 2]; - reader - .read_exact(&mut bytes) - .await - .map_err(|e| miette!("malformed websocket extended length: {e}"))?; + read_exact_for_assembly(reader, &mut bytes, assembly).await?; raw_header.extend_from_slice(&bytes); let len = u64::from(u16::from_be_bytes(bytes)); if len < 126 { - return Err(miette!( - "websocket frame uses non-minimal 16-bit extended length" + return Err(FrameError::protocol( + FrameFailureClass::InvalidLength, + miette!("websocket frame uses non-minimal 16-bit extended length"), )); } len } 127 => { let mut bytes = [0u8; 8]; - reader - .read_exact(&mut bytes) - .await - .map_err(|e| miette!("malformed websocket extended length: {e}"))?; + read_exact_for_assembly(reader, &mut bytes, assembly).await?; if bytes[0] & 0x80 != 0 { - return Err(miette!("websocket frame uses non-canonical 64-bit length")); + return Err(FrameError::protocol( + FrameFailureClass::InvalidLength, + miette!("websocket frame uses non-canonical 64-bit length"), + )); } raw_header.extend_from_slice(&bytes); let len = u64::from_be_bytes(bytes); if u16::try_from(len).is_ok() { - return Err(miette!( - "websocket frame uses non-minimal 64-bit extended length" + return Err(FrameError::protocol( + FrameFailureClass::InvalidLength, + miette!("websocket frame uses non-minimal 64-bit extended length"), )); } len @@ -347,10 +891,7 @@ async fn read_frame_header(reader: &mut R) -> Result Result<()> { +) -> FrameResult<()> { if !valid_rsv_bits(frame, fragments, compression) { - return Err(miette!( - "websocket frame has unsupported RSV bits or extension state" + return Err(FrameError::protocol( + FrameFailureClass::RsvBits, + miette!("websocket frame has unsupported RSV bits or extension state"), )); } if !frame.masked { - return Err(miette!("websocket client frame is not masked")); + return Err(FrameError::protocol( + FrameFailureClass::UnmaskedClientFrame, + miette!("websocket client frame is not masked"), + )); } if !matches!( frame.opcode, @@ -390,34 +935,49 @@ fn validate_frame_header( | OPCODE_PING | OPCODE_PONG ) { - return Err(miette!("websocket frame uses reserved opcode")); + return Err(FrameError::protocol( + FrameFailureClass::ReservedOpcode, + miette!("websocket frame uses reserved opcode"), + )); } if matches!(frame.opcode, OPCODE_CLOSE | OPCODE_PING | OPCODE_PONG) { if !frame.fin { - return Err(miette!("websocket control frame is fragmented")); + return Err(FrameError::protocol( + FrameFailureClass::InvalidControlFrame, + miette!("websocket control frame is fragmented"), + )); } if frame.payload_len > 125 { - return Err(miette!("websocket control frame exceeds 125 bytes")); + return Err(FrameError::protocol( + FrameFailureClass::InvalidControlFrame, + miette!("websocket control frame exceeds 125 bytes"), + )); } } if matches!(frame.opcode, OPCODE_TEXT | OPCODE_BINARY) && !matches!(fragments, FragmentState::None) { - return Err(miette!( - "websocket data frame started before previous fragmented message completed" + return Err(FrameError::protocol( + FrameFailureClass::InvalidFragmentation, + miette!("websocket data frame started before previous fragmented message completed"), )); } if matches!(frame.opcode, OPCODE_CONTINUATION) && matches!(fragments, FragmentState::None) { - return Err(miette!( - "websocket continuation frame without active fragmented message" + return Err(FrameError::protocol( + FrameFailureClass::InvalidFragmentation, + miette!("websocket continuation frame without active fragmented message"), )); } if (frame.opcode == OPCODE_BINARY - || (frame.opcode == OPCODE_CONTINUATION && matches!(fragments, FragmentState::Binary))) + || (frame.opcode == OPCODE_CONTINUATION + && matches!(fragments, FragmentState::Binary { .. }))) && frame.payload_len > MAX_RAW_FRAME_PAYLOAD_BYTES { - return Err(miette!( - "websocket binary frame exceeds {MAX_RAW_FRAME_PAYLOAD_BYTES} byte relay limit" + return Err(FrameError::protocol( + FrameFailureClass::InvalidLength, + miette!( + "websocket binary frame exceeds {MAX_RAW_FRAME_PAYLOAD_BYTES} byte relay limit" + ), )); } Ok(()) @@ -440,82 +1000,195 @@ fn valid_rsv_bits( async fn read_masked_payload( reader: &mut R, frame: &FrameHeader, -) -> Result> { - let payload_len = usize::try_from(frame.payload_len) - .map_err(|_| miette!("websocket text frame is too large to buffer"))?; + assembly: Option<&TextMessageAssembly>, +) -> FrameResult> { + let payload_len = usize::try_from(frame.payload_len).map_err(|_| { + FrameError::message_too_big(miette!("websocket text frame is too large to buffer")) + })?; if payload_len > MAX_TEXT_MESSAGE_BYTES { - return Err(miette!( + return Err(FrameError::message_too_big(miette!( "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" - )); + ))); } let mut payload = vec![0u8; payload_len]; - reader - .read_exact(&mut payload) - .await - .map_err(|e| miette!("malformed websocket payload: {e}"))?; - let mask_key = frame - .mask_key - .ok_or_else(|| miette!("websocket client frame is not masked"))?; + read_exact_for_assembly(reader, &mut payload, assembly).await?; + let mask_key = frame.mask_key.ok_or_else(|| { + FrameError::protocol( + FrameFailureClass::UnmaskedClientFrame, + miette!("websocket client frame is not masked"), + ) + })?; apply_mask(&mut payload, mask_key); Ok(payload) } -fn append_text_fragment(buffer: &mut Vec, next: Vec) -> Result<()> { - let new_len = buffer - .len() - .checked_add(next.len()) - .ok_or_else(|| miette!("websocket text message length overflow"))?; +fn append_text_fragment(buffer: &mut Vec, next: Vec) -> FrameResult<()> { + let new_len = buffer.len().checked_add(next.len()).ok_or_else(|| { + FrameError::message_too_big(miette!("websocket text message length overflow")) + })?; if new_len > MAX_TEXT_MESSAGE_BYTES { - return Err(miette!( + return Err(FrameError::message_too_big(miette!( "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" - )); + ))); } buffer.extend_from_slice(&next); Ok(()) } +fn ensure_generation_current( + host: &str, + port: u16, + options: &RelayOptions<'_>, +) -> WebSocketRelayResult<()> { + ensure_generation_guard_current(host, port, options.policy_name, options.generation_guard) +} + +fn ensure_generation_guard_current( + host: &str, + port: u16, + policy_name: &str, + generation_guard: Option<&PolicyGenerationGuard>, +) -> WebSocketRelayResult<()> { + let Some(guard) = generation_guard else { + return Ok(()); + }; + guard.ensure_current().map_err(|error| { + crate::l7::relay::emit_policy_reload(guard, host, port, policy_name); + terminate(WebSocketTerminationCause::PolicyReload, error) + }) +} + #[allow(clippy::too_many_arguments)] async fn relay_text_payload( writer: &mut W, frame: &FrameHeader, payload: Vec, + _assembly_admission: WebSocketAssemblyAdmission, + admission: Option, force_reframe: bool, compressed: bool, host: &str, port: u16, - options: &RelayOptions<'_>, -) -> Result<()> { + options: &mut RelayOptions<'_>, +) -> WebSocketRelayResult<()> { + ensure_generation_current(host, port, options)?; let message_payload = if compressed { - decompress_permessage_deflate(&payload)? + decompress_permessage_deflate(&payload).map_err(|error| match error { + WebSocketDecompressionError::MessageTooBig(error) => { + WebSocketTermination::from(FrameError::message_too_big(error)) + } + WebSocketDecompressionError::Protocol(error) => WebSocketTermination::from( + FrameError::protocol(FrameFailureClass::ProtocolError, error), + ), + })? } else { payload }; - let mut text = String::from_utf8(message_payload) - .map_err(|_| miette!("websocket text message is not valid UTF-8"))?; + let mut text = String::from_utf8(message_payload).map_err(|_| { + WebSocketTermination::from(FrameError::invalid_utf8(miette!( + "websocket text message is not valid UTF-8" + ))) + })?; + + // 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() { + inspect_websocket_text_message(host, port, options.policy_name, inspector, &text)?; + } + ensure_generation_current(host, port, options)?; + + let mut middleware_transformed = false; + if let Some(session) = options.middleware_session.as_mut() { + let admission = admission.ok_or_else(|| { + terminate( + WebSocketTerminationCause::MiddlewareFailure, + miette!("websocket middleware message missing work admission"), + ) + })?; + let outcome = session + .evaluate_text_admitted(text.into_bytes(), admission) + .await; + ensure_generation_current(host, port, options)?; + if let Some(ctx) = options.middleware_context { + crate::l7::middleware::emit_websocket_message_events(ctx, &outcome); + } + if !outcome.allowed { + if outcome.platform_oversize { + return Err(FrameError::message_too_big(miette!( + "websocket message over middleware platform capacity" + )) + .into()); + } + let cause = if outcome.denial.is_some() { + WebSocketTerminationCause::MiddlewareDenial + } else { + WebSocketTerminationCause::MiddlewareFailure + }; + return Err(terminate( + cause, + miette!("websocket middleware denied message: {}", outcome.reason), + )); + } + middleware_transformed = outcome + .invocations + .iter() + .any(|invocation| invocation.transformed); + text = String::from_utf8(outcome.payload).map_err(|_| { + terminate( + WebSocketTerminationCause::MiddlewareFailure, + miette!("websocket middleware returned invalid UTF-8"), + ) + })?; + } + + if middleware_transformed && let Some(inspector) = options.inspector.as_ref() { + inspect_websocket_text_message(host, port, options.policy_name, inspector, &text)?; + } + ensure_generation_current(host, port, options)?; + let replacements = if let Some(resolver) = options.resolver { resolver .rewrite_websocket_text_placeholders(&mut text) - .map_err(|_| miette!("websocket credential placeholder resolution failed"))? + .map_err(|_| { + terminate( + WebSocketTerminationCause::MiddlewareFailure, + miette!("websocket credential placeholder resolution failed"), + ) + })? } else { 0 }; + ensure_generation_current(host, port, options)?; - if let Some(inspector) = options.inspector.as_ref() { - inspect_websocket_text_message(host, port, options.policy_name, inspector, &text)?; - } - - if replacements == 0 && !force_reframe && !compressed { - writer - .write_all(&frame.raw_header) - .await - .into_diagnostic()?; + if replacements == 0 && !middleware_transformed && !force_reframe && !compressed { + writer.write_all(&frame.raw_header).await.map_err(|error| { + terminate( + WebSocketTerminationCause::PeerDisconnect, + miette!("websocket upstream write failed: {error}"), + ) + })?; + ensure_generation_current(host, port, options)?; let mut payload = text.into_bytes(); - let mask_key = frame - .mask_key - .ok_or_else(|| miette!("websocket client frame is not masked"))?; + let mask_key = frame.mask_key.ok_or_else(|| { + WebSocketTermination::from(FrameError::protocol( + FrameFailureClass::UnmaskedClientFrame, + miette!("websocket client frame is not masked"), + )) + })?; apply_mask(&mut payload, mask_key); - writer.write_all(&payload).await.into_diagnostic()?; - writer.flush().await.into_diagnostic()?; + writer.write_all(&payload).await.map_err(|error| { + terminate( + WebSocketTerminationCause::PeerDisconnect, + miette!("websocket upstream write failed: {error}"), + ) + })?; + writer.flush().await.map_err(|error| { + terminate( + WebSocketTerminationCause::PeerDisconnect, + miette!("websocket upstream flush failed: {error}"), + ) + })?; return Ok(()); } @@ -523,10 +1196,33 @@ async fn relay_text_payload( emit_rewrite_event(host, port, options.policy_name, replacements); } if compressed { - let compressed_payload = compress_permessage_deflate(text.as_bytes())?; - return write_masked_frame_with_rsv(writer, OPCODE_TEXT, 0x40, &compressed_payload).await; + let compressed_payload = compress_permessage_deflate(text.as_bytes()).map_err(|error| { + WebSocketTermination::from(FrameError::protocol( + FrameFailureClass::ProtocolError, + error, + )) + })?; + return write_masked_text_frame_guarded( + writer, + 0x40, + &compressed_payload, + host, + port, + options.policy_name, + options.generation_guard, + ) + .await; } - write_masked_frame(writer, OPCODE_TEXT, text.as_bytes()).await + write_masked_text_frame_guarded( + writer, + 0, + text.as_bytes(), + host, + port, + options.policy_name, + options.generation_guard, + ) + .await } fn inspect_websocket_text_message( @@ -535,7 +1231,7 @@ fn inspect_websocket_text_message( policy_name: &str, inspector: &InspectionOptions<'_>, text: &str, -) -> Result<()> { +) -> WebSocketRelayResult<()> { if inspector.graphql_policy { return inspect_graphql_websocket_message(host, port, policy_name, inspector, text); } @@ -547,7 +1243,8 @@ fn inspect_websocket_text_message( graphql: None, jsonrpc: None, }; - let (allowed, reason) = evaluate_l7_request(inspector.engine, inspector.ctx, &request_info)?; + let (allowed, reason) = evaluate_l7_request(inspector.engine, inspector.ctx, &request_info) + .map_err(|error| terminate(WebSocketTerminationCause::PolicyReload, error))?; let decision = match (allowed, inspector.enforcement) { (true, _) => "allow", (false, EnforcementMode::Audit) => "audit", @@ -563,7 +1260,10 @@ fn inspect_websocket_text_message( None, ); if !allowed && inspector.enforcement == EnforcementMode::Enforce { - return Err(miette!("websocket text message denied by policy")); + return Err(terminate( + WebSocketTerminationCause::PolicyDenial, + miette!("websocket text message denied by policy"), + )); } Ok(()) } @@ -574,7 +1274,7 @@ fn inspect_graphql_websocket_message( policy_name: &str, inspector: &InspectionOptions<'_>, text: &str, -) -> Result<()> { +) -> WebSocketRelayResult<()> { match classify_graphql_websocket_message(text) { GraphqlWebSocketMessage::Control { message_type } => { let request_info = L7RequestInfo { @@ -614,7 +1314,8 @@ fn inspect_graphql_websocket_message( let (allowed, reason) = if let Some(reason) = parse_error_reason { (false, reason) } else { - evaluate_l7_request(inspector.engine, inspector.ctx, &request_info)? + evaluate_l7_request(inspector.engine, inspector.ctx, &request_info) + .map_err(|error| terminate(WebSocketTerminationCause::PolicyReload, error))? }; let decision = match (allowed, inspector.enforcement) { (_, _) if force_deny => "deny", @@ -633,7 +1334,10 @@ fn inspect_graphql_websocket_message( Some(&graphql), ); if (!allowed && inspector.enforcement == EnforcementMode::Enforce) || force_deny { - return Err(miette!("websocket GraphQL message denied by policy")); + return Err(terminate( + WebSocketTerminationCause::PolicyDenial, + miette!("websocket GraphQL message denied by policy"), + )); } Ok(()) } @@ -728,24 +1432,29 @@ async fn relay_control_frame( reader: &mut R, writer: &mut W, frame: &FrameHeader, -) -> Result<()> + assembly: Option<&TextMessageAssembly>, +) -> FrameResult<()> where R: AsyncRead + Unpin, W: AsyncWrite + Unpin, { - let raw_payload_len = usize::try_from(frame.payload_len) - .map_err(|_| miette!("websocket control frame payload length overflow"))?; + let raw_payload_len = usize::try_from(frame.payload_len).map_err(|_| { + FrameError::protocol( + FrameFailureClass::InvalidControlFrame, + miette!("websocket control frame payload length overflow"), + ) + })?; let mut raw_payload = vec![0u8; raw_payload_len]; - reader - .read_exact(&mut raw_payload) - .await - .map_err(|e| miette!("malformed websocket control payload: {e}"))?; + read_exact_for_assembly(reader, &mut raw_payload, assembly).await?; if frame.opcode == OPCODE_CLOSE { let mut payload = raw_payload.clone(); - let mask_key = frame - .mask_key - .ok_or_else(|| miette!("websocket client frame is not masked"))?; + let mask_key = frame.mask_key.ok_or_else(|| { + FrameError::protocol( + FrameFailureClass::UnmaskedClientFrame, + miette!("websocket client frame is not masked"), + ) + })?; apply_mask(&mut payload, mask_key); validate_close_payload(&payload)?; } @@ -753,16 +1462,23 @@ where writer .write_all(&frame.raw_header) .await - .into_diagnostic()?; - writer.write_all(&raw_payload).await.into_diagnostic()?; - writer.flush().await.into_diagnostic()?; + .map_err(|error| FrameError::peer_io("websocket upstream write failed", error))?; + writer + .write_all(&raw_payload) + .await + .map_err(|error| FrameError::peer_io("websocket upstream write failed", error))?; + writer + .flush() + .await + .map_err(|error| FrameError::peer_io("websocket upstream flush failed", error))?; Ok(()) } -fn validate_close_payload(payload: &[u8]) -> Result<()> { +fn validate_close_payload(payload: &[u8]) -> FrameResult<()> { if payload.len() == 1 { - return Err(miette!( - "websocket close frame payload cannot be exactly one byte" + return Err(FrameError::protocol( + FrameFailureClass::InvalidCloseFrame, + miette!("websocket close frame payload cannot be exactly one byte"), )); } if payload.len() < 2 { @@ -771,10 +1487,15 @@ fn validate_close_payload(payload: &[u8]) -> Result<()> { let code = u16::from_be_bytes([payload[0], payload[1]]); if !valid_close_code(code) { - return Err(miette!("websocket close frame uses invalid close code")); + return Err(FrameError::protocol( + FrameFailureClass::InvalidCloseFrame, + miette!("websocket close frame uses invalid close code"), + )); } if std::str::from_utf8(&payload[2..]).is_err() { - return Err(miette!("websocket close frame reason is not valid UTF-8")); + return Err(FrameError::invalid_utf8(miette!( + "websocket close frame reason is not valid UTF-8" + ))); } Ok(()) } @@ -787,7 +1508,7 @@ async fn copy_raw_frame_payload( reader: &mut R, writer: &mut W, frame: &FrameHeader, -) -> Result<()> +) -> FrameResult<()> where R: AsyncRead + Unpin, W: AsyncWrite + Unpin, @@ -795,21 +1516,32 @@ where writer .write_all(&frame.raw_header) .await - .into_diagnostic()?; + .map_err(|error| FrameError::peer_io("websocket upstream write failed", error))?; let mut remaining = frame.payload_len; let mut buf = [0u8; COPY_BUF_SIZE]; while remaining > 0 { let to_read = usize::try_from(remaining) .unwrap_or(buf.len()) .min(buf.len()); - let n = reader.read(&mut buf[..to_read]).await.into_diagnostic()?; + let n = reader + .read(&mut buf[..to_read]) + .await + .map_err(|error| FrameError::peer_io("websocket client read failed", error))?; if n == 0 { - return Err(miette!("websocket payload ended before declared length")); + return Err(FrameError::peer_disconnect(miette!( + "websocket payload ended before declared length" + ))); } - writer.write_all(&buf[..n]).await.into_diagnostic()?; + writer + .write_all(&buf[..n]) + .await + .map_err(|error| FrameError::peer_io("websocket upstream write failed", error))?; remaining -= n as u64; } - writer.flush().await.into_diagnostic()?; + writer + .flush() + .await + .map_err(|error| FrameError::peer_io("websocket upstream flush failed", error))?; Ok(()) } @@ -821,12 +1553,73 @@ async fn write_masked_frame( write_masked_frame_with_rsv(writer, opcode, 0, payload).await } +pub(super) async fn write_masked_close( + writer: &mut W, + payload: &[u8], +) -> Result<()> { + write_masked_frame(writer, OPCODE_CLOSE, payload).await +} + +pub(super) async fn write_unmasked_close( + writer: &mut W, + payload: &[u8], +) -> Result<()> { + let payload_len = u8::try_from(payload.len()) + .map_err(|_| miette!("websocket close payload exceeds 125 bytes"))?; + writer + .write_all(&[0x80 | OPCODE_CLOSE, payload_len]) + .await + .into_diagnostic()?; + writer.write_all(payload).await.into_diagnostic()?; + writer.flush().await.into_diagnostic()?; + Ok(()) +} + async fn write_masked_frame_with_rsv( writer: &mut W, opcode: u8, rsv: u8, payload: &[u8], ) -> Result<()> { + let (header, masked) = masked_frame_parts(opcode, rsv, payload); + writer.write_all(&header).await.into_diagnostic()?; + writer.write_all(&masked).await.into_diagnostic()?; + writer.flush().await.into_diagnostic()?; + Ok(()) +} + +async fn write_masked_text_frame_guarded( + writer: &mut W, + rsv: u8, + payload: &[u8], + host: &str, + port: u16, + policy_name: &str, + generation_guard: Option<&PolicyGenerationGuard>, +) -> WebSocketRelayResult<()> { + let (header, masked) = masked_frame_parts(OPCODE_TEXT, rsv, payload); + writer.write_all(&header).await.map_err(|error| { + terminate( + WebSocketTerminationCause::PeerDisconnect, + miette!("websocket upstream write failed: {error}"), + ) + })?; + ensure_generation_guard_current(host, port, policy_name, generation_guard)?; + writer.write_all(&masked).await.map_err(|error| { + terminate( + WebSocketTerminationCause::PeerDisconnect, + miette!("websocket upstream write failed: {error}"), + ) + })?; + writer.flush().await.map_err(|error| { + terminate( + WebSocketTerminationCause::PeerDisconnect, + miette!("websocket upstream flush failed: {error}"), + ) + }) +} + +fn masked_frame_parts(opcode: u8, rsv: u8, payload: &[u8]) -> (Vec, Vec) { let mut header = Vec::with_capacity(14); header.push(0x80 | rsv | opcode); match payload.len() { @@ -849,13 +1642,12 @@ async fn write_masked_frame_with_rsv( let mut masked = payload.to_vec(); apply_mask(&mut masked, mask_key); - writer.write_all(&header).await.into_diagnostic()?; - writer.write_all(&masked).await.into_diagnostic()?; - writer.flush().await.into_diagnostic()?; - Ok(()) + (header, masked) } -fn decompress_permessage_deflate(payload: &[u8]) -> Result> { +fn decompress_permessage_deflate( + payload: &[u8], +) -> std::result::Result, WebSocketDecompressionError> { let mut decoder = Decompress::new(false); let mut input = Vec::with_capacity(payload.len() + 4); input.extend_from_slice(payload); @@ -868,18 +1660,30 @@ fn decompress_permessage_deflate(payload: &[u8]) -> Result> { let before_out = decoder.total_out(); let status = decoder .decompress(&input[input_pos..], &mut scratch, FlushDecompress::Sync) - .map_err(|e| miette!("websocket permessage-deflate decompression failed: {e}"))?; - let read = usize::try_from(decoder.total_in() - before_in) - .map_err(|_| miette!("websocket permessage-deflate input length overflow"))?; - let written = usize::try_from(decoder.total_out() - before_out) - .map_err(|_| miette!("websocket permessage-deflate output length overflow"))?; - input_pos = input_pos - .checked_add(read) - .ok_or_else(|| miette!("websocket permessage-deflate input length overflow"))?; + .map_err(|e| { + WebSocketDecompressionError::Protocol(miette!( + "websocket permessage-deflate decompression failed: {e}" + )) + })?; + let read = usize::try_from(decoder.total_in() - before_in).map_err(|_| { + WebSocketDecompressionError::Protocol(miette!( + "websocket permessage-deflate input length overflow" + )) + })?; + let written = usize::try_from(decoder.total_out() - before_out).map_err(|_| { + WebSocketDecompressionError::Protocol(miette!( + "websocket permessage-deflate output length overflow" + )) + })?; + input_pos = input_pos.checked_add(read).ok_or_else(|| { + WebSocketDecompressionError::Protocol(miette!( + "websocket permessage-deflate input length overflow" + )) + })?; if out.len().saturating_add(written) > MAX_TEXT_MESSAGE_BYTES { - return Err(miette!( + return Err(WebSocketDecompressionError::MessageTooBig(miette!( "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" - )); + ))); } out.extend_from_slice(&scratch[..written]); if matches!(status, Status::StreamEnd) { @@ -889,9 +1693,9 @@ fn decompress_permessage_deflate(payload: &[u8]) -> Result> { break; } if read == 0 && written == 0 { - return Err(miette!( + return Err(WebSocketDecompressionError::Protocol(miette!( "websocket permessage-deflate decompression did not make progress" - )); + ))); } } Ok(out) @@ -1048,44 +1852,60 @@ fn graphql_log_summary(info: &crate::l7::graphql::GraphqlRequestInfo) -> String format!(" graphql_ops={}", ops.join(";")) } -fn protocol_failure_class(error: &miette::Report) -> &'static str { - let msg = error.to_string().to_ascii_lowercase(); - if msg.contains("credential") { - "credential_resolution_failed" - } else if msg.contains("utf-8") { - "invalid_utf8" - } else if msg.contains("close frame") || msg.contains("after close") { - "invalid_close_frame" - } else if msg.contains("control frame") { - "invalid_control_frame" - } else if msg.contains("length") - || msg.contains("too large") - || msg.contains("exceeds") - || msg.contains("overflow") - { - "invalid_length" - } else if msg.contains("continuation") || msg.contains("fragmented") { - "invalid_fragmentation" - } else if msg.contains("reserved opcode") { - "reserved_opcode" - } else if msg.contains("not masked") { - "unmasked_client_frame" - } else if msg.contains("rsv") { - "rsv_bits" - } else if msg.contains("malformed") { - "malformed_frame" - } else { - "protocol_error" +fn observe_termination( + host: &str, + port: u16, + policy_name: &str, + termination: &WebSocketTermination, +) { + if let Some(failure_class) = termination.failure_class { + emit_protocol_failure(host, port, policy_name, failure_class); } + if termination.cause == WebSocketTerminationCause::CapacityExhausted { + ocsf_emit!( + NetworkActivityBuilder::new(openshell_ocsf::ctx::ctx()) + .activity(ActivityId::Open) + .action(ActionId::Denied) + .disposition(DispositionId::Blocked) + .severity(SeverityId::Medium) + .status(StatusId::Failure) + .dst_endpoint(Endpoint::from_domain(host, port)) + .firewall_rule(policy_name, "l7-websocket") + .status_detail("assembly_capacity_exhausted") + .message(format!( + "WebSocket text assembly capacity exhausted [host:{host} port:{port}]" + )) + .build() + ); + } +} + +fn emit_protocol_failure( + host: &str, + port: u16, + policy_name: &str, + failure_class: FrameFailureClass, +) { + ocsf_emit!(protocol_failure_event( + host, + port, + policy_name, + failure_class + )); } -fn emit_protocol_failure(host: &str, port: u16, policy_name: &str, failure_class: &str) { +fn protocol_failure_event( + host: &str, + port: u16, + policy_name: &str, + failure_class: FrameFailureClass, +) -> openshell_ocsf::OcsfEvent { let policy_name = if policy_name.is_empty() { "-" } else { policy_name }; - let event = NetworkActivityBuilder::new(openshell_ocsf::ctx::ctx()) + NetworkActivityBuilder::new(openshell_ocsf::ctx::ctx()) .activity(ActivityId::Open) .action(ActionId::Denied) .disposition(DispositionId::Blocked) @@ -1094,9 +1914,8 @@ fn emit_protocol_failure(host: &str, port: u16, policy_name: &str, failure_class .dst_endpoint(Endpoint::from_domain(host, port)) .firewall_rule(policy_name, "l7-websocket") .message(protocol_failure_message(host, port)) - .status_detail(failure_class) - .build(); - ocsf_emit!(event); + .status_detail(failure_class.as_str()) + .build() } fn protocol_failure_message(host: &str, port: u16) -> String { @@ -1108,9 +1927,28 @@ mod tests { use super::*; use crate::l7::relay::L7EvalContext; use crate::opa::{NetworkInput, OpaEngine}; + use openshell_core::proto::middleware::v1::supervisor_middleware_server::{ + SupervisorMiddleware, SupervisorMiddlewareServer, + }; + use openshell_core::secrets::SecretResolver; use std::path::PathBuf; use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio_stream::wrappers::{ReceiverStream, TcpListenerStream}; + use tonic::{Request, Response, Status}; + + async fn test_assembly_admission() -> WebSocketAssemblyAdmission { + match WebSocketAssemblyBudget::new(1, 0) + .reserve() + .await + .expect("reserve test assembly") + { + WebSocketAssemblyAdmissionOutcome::Admitted(admission) => admission, + WebSocketAssemblyAdmissionOutcome::QueueExhausted => { + panic!("fresh test assembly budget must admit") + } + } + } const TEST_POLICY: &str = include_str!("../../data/sandbox-policy.rego"); const GRAPHQL_WS_POLICY: &str = r#" @@ -1137,6 +1975,54 @@ network_policies: - { path: /usr/bin/node } "#; + #[test] + fn termination_causes_map_to_protocol_close_codes_and_session_reasons() { + use openshell_core::proto::WebSocketSessionEndReason as EndReason; + + assert_eq!( + WebSocketTerminationCause::InvalidUtf8.close_code(), + Some(1007) + ); + assert_eq!( + WebSocketTerminationCause::ProtocolError.close_code(), + Some(1002) + ); + assert_eq!( + WebSocketTerminationCause::MessageTooBig.close_code(), + Some(1009) + ); + assert_eq!( + WebSocketTerminationCause::MiddlewareDenial.close_code(), + Some(1008) + ); + assert_eq!( + WebSocketTerminationCause::PolicyReload.close_code(), + Some(1012) + ); + assert_eq!(WebSocketTerminationCause::PeerDisconnect.close_code(), None); + + assert_eq!( + WebSocketTerminationCause::InvalidUtf8.session_end_reason(), + EndReason::ProtocolError + ); + assert_eq!( + WebSocketTerminationCause::MiddlewareDenial.session_end_reason(), + EndReason::MiddlewareDenial + ); + assert_eq!( + WebSocketTerminationCause::PolicyDenial.session_end_reason(), + EndReason::PolicyDenial + ); + assert_eq!( + WebSocketTerminationCause::MiddlewareFailure.session_end_reason(), + EndReason::MiddlewareFailure + ); + assert_eq!( + WebSocketTerminationCause::PolicyReload.session_end_reason(), + EndReason::PolicyReload + ); + } + fn resolver() -> (HashMap, SecretResolver) { let (child_env, resolver) = SecretResolver::from_provider_env( std::iter::once(("DISCORD_BOT_TOKEN".to_string(), "real-token".to_string())).collect(), @@ -1208,6 +2094,36 @@ network_policies: frame } + fn test_frame_header(fin: bool, opcode: u8, payload_len: u64) -> FrameHeader { + FrameHeader { + fin, + rsv: 0, + opcode, + masked: true, + payload_len, + mask_key: Some([0x37, 0xfa, 0x21, 0x3d]), + raw_header: Vec::new(), + } + } + + async fn wait_for_stalled_task(task: &tokio::task::JoinHandle) { + tokio::task::yield_now().await; + assert!(!task.is_finished(), "fixture task must be pending"); + } + + fn relay_options_with_budget(budget: WebSocketAssemblyBudget) -> RelayOptions<'static> { + RelayOptions { + policy_name: "test-policy", + assembly_budget: budget, + resolver: None, + generation_guard: None, + inspector: None, + compression: WebSocketCompression::None, + middleware_session: None, + middleware_context: None, + } + } + fn close_payload(code: u16, reason: &[u8]) -> Vec { let mut payload = Vec::with_capacity(2 + reason.len()); payload.extend_from_slice(&code.to_be_bytes()); @@ -1223,25 +2139,141 @@ network_policies: client_write.write_all(&input).await.unwrap(); drop(client_write); - let options = RelayOptions { + let mut options = RelayOptions { policy_name: "test-policy", + assembly_budget: WebSocketAssemblyBudget::default(), resolver: Some(&resolver), + generation_guard: None, inspector: None, compression: WebSocketCompression::None, + middleware_session: None, + middleware_context: None, }; let result = relay_client_to_server( &mut relay_read, &mut relay_write, "gateway.example.test", 443, - &options, + &mut options, ) .await; drop(relay_write); let mut output = Vec::new(); upstream_read.read_to_end(&mut output).await.unwrap(); - result.map(|()| output) + result + .map(|_| output) + .map_err(|termination| termination.error) + } + + #[tokio::test] + async fn parsed_relays_share_assembly_budget_without_middleware() { + let budget = WebSocketAssemblyBudget::new(1, 0); + let held_frame = masked_frame(true, OPCODE_TEXT, b"held"); + let (mut first_client, mut first_reader) = tokio::io::duplex(64); + let (mut first_writer, _first_upstream) = tokio::io::duplex(64); + first_client + .write_all(&held_frame[..6]) + .await + .expect("send frame header without payload"); + let first_budget = budget.clone(); + let first = tokio::spawn(async move { + let mut options = relay_options_with_budget(first_budget); + relay_client_to_server( + &mut first_reader, + &mut first_writer, + "gateway.example.test", + 443, + &mut options, + ) + .await + }); + wait_for_stalled_task(&first).await; + + let rejected_frame = masked_frame(true, OPCODE_TEXT, b"rejected"); + let (mut second_client, mut second_reader) = tokio::io::duplex(64); + let (mut second_writer, mut second_upstream) = tokio::io::duplex(64); + second_client + .write_all(&rejected_frame) + .await + .expect("send rejected frame"); + drop(second_client); + let mut second_options = relay_options_with_budget(budget.clone()); + let rejected = relay_client_to_server( + &mut second_reader, + &mut second_writer, + "gateway.example.test", + 443, + &mut second_options, + ) + .await + .expect_err("full assembly budget must shed"); + assert_eq!(rejected.cause, WebSocketTerminationCause::CapacityExhausted); + drop(second_writer); + let mut rejected_output = Vec::new(); + second_upstream + .read_to_end(&mut rejected_output) + .await + .expect("read rejected upstream"); + assert!(rejected_output.is_empty()); + + drop(first_client); + first + .await + .expect("join held relay") + .expect_err("incomplete message must end on disconnect"); + + let admitted_frame = masked_frame(true, OPCODE_TEXT, b"admitted"); + let (mut third_client, mut third_reader) = tokio::io::duplex(64); + let (mut third_writer, mut third_upstream) = tokio::io::duplex(64); + third_client + .write_all(&admitted_frame) + .await + .expect("send admitted frame"); + drop(third_client); + let mut third_options = relay_options_with_budget(budget); + relay_client_to_server( + &mut third_reader, + &mut third_writer, + "gateway.example.test", + 443, + &mut third_options, + ) + .await + .expect("released assembly budget must admit"); + drop(third_writer); + let mut admitted_output = Vec::new(); + third_upstream + .read_to_end(&mut admitted_output) + .await + .expect("read admitted upstream"); + assert_eq!(admitted_output, admitted_frame); + } + + #[tokio::test] + async fn assembly_budget_bounds_waiters_and_recovers_capacity() { + let budget = WebSocketAssemblyBudget::new(1, 1); + let first = match budget.reserve().await.expect("reserve active assembly") { + WebSocketAssemblyAdmissionOutcome::Admitted(admission) => admission, + WebSocketAssemblyAdmissionOutcome::QueueExhausted => { + panic!("fresh budget must admit") + } + }; + let waiting_budget = budget.clone(); + let waiting = tokio::spawn(async move { waiting_budget.reserve().await }); + wait_for_stalled_task(&waiting).await; + assert!(matches!( + budget.reserve().await.expect("attempt shed admission"), + WebSocketAssemblyAdmissionOutcome::QueueExhausted + )); + drop(first); + assert!(matches!( + waiting + .await + .expect("join waiting admission") + .expect("waiting admission result"), + WebSocketAssemblyAdmissionOutcome::Admitted(_) + )); } async fn run_client_to_server_with_graphql_policy( @@ -1281,9 +2313,11 @@ network_policies: client_write.write_all(&input).await.unwrap(); drop(client_write); - let options = RelayOptions { + let mut options = RelayOptions { policy_name: "graphql_ws", + assembly_budget: WebSocketAssemblyBudget::default(), resolver, + generation_guard: Some(tunnel_engine.generation_guard()), inspector: Some(InspectionOptions { engine: &tunnel_engine, ctx: &ctx, @@ -1293,20 +2327,24 @@ network_policies: graphql_policy: true, }), compression: WebSocketCompression::None, + middleware_session: None, + middleware_context: None, }; let result = relay_client_to_server( &mut relay_read, &mut relay_write, "realtime.graphql.test", 443, - &options, + &mut options, ) .await; drop(relay_write); let mut output = Vec::new(); upstream_read.read_to_end(&mut output).await.unwrap(); - result.map(|()| output) + result + .map(|_| output) + .map_err(|termination| termination.error) } async fn run_client_to_server_compressed(input: Vec) -> Result> { @@ -1317,25 +2355,31 @@ network_policies: client_write.write_all(&input).await.unwrap(); drop(client_write); - let options = RelayOptions { + let mut options = RelayOptions { policy_name: "test-policy", + assembly_budget: WebSocketAssemblyBudget::default(), resolver: Some(&resolver), + generation_guard: None, inspector: None, compression: WebSocketCompression::PermessageDeflate, + middleware_session: None, + middleware_context: None, }; let result = relay_client_to_server( &mut relay_read, &mut relay_write, "gateway.example.test", 443, - &options, + &mut options, ) .await; drop(relay_write); let mut output = Vec::new(); upstream_read.read_to_end(&mut output).await.unwrap(); - result.map(|()| output) + result + .map(|_| output) + .map_err(|termination| termination.error) } fn decode_masked_text_frame(frame: &[u8]) -> String { @@ -1411,6 +2455,202 @@ network_policies: frame } + #[tokio::test(start_paused = true)] + async fn stalled_initial_text_payload_releases_middleware_admission() { + let runner = openshell_supervisor_middleware::ChainRunner::default(); + let admission = runner + .reserve_middleware_work_admission() + .await + .expect("reserve middleware work"); + let (client, mut reader) = tokio::io::duplex(8); + let frame = test_frame_header(true, OPCODE_TEXT, 4); + let task = tokio::spawn(async move { + let mut assembly = + TextMessageAssembly::new(false, test_assembly_admission().await, Some(admission)); + assembly.read_payload(&mut reader, &frame).await + }); + + wait_for_stalled_task(&task).await; + tokio::time::advance(TEXT_MESSAGE_ASSEMBLY_IDLE_TIMEOUT + StdDuration::from_millis(1)) + .await; + let error = task + .await + .expect("join stalled assembly") + .expect_err("initial payload must time out"); + assert!(error.to_string().contains("assembly idle timeout")); + + runner + .reserve_middleware_work_admission() + .await + .expect("timed-out initial payload releases admission"); + drop(client); + } + + #[tokio::test(start_paused = true)] + async fn stalled_continuation_payload_releases_middleware_admission() { + let runner = openshell_supervisor_middleware::ChainRunner::default(); + let admission = runner + .reserve_middleware_work_admission() + .await + .expect("reserve middleware work"); + let (mut client, mut reader) = tokio::io::duplex(8); + client + .write_all(&[0x37]) + .await + .expect("send partial continuation payload"); + let frame = test_frame_header(true, OPCODE_CONTINUATION, 2); + let task = tokio::spawn(async move { + let mut assembly = + TextMessageAssembly::new(false, test_assembly_admission().await, Some(admission)); + assembly.payload.extend_from_slice(b"first"); + assembly.add_fragment().expect("continuation within limit"); + assembly.read_payload(&mut reader, &frame).await + }); + + wait_for_stalled_task(&task).await; + tokio::time::advance(TEXT_MESSAGE_ASSEMBLY_IDLE_TIMEOUT + StdDuration::from_millis(1)) + .await; + let error = task + .await + .expect("join stalled assembly") + .expect_err("continuation payload must time out"); + assert!(error.to_string().contains("assembly idle timeout")); + + runner + .reserve_middleware_work_admission() + .await + .expect("timed-out continuation releases admission"); + drop(client); + } + + #[tokio::test(start_paused = true)] + async fn missing_continuation_header_releases_middleware_admission() { + let runner = openshell_supervisor_middleware::ChainRunner::default(); + let admission = runner + .reserve_middleware_work_admission() + .await + .expect("reserve middleware work"); + let (client, mut reader) = tokio::io::duplex(8); + let task = tokio::spawn(async move { + let mut assembly = + TextMessageAssembly::new(false, test_assembly_admission().await, Some(admission)); + assembly.payload.extend_from_slice(b"first"); + let result = read_frame_header(&mut reader, Some(&assembly)).await; + drop(assembly); + result + }); + + wait_for_stalled_task(&task).await; + tokio::time::advance(TEXT_MESSAGE_ASSEMBLY_IDLE_TIMEOUT + StdDuration::from_millis(1)) + .await; + let error = task + .await + .expect("join stalled assembly") + .expect_err("missing continuation header must time out"); + assert!(error.to_string().contains("assembly idle timeout")); + + runner + .reserve_middleware_work_admission() + .await + .expect("missing continuation releases admission"); + drop(client); + } + + #[tokio::test(start_paused = true)] + async fn text_assembly_total_deadline_does_not_reset_on_input_progress() { + let runner = openshell_supervisor_middleware::ChainRunner::default(); + let admission = runner + .reserve_middleware_work_admission() + .await + .expect("reserve middleware work"); + let (mut client, mut reader) = tokio::io::duplex(8); + let frame = test_frame_header(true, OPCODE_TEXT, 1_024); + let task = tokio::spawn(async move { + let mut assembly = + TextMessageAssembly::new(false, test_assembly_admission().await, Some(admission)); + assembly.read_payload(&mut reader, &frame).await + }); + wait_for_stalled_task(&task).await; + + let progress_interval = TEXT_MESSAGE_ASSEMBLY_IDLE_TIMEOUT / 2; + let mut elapsed = StdDuration::ZERO; + while elapsed + progress_interval < TEXT_MESSAGE_ASSEMBLY_TOTAL_TIMEOUT { + tokio::time::advance(progress_interval).await; + elapsed += progress_interval; + client + .write_all(&[0x37]) + .await + .expect("make progress before idle timeout"); + tokio::task::yield_now().await; + assert!(!task.is_finished(), "idle deadline must reset on progress"); + } + let until_total_timeout = TEXT_MESSAGE_ASSEMBLY_TOTAL_TIMEOUT + .checked_sub(elapsed) + .expect("progress schedule stays within total timeout"); + tokio::time::advance(until_total_timeout + StdDuration::from_millis(1)).await; + let error = task + .await + .expect("join total assembly timeout") + .expect_err("total assembly deadline must not reset"); + assert!(error.to_string().contains("assembly total timeout")); + + runner + .reserve_middleware_work_admission() + .await + .expect("total timeout releases admission"); + } + + #[tokio::test(start_paused = true)] + async fn stalled_assemblies_release_the_saturated_shared_work_budget() { + let runner = openshell_supervisor_middleware::ChainRunner::default(); + let mut clients = Vec::new(); + let mut stalled = Vec::new(); + for _ in 0..openshell_supervisor_middleware::MAX_CONCURRENT_MIDDLEWARE_WORK { + let admission = runner + .reserve_middleware_work_admission() + .await + .expect("fill shared middleware work budget"); + let (client, mut reader) = tokio::io::duplex(8); + clients.push(client); + let frame = test_frame_header(true, OPCODE_TEXT, 4); + stalled.push(tokio::spawn(async move { + let mut assembly = TextMessageAssembly::new( + false, + test_assembly_admission().await, + Some(admission), + ); + assembly.read_payload(&mut reader, &frame).await + })); + } + tokio::task::yield_now().await; + assert!(stalled.iter().all(|task| !task.is_finished())); + + let later_runner = runner.clone(); + let later_work = + tokio::spawn(async move { later_runner.reserve_middleware_work_admission().await }); + wait_for_stalled_task(&later_work).await; + + tokio::time::advance(TEXT_MESSAGE_ASSEMBLY_IDLE_TIMEOUT + StdDuration::from_millis(1)) + .await; + for task in stalled { + let error = task + .await + .expect("join stalled assembly") + .expect_err("stalled assembly must time out"); + assert!(error.to_string().contains("assembly idle timeout")); + } + + let admission = later_work + .await + .expect("join later middleware work") + .expect("later middleware work acquires released capacity"); + assert!( + admission.saturated(), + "later work must have waited behind the saturated budget" + ); + drop(clients); + } + #[test] fn classifies_graphql_transport_ws_subscribe_operation() { let message = r#"{"type":"subscribe","id":"1","payload":{"query":"subscription NewMessages { messageAdded }"}}"#; @@ -1556,9 +2796,13 @@ network_policies: 443, RelayOptions { policy_name: "test-policy", + assembly_budget: WebSocketAssemblyBudget::default(), resolver: Some(&resolver), + generation_guard: None, inspector: None, compression: WebSocketCompression::None, + middleware_session: None, + middleware_context: None, }, ) .await @@ -1583,6 +2827,1362 @@ network_policies: let _ = tokio::time::timeout(std::time::Duration::from_secs(2), relay).await; } + #[derive(Debug)] + enum ObservedWebSocketRequest { + SessionStart, + Message { sequence: u64, payload: Vec }, + SessionEnd(openshell_core::proto::WebSocketSessionEndReason), + } + + #[derive(Clone, Default)] + struct OpenAiWebSocketRedactor { + observed: Option>, + close_on_first_message: bool, + message_received: Option>, + release_message: Option>, + } + + #[tonic::async_trait] + impl SupervisorMiddleware for OpenAiWebSocketRedactor { + type EvaluateWebSocketSessionStream = + openshell_supervisor_middleware::WebSocketResponseStream; + + async fn describe( + &self, + _request: Request<()>, + ) -> std::result::Result, Status> + { + use openshell_core::proto::{ + MiddlewareBinding, MiddlewareManifest, SupervisorMiddlewareOperation, + SupervisorMiddlewarePhase, + }; + Ok(Response::new(MiddlewareManifest { + name: "test/openai-websocket-redactor".into(), + service_version: "test".into(), + bindings: vec![MiddlewareBinding { + operation: SupervisorMiddlewareOperation::WebsocketMessage as i32, + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + max_payload_bytes: openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES + as u64, + timeout: "1s".into(), + }], + })) + } + + async fn validate_config( + &self, + _request: Request, + ) -> std::result::Result, Status> + { + Ok(Response::new( + openshell_core::proto::ValidateConfigResponse { + valid: true, + reason: String::new(), + }, + )) + } + + async fn evaluate_http_request( + &self, + _request: Request, + ) -> std::result::Result, Status> + { + Err(Status::unimplemented("WebSocket-only test middleware")) + } + + async fn evaluate_web_socket_session( + &self, + request: Request>, + ) -> std::result::Result, Status> { + use openshell_core::proto::{ + Decision, WebSocketMessageResult, WebSocketPreflightAction, + WebSocketPreflightDecision, WebSocketSessionEventResult, web_socket_session_event, + web_socket_session_event_result, + }; + let mut requests = request.into_inner(); + let observed = self.observed.clone(); + let close_on_first_message = self.close_on_first_message; + let message_received = self.message_received.clone(); + let release_message = self.release_message.clone(); + let (responses_tx, responses_rx) = tokio::sync::mpsc::channel(4); + tokio::spawn(async move { + while let Ok(Some(request)) = requests.message().await { + let response = match request.event { + Some(web_socket_session_event::Event::Preflight(_)) => { + Some(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::PreflightDecision( + WebSocketPreflightDecision { + action: WebSocketPreflightAction::Inspect as i32, + ..Default::default() + }, + ), + ), + }) + } + Some(web_socket_session_event::Event::Message(message)) => { + if let Some(observed) = &observed { + let _ = observed.send(ObservedWebSocketRequest::Message { + sequence: message.sequence, + payload: message.payload.clone(), + }); + } + if close_on_first_message { + break; + } + if let Some(received) = &message_received { + received.notify_one(); + } + if let Some(release) = &release_message { + release.notified().await; + } + let text = + String::from_utf8(message.payload).expect("OpenAI event UTF-8"); + let deny = text.contains("deny-me"); + let replacement = + text.replace("customer-secret", "[REDACTED]").into_bytes(); + Some(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::MessageResult( + WebSocketMessageResult { + sequence: message.sequence, + decision: if deny { + Decision::Deny as i32 + } else { + Decision::Allow as i32 + }, + replacement: if deny { + Vec::new() + } else { + replacement + }, + has_replacement: !deny, + reason_code: if deny { + "blocked".into() + } else { + "redacted".into() + }, + ..Default::default() + }, + ), + ), + }) + } + Some(web_socket_session_event::Event::SessionStart(_)) => { + if let Some(observed) = &observed { + let _ = observed.send(ObservedWebSocketRequest::SessionStart); + } + None + } + Some(web_socket_session_event::Event::SessionEnd(end)) => { + if let Some(observed) = &observed + && let Ok(reason) = + openshell_core::proto::WebSocketSessionEndReason::try_from( + end.reason, + ) + { + let _ = observed.send(ObservedWebSocketRequest::SessionEnd(reason)); + } + None + } + _ => None, + }; + if let Some(response) = response + && responses_tx.send(Ok(response)).await.is_err() + { + break; + } + } + }); + Ok(Response::new(Box::pin(ReceiverStream::new(responses_rx)))) + } + } + + async fn recording_middleware_session( + scheme: &str, + ) -> ( + openshell_supervisor_middleware::WebSocketSession, + tokio::sync::mpsc::UnboundedReceiver, + tokio::sync::oneshot::Sender<()>, + tokio::task::JoinHandle>, + ) { + recording_middleware_session_with_controls(scheme, None, None).await + } + + async fn recording_middleware_session_with_controls( + scheme: &str, + message_received: Option>, + release_message: Option>, + ) -> ( + openshell_supervisor_middleware::WebSocketSession, + tokio::sync::mpsc::UnboundedReceiver, + tokio::sync::oneshot::Sender<()>, + tokio::task::JoinHandle>, + ) { + use openshell_core::proto::SupervisorMiddlewareService; + use openshell_supervisor_middleware::{ChainEntry, MiddlewareRegistry, OnError}; + + let (observed_tx, observed_rx) = tokio::sync::mpsc::unbounded_channel(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket middleware"); + let address = listener.local_addr().expect("middleware address"); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tonic::transport::Server::builder() + .add_service(SupervisorMiddlewareServer::new(OpenAiWebSocketRedactor { + observed: Some(observed_tx), + message_received, + release_message, + ..Default::default() + })) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }); + let server_task = tokio::spawn(server); + let registry = MiddlewareRegistry::connect_services( + Vec::new(), + vec![SupervisorMiddlewareService { + name: "openai-redactor".into(), + grpc_endpoint: format!("http://{address}"), + max_payload_bytes: openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES + as u64, + timeout: "2s".into(), + }], + ) + .await + .expect("connect middleware"); + let runner = openshell_supervisor_middleware::ChainRunner::from_registry(registry); + let preflight = runner + .preflight_websocket( + &[ChainEntry { + name: "redact-openai".into(), + implementation: "openai-redactor".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailClosed, + }], + openshell_supervisor_middleware::WebSocketPreflightInput { + session_id: "session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: scheme.into(), + host: "api.openai.com".into(), + port: if scheme == "wss" { 443 } else { 80 }, + path: "/v1/responses".into(), + requested_subprotocols: Vec::new(), + }, + ) + .await + .expect("preflight"); + + ( + preflight.session.expect("middleware inspects session"), + observed_rx, + shutdown_tx, + server_task, + ) + } + + async fn assert_invalid_close_termination(payload: Vec, expected_close_code: u16) { + let (mut session, mut observed, shutdown_tx, server_task) = + recording_middleware_session("wss").await; + assert!(session.start("").await.allowed); + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionStart) + )); + + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let relay = tokio::spawn(async move { + relay_with_options( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + 443, + RelayOptions { + policy_name: "rest-api", + assembly_budget: WebSocketAssemblyBudget::default(), + resolver: None, + generation_guard: None, + inspector: None, + compression: WebSocketCompression::None, + middleware_session: Some(session), + middleware_context: None, + }, + ) + .await + }); + + client_app + .write_all(&masked_frame(true, OPCODE_CLOSE, &payload)) + .await + .expect("send invalid close frame"); + client_app.flush().await.expect("flush invalid close frame"); + + let upstream_close = read_one_frame(&mut upstream_app).await; + assert_eq!(upstream_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes( + decode_masked_payload(&upstream_close)[..2] + .try_into() + .expect("upstream close code"), + ), + expected_close_code + ); + let client_close = read_one_frame(&mut client_app).await; + assert_eq!(client_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes(client_close[2..4].try_into().expect("client close code")), + expected_close_code + ); + + relay + .await + .expect("join relay") + .expect_err("invalid close frame must terminate relay"); + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionEnd( + openshell_core::proto::WebSocketSessionEndReason::ProtocolError, + )) + )); + assert!( + observed.try_recv().is_err(), + "control-frame failure must not invoke middleware" + ); + + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + #[tokio::test] + async fn invalid_close_frames_close_both_peers_and_report_protocol_session_end() { + assert_invalid_close_termination(vec![0x03], 1002).await; + assert_invalid_close_termination(close_payload(1005, b""), 1002).await; + assert_invalid_close_termination(close_payload(1000, &[0xff]), 1007).await; + } + + #[test] + fn invalid_close_frames_have_stable_telemetry_classes() { + let cases = [ + ( + vec![0x03], + WebSocketTerminationCause::ProtocolError, + FrameFailureClass::InvalidCloseFrame, + 1002, + ), + ( + close_payload(1005, b""), + WebSocketTerminationCause::ProtocolError, + FrameFailureClass::InvalidCloseFrame, + 1002, + ), + ( + close_payload(1000, &[0xff]), + WebSocketTerminationCause::InvalidUtf8, + FrameFailureClass::InvalidUtf8, + 1007, + ), + ]; + + for (payload, expected_cause, expected_class, expected_close_code) in cases { + let frame_error = + validate_close_payload(&payload).expect_err("close payload must be invalid"); + let termination = WebSocketTermination::from(frame_error); + assert_eq!(termination.cause, expected_cause); + assert_eq!(termination.failure_class, Some(expected_class)); + assert_eq!(termination.cause.close_code(), Some(expected_close_code)); + + let event = + protocol_failure_event("gateway.example.test", 443, "test-policy", expected_class); + assert_eq!( + event.base().status_detail.as_deref(), + Some(expected_class.as_str()) + ); + } + } + + #[tokio::test] + async fn partial_control_payload_eof_is_peer_disconnect_without_generated_close() { + let (mut session, mut observed, shutdown_tx, server_task) = + recording_middleware_session("wss").await; + assert!(session.start("").await.allowed); + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionStart) + )); + + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let relay = tokio::spawn(async move { + relay_with_options( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + 443, + RelayOptions { + policy_name: "rest-api", + assembly_budget: WebSocketAssemblyBudget::default(), + resolver: None, + generation_guard: None, + inspector: None, + compression: WebSocketCompression::None, + middleware_session: Some(session), + middleware_context: None, + }, + ) + .await + }); + + client_app + .write_all(&[0x80 | OPCODE_CLOSE, 0x80 | 2, 0x37, 0xfa, 0x21, 0x3d]) + .await + .expect("send truncated close frame"); + drop(client_app); + + let mut upstream_output = Vec::new(); + upstream_app + .read_to_end(&mut upstream_output) + .await + .expect("read upstream shutdown"); + assert!( + upstream_output.is_empty(), + "peer EOF must not generate an upstream close frame" + ); + relay + .await + .expect("join relay") + .expect_err("truncated payload must end relay"); + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionEnd( + openshell_core::proto::WebSocketSessionEndReason::PeerDisconnect, + )) + )); + + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + #[tokio::test] + async fn denied_websocket_session_start_reports_middleware_failure_before_close() { + let (session, mut observed, shutdown_tx, server_task) = + recording_middleware_session("wss").await; + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + + crate::l7::relay::handle_upgrade( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + 443, + crate::l7::relay::UpgradeRelayOptions { + websocket_request: true, + assembly_budget: Some(WebSocketAssemblyBudget::default()), + middleware_session: Some(session), + selected_subprotocol: Some("x".repeat(257)), + ..Default::default() + }, + ) + .await + .expect("denied session start closes cleanly"); + + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionEnd( + openshell_core::proto::WebSocketSessionEndReason::MiddlewareFailure, + )) + )); + assert!( + observed.try_recv().is_err(), + "session start failure must send exactly one terminal event" + ); + + let client_close = read_one_frame(&mut client_app).await; + assert_eq!(client_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes(client_close[2..4].try_into().expect("client close code")), + 1008 + ); + let upstream_close = read_one_frame(&mut upstream_app).await; + assert_eq!(upstream_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes( + decode_masked_payload(&upstream_close)[..2] + .try_into() + .expect("upstream close code"), + ), + 1008 + ); + + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + #[tokio::test] + async fn selected_middleware_reports_binary_gap_and_inspects_later_text_at_next_sequence() { + let (session, mut observed, shutdown_tx, server_task) = + recording_middleware_session("wss").await; + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let relay = tokio::spawn(async move { + crate::l7::relay::handle_upgrade( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + 443, + crate::l7::relay::UpgradeRelayOptions { + websocket_request: true, + assembly_budget: Some(WebSocketAssemblyBudget::default()), + ctx: Some(&L7EvalContext { + host: "api.openai.com".into(), + port: 443, + policy_name: "rest-api".into(), + ..Default::default() + }), + policy_name: "rest-api".into(), + middleware_session: Some(session), + ..Default::default() + }, + ) + .await + }); + + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionStart) + )); + + let binary = masked_frame(true, OPCODE_BINARY, &[0, 1, 2, 3, 255]); + let text = masked_frame(true, OPCODE_TEXT, br#"{"type":"response.create"}"#); + client_app + .write_all(&binary) + .await + .expect("send binary message"); + client_app + .write_all(&text) + .await + .expect("send text message"); + + assert_eq!(read_one_frame(&mut upstream_app).await, binary); + assert_eq!( + decode_masked_text_frame(&read_one_frame(&mut upstream_app).await), + r#"{"type":"response.create"}"# + ); + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::Message { + sequence: 2, + payload, + }) if payload == br#"{"type":"response.create"}"# + )); + + drop(client_app); + drop(upstream_app); + tokio::time::timeout(std::time::Duration::from_secs(2), relay) + .await + .expect("relay should stop") + .expect("join relay") + .expect("relay result"); + + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + async fn assert_middleware_only_relay_observes_reload(scheme: &str, fragmented: bool) { + use openshell_supervisor_middleware::MiddlewareRegistry; + + let (session, mut observed, shutdown_tx, server_task) = + recording_middleware_session(scheme).await; + let engine = Arc::new( + OpaEngine::from_strings(TEST_POLICY, "network_policies: {}\n").expect("test policy"), + ); + let generation_guard = engine + .generation_guard(engine.current_generation()) + .expect("generation guard"); + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let port = if scheme == "wss" { 443 } else { 80 }; + let relay = tokio::spawn(async move { + crate::l7::relay::handle_upgrade( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + port, + crate::l7::relay::UpgradeRelayOptions { + websocket_request: true, + assembly_budget: Some(WebSocketAssemblyBudget::default()), + generation_guard: Some(&generation_guard), + ctx: Some(&L7EvalContext { + host: "api.openai.com".into(), + port, + policy_name: "rest-api".into(), + ..Default::default() + }), + policy_name: "rest-api".into(), + middleware_session: Some(session), + ..Default::default() + }, + ) + .await + }); + + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionStart) + )); + let stale_frame = if fragmented { + masked_frame(true, OPCODE_CONTINUATION, b"message") + } else { + masked_frame(true, OPCODE_TEXT, b"stale-message") + }; + if fragmented { + client_app + .write_all(&masked_frame(false, OPCODE_TEXT, b"stale-")) + .await + .expect("send initial fragment"); + } else { + client_app + .write_all(&stale_frame[..7]) + .await + .expect("send frame header and partial payload"); + } + + engine + .replace_middleware_registry(MiddlewareRegistry::default()) + .expect("invalidate generation"); + let stale_input = if fragmented { + &stale_frame[..6] + } else { + &stale_frame[7..] + }; + client_app + .write_all(stale_input) + .await + .expect("send stale data frame"); + + let upstream_close = read_one_frame(&mut upstream_app).await; + assert_eq!(upstream_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes( + decode_masked_payload(&upstream_close)[..2] + .try_into() + .expect("upstream close code"), + ), + 1012 + ); + let client_close = read_one_frame(&mut client_app).await; + assert_eq!(client_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes(client_close[2..4].try_into().expect("client close code")), + 1012 + ); + let error = relay + .await + .expect("join relay") + .expect_err("stale generation must terminate relay"); + assert!(error.to_string().contains("policy generation is stale")); + match observed.recv().await { + Some(ObservedWebSocketRequest::SessionEnd( + openshell_core::proto::WebSocketSessionEndReason::PolicyReload, + )) => {} + Some(ObservedWebSocketRequest::Message { payload, .. }) => { + panic!("stale message leaked {} bytes to middleware", payload.len()); + } + other => panic!("unexpected middleware lifecycle event: {other:?}"), + } + assert!( + observed.try_recv().is_err(), + "stale data must not reach middleware" + ); + + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + #[tokio::test] + async fn middleware_only_rest_wss_relay_observes_policy_reload() { + assert_middleware_only_relay_observes_reload("wss", false).await; + } + + #[tokio::test] + async fn middleware_only_rest_ws_relay_stops_fragment_after_policy_reload() { + assert_middleware_only_relay_observes_reload("ws", true).await; + } + + #[tokio::test] + async fn reload_during_middleware_evaluation_blocks_credentials_and_payload() { + use openshell_supervisor_middleware::MiddlewareRegistry; + + let message_received = Arc::new(tokio::sync::Notify::new()); + let release_message = Arc::new(tokio::sync::Notify::new()); + let (session, mut observed, shutdown_tx, server_task) = + recording_middleware_session_with_controls( + "wss", + Some(Arc::clone(&message_received)), + Some(Arc::clone(&release_message)), + ) + .await; + let engine = Arc::new( + OpaEngine::from_strings(TEST_POLICY, "network_policies: {}\n").expect("test policy"), + ); + let generation_guard = engine + .generation_guard(engine.current_generation()) + .expect("generation guard"); + let assembly_budget = engine.websocket_assembly_budget(); + let (child_env, resolver) = resolver(); + let placeholder = child_env + .get("DISCORD_BOT_TOKEN") + .expect("credential placeholder") + .clone(); + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let relay = tokio::spawn(async move { + crate::l7::relay::handle_upgrade( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + 443, + crate::l7::relay::UpgradeRelayOptions { + websocket_request: true, + websocket: crate::l7::relay::WebSocketUpgradeBehavior { + credential_rewrite: true, + ..Default::default() + }, + assembly_budget: Some(assembly_budget), + secret_resolver: Some(Arc::new(resolver)), + generation_guard: Some(&generation_guard), + policy_name: "rest-api".into(), + middleware_session: Some(session), + ..Default::default() + }, + ) + .await + }); + + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionStart) + )); + let payload = format!(r#"{{"authorization":"{placeholder}"}}"#); + client_app + .write_all(&masked_frame(true, OPCODE_TEXT, payload.as_bytes())) + .await + .expect("send credential-bearing message"); + message_received.notified().await; + + engine + .replace_middleware_registry(MiddlewareRegistry::default()) + .expect("invalidate generation"); + release_message.notify_one(); + + let upstream_close = read_one_frame(&mut upstream_app).await; + assert_eq!(upstream_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes( + decode_masked_payload(&upstream_close)[..2] + .try_into() + .expect("upstream close code"), + ), + 1012 + ); + assert!( + !String::from_utf8_lossy(&upstream_close).contains("real-token"), + "resolved credential must not reach upstream" + ); + let error = relay + .await + .expect("join relay") + .expect_err("stale generation must terminate relay"); + assert!(error.to_string().contains("policy generation is stale")); + + let mut end_reason = None; + while let Some(event) = observed.recv().await { + if let ObservedWebSocketRequest::SessionEnd(reason) = event { + end_reason = Some(reason); + break; + } + } + assert_eq!( + end_reason, + Some(openshell_core::proto::WebSocketSessionEndReason::PolicyReload) + ); + + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + #[tokio::test] + async fn pre_upgrade_reload_finalizes_session_as_policy_reload() { + use openshell_supervisor_middleware::MiddlewareRegistry; + + let (session, mut observed, shutdown_tx, server_task) = + recording_middleware_session("wss").await; + let engine = + OpaEngine::from_strings(TEST_POLICY, "network_policies: {}\n").expect("test policy"); + let generation_guard = engine + .generation_guard(engine.current_generation()) + .expect("generation guard"); + engine + .replace_middleware_registry(MiddlewareRegistry::default()) + .expect("invalidate generation"); + let mut session = Some(session); + + crate::l7::relay::finalize_websocket_pre_upgrade( + &mut session, + &generation_guard, + "api.openai.com", + 443, + "rest-api", + Ok(()), + ) + .await + .expect_err("reload before upgrade must terminate"); + assert!(session.is_none()); + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionEnd( + openshell_core::proto::WebSocketSessionEndReason::PolicyReload + )) + )); + + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + #[tokio::test] + async fn graphql_policy_denial_reports_policy_denial_to_middleware() { + let (mut session, mut observed, shutdown_tx, server_task) = + recording_middleware_session("wss").await; + assert!(session.start("").await.allowed); + assert!(matches!( + observed.recv().await, + Some(ObservedWebSocketRequest::SessionStart) + )); + + let engine = OpaEngine::from_strings(TEST_POLICY, GRAPHQL_WS_POLICY) + .expect("GraphQL WebSocket policy should load"); + let network_input = NetworkInput { + host: "realtime.graphql.test".into(), + port: 443, + binary_path: PathBuf::from("/usr/bin/node"), + binary_sha256: "unused".into(), + ancestors: vec![], + cmdline_paths: vec![], + }; + let generation = engine + .evaluate_network_action_with_generation(&network_input) + .expect("network action should evaluate") + .1; + let tunnel_engine = engine + .clone_engine_for_tunnel(generation) + .expect("tunnel engine"); + let ctx = L7EvalContext { + host: "realtime.graphql.test".into(), + port: 443, + policy_name: "graphql_ws".into(), + binary_path: "/usr/bin/node".into(), + ..Default::default() + }; + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let relay = tokio::spawn(async move { + relay_with_options( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "realtime.graphql.test", + 443, + RelayOptions { + policy_name: "graphql_ws", + assembly_budget: WebSocketAssemblyBudget::default(), + resolver: None, + generation_guard: Some(tunnel_engine.generation_guard()), + inspector: Some(InspectionOptions { + engine: &tunnel_engine, + ctx: &ctx, + enforcement: EnforcementMode::Enforce, + target: "/graphql".into(), + query_params: HashMap::new(), + graphql_policy: true, + }), + compression: WebSocketCompression::None, + middleware_session: Some(session), + middleware_context: Some(&ctx), + }, + ) + .await + }); + + let payload = + br#"{"type":"subscribe","id":"1","payload":{"query":"query Admin { adminAuditLog }"}}"#; + client_app + .write_all(&masked_frame(true, OPCODE_TEXT, payload)) + .await + .expect("send policy-denied message"); + + let upstream_close = read_one_frame(&mut upstream_app).await; + assert_eq!( + u16::from_be_bytes( + decode_masked_payload(&upstream_close)[..2] + .try_into() + .expect("upstream close code"), + ), + 1008 + ); + let client_close = read_one_frame(&mut client_app).await; + assert_eq!( + u16::from_be_bytes(client_close[2..4].try_into().expect("client close code")), + 1008 + ); + let error = relay + .await + .expect("join relay") + .expect_err("policy denial must terminate relay"); + assert!( + error + .to_string() + .contains("websocket GraphQL message denied") + ); + match observed.recv().await { + Some(ObservedWebSocketRequest::SessionEnd( + openshell_core::proto::WebSocketSessionEndReason::PolicyDenial, + )) => {} + Some(ObservedWebSocketRequest::Message { payload, .. }) => { + panic!( + "built-in policy denial leaked {} bytes to middleware", + payload.len() + ); + } + other => panic!("unexpected middleware lifecycle event: {other:?}"), + } + + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware server") + .expect("middleware server"); + } + + #[tokio::test] + async fn parsed_relay_sends_redacted_openai_event_to_upstream() { + use openshell_core::proto::SupervisorMiddlewareService; + use openshell_supervisor_middleware::{ChainEntry, MiddlewareRegistry, OnError}; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket middleware"); + let address = listener.local_addr().expect("middleware address"); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tonic::transport::Server::builder() + .add_service(SupervisorMiddlewareServer::new( + OpenAiWebSocketRedactor::default(), + )) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }); + let server_task = tokio::spawn(server); + let registry = MiddlewareRegistry::connect_services( + Vec::new(), + vec![SupervisorMiddlewareService { + name: "openai-redactor".into(), + grpc_endpoint: format!("http://{address}"), + max_payload_bytes: openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES + as u64, + timeout: "2s".into(), + }], + ) + .await + .expect("connect middleware"); + let runner = openshell_supervisor_middleware::ChainRunner::from_registry(registry); + let preflight = runner + .preflight_websocket( + &[ChainEntry { + name: "redact-openai".into(), + implementation: "openai-redactor".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailClosed, + }], + openshell_supervisor_middleware::WebSocketPreflightInput { + session_id: "session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: Vec::new(), + }, + ) + .await + .expect("preflight"); + let mut session = preflight.session.expect("middleware inspects session"); + assert!(session.start("").await.allowed); + + let original = br#"{"type":"response.create","response":{"input":"customer-secret"}}"#; + let client_frame = masked_frame(true, OPCODE_TEXT, original); + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let relay = tokio::spawn(async move { + relay_with_options( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + 443, + RelayOptions { + policy_name: "openai", + assembly_budget: WebSocketAssemblyBudget::default(), + resolver: None, + generation_guard: None, + inspector: None, + compression: WebSocketCompression::None, + middleware_session: Some(session), + middleware_context: None, + }, + ) + .await + }); + + client_app + .write_all(&client_frame) + .await + .expect("send event"); + client_app.flush().await.expect("flush event"); + let upstream_frame = tokio::time::timeout( + std::time::Duration::from_secs(2), + read_one_frame(&mut upstream_app), + ) + .await + .expect("upstream receives event"); + let upstream_text = decode_masked_text_frame(&upstream_frame); + assert!(upstream_text.contains("[REDACTED]")); + assert!(!upstream_text.contains("customer-secret")); + + let denied = br#"{"type":"response.create","response":{"input":"deny-me"}}"#; + client_app + .write_all(&masked_frame(true, OPCODE_TEXT, denied)) + .await + .expect("send denied event"); + client_app.flush().await.expect("flush denied event"); + let upstream_close = tokio::time::timeout( + std::time::Duration::from_secs(2), + read_one_frame(&mut upstream_app), + ) + .await + .expect("upstream receives close"); + assert_eq!(upstream_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes( + decode_masked_payload(&upstream_close)[..2] + .try_into() + .expect("close code"), + ), + 1008 + ); + let client_close = tokio::time::timeout( + std::time::Duration::from_secs(2), + read_one_frame(&mut client_app), + ) + .await + .expect("client receives close"); + assert_eq!(client_close[0] & 0x0f, OPCODE_CLOSE); + assert_eq!( + u16::from_be_bytes(client_close[2..4].try_into().expect("close code")), + 1008 + ); + + drop(client_app); + drop(upstream_app); + let _ = tokio::time::timeout(std::time::Duration::from_secs(2), relay).await; + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware") + .expect("serve middleware"); + } + + #[tokio::test] + async fn fully_disabled_session_bypasses_saturated_work_budget_without_rpc() { + use openshell_core::proto::SupervisorMiddlewareService; + use openshell_supervisor_middleware::{ChainEntry, MiddlewareRegistry, OnError}; + + let (observed_tx, mut observed_rx) = tokio::sync::mpsc::unbounded_channel(); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind WebSocket middleware"); + let address = listener.local_addr().expect("middleware address"); + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel(); + let server = tonic::transport::Server::builder() + .add_service(SupervisorMiddlewareServer::new(OpenAiWebSocketRedactor { + observed: Some(observed_tx), + close_on_first_message: true, + ..Default::default() + })) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }); + let server_task = tokio::spawn(server); + let registry = MiddlewareRegistry::connect_services( + Vec::new(), + vec![SupervisorMiddlewareService { + name: "openai-redactor".into(), + grpc_endpoint: format!("http://{address}"), + max_payload_bytes: openshell_supervisor_middleware::MAX_MIDDLEWARE_PAYLOAD_BYTES + as u64, + timeout: "2s".into(), + }], + ) + .await + .expect("connect middleware"); + let runner = openshell_supervisor_middleware::ChainRunner::from_registry(registry); + let preflight = runner + .preflight_websocket( + &[ChainEntry { + name: "redact-openai".into(), + implementation: "openai-redactor".into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailOpen, + }], + openshell_supervisor_middleware::WebSocketPreflightInput { + session_id: "disabled-session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: Vec::new(), + }, + ) + .await + .expect("preflight"); + let mut session = preflight.session.expect("middleware inspects session"); + assert!(session.start("").await.allowed); + assert!(matches!( + observed_rx.recv().await, + Some(ObservedWebSocketRequest::SessionStart) + )); + + let failed = session + .evaluate_text(br#"{"type":"response.create"}"#.to_vec()) + .await; + assert!(failed.allowed); + assert!(failed.invocations[0].stage_disabled); + assert!(matches!( + observed_rx.recv().await, + Some(ObservedWebSocketRequest::Message { .. }) + )); + + let mut occupied_work = Vec::new(); + for _ in 0..openshell_supervisor_middleware::MAX_CONCURRENT_MIDDLEWARE_WORK { + occupied_work.push( + runner + .reserve_middleware_work_admission() + .await + .expect("fill middleware work budget"), + ); + } + + let original = r#"{"type":"response.cancel","reason":"keep-original"}"#; + let client_frame = masked_frame(true, OPCODE_TEXT, original.as_bytes()); + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let relay = tokio::spawn(async move { + relay_with_options( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + 443, + RelayOptions { + policy_name: "openai", + assembly_budget: WebSocketAssemblyBudget::default(), + resolver: None, + generation_guard: None, + inspector: None, + compression: WebSocketCompression::None, + middleware_session: Some(session), + middleware_context: None, + }, + ) + .await + }); + + client_app + .write_all(&client_frame) + .await + .expect("send bypassed event"); + client_app.flush().await.expect("flush bypassed event"); + let upstream_frame = tokio::time::timeout( + std::time::Duration::from_secs(2), + read_one_frame(&mut upstream_app), + ) + .await + .expect("fully disabled session bypasses saturated work"); + assert_eq!(decode_masked_text_frame(&upstream_frame), original); + assert!( + observed_rx.try_recv().is_err(), + "fully disabled session must not make another middleware RPC" + ); + + drop(occupied_work); + drop(client_app); + drop(upstream_app); + tokio::time::timeout(std::time::Duration::from_secs(2), relay) + .await + .expect("relay finishes after disconnect") + .expect("join relay") + .expect("relay"); + let _ = shutdown_tx.send(()); + server_task + .await + .expect("join middleware") + .expect("serve middleware"); + } + + #[tokio::test] + async fn parsed_relay_uses_builtin_regex_for_complete_client_text_messages() { + use openshell_supervisor_middleware::{ChainEntry, MiddlewareRegistry, OnError}; + + let registry = MiddlewareRegistry::connect_services( + openshell_supervisor_middleware_builtins::services(), + Vec::new(), + ) + .await + .expect("connect built-in middleware"); + let runner = openshell_supervisor_middleware::ChainRunner::from_registry(registry); + let preflight = runner + .preflight_websocket( + &[ChainEntry { + name: "regex-redactor".into(), + implementation: openshell_supervisor_middleware_builtins::BUILTIN_REGEX.into(), + order: 0, + config: prost_types::Struct::default(), + on_error: OnError::FailClosed, + }], + openshell_supervisor_middleware::WebSocketPreflightInput { + session_id: "builtin-regex-session".into(), + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: Vec::new(), + }, + ) + .await + .expect("preflight"); + let mut session = preflight.session.expect("built-in inspects session"); + assert!(session.start("").await.allowed); + + let original = br#"{"type":"response.create","response":{"input":"sk-ABCDEFGHIJKLMNOP"}}"#; + let (mut client_app, mut relay_client) = tokio::io::duplex(4096); + let (mut relay_upstream, mut upstream_app) = tokio::io::duplex(4096); + let relay = tokio::spawn(async move { + relay_with_options( + &mut relay_client, + &mut relay_upstream, + Vec::new(), + "api.openai.com", + 443, + RelayOptions { + policy_name: "openai", + assembly_budget: WebSocketAssemblyBudget::default(), + resolver: None, + generation_guard: None, + inspector: None, + compression: WebSocketCompression::None, + middleware_session: Some(session), + middleware_context: None, + }, + ) + .await + }); + + let split = original + .windows(b"sk-ABC".len()) + .position(|window| window == b"sk-ABC") + .expect("token prefix") + + b"sk-ABC".len(); + client_app + .write_all(&masked_frame(false, OPCODE_TEXT, &original[..split])) + .await + .expect("send initial event fragment"); + client_app + .write_all(&masked_frame(true, OPCODE_CONTINUATION, &original[split..])) + .await + .expect("send final event fragment"); + let upstream_frame = tokio::time::timeout( + std::time::Duration::from_secs(2), + read_one_frame(&mut upstream_app), + ) + .await + .expect("upstream receives event"); + let upstream_text = decode_masked_text_frame(&upstream_frame); + assert_eq!( + upstream_text, + r#"{"type":"response.create","response":{"input":"[REDACTED]"}}"# + ); + + drop(client_app); + drop(upstream_app); + tokio::time::timeout(std::time::Duration::from_secs(2), relay) + .await + .expect("relay finishes after disconnect") + .expect("join relay") + .expect("relay"); + } + + #[tokio::test] + #[ignore = "PR 2 adds return-path inspection and completes this full-duplex fixture"] + async fn pr2_full_duplex_external_middleware_vertical_slice() { + // PR 2 should extend the real relay fixture above with controllable + // server-to-client transforms plus slow, hanging, closed, duplicate, + // missing, out-of-order, and oversized middleware responses. + let deferred_faults = [ + "slow", + "hanging", + "closed", + "duplicate", + "missing", + "out-of-order", + "oversized", + ]; + assert_eq!(deferred_faults.len(), 7); + } + #[tokio::test] async fn graphql_websocket_policy_allows_subscription_operation() { let payload = r#"{"type":"subscribe","id":"1","payload":{"query":"subscription NewMessages { messageAdded }"}}"#; diff --git a/crates/openshell-supervisor-network/src/opa.rs b/crates/openshell-supervisor-network/src/opa.rs index f0654c287d..8c3425af95 100644 --- a/crates/openshell-supervisor-network/src/opa.rs +++ b/crates/openshell-supervisor-network/src/opa.rs @@ -123,6 +123,7 @@ pub struct OpaEngine { engine: Mutex, generation: Arc, middleware_runner: RwLock, + websocket_assembly_budget: crate::l7::websocket::WebSocketAssemblyBudget, } /// Generation guard captured when an HTTP tunnel or request path starts. @@ -163,6 +164,7 @@ pub struct TunnelPolicyEngine { engine: Mutex, generation_guard: PolicyGenerationGuard, middleware_runner: ChainRunner, + websocket_assembly_budget: crate::l7::websocket::WebSocketAssemblyBudget, } impl TunnelPolicyEngine { @@ -190,6 +192,12 @@ impl TunnelPolicyEngine { &self.middleware_runner } + pub(crate) fn websocket_assembly_budget( + &self, + ) -> crate::l7::websocket::WebSocketAssemblyBudget { + self.websocket_assembly_budget.clone() + } + /// Query the ordered middleware chain for a destination within this tunnel. pub fn query_middleware_chain(&self, input: &NetworkInput) -> Result> { let mut engine = self @@ -201,6 +209,12 @@ impl TunnelPolicyEngine { } impl OpaEngine { + pub(crate) fn websocket_assembly_budget( + &self, + ) -> crate::l7::websocket::WebSocketAssemblyBudget { + self.websocket_assembly_budget.clone() + } + /// Load policy from a `.rego` rules file and data from a YAML file. /// /// Preprocesses the YAML data to expand access presets and validate L7 config. @@ -236,6 +250,7 @@ impl OpaEngine { engine: Mutex::new(engine), generation: Arc::new(AtomicU64::new(0)), middleware_runner: RwLock::new(ChainRunner::default()), + websocket_assembly_budget: crate::l7::websocket::WebSocketAssemblyBudget::default(), }) } @@ -291,6 +306,7 @@ impl OpaEngine { engine: Mutex::new(engine), generation: Arc::new(AtomicU64::new(0)), middleware_runner: RwLock::new(ChainRunner::default()), + websocket_assembly_budget: crate::l7::websocket::WebSocketAssemblyBudget::default(), }) } @@ -370,6 +386,7 @@ impl OpaEngine { engine: Mutex::new(engine), generation: Arc::new(AtomicU64::new(0)), middleware_runner: RwLock::new(ChainRunner::default()), + websocket_assembly_budget: crate::l7::websocket::WebSocketAssemblyBudget::default(), }) } @@ -539,8 +556,6 @@ impl OpaEngine { .engine .into_inner() .map_err(|_| miette::miette!("lock poisoned on new engine"))?; - let new_runner = ChainRunner::from_registry(registry); - // Match clone_engine_for_tunnel's lock order (engine, then runner) so // readers can observe only the old pair or the new pair. let mut engine = self @@ -551,6 +566,7 @@ impl OpaEngine { .middleware_runner .write() .map_err(|_| miette::miette!("middleware runner lock poisoned"))?; + let new_runner = runner.with_replacement_registry(registry); *engine = new_engine; *runner = new_runner; self.generation.fetch_add(1, Ordering::AcqRel); @@ -569,7 +585,7 @@ impl OpaEngine { .middleware_runner .write() .map_err(|_| miette::miette!("middleware runner lock poisoned"))?; - *runner = ChainRunner::from_registry(registry); + *runner = runner.with_replacement_registry(registry); self.generation.fetch_add(1, Ordering::AcqRel); Ok(()) } @@ -764,6 +780,7 @@ impl OpaEngine { current_generation: Arc::clone(&self.generation), }, middleware_runner: self.middleware_runner()?, + websocket_assembly_budget: self.websocket_assembly_budget(), }) } } @@ -3047,6 +3064,7 @@ network_policies: engine: Mutex::new(rego), generation: Arc::new(AtomicU64::new(0)), middleware_runner: RwLock::new(ChainRunner::default()), + websocket_assembly_budget: crate::l7::websocket::WebSocketAssemblyBudget::default(), }; let input = l7_websocket_graphql_input( "realtime.graphql.com", @@ -6809,6 +6827,142 @@ network_policies: assert!(described[0].is_resolved()); } + #[tokio::test] + async fn middleware_session_budget_survives_atomic_registry_reload() { + let proto = test_proto(); + let engine = OpaEngine::from_proto(&proto).expect("initial load should succeed"); + let registry = MiddlewareRegistry::connect_services( + openshell_supervisor_middleware_builtins::services(), + Vec::new(), + ) + .await + .expect("built-in registry"); + engine + .replace_middleware_registry(registry) + .expect("install registry"); + let old_runner = engine.middleware_runner().expect("old runner"); + let entry = ChainEntry { + name: "regex".into(), + implementation: openshell_supervisor_middleware_builtins::BUILTIN_REGEX.into(), + order: 0, + config: prost_types::Struct { + fields: std::iter::once(( + "mode".into(), + prost_types::Value { + kind: Some(prost_types::value::Kind::StringValue("redact".into())), + }, + )) + .collect(), + }, + on_error: openshell_supervisor_middleware::OnError::FailClosed, + }; + let preflight_input = + |session_id: String| openshell_supervisor_middleware::WebSocketPreflightInput { + session_id, + request_id: "request".into(), + sandbox_id: "sandbox".into(), + scheme: "wss".into(), + host: "api.openai.com".into(), + port: 443, + path: "/v1/responses".into(), + requested_subprotocols: Vec::new(), + }; + let mut old_sessions = Vec::new(); + for index in 0..openshell_supervisor_middleware::MAX_CONCURRENT_MIDDLEWARE_SESSIONS { + let preflight = old_runner + .preflight_websocket( + std::slice::from_ref(&entry), + preflight_input(format!("old-generation-{index}")), + ) + .await + .expect("admit old-generation session"); + old_sessions.push(preflight.session.expect("built-in inspects session")); + } + + let replacement = MiddlewareRegistry::connect_services( + openshell_supervisor_middleware_builtins::services(), + Vec::new(), + ) + .await + .expect("replacement registry"); + engine + .reload_policy_and_middleware_from_proto_with_pid(&proto, 0, replacement) + .expect("atomic policy and registry reload"); + let current_runner = engine.middleware_runner().expect("current runner"); + let overflow = current_runner + .preflight_websocket( + std::slice::from_ref(&entry), + preflight_input("new-generation-overflow".into()), + ) + .await + .expect("capacity exhaustion is a typed preflight outcome"); + assert!(!overflow.allowed); + assert!(overflow.session_capacity_exhausted); + + old_sessions + .pop() + .expect("old-generation session") + .end(openshell_core::proto::WebSocketSessionEndReason::PolicyReload) + .await; + let admitted = current_runner + .preflight_websocket( + std::slice::from_ref(&entry), + preflight_input("new-generation-admitted".into()), + ) + .await + .expect("released capacity is reusable after reload"); + assert!(admitted.allowed); + assert!(admitted.session.is_some()); + } + + #[tokio::test] + async fn websocket_assembly_budget_survives_policy_reload() { + use crate::l7::websocket::{ + MAX_CONCURRENT_WEBSOCKET_ASSEMBLIES, WebSocketAssemblyAdmissionOutcome, + }; + + let proto = test_proto(); + let engine = OpaEngine::from_proto(&proto).expect("initial policy"); + let old_tunnel = engine + .clone_engine_for_tunnel(engine.current_generation()) + .expect("old tunnel"); + let old_budget = old_tunnel.websocket_assembly_budget(); + let mut old_admissions = Vec::new(); + for _ in 0..MAX_CONCURRENT_WEBSOCKET_ASSEMBLIES { + match old_budget.reserve().await.expect("reserve old assembly") { + WebSocketAssemblyAdmissionOutcome::Admitted(admission) => { + old_admissions.push(admission); + } + WebSocketAssemblyAdmissionOutcome::QueueExhausted => { + panic!("active assembly capacity exhausted too early") + } + } + } + + engine + .reload_from_proto_with_pid(&proto, 0) + .expect("policy reload"); + let new_tunnel = engine + .clone_engine_for_tunnel(engine.current_generation()) + .expect("new tunnel"); + let new_budget = new_tunnel.websocket_assembly_budget(); + let waiting = tokio::spawn(async move { new_budget.reserve().await }); + tokio::task::yield_now().await; + assert!( + !waiting.is_finished(), + "new generation must share old generation assembly capacity" + ); + + drop(old_admissions.pop()); + assert!(matches!( + waiting + .await + .expect("join waiting assembly") + .expect("waiting assembly result"), + WebSocketAssemblyAdmissionOutcome::Admitted(_) + )); + } + #[tokio::test] async fn policy_only_reload_keeps_connected_middleware_registry() { let proto = test_proto(); diff --git a/crates/openshell-supervisor-network/src/proxy.rs b/crates/openshell-supervisor-network/src/proxy.rs index 6e9c48220b..f15f88d30a 100644 --- a/crates/openshell-supervisor-network/src/proxy.rs +++ b/crates/openshell-supervisor-network/src/proxy.rs @@ -4841,47 +4841,6 @@ async fn handle_forward_proxy( return Ok(()); } - // 6. Connect upstream. Plain-HTTP requests always dial the destination - // directly: only TLS (CONNECT) tunnels chain through the corporate - // proxy, since plain-HTTP forwarding would need absolute-form requests - // rather than a CONNECT tunnel. - let dial_result = TcpStream::connect(addrs.as_slice()).await; - let mut upstream = match dial_result { - Ok(s) => s, - Err(e) => { - let event = HttpActivityBuilder::new(openshell_ocsf::ctx::ctx()) - .activity(ActivityId::Fail) - .severity(SeverityId::Low) - .status(StatusId::Failure) - .http_request(HttpRequest::new( - method, - OcsfUrl::new("http", &host_lc, &path, port), - )) - .dst_endpoint(Endpoint::from_domain(&host_lc, port)) - .src_endpoint(Endpoint::from_ip(workload_addr.ip(), workload_addr.port())) - .actor_process( - Process::from_bypass(&binary_str, &pid_str, &ancestors_str) - .with_cmd_line(&cmdline_str), - ) - .message(format!( - "FORWARD upstream connect failed for {host_lc}:{port}: {e}" - )) - .build(); - ocsf_emit!(event); - respond( - client, - &build_json_error_response( - 502, - "Bad Gateway", - "upstream_unreachable", - &format!("connection to {host_lc}:{port} failed"), - ), - ) - .await?; - return Ok(()); - } - }; - let middleware_path = path.split_once('?').map_or(path.as_str(), |(path, _)| path); let middleware_input = crate::opa::NetworkInput { host: host_lc.clone(), @@ -4915,6 +4874,7 @@ async fn handle_forward_proxy( .await?; return Ok(()); } + let websocket_chain = forward_websocket_request.then(|| chain.clone()); if !chain.is_empty() { let middleware_runner = opa_engine.middleware_runner()?; let request = crate::l7::rest::request_from_buffered_http( @@ -4949,9 +4909,62 @@ async fn handle_forward_proxy( respond(client, &response).await?; return Ok(()); } + crate::l7::middleware::MiddlewareApplyResult::AdmissionExhausted => { + emit_activity_simple(activity_tx, true, "middleware"); + let response = build_middleware_unavailable_response(&l7_ctx.policy_name); + respond(client, &response).await?; + return Ok(()); + } }; } + let mut middleware_session = if let Some(chain) = websocket_chain.as_deref() { + let request = crate::l7::rest::request_from_buffered_http( + method, + middleware_path, + &upstream_target, + forward_request_bytes.clone(), + )?; + let middleware_runner = opa_engine.middleware_runner()?; + let preflight = crate::l7::relay::websocket_middleware_preflight( + &request, + chain, + &middleware_runner, + &l7_ctx, + "ws", + ) + .await; + let preflight = match preflight { + Ok(preflight) => preflight, + Err(error) => { + warn!(error = %error, "Plaintext WebSocket middleware preflight failed"); + respond( + client, + &build_json_error_response( + 502, + "Bad Gateway", + "middleware_failed", + "WebSocket middleware preflight failed", + ), + ) + .await?; + return Ok(()); + } + }; + crate::l7::middleware::emit_websocket_preflight_events(&l7_ctx, &preflight); + if preflight.terminal_reason.is_some() { + let response = preflight.denial.as_ref().map_or_else( + || build_middleware_failure_response(&l7_ctx.policy_name), + |denial| build_middleware_deny_response(&l7_ctx.policy_name, denial), + ); + respond(client, &response).await?; + return Ok(()); + } + preflight.session + } else { + None + }; + forward_request_bytes = match inject_token_grant_for_forward_request( method, &upstream_target, @@ -4968,6 +4981,11 @@ async fn handle_forward_proxy( error = %e, "token grant failed in forward proxy" ); + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::Cancellation) + .await; + } respond( client, &build_json_error_response( @@ -4999,6 +5017,11 @@ async fn handle_forward_proxy( error = %e, "credential injection failed in forward proxy" ); + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::Cancellation) + .await; + } respond( client, &build_json_error_response( @@ -5023,6 +5046,86 @@ async fn handle_forward_proxy( "Forward proxy rejected request because policy changed before relay" ); emit_l7_tunnel_close_after_policy_change(&host_lc, port, e); + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::PolicyReload) + .await; + } + respond( + client, + &build_json_error_response( + 403, + "Forbidden", + "policy_denied", + &format!("{method} {host_lc}:{port}{path} not permitted by policy"), + ), + ) + .await?; + return Ok(()); + } + + // Plain-HTTP requests dial the destination directly: only TLS (CONNECT) + // tunnels chain through the corporate proxy, since plain-HTTP forwarding + // would need absolute-form requests rather than a CONNECT tunnel. Dial + // only after every local authorization and transformation step so a + // rejected WebSocket preflight cannot contact the destination. + let dial_result = TcpStream::connect(addrs.as_slice()).await; + let mut upstream = match dial_result { + Ok(s) => s, + Err(e) => { + let event = HttpActivityBuilder::new(openshell_ocsf::ctx::ctx()) + .activity(ActivityId::Fail) + .severity(SeverityId::Low) + .status(StatusId::Failure) + .http_request(HttpRequest::new( + method, + OcsfUrl::new("http", &host_lc, &path, port), + )) + .dst_endpoint(Endpoint::from_domain(&host_lc, port)) + .src_endpoint(Endpoint::from_ip(workload_addr.ip(), workload_addr.port())) + .actor_process( + Process::from_bypass(&binary_str, &pid_str, &ancestors_str) + .with_cmd_line(&cmdline_str), + ) + .message(format!( + "FORWARD upstream connect failed for {host_lc}:{port}: {e}" + )) + .build(); + ocsf_emit!(event); + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::UpstreamRejected) + .await; + } + respond( + client, + &build_json_error_response( + 502, + "Bad Gateway", + "upstream_unreachable", + &format!("connection to {host_lc}:{port} failed"), + ), + ) + .await?; + return Ok(()); + } + }; + + if let Err(e) = forward_generation_guard.ensure_current() { + warn!( + host = %host_lc, + port, + captured_generation = forward_generation_guard.captured_generation(), + current_generation = forward_generation_guard.current_generation(), + error = %e, + "Forward proxy rejected request because policy changed during upstream connect" + ); + emit_l7_tunnel_close_after_policy_change(&host_lc, port, e); + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::PolicyReload) + .await; + } respond( client, &build_json_error_response( @@ -5035,7 +5138,8 @@ async fn handle_forward_proxy( .await?; return Ok(()); } - let outcome = relay_rewritten_forward_request( + + let outcome_result = relay_rewritten_forward_request( method, &path, rewritten, @@ -5048,6 +5152,15 @@ async fn handle_forward_proxy( request_body_credential_rewrite, }, ) + .await; + let outcome = crate::l7::relay::finalize_websocket_pre_upgrade( + &mut middleware_session, + &forward_generation_guard, + &host_lc, + port, + policy_str, + outcome_result, + ) .await?; // The request has now survived middleware, token grant, credential @@ -5077,39 +5190,55 @@ async fn handle_forward_proxy( } emit_forward_success_activity(activity_tx, l7_activity_pending); - if let crate::l7::provider::RelayOutcome::Upgraded { - overflow, - websocket_permessage_deflate, - } = outcome - { - let mut upgrade_options = if let (Some(config), Some(engine)) = ( - forward_upgrade_config.as_ref(), - forward_tunnel_engine.as_ref(), - ) { - crate::l7::relay::upgrade_options( - config, - &l7_ctx, - forward_websocket_request, - &forward_upgrade_target, - &forward_upgrade_query_params, - Some(engine), - ) - } else { - crate::l7::relay::UpgradeRelayOptions { - websocket_request: forward_websocket_request, - ..Default::default() + match outcome { + crate::l7::provider::RelayOutcome::Reusable + | crate::l7::provider::RelayOutcome::Consumed => { + if let Some(session) = middleware_session.take() { + session + .end(openshell_core::proto::WebSocketSessionEndReason::UpstreamRejected) + .await; } - }; - upgrade_options.websocket.permessage_deflate = websocket_permessage_deflate; - crate::l7::relay::handle_upgrade( - client, - &mut upstream, + } + crate::l7::provider::RelayOutcome::Upgraded { overflow, - &host_lc, - port, - upgrade_options, - ) - .await?; + websocket_permessage_deflate, + websocket_subprotocol, + } => { + let mut upgrade_options = if let (Some(config), Some(engine)) = ( + forward_upgrade_config.as_ref(), + forward_tunnel_engine.as_ref(), + ) { + crate::l7::relay::upgrade_options( + config, + &l7_ctx, + forward_websocket_request, + &forward_upgrade_target, + &forward_upgrade_query_params, + Some(engine), + ) + } else { + crate::l7::relay::UpgradeRelayOptions { + websocket_request: forward_websocket_request, + ctx: Some(&l7_ctx), + policy_name: l7_ctx.policy_name.clone(), + ..Default::default() + } + }; + upgrade_options.generation_guard = Some(&forward_generation_guard); + upgrade_options.assembly_budget = Some(opa_engine.websocket_assembly_budget()); + upgrade_options.websocket.permessage_deflate = websocket_permessage_deflate; + upgrade_options.middleware_session = middleware_session.take(); + upgrade_options.selected_subprotocol = websocket_subprotocol; + crate::l7::relay::handle_upgrade( + client, + &mut upstream, + overflow, + &host_lc, + port, + upgrade_options, + ) + .await?; + } } Ok(()) @@ -5231,6 +5360,14 @@ fn build_middleware_deny_response( } fn build_middleware_failure_response(policy_name: &str) -> Vec { + build_middleware_platform_response(policy_name, "403 Forbidden") +} + +fn build_middleware_unavailable_response(policy_name: &str) -> Vec { + build_middleware_platform_response(policy_name, "503 Service Unavailable") +} + +fn build_middleware_platform_response(policy_name: &str, status: &str) -> Vec { let body = serde_json::json!({ "error": "middleware_failed", "detail": "Request could not be processed by configured middleware", @@ -5238,7 +5375,7 @@ fn build_middleware_failure_response(policy_name: &str) -> Vec { }); let body_str = body.to_string(); format!( - "HTTP/1.1 403 Forbidden\r\n\ + "HTTP/1.1 {status}\r\n\ Content-Type: application/json\r\n\ Content-Length: {}\r\n\ Connection: close\r\n\ @@ -5327,6 +5464,104 @@ mod tests { use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; + struct DenyWebSocketPreflight; + + #[tonic::async_trait] + impl openshell_core::middleware::SupervisorMiddlewareEndpoint for DenyWebSocketPreflight { + async fn describe( + &self, + _request: tonic::Request<()>, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + Ok(tonic::Response::new( + openshell_core::proto::MiddlewareManifest { + name: "test/deny-websocket".into(), + service_version: "test".into(), + bindings: vec![openshell_core::proto::MiddlewareBinding { + operation: + openshell_core::proto::SupervisorMiddlewareOperation::WebsocketMessage + as i32, + phase: openshell_core::proto::SupervisorMiddlewarePhase::PreCredentials + as i32, + max_payload_bytes: 1024, + timeout: "1s".into(), + }], + }, + )) + } + + async fn validate_config( + &self, + _request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + Ok(tonic::Response::new( + openshell_core::proto::ValidateConfigResponse { + valid: true, + reason: String::new(), + }, + )) + } + + async fn evaluate_http_request( + &self, + _request: tonic::Request, + ) -> std::result::Result< + tonic::Response, + tonic::Status, + > { + Err(tonic::Status::unimplemented("WebSocket-only test service")) + } + + async fn open_websocket_session( + &self, + mut receiver: mpsc::Receiver, + ) -> std::result::Result + { + let (responses, response_stream) = mpsc::channel(1); + tokio::spawn(async move { + while let Some(event) = receiver.recv().await { + if matches!( + event.event, + Some(openshell_core::proto::web_socket_session_event::Event::Preflight(_)) + ) { + let _ = responses + .send(Ok(openshell_core::proto::WebSocketSessionEventResult { + result: Some( + openshell_core::proto::web_socket_session_event_result::Result::PreflightDecision( + openshell_core::proto::WebSocketPreflightDecision { + action: openshell_core::proto::WebSocketPreflightAction::Deny as i32, + reason: "test denial".into(), + reason_code: "test_denial".into(), + ..Default::default() + }, + ), + ), + })) + .await; + break; + } + } + }); + Ok(Box::pin(tokio_stream::wrappers::ReceiverStream::new( + response_stream, + ))) + } + } + + fn non_loopback_test_ipv4() -> Option { + let socket = std::net::UdpSocket::bind("0.0.0.0:0").ok()?; + socket.connect("192.0.2.1:9").ok()?; + match socket.local_addr().ok()?.ip() { + IpAddr::V4(ip) if !ip.is_loopback() && !ip.is_unspecified() => Some(ip), + _ => None, + } + } + async fn drive_raw_request_through_handler(raw: Vec) -> Vec { let policy = include_str!("../data/sandbox-policy.rego"); let data = r#" @@ -5385,6 +5620,117 @@ network_policies: {} } } + #[tokio::test] + async fn plaintext_websocket_preflight_denial_does_not_connect_upstream() { + if !cfg!(target_os = "linux") { + eprintln!("skipping: handler identity binding requires /proc (Linux)"); + return; + } + let Some(upstream_ip) = non_loopback_test_ipv4() else { + eprintln!("skipping: no routable non-loopback IPv4 test address"); + return; + }; + + let upstream_listener = TcpListener::bind((upstream_ip, 0)) + .await + .expect("bind upstream listener"); + let upstream_port = upstream_listener.local_addr().unwrap().port(); + let executable = std::env::current_exe().expect("current executable"); + let data = format!( + r#" +network_middlewares: + deny-upgrade: + middleware: test/deny-websocket + on_error: fail_closed + endpoints: + include: ["{upstream_ip}"] +network_policies: + allow-upstream: + name: allow-upstream + endpoints: + - host: "{upstream_ip}" + port: {upstream_port} + binaries: + - {{ path: "{executable}" }} +"#, + executable = executable.display(), + ); + let engine = Arc::new( + OpaEngine::from_strings(include_str!("../data/sandbox-policy.rego"), &data) + .expect("load policy"), + ); + let registry = openshell_supervisor_middleware::MiddlewareRegistry::connect_services( + vec![Arc::new(DenyWebSocketPreflight)], + Vec::new(), + ) + .await + .expect("connect test middleware"); + engine + .replace_middleware_registry(registry) + .expect("install test middleware"); + + let proxy_listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind proxy listener"); + let proxy_address = proxy_listener.local_addr().unwrap(); + let target = format!("http://{upstream_ip}:{upstream_port}/ws"); + let request = format!( + "GET {target} HTTP/1.1\r\nHost: {upstream_ip}:{upstream_port}\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Version: 13\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n\r\n" + ); + let client = tokio::spawn(async move { + let mut socket = TcpStream::connect(proxy_address) + .await + .expect("connect proxy"); + let mut response = Vec::new(); + socket + .read_to_end(&mut response) + .await + .expect("read proxy response"); + response + }); + let (mut proxy_connection, _) = proxy_listener.accept().await.unwrap(); + + tokio::time::timeout( + std::time::Duration::from_secs(2), + handle_forward_proxy( + "GET", + &target, + request.as_bytes(), + request.len(), + &mut proxy_connection, + engine, + Arc::new(BinaryIdentityCache::new()), + Arc::new(AtomicU32::new(std::process::id())), + None, + AgentProposals::default(), + Arc::new(None), + None, + None, + None, + None, + ), + ) + .await + .expect("denied preflight must complete without an upstream response") + .expect("handle denied plaintext WebSocket upgrade"); + drop(proxy_connection); + + let response = String::from_utf8(client.await.unwrap()).expect("UTF-8 response"); + assert!( + response.contains("\"error\":\"middleware_denied\""), + "preflight must deny the upgrade: {response}" + ); + assert!( + tokio::time::timeout( + std::time::Duration::from_millis(100), + upstream_listener.accept() + ) + .await + .is_err(), + "denied preflight must not establish an upstream connection" + ); + } + #[tokio::test] async fn malformed_request_lines_are_rejected_before_connect_or_forward_dispatch() { for host in ["api.example.com", "unmatched.example.com"] { @@ -5450,6 +5796,28 @@ network_policies: {} assert!(body.get("agent_guidance").is_none()); } + #[test] + fn middleware_unavailable_response_is_complete_and_has_no_retry_hint() { + let response = build_middleware_unavailable_response("api-policy"); + let response = String::from_utf8(response).expect("UTF-8 error response"); + assert!(response.starts_with("HTTP/1.1 503 Service Unavailable\r\n")); + assert!(!response.to_ascii_lowercase().contains("retry-after")); + let (headers, body) = response.split_once("\r\n\r\n").expect("HTTP response"); + let content_length = headers + .lines() + .find_map(|line| { + line.strip_prefix("Content-Length: ") + .and_then(|value| value.parse::().ok()) + }) + .expect("Content-Length"); + assert_eq!(content_length, body.len()); + let body: serde_json::Value = serde_json::from_str(body).expect("JSON response"); + assert_eq!(body["error"], "middleware_failed"); + assert_eq!(body["policy"], "api-policy"); + assert!(body.get("middleware").is_none()); + assert!(body.get("reason_code").is_none()); + } + #[test] fn policy_deny_response_includes_reason() { let response = build_json_error_response_with_reason( @@ -5869,7 +6237,7 @@ network_policies: secret_resolver: None, ..Default::default() }; - let runner = openshell_supervisor_middleware::ChainRunner::new( + let runner = openshell_supervisor_middleware::ChainRunner::from_endpoint( openshell_supervisor_middleware_builtins::services() .into_iter() .next() @@ -5907,6 +6275,9 @@ network_policies: crate::l7::middleware::MiddlewareApplyResult::Allowed(_) => { panic!("policy-invalid transformed request must be denied") } + crate::l7::middleware::MiddlewareApplyResult::AdmissionExhausted => { + panic!("test middleware work admission must be available") + } } } @@ -6192,6 +6563,14 @@ network_policies: frame } + fn masked_close_code(frame: &[u8]) -> u16 { + assert_eq!(frame[0] & 0x0f, 0x08, "expected a close frame"); + assert_eq!(frame[1] & 0x7f, 2, "expected a two-byte close code"); + assert_ne!(frame[1] & 0x80, 0, "client-to-upstream close is masked"); + let decoded = [frame[6] ^ frame[2], frame[7] ^ frame[3]]; + u16::from_be_bytes(decoded) + } + async fn forward_websocket_denied_after_upgrade( config: crate::l7::L7EndpointConfig, tunnel_engine: crate::opa::TunnelPolicyEngine, @@ -6243,6 +6622,7 @@ network_policies: if let crate::l7::provider::RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + .. } = outcome { let mut options = crate::l7::relay::upgrade_options( @@ -6348,6 +6728,7 @@ network_policies: ); assert!(options.websocket.credential_rewrite); assert!(options.secret_resolver.is_some()); + assert!(options.generation_guard.is_some()); assert!(options.engine.is_some()); assert!(options.ctx.is_some()); assert!(matches!( @@ -6388,6 +6769,41 @@ network_policies: )); } + #[test] + fn rest_websocket_upgrade_carries_guard_without_message_inspector() { + let engine = OpaEngine::from_strings( + include_str!("../data/sandbox-policy.rego"), + "network_policies: {}\n", + ) + .expect("test policy"); + let tunnel_engine = engine + .clone_engine_for_tunnel(engine.current_generation()) + .expect("tunnel engine"); + let ctx = crate::l7::relay::L7EvalContext { + host: "gateway.example.test".into(), + port: 443, + policy_name: "rest_api".into(), + binary_path: "/usr/bin/node".into(), + ..Default::default() + }; + let config = websocket_l7_config(crate::l7::L7Protocol::Rest, false); + let options = crate::l7::relay::upgrade_options( + &config, + &ctx, + true, + "/ws", + &std::collections::HashMap::new(), + Some(&tunnel_engine), + ); + + assert!(matches!( + options.websocket.message_policy, + crate::l7::relay::WebSocketMessagePolicy::None + )); + assert!(options.generation_guard.is_some()); + assert!(options.engine.is_some()); + } + #[tokio::test] async fn forward_websocket_upgrade_blocks_text_frame_by_policy() { let data = r#" @@ -6426,9 +6842,10 @@ network_policies: .await; assert!(err.to_string().contains("websocket text message denied")); - assert!( - leaked.is_empty(), - "denied forward-proxy WebSocket text frames must not reach upstream" + assert_eq!( + masked_close_code(&leaked), + 1008, + "only a policy close, not the denied text frame, may reach upstream" ); } @@ -6479,9 +6896,10 @@ network_policies: .await; assert!(err.to_string().contains("websocket GraphQL message denied")); - assert!( - leaked.is_empty(), - "denied forward-proxy GraphQL WebSocket operations must not reach upstream" + assert_eq!( + masked_close_code(&leaked), + 1008, + "only a policy close, not the denied GraphQL operation, may reach upstream" ); } @@ -8793,7 +9211,7 @@ network_policies: secret_resolver: None, ..Default::default() }; - let runner = openshell_supervisor_middleware::ChainRunner::new( + let runner = openshell_supervisor_middleware::ChainRunner::from_endpoint( openshell_supervisor_middleware_builtins::services() .into_iter() .next() diff --git a/docs/extensibility/supervisor-middleware.mdx b/docs/extensibility/supervisor-middleware.mdx index f3bdcac9bb..452244ebed 100644 --- a/docs/extensibility/supervisor-middleware.mdx +++ b/docs/extensibility/supervisor-middleware.mdx @@ -3,11 +3,11 @@ # SPDX-License-Identifier: Apache-2.0 title: "Supervisor Middleware" sidebar-title: "Supervisor Middleware" -description: "Configure and operate built-in and operator-run middleware for sandbox HTTP requests." +description: "Configure and operate built-in and operator-run middleware for sandbox HTTP requests and WebSocket messages." keywords: "Generative AI, Cybersecurity, AI Agents, Supervisor Middleware, Extensibility, Request Filtering" --- -Supervisor middleware adds ordered request-processing stages to allowed HTTP egress. Middleware runs after network and L7 policy admit a request and before OpenShell injects provider credentials. A stage can allow or deny the request, replace its body, add approved headers, and report audit-safe findings. +Supervisor middleware adds ordered processing stages to allowed HTTP and WebSocket egress. Middleware runs after network and L7 policy admit traffic and before OpenShell injects provider credentials. A stage can allow or deny an HTTP request or client WebSocket text message, replace its payload, add approved HTTP headers, and report audit-safe findings. Middleware selection is independent of the network policy rule that admitted the request. OpenShell matches middleware by destination host, so the same middleware applies consistently across broad, specific, user-authored, and provider-derived network policies. @@ -22,6 +22,15 @@ For each inspected HTTP request, the supervisor: 5. Re-checks body-aware protocol policy (GraphQL, JSON-RPC, MCP) after each stage that replaces the body. Every middleware receives a payload the policy admits, and a transformation cannot smuggle a denied or unparseable operation to a later stage or the upstream. 6. Applies allowed transformations, injects provider credentials, and forwards the request. +For an RFC 6455 upgrade over `ws://` or `wss://`, the supervisor first finds every host-matched attachment, then selects only implementations that advertise `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`. It opens one ordered, phase-specific `EvaluateWebSocketSession` stream per selected stage. OpenShell sends `WebSocketSessionEvent` values, while the service returns `WebSocketSessionEventResult` values only for preflight and message events; session start and end are notifications. Future upstream-to-client inspection uses the same RPC with `PRE_RETURN`; an implementation that advertises both phases receives two independent streams for the WebSocket session. An attachment without the selected binding can still inspect the HTTP upgrade request when it advertises the HTTP binding, but it is not a failed WebSocket stage. OpenShell allows post-upgrade traffic and emits an informational `binding_not_selected` coverage event for that attachment. + +1. A preflight before the upgrade is sent upstream. The stage chooses `INSPECT`, voluntary `SKIP`, or authoritative `DENY` and may return a bounded diagnostic reason, stable reason code, findings, and metadata. OpenShell runs selected preflights concurrently; any `DENY` rejects the upgrade regardless of `on_error`. +2. A session-start event after the upstream accepts the upgrade, including the negotiated subprotocol. +3. Complete client-to-upstream text messages in sequence order. OpenShell reassembles fragmented messages and decompresses negotiated `permessage-deflate` messages before evaluation. +4. A best-effort session-end event when the stage stream remains writable. OpenShell attempts at most one terminal event for each opened stream, including streams opened during a preflight that rejects the upgrade before session start. + +Allowed replacements are re-framed, re-compressed when required, and forwarded. Binary messages, control frames, and upstream-to-client traffic remain uninspected. Binary messages pass through under both `on_error` modes. For each active selected stage, OpenShell emits an informational `unsupported_message_type` coverage event and advances the session-global sequence; the next text message can therefore reach the stage with a valid sequence gap. The network supervisor reserves process-wide assembly capacity before buffering every parsed WebSocket text message. Active middleware sessions additionally reserve shared middleware capacity before buffering WebSocket text, and HTTP middleware reserves the same capacity before buffering request bodies; at most 32 evaluations run and 64 additional unbuffered callers wait for capacity. When both middleware bounds are full, OpenShell sheds an HTTP request with `503 Service Unavailable` before reading its body. Persistent middleware streams use a separate process-wide budget of 32 sessions. WebSocket session admission does not wait: if the budget is full, OpenShell applies each selected config's `on_error` behavior before opening a stream. + Because each transformed body is re-checked before the next stage runs, a middleware hook always receives a request that satisfies the sandbox policy. A stage whose output the policy rejects stops the chain; under `enforcement: audit` the rejection is logged and the request proceeds. If post-transformation policy evaluation itself fails, OpenShell denies the request and emits a high-severity detection finding. This failure is separate from middleware `on_error` because the middleware completed successfully; the sandbox policy could not validate its output. @@ -30,14 +39,14 @@ Middleware receives the request before credential injection. Operator-run servic ## Choose a Middleware Type -| Type | Registration | Body limit | Deployment | +| Type | Registration | Payload limit | Deployment | | --- | --- | --- | --- | | Built-in | None | Defined by OpenShell | Runs inside the supervisor | | 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. -Operator-run services expose bindings for supported operation and phase pairs. A binding is identified by its operation and phase; V1 supports only `HttpRequest/pre_credentials`. Policies attach the complete middleware by its operator-owned gateway registration name. +Operator-run services expose bindings for supported operation and phase pairs. A binding is identified by its operation and phase. V1 supports `HttpRequest/pre_credentials` and `WebSocketMessage/pre_credentials`; a service may expose either or both. Policies attach the complete middleware by its operator-owned gateway registration name. ## Register a Middleware Service @@ -47,7 +56,7 @@ Start an operator-run service before starting the gateway, then add a registrati [[openshell.supervisor.middleware]] name = "local-content-guard" grpc_endpoint = "http://host.openshell.internal:50051" -max_body_bytes = 262144 +max_payload_bytes = 262144 timeout = "500ms" ``` @@ -55,10 +64,10 @@ timeout = "500ms" | --- | --- | | `name` | Operator-owned registration name used by policy attachments and diagnostics. Names must be unique, and `openshell/` is reserved for built-ins. | | `grpc_endpoint` | Service address reachable from both the gateway and sandbox supervisors. Supports plaintext `http://` and TLS `https://` with platform trust roots. | -| `max_body_bytes` | Operator limit applied to every binding exposed by the service, up to the 4 MiB platform maximum. | +| `max_payload_bytes` | Shared operator limit applied to inspectable logical payloads across every binding exposed by the service, up to the 4 MiB platform maximum. It caps HTTP bodies and complete WebSocket text messages. | | `timeout` | Optional service-wide RPC timeout using an integer with an `ms` or `s` suffix. Defaults to `500ms`; valid values range from `10ms` through `30s`. | -Each binding returned by `Describe` may advertise a shorter `timeout` using the same syntax and bounds. The operator-configured service timeout is a ceiling: OpenShell uses the smaller of the binding and service values. An omitted binding timeout inherits the service setting, and an omitted service setting uses the 500 ms platform default. OpenShell rejects an invalid timeout before accepting the manifest. The service timeout applies to `Describe`, while the effective binding timeout applies to `ValidateConfig` and `EvaluateHttpRequest`. +Each binding returned by `Describe` may advertise a shorter `timeout` using the same syntax and bounds. The operator-configured service timeout is a ceiling: OpenShell uses the smaller of the binding and service values. An omitted binding timeout inherits the service setting, and an omitted service setting uses the 500 ms platform default. OpenShell rejects an invalid timeout before accepting the manifest. The service timeout applies to `Describe`; the effective binding timeout applies to `ValidateConfig`, `EvaluateHttpRequest`, WebSocket preflight, and each WebSocket message. WebSocket streams have no connection-wide deadline. The gateway connects to every registered service and verifies its capabilities before accepting traffic. Gateway startup fails when a service is unavailable, reports an invalid capability, or exposes more than one binding for the same operation and phase. The manifest `name` is diagnostic metadata and does not need to match the operator registration name. Operator-run registration names cannot claim the reserved `openshell/` namespace. @@ -92,16 +101,18 @@ See [Policy Schema](/reference/policy-schema#network-middleware) for the complet ## Configure Failure Behavior -`on_error` controls what happens when middleware is unavailable, rejects its configuration, returns an invalid result, or exceeds a body limit. +`on_error` controls what happens after an operation binding is selected and middleware is unavailable, rejects its configuration, returns an invalid result, or exceeds the selected binding's payload limit. It does not turn an unadvertised operation or an unsupported WebSocket message class into a middleware failure. | Value | Behavior | | --- | --- | -| `fail_closed` | Denies the request when the middleware stage fails. This is the default. | -| `fail_open` | Skips the failed stage and continues the request through the remaining chain. | +| `fail_closed` | Denies the HTTP request or closes the WebSocket when the stage fails. This is the default. | +| `fail_open` | Skips the failed HTTP stage. For a broken WebSocket stage stream, disables that stage for the rest of the connection and continues the remaining chain. | -Use `fail_open` only when bypassing the middleware preserves the intended security policy. OpenShell emits a detection finding when a failed stage is bypassed. +Use `fail_open` only when bypassing the middleware preserves the intended security policy. OpenShell emits a detection finding when a failed stage is bypassed and a separate state-change finding when a WebSocket stage is disabled for the session. -An explicit deny decision always stops the chain and denies the request, regardless of `on_error`. The HTTP response uses `error: middleware_denied`, identifies the policy-local middleware config, and omits policy-advisor remediation because the network and L7 allow rules already matched. OpenShell never copies the free-form middleware `reason` into the response or security logs. A service can instead return an optional stable `reason_code`: 1–64 bytes, starting with a lowercase ASCII letter and containing only lowercase ASCII letters, digits, and underscores. Invalid codes make the result a middleware failure governed by `on_error`. +Capability coverage is separate from failure handling. A host-matched HTTP-only attachment does not join the WebSocket chain, regardless of `on_error`. Binary messages are outside the V1 text-message binding and pass through even when a selected stage is `fail_closed`. OpenShell records both states as informational coverage events so operators do not mistake pass-through traffic for inspected traffic. If a deployment requires all WebSocket message classes to be inspected, V1 cannot express that requirement. + +An explicit deny decision always stops the chain and denies the request or WebSocket upgrade, regardless of `on_error`. A WebSocket preflight `DENY` is a successful policy decision, not a middleware failure; OpenShell rejects the upgrade before upstream contact and ends each still-writable stream opened by a successful preflight decision with `MIDDLEWARE_DENIAL`. The HTTP response uses `error: middleware_denied`, identifies the policy-local middleware config, and omits policy-advisor remediation because the network and L7 allow rules already matched. OpenShell never copies the free-form middleware `reason` into the response or security logs. HTTP results, WebSocket preflight decisions, and WebSocket message results can instead return an optional stable `reason_code`: 1–64 bytes, starting with a lowercase ASCII letter and containing only lowercase ASCII letters, digits, and underscores. Invalid codes make the result a middleware failure governed by `on_error`. Preflight findings and metadata use the same bounds and audit-safe handling as message results. ```json { @@ -117,19 +128,21 @@ A failed `fail_closed` stage uses `error: middleware_failed` and a platform-owne Middleware decisions are enforced regardless of the endpoint's `enforcement` mode. `enforcement: audit` applies to an endpoint's network and L7 policy rules and does not bypass middleware: a middleware deny, or a failed `fail_closed` stage, blocks the request even on an audit endpoint. A middleware service that needs to observe traffic without blocking should return an allow decision with findings, which OpenShell emits as detection findings. -## Set Body Limits +## Set Payload Limits -Every middleware binding declares the largest request or replacement body it supports. +Every middleware binding declares the largest logical payload or replacement it supports through `max_payload_bytes`. For `HTTP_REQUEST`, that payload is one request body. For `WEBSOCKET_MESSAGE`, it is one complete message rather than the whole session. - Built-in middleware uses its OpenShell-defined limit. -- Each operator-run registration sets `max_body_bytes` no higher than the service capability. +- Each operator-run registration sets one `max_payload_bytes` ceiling no higher than any binding's advertised `max_payload_bytes` capability. - A selected chain buffers using its largest stage limit, so every stage that can process the body receives it. - The same per-stage limit applies to request bodies and replacement bodies. -The gateway rejects a registration whose operator limit exceeds the service capability or the 4 MiB platform maximum instead of silently clamping it. OpenShell also bounds the non-body protobuf components: 64 KiB for service config, 4 KiB for request context, 32 KiB for the target, and 128 request header lines totaling at most 64 KiB encoded. Results allow a 4 KiB discarded free-form reason, a 64-byte validated reason code, 64 header mutations totaling at most 64 KiB encoded, 32 findings of at most 4 KiB encoded each, and 64 metadata entries totaling at most 32 KiB. Middleware gRPC servers should configure request and response message limits to at least 4 MiB plus 293 KiB so every platform-valid envelope fits. +The gateway rejects a registration whose operator limit exceeds the service capability or the 4 MiB platform maximum instead of silently clamping it. OpenShell also bounds the non-payload protobuf components: 64 KiB for service config, 4 KiB for request context, 32 KiB for the target, and 128 request header lines totaling at most 64 KiB encoded. Results allow a 4 KiB discarded free-form reason, a 64-byte validated reason code, 64 header mutations totaling at most 64 KiB encoded, 32 findings of at most 4 KiB encoded each, and 64 metadata entries totaling at most 32 KiB. Middleware gRPC servers should configure request and response message limits to at least 4 MiB plus 293 KiB so every platform-valid envelope fits. At request time, exceeding a selected stage's limit is a middleware failure for that stage alone and follows that config's `on_error` behavior; other stages in the chain still run against their own limits. OpenShell can apply `fail_open` to an oversized `Content-Length` before consuming body bytes. A chunked body can cross the limit only after bytes have been consumed, so OpenShell denies that request because it cannot safely resume the original stream. +For a WebSocket binding, `max_payload_bytes` covers complete client text messages and replacements. Exceeding a selected stage's effective text-message limit follows that stage's `on_error`. The 4 MiB parsed-text platform cap and other protocol-safety limits are independent of middleware failure policy. Binary messages are not delivered to middleware, so the operator ceiling does not become a binary relay limit; individual raw binary frames retain the 16 MiB relay-safety bound. Oversized parsed text closes the connection with code `1009`; invalid UTF-8 uses `1007`; protocol errors use `1002`; middleware or policy denials use `1008`; and policy reload uses `1012`. + ## Mutate Request Headers A middleware result can return ordered header mutations before OpenShell injects credentials. A `write` mutation adds a value when the case-insensitive header name is absent and selects one behavior when it is already present: @@ -164,6 +177,8 @@ Middleware activity is emitted through OpenShell's OCSF logging: - A denied invocation records a platform-owned reason derived from the policy-local config name and optional validated reason code. OpenShell does not record service-provided free-form reason text. - A bypass under `fail_open` emits a detection finding. - A required stage that fails closed emits a high-severity detection finding. +- A host-matched attachment without a WebSocket binding emits an informational `binding_not_selected` coverage event. +- A binary message encountered by an active WebSocket stage emits an informational `unsupported_message_type` coverage event with message type, sequence, and byte count. It is not reported as an invocation or failure. - Built-in findings include their type, label, and aggregate count. Operator-run findings use the operator-owned registration name and a platform label plus the aggregate count; OpenShell does not log service-provided finding text or diagnostic metadata. A stage can return at most 32 findings. Exceeding the per-stage cap is an invalid response handled through `on_error`. A maximum 10-stage chain retains and emits up to 320 findings without silently dropping findings from later stages. - Registry reload success and failure are emitted as configuration state changes. @@ -171,8 +186,10 @@ See [Logging](/observability/logging) for log access and [OCSF JSON Export](/obs ## Current Limitations -- Middleware applies only to HTTP requests parsed by the supervisor. Traffic the supervisor cannot inspect, such as HTTP/2 prior knowledge or non-HTTP TCP protocols, is denied on hosts matched by a fail-closed middleware, and relayed with a detection finding when every matching middleware is `fail_open`. -- The typed operation and phase are `HTTP_REQUEST/PRE_CREDENTIALS`. +- Middleware applies only through operation bindings advertised by each implementation. For protocols that have no supported middleware operation at all, such as HTTP/2 prior knowledge or non-HTTP TCP, the existing uninspectable-traffic gate denies a host match containing `fail_closed` and relays an all-`fail_open` match with a detection finding. +- The typed operation and phase pairs are `HTTP_REQUEST/PRE_CREDENTIALS` and `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`. +- A host match does not imply every advertised operation: an HTTP-only attachment can inspect the upgrade GET, then post-upgrade traffic passes with `binding_not_selected` coverage. +- The V1 WebSocket binding inspects complete client text messages only. Binary messages pass with `unsupported_message_type` coverage for active stages; control frames and upstream-to-client messages remain outside the middleware operation. - Selection uses destination host include and exclude patterns. - A fail-closed middleware cannot cover `tls: skip` endpoints because OpenShell cannot inspect that traffic. An all-`fail_open` match may cover the endpoint; OpenShell bypasses the middleware and emits a detection finding. - Operator-run services support plaintext `http://` and TLS `https://` endpoints. HTTPS certificates must chain to a CA in the platform trust store. diff --git a/docs/observability/logging.mdx b/docs/observability/logging.mdx index 1ab0b5bfed..c3998629e5 100644 --- a/docs/observability/logging.mdx +++ b/docs/observability/logging.mdx @@ -137,6 +137,21 @@ OCSF HTTP:POST [INFO] ALLOWED POST http://httpbin.org:443/anything [policy:httpb OCSF FINDING:CREATE [MED] "configured content matched" [type:content_guard.match count:1 middleware:prototype-content-guard] ``` +WebSocket middleware emits one safe event per preflight, session-start, or client text-message decision. The event includes the policy-local config, registered implementation, sequence, byte counts, transformation flag, and validated reason code. It never includes the message payload or service-provided free-form reason: + +```text +OCSF NET:OTHER [INFO] WEBSOCKET_MIDDLEWARE allow config=api-redactor implementation=openshell/regex sequence=3 input_bytes=128 replacement_bytes=96 transformed=true reason_code=- +``` + +Coverage events are separate from invocation decisions. `binding_not_selected` means a host-matched attachment did not advertise the WebSocket operation. `unsupported_message_type` means an active WebSocket stage encountered a binary message, which V1 passes through without inspection. Both are informational, do not apply `on_error`, and never claim the traffic was inspected: + +```text +OCSF NET:OTHER [INFO] WEBSOCKET_MIDDLEWARE_COVERAGE state=binding_not_selected config=http-dlp implementation=example/http-dlp sequence=- message_type=- input_bytes=0 +OCSF NET:OTHER [INFO] WEBSOCKET_MIDDLEWARE_COVERAGE state=unsupported_message_type config=api-redactor implementation=openshell/regex sequence=4 message_type=binary input_bytes=512 +``` + +A fail-open stream error emits both a middleware failure and `openshell.middleware.websocket_stage_disabled`. The latter records that OpenShell will bypass that stage for later messages on the same connection. Waiting for saturated admission capacity also emits a detection finding without payload content. When both active capacity and the bounded wait queue are full, HTTP work is rejected before its payload is buffered, returns `503 Service Unavailable`, and emits `openshell.middleware.admission_exhausted`. + Proxy and SSH servers ready: ```text @@ -167,7 +182,7 @@ OCSF CONFIG:LOADED [INFO] Policy reloaded successfully [policy_hash:0cc0c2b52557 Denied `NET:` and `HTTP:` events carry a `[reason:...]` suffix that surfaces the decision detail from the event's `status_detail` field. The reason helps distinguish between policy misses, SSRF hardening, and L7 enforcement without inspecting the full OCSF JSONL record. -For supervisor middleware denials, `status_detail` contains a platform-owned reason derived from the policy-local middleware config name and optional validated reason code. Middleware failure details also use platform-owned error codes. OpenShell does not copy per-request service text into logs. +For supervisor middleware denials, `status_detail` contains a platform-owned reason derived from the policy-local middleware config name and optional validated reason code. Middleware failure details also use platform-owned error codes. OpenShell does not copy per-request service text or WebSocket message content into logs. Common reason phrases emitted by the sandbox include: diff --git a/docs/reference/gateway-config.mdx b/docs/reference/gateway-config.mdx index 27a1d3d45a..1e3c77c48c 100644 --- a/docs/reference/gateway-config.mdx +++ b/docs/reference/gateway-config.mdx @@ -116,7 +116,7 @@ provider_profile_sources = [ [[openshell.supervisor.middleware]] name = "local-content-guard" grpc_endpoint = "http://host.openshell.internal:50051" -max_body_bytes = 262144 +max_payload_bytes = 262144 timeout = "500ms" # Gateway listener TLS (distinct from the per-driver guest_tls_*). @@ -183,21 +183,21 @@ Register operator-run supervisor middleware services with one or more `[[openshe [[openshell.supervisor.middleware]] name = "local-content-guard" grpc_endpoint = "http://host.openshell.internal:50051" -max_body_bytes = 262144 +max_payload_bytes = 262144 timeout = "500ms" ``` -Each service implements the supervisor middleware gRPC contract and exposes bindings through `Describe`. Policies reference the operator-owned registration `name`, attaching the complete middleware and all of its bindings. Bindings are identified by operation and phase. A manifest may expose at most one binding for each operation and phase pair; V1 supports only `HttpRequest/pre_credentials`, so a service currently exposes one binding. Registration names must be unique, and operator-run registrations cannot claim the reserved `openshell/` namespace. The service-reported manifest name is diagnostic metadata and does not need to match the registration name. +Each service implements the supervisor middleware gRPC contract and exposes bindings through `Describe`. Policies reference the operator-owned registration `name`, attaching the complete middleware and all of its bindings. Bindings are identified by operation and phase. A manifest may expose at most one binding for each operation and phase pair. V1 supports `HttpRequest/pre_credentials` and `WebSocketMessage/pre_credentials`, so a service can inspect HTTP, WebSocket, or both. Registration names must be unique, and operator-run registrations cannot claim the reserved `openshell/` namespace. The service-reported manifest name is diagnostic metadata and does not need to match the registration name. The gateway connects to every registered service and validates `Describe` before it starts. The service must therefore be running before the gateway. Policy creation and full policy updates call `ValidateConfig`; an unavailable service or invalid middleware configuration rejects the policy before persistence. -`max_body_bytes` is the operator limit for every binding exposed by the service. It must be greater than zero, no larger than each binding's advertised limit, and no larger than the 4 MiB platform maximum. OpenShell rejects an oversized value instead of silently clamping it. Middleware gRPC servers should allow messages of at least 4 MiB plus 293 KiB so a maximum-size body and its protobuf envelope fit on the transport. +`max_payload_bytes` is the shared operator limit for inspectable logical payloads across every binding exposed by the service. It caps HTTP request and replacement bodies as well as complete WebSocket text messages and replacements. The value must be greater than zero, no larger than each binding's advertised `max_payload_bytes` capability, and no larger than the 4 MiB platform maximum. OpenShell rejects oversized values instead of silently clamping them. Binary WebSocket messages are not exposed to V1 middleware, so this field does not limit binary pass-through. Middleware gRPC servers should allow messages of at least 4 MiB plus 293 KiB so a maximum-size payload and its protobuf envelope fit on the transport. -`timeout` is the operator-configured service-wide RPC timeout. It accepts the same compact duration syntax as gateway interceptors: an integer followed by `ms` or `s`, such as `500ms` or `2s`. Values must be between `10ms` and `30s`, inclusive. Omit the field to use the 500 ms platform default. A binding may advertise a shorter `timeout` in the `Describe` manifest, but it cannot extend the operator-configured deadline; OpenShell uses the smaller value. OpenShell validates both levels before accepting the service. The effective timeout covers `ValidateConfig` and `EvaluateHttpRequest`; `Describe` uses the service timeout because binding metadata is not available yet. +`timeout` is the operator-configured service-wide RPC timeout. It accepts the same compact duration syntax as gateway interceptors: an integer followed by `ms` or `s`, such as `500ms` or `2s`. Values must be between `10ms` and `30s`, inclusive. Omit the field to use the 500 ms platform default. A binding may advertise a shorter `timeout` in the `Describe` manifest, but it cannot extend the operator-configured deadline; OpenShell uses the smaller value. OpenShell validates both levels before accepting the service. The effective timeout covers `ValidateConfig`, `EvaluateHttpRequest`, WebSocket preflight, and each WebSocket message. `Describe` uses the service timeout because binding metadata is not available yet. An accepted WebSocket stream has no connection-wide RPC deadline. The service `grpc_endpoint` currently supports plaintext `http://` and TLS `https://` using the platform trust store. Custom trust roots, client authentication, health checks, and runtime registration are not currently supported. The endpoint must be reachable from both the gateway and sandbox supervisors; use `host.openshell.internal` or another shared address that can be resolved in both places. -See [Supervisor Middleware](/extensibility/supervisor-middleware) for selection, failure, body-limit, and operational guidance. +See [Supervisor Middleware](/extensibility/supervisor-middleware) for selection, failure, payload-limit, and operational guidance. ## Gateway Interceptors diff --git a/docs/reference/policy-schema.mdx b/docs/reference/policy-schema.mdx index 9c4e877eaf..9047f47cf8 100644 --- a/docs/reference/policy-schema.mdx +++ b/docs/reference/policy-schema.mdx @@ -30,7 +30,7 @@ network_middlewares: { ... } | `landlock` | object | No | Static | Configures Landlock LSM enforcement behavior. | | `process` | object | No | Static | Sets the user and group the agent process runs as. | | `network_policies` | map | No | Dynamic | Declares which binaries can reach which network endpoints. | -| `network_middlewares` | map | No | Dynamic | Selects ordered HTTP request middleware by destination host. | +| `network_middlewares` | map | No | Dynamic | Attaches ordered middleware by destination host; each implementation's manifest selects its supported HTTP and WebSocket operations. | Static fields are set at sandbox creation time. Changing them requires destroying and recreating the sandbox. Dynamic fields can be updated on a running sandbox with `openshell policy update` for incremental merges or `openshell policy set` for full replacement, and take effect without restarting. @@ -497,7 +497,7 @@ Identifies an executable that is permitted to use the associated endpoints. **Category:** Dynamic -A map of up to 10 middleware configs selected after network and L7 policy admit an HTTP request. Each map key is the stable policy-local config identity. Middleware selection is independent of the network policy entry that admitted the request. Every matching config runs once by ascending `order` before provider credential injection. Order values must be unique across the policy, and runtime selection also enforces the 10-stage maximum. +A map of up to 10 middleware configs selected after network and L7 policy admit an HTTP request or WebSocket upgrade. Each map key is the stable policy-local config identity. Middleware selection is independent of the network policy entry that admitted the traffic. Every matching config runs once by ascending `order` before provider credential injection. WebSocket-capable bindings continue on client text messages after the upgrade. Order values must be unique across the policy, and runtime selection also enforces the 10-stage maximum. ```yaml showLineNumbers={false} network_middlewares: @@ -519,10 +519,10 @@ network_middlewares: | `middleware` | string | Yes | Built-in middleware name or operator-owned registration name. `openshell/` is reserved for built-ins. | | `order` | integer | No | Execution priority. Lower values run first, and values must be unique across the policy. Defaults to `0`; therefore, policies with multiple configs normally specify it explicitly. | | `config` | object | No | Implementation-owned configuration validated by the selected middleware. | -| `on_error` | string | No | `fail_closed` denies the request when the stage fails; `fail_open` skips the failed stage. Defaults to `fail_closed`. | +| `on_error` | string | No | Applies only after an advertised operation binding is selected. `fail_closed` denies the HTTP request or closes the WebSocket when that stage fails; `fail_open` skips a failed HTTP stage or disables a broken WebSocket stage for the rest of that connection. Defaults to `fail_closed`. | | `endpoints` | object | Yes | Host selector with required non-empty `include` and optional `exclude` lists, limited to 32 combined patterns. Exclusions take precedence. | -Host selectors use the same case-insensitive exact and DNS glob semantics as network endpoints: `*` matches exactly one DNS label and `**` matches one or more labels, so `**.example.com` covers subdomains but not `example.com` itself. Brace alternates are rejected at validation. Middleware runs only on HTTP requests the supervisor parses. A fail-closed selector that can cover a `tls: skip` endpoint is rejected because OpenShell cannot inspect that traffic. An all-`fail_open` match may cover the endpoint; the supervisor bypasses the middleware and emits a detection finding. +Host selectors use the same case-insensitive exact and DNS glob semantics as network endpoints: `*` matches exactly one DNS label and `**` matches one or more labels, so `**.example.com` covers subdomains but not `example.com` itself. Brace alternates are rejected at validation. A matching attachment joins only the operation chains its implementation advertises. An HTTP-only attachment may inspect a WebSocket upgrade GET without joining the post-upgrade chain; OpenShell permits the messages and records `binding_not_selected` coverage regardless of `on_error`. WebSocket bindings inspect complete client text messages. Binary messages pass with `unsupported_message_type` coverage for active stages. A fail-closed selector that can cover a `tls: skip` endpoint is rejected because OpenShell cannot inspect that traffic through any operation. An all-`fail_open` match may cover the endpoint; the supervisor bypasses the middleware and emits a detection finding. See [Supervisor Middleware](/extensibility/supervisor-middleware) for registration, failure behavior, body limits, and operational guidance. diff --git a/docs/sandboxes/policies.mdx b/docs/sandboxes/policies.mdx index 4c2e12dacf..49688bf0f7 100644 --- a/docs/sandboxes/policies.mdx +++ b/docs/sandboxes/policies.mdx @@ -62,7 +62,7 @@ network_middlewares: Static sections are locked at sandbox creation. Changing them requires destroying and recreating the sandbox. Dynamic sections can be updated on a running sandbox with `openshell policy update` for incremental merges or `openshell policy set` for full replacement, and take effect without restarting. When a hot reload changes rules on an active HTTP L7 endpoint, existing keep-alive tunnels are closed before forwarding another parsed request. Credential-injection-only HTTP passthrough tunnels use the same reload boundary. Most HTTP clients reconnect automatically, and the next request is evaluated against the current policy. -Raw streams are connection-scoped and outside L7 live-reload guarantees. This includes `tls: skip`, non-HTTP TCP payloads, HTTP upgrades such as WebSocket, and long-lived response streams such as SSE. A reload applies to the next connection or next parsed HTTP request; it does not interrupt an already-forwarded raw stream. Use `protocol: websocket` when policy should stay attached to the RFC 6455 upgrade and client text messages after the allowed upgrade. Add `websocket_credential_rewrite: true` only when the relay should rewrite credential placeholders in client-to-server WebSocket text messages. Add `request_body_credential_rewrite: true` only on inspected REST endpoints that need OpenShell to rewrite placeholders in supported text request bodies. +Raw streams are connection-scoped and outside L7 live-reload guarantees. This includes `tls: skip`, non-HTTP TCP payloads, uninspected HTTP upgrades, and long-lived response streams such as SSE. A reload applies to the next connection or next parsed HTTP request. A parsed WebSocket relay closes with code `1012` when its attached policy generation becomes stale. Use `protocol: websocket` when policy should stay attached to the RFC 6455 upgrade and client text messages after the allowed upgrade. Add `websocket_credential_rewrite: true` only when the relay should rewrite credential placeholders in client-to-server WebSocket text messages. Add `request_body_credential_rewrite: true` only on inspected REST endpoints that need OpenShell to rewrite placeholders in supported text request bodies. | Section | Type | Description | |---|---|---| @@ -70,11 +70,11 @@ Raw streams are connection-scoped and outside L7 live-reload guarantees. This in | `landlock` | Static | Configures Landlock LSM enforcement behavior. Set `compatibility` to `best_effort` (skip individual inaccessible paths while applying remaining rules) or `hard_requirement` (fail if any path is inaccessible or the required kernel ABI is unavailable). Refer to the [Policy Schema Reference](/reference/policy-schema#landlock) for the full behavior table. | | `process` | Static | Optionally overrides the OS-level identity for the agent process. Explicit values must be `sandbox` or numeric UID/GID values in the allowed sandbox range. Docker and Podman may use named identities or non-root system IDs only through per-field OCI `USER` fallback; Kubernetes uses its platform-selected numeric identity. Root identities are always rejected. The agent also runs with seccomp filters that block dangerous system calls. | | `network_policies` | Dynamic | Controls network access for ordinary outbound traffic from the sandbox. Each block has a name, a list of endpoints (host, port, protocol, and optional rules), and a list of binaries allowed to use those endpoints.
Every outbound connection except `https://inference.local` goes through the proxy, which queries the [policy engine](/about/how-it-works#core-components) with the destination and calling binary. A connection is allowed only when both match an entry in the same policy block.
For endpoints with `protocol: rest`, the proxy auto-detects TLS and terminates it so each HTTP request can be checked against that endpoint's `rules` (method and path). For endpoints with `protocol: websocket`, the proxy validates the RFC 6455 upgrade and evaluates `GET` rules for the handshake plus either `WEBSOCKET_TEXT` rules for raw client text messages or GraphQL operation rules for GraphQL-over-WebSocket messages. Set `websocket_credential_rewrite: true` only when a WebSocket or REST compatibility endpoint must keep placeholder credentials in sandbox-owned text frames and resolve them at the OpenShell relay boundary.
Endpoints without `protocol` allow the TCP stream through without inspecting payloads.
If no endpoint matches, the connection is denied. Configure managed inference separately through [Inference Routing](/sandboxes/inference-routing). | -| `network_middlewares` | Dynamic | Declares keyed HTTP request middleware configs. After network and L7 policy admit a request, OpenShell matches each config's host selectors independently and runs matching entries by their unique ascending `order` before credential injection. | +| `network_middlewares` | Dynamic | Declares keyed HTTP and WebSocket middleware configs. After network and L7 policy admit a request or upgrade, OpenShell matches each config's host selectors independently and runs matching entries by their unique ascending `order` before credential injection. WebSocket-capable entries continue on complete client text messages. | ## Supervisor Middleware -Supervisor middleware can inspect, deny, or replace admitted HTTP request bodies before provider credentials are injected. Middleware selection is independent of the `network_policies` rule that admitted the request: each keyed `network_middlewares` entry matches the destination host through `endpoints.include` and `endpoints.exclude`. +Supervisor middleware can inspect, deny, or replace admitted HTTP request bodies and client WebSocket text messages before provider credentials are injected. Middleware selection is independent of the `network_policies` rule that admitted the traffic: each keyed `network_middlewares` entry matches the destination host through `endpoints.include` and `endpoints.exclude`. ```yaml network_middlewares: @@ -92,9 +92,9 @@ network_middlewares: Matching entries run once each by ascending `order`; lower values run first, and duplicate order values are rejected. The default order is `0`, so policies with multiple entries normally set it explicitly. Map keys are structurally unique. An optional `name` provides a human-readable label, defaults to the map key, and does not change the key used as the config identity. Different keys may use the same implementation and run as distinct stages. `exclude` takes precedence over `include`. -`openshell/regex` is an example built into the supervisor. It applies fixed regular expressions as a best-effort UTF-8 text transformation, without guarantees that sensitive values will be detected or fully removed. Custom expressions are not configurable yet. Operator-run middleware must be registered by name before a policy can reference it. The gateway validates implementation-owned config before accepting the policy. +`openshell/regex` is an example built into the supervisor. It applies fixed regular expressions to UTF-8 HTTP request bodies and complete client-to-upstream WebSocket text messages before credential injection. The initial pattern recognizes `sk-` tokens. This is a best-effort text transformation without guarantees that sensitive values will be detected or fully removed; it does not inspect binary or upstream-to-client WebSocket messages. Custom expressions are not configurable yet. Operator-run middleware must be registered by name before a policy can reference it. The gateway validates implementation-owned config before accepting the policy. -`on_error` defaults to `fail_closed`. Use `fail_open` only when skipping a failed middleware is acceptable. Middleware applies only to HTTP traffic the supervisor can parse and inspect. Policy validation rejects a fail-closed selector that can cover a `tls: skip` endpoint. An all-`fail_open` match may cover the endpoint; the supervisor bypasses the middleware and emits a detection finding. +`on_error` defaults to `fail_closed`. Use `fail_open` only when skipping a selected stage that fails is acceptable. For WebSocket streams, a broken fail-open stage is disabled for the rest of that connection and OpenShell emits a state-change finding. A host-matched attachment joins only operation chains advertised by its implementation. An HTTP-only attachment may inspect the WebSocket upgrade GET, but post-upgrade messages pass with informational `binding_not_selected` coverage under either error mode. Binary messages also pass without middleware inspection and produce `unsupported_message_type` coverage for active WebSocket stages. Upstream-to-client messages remain uninspected. Policy validation rejects a fail-closed selector that can cover a `tls: skip` endpoint. An all-`fail_open` match may cover the endpoint; the supervisor bypasses the middleware and emits a detection finding. See [Supervisor Middleware](/extensibility/supervisor-middleware) for registration, chain ordering, body limits, failure behavior, and operations. diff --git a/examples/supervisor-middleware-content-guard/Cargo.lock b/examples/supervisor-middleware-content-guard/Cargo.lock index f397d82dee..dc3eff06c4 100644 --- a/examples/supervisor-middleware-content-guard/Cargo.lock +++ b/examples/supervisor-middleware-content-guard/Cargo.lock @@ -864,6 +864,7 @@ dependencies = [ "serde_json", "thiserror", "tokio", + "tokio-stream", "tonic", "tonic-prost", "tonic-prost-build", @@ -879,6 +880,7 @@ dependencies = [ "openshell-core", "prost-types", "tokio", + "tokio-stream", "tonic", ] diff --git a/examples/supervisor-middleware-content-guard/Cargo.toml b/examples/supervisor-middleware-content-guard/Cargo.toml index eceaeac509..ef0d1f47c3 100644 --- a/examples/supervisor-middleware-content-guard/Cargo.toml +++ b/examples/supervisor-middleware-content-guard/Cargo.toml @@ -17,6 +17,7 @@ clap = { version = "4.5", features = ["derive"] } openshell-core = { path = "../../crates/openshell-core", default-features = false } prost-types = "0.14" tokio = { version = "1.43", features = ["macros", "rt-multi-thread"] } +tokio-stream = "0.1" tonic = { version = "0.14", features = ["transport"] } [[bin]] diff --git a/examples/supervisor-middleware-content-guard/README.md b/examples/supervisor-middleware-content-guard/README.md index 10f76effa7..53eda94210 100644 --- a/examples/supervisor-middleware-content-guard/README.md +++ b/examples/supervisor-middleware-content-guard/README.md @@ -8,10 +8,10 @@ SPDX-License-Identifier: Apache-2.0 > [!WARNING] > Supervisor middleware is a research preview. Its policy and service contracts may change without compatibility guarantees. Use it only to prototype and evaluate middleware integrations. -This example implements an operator-run supervisor middleware service. It scans UTF-8 HTTP request bodies for configured literal strings, then either replaces every match or denies the request. Findings report only aggregate counts and never include configured terms or request content. +This example implements an operator-run supervisor middleware service. It scans UTF-8 HTTP request bodies and complete client-to-upstream WebSocket text messages for configured literal strings, then either replaces every match or denies the request or message. Findings report only aggregate counts and never include configured terms or inspected content. > [!WARNING] -> This intentionally simple implementation demonstrates the supervisor middleware service contract. It is not a complete or reliable content guard and must not be used as a security control. It handles only UTF-8 request bodies and case-sensitive literal terms, merges overlapping literal match ranges before redaction, and does not address the encodings, transformations, normalization, streaming, or adversarial inputs that a production content guard must handle. +> This intentionally simple implementation demonstrates the supervisor middleware service contract. It is not a complete or reliable content guard and must not be used as a security control. It handles only UTF-8 HTTP request bodies and WebSocket text messages with case-sensitive literal terms, merges overlapping literal match ranges before redaction, and does not address encodings, transformations, normalization, binary WebSocket messages, upstream-to-client messages, or adversarial inputs that a production content guard must handle. ## Prerequisites @@ -54,7 +54,7 @@ Add the service registration to your local gateway TOML: [[openshell.supervisor.middleware]] name = "content-guard-example" grpc_endpoint = "http://host.openshell.internal:50051" -max_body_bytes = 262144 +max_payload_bytes = 262144 timeout = "500ms" ``` @@ -84,6 +84,12 @@ curl -sS https://httpbin.org/anything \ The echoed JSON body contains `[FILTERED]` instead of the configured term. +## WebSocket behavior + +For a selected WebSocket upgrade, the service accepts preflight, waits for the session-start notification, and evaluates each complete client-to-upstream text message. Redact mode returns a replacement message, while deny mode returns `content_match` and OpenShell closes the session according to middleware policy. Session-start and session-end events are notifications and do not produce results. + +The service advertises a 256 KiB limit for complete WebSocket text messages. OpenShell does not send binary messages, control frames, or upstream-to-client messages to this binding. The smoke script exercises the HTTP path; the example's unit tests cover the WebSocket lifecycle and both redact and deny results. + ## Configuration | Field | Required | Description | @@ -101,4 +107,4 @@ config: - prototype-secret ``` -The implementation supports only `HttpRequest/pre_credentials`, advertises a 256 KiB body limit, and inherits the service-wide RPC timeout. The gateway registration may set a smaller body limit. A binding can advertise a shorter timeout, but it cannot extend the operator-configured timeout. +The implementation supports `HTTP_REQUEST/PRE_CREDENTIALS` and `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`, advertises a 256 KiB limit for each operation, and inherits the service-wide RPC timeout. The gateway registration's `max_payload_bytes` may set a smaller shared limit. A binding can advertise a shorter timeout, but it cannot extend the operator-configured timeout. diff --git a/examples/supervisor-middleware-content-guard/smoke.sh b/examples/supervisor-middleware-content-guard/smoke.sh index aab026e983..509ffd403d 100755 --- a/examples/supervisor-middleware-content-guard/smoke.sh +++ b/examples/supervisor-middleware-content-guard/smoke.sh @@ -214,7 +214,7 @@ ttl_secs = 0 [[openshell.supervisor.middleware]] name = "content-guard-example" grpc_endpoint = "http://$SERVICE_HOST:$MIDDLEWARE_PORT" -max_body_bytes = 262144 +max_payload_bytes = 262144 timeout = "500ms" EOF } diff --git a/examples/supervisor-middleware-content-guard/src/main.rs b/examples/supervisor-middleware-content-guard/src/main.rs index cf36a4cb0c..cea1a7c270 100644 --- a/examples/supervisor-middleware-content-guard/src/main.rs +++ b/examples/supervisor-middleware-content-guard/src/main.rs @@ -6,23 +6,27 @@ use std::net::SocketAddr; use std::ops::Range; use clap::Parser; +use openshell_core::middleware::WebSocketResponseStream; use openshell_core::proto::middleware::v1::supervisor_middleware_server::{ SupervisorMiddleware, SupervisorMiddlewareServer, }; use openshell_core::proto::{ Decision, Finding, HttpRequestEvaluation, HttpRequestResult, MiddlewareBinding, MiddlewareManifest, SupervisorMiddlewareOperation, SupervisorMiddlewarePhase, - ValidateConfigRequest, ValidateConfigResponse, + ValidateConfigRequest, ValidateConfigResponse, WebSocketMessage, WebSocketMessageResult, + WebSocketMessageType, WebSocketPreflightAction, WebSocketPreflightDecision, + WebSocketSessionEvent, WebSocketSessionEventResult, web_socket_session_event, + web_socket_session_event_result, }; use prost_types::Struct; use prost_types::value::Kind; +use tokio_stream::{Stream, StreamExt}; use tonic::transport::Server; use tonic::{Request, Response, Status}; const MANIFEST_NAME: &str = "example/content-guard-service"; -const OPERATION: SupervisorMiddlewareOperation = SupervisorMiddlewareOperation::HttpRequest; const PHASE: SupervisorMiddlewarePhase = SupervisorMiddlewarePhase::PreCredentials; -const MAX_BODY_BYTES: u64 = 256 * 1024; +const MAX_PAYLOAD_BYTES: u64 = 256 * 1024; const DEFAULT_REPLACEMENT: &str = "[REDACTED]"; #[derive(Debug, Parser)] @@ -113,8 +117,107 @@ fn optional_string_field<'a>(config: &'a Struct, name: &str) -> Result(mut events: S) -> WebSocketResponseStream + where + S: Stream> + Send + Unpin + 'static, + { + let (results_tx, results_rx) = tokio::sync::mpsc::channel(4); + tokio::spawn(async move { + let mut config = None; + let mut started = false; + let mut sequence_lower_bound = Some(1_u64); + + while let Some(event) = events.next().await { + let event = match event { + Ok(event) => event, + Err(error) => { + let _ = results_tx.send(Err(error)).await; + break; + } + }; + let result = match event.event { + Some(web_socket_session_event::Event::Preflight(preflight)) + if config.is_none() && !started => + { + if let Err(error) = validate_phase(preflight.phase) { + Err(Status::invalid_argument(error)) + } else { + match GuardConfig::parse(preflight.config.as_ref()) { + Ok(selected_config) => { + config = Some(selected_config); + Ok(Some(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::PreflightDecision( + WebSocketPreflightDecision { + action: WebSocketPreflightAction::Inspect + as i32, + ..Default::default() + }, + ), + ), + })) + } + Err(error) => Err(Status::invalid_argument(error)), + } + } + } + Some(web_socket_session_event::Event::SessionStart(_)) + if config.is_some() && !started => + { + started = true; + Ok(None) + } + Some(web_socket_session_event::Event::Message(message)) if started => { + if let Err(error) = advance_sequence_lower_bound( + &mut sequence_lower_bound, + message.sequence, + ) { + Err(error) + } else { + let selected_config = + config.as_ref().expect("started stream has config"); + evaluate_websocket_message(selected_config, &message).map(|result| { + Some(WebSocketSessionEventResult { + result: Some( + web_socket_session_event_result::Result::MessageResult( + result, + ), + ), + }) + }) + } + } + Some(web_socket_session_event::Event::SessionEnd(_)) if config.is_some() => { + break; + } + _ => Err(Status::failed_precondition( + "invalid content guard WebSocket session lifecycle", + )), + }; + + match result { + Ok(Some(result)) => { + if results_tx.send(Ok(result)).await.is_err() { + break; + } + } + Ok(None) => {} + Err(error) => { + let _ = results_tx.send(Err(error)).await; + break; + } + } + } + }); + Box::pin(tokio_stream::wrappers::ReceiverStream::new(results_rx)) + } +} + #[tonic::async_trait] impl SupervisorMiddleware for ContentGuard { + type EvaluateWebSocketSessionStream = WebSocketResponseStream; + async fn describe( &self, _request: Request<()>, @@ -122,12 +225,20 @@ impl SupervisorMiddleware for ContentGuard { Ok(Response::new(MiddlewareManifest { name: MANIFEST_NAME.into(), service_version: env!("CARGO_PKG_VERSION").into(), - bindings: vec![MiddlewareBinding { - operation: OPERATION as i32, - phase: PHASE as i32, - max_body_bytes: MAX_BODY_BYTES, - timeout: String::new(), - }], + bindings: vec![ + MiddlewareBinding { + operation: SupervisorMiddlewareOperation::HttpRequest as i32, + phase: PHASE as i32, + max_payload_bytes: MAX_PAYLOAD_BYTES, + timeout: String::new(), + }, + MiddlewareBinding { + operation: SupervisorMiddlewareOperation::WebsocketMessage as i32, + phase: PHASE as i32, + max_payload_bytes: MAX_PAYLOAD_BYTES, + timeout: String::new(), + }, + ], })) } @@ -161,6 +272,13 @@ impl SupervisorMiddleware for ContentGuard { .map_err(|_| Status::invalid_argument("content guard requires a UTF-8 body"))?; Ok(Response::new(evaluate(&config, &body))) } + + async fn evaluate_web_socket_session( + &self, + request: Request>, + ) -> Result, Status> { + Ok(Response::new(Self::websocket_stream(request.into_inner()))) + } } fn validate_phase(phase: i32) -> Result<(), String> { @@ -209,7 +327,7 @@ fn evaluate(config: &GuardConfig, body: &str) -> HttpRequestResult { }, Mode::Deny => HttpRequestResult { decision: Decision::Deny as i32, - reason: "request body matched configured content".into(), + reason: "payload matched configured content".into(), body: Vec::new(), has_body: false, header_mutations: Vec::new(), @@ -220,6 +338,56 @@ fn evaluate(config: &GuardConfig, body: &str) -> HttpRequestResult { } } +fn evaluate_websocket_message( + config: &GuardConfig, + message: &WebSocketMessage, +) -> Result { + if message.message_type != WebSocketMessageType::Text as i32 { + return Err(Status::invalid_argument( + "content guard supports only WebSocket text messages", + )); + } + let payload_bytes = u64::try_from(message.payload.len()).map_err(|_| { + Status::invalid_argument("WebSocket text message length is not representable") + })?; + if payload_bytes > MAX_PAYLOAD_BYTES { + return Err(Status::invalid_argument(format!( + "WebSocket text message exceeds {MAX_PAYLOAD_BYTES} bytes" + ))); + } + let payload = std::str::from_utf8(&message.payload) + .map_err(|_| Status::invalid_argument("content guard requires a UTF-8 message"))?; + let result = evaluate(config, payload); + Ok(WebSocketMessageResult { + sequence: message.sequence, + decision: result.decision, + replacement: result.body, + has_replacement: result.has_body, + reason: result.reason, + findings: result.findings, + metadata: result.metadata, + reason_code: result.reason_code, + }) +} + +fn advance_sequence_lower_bound( + lower_bound: &mut Option, + sequence: u64, +) -> Result<(), Status> { + let Some(current_lower_bound) = *lower_bound else { + return Err(Status::invalid_argument( + "WebSocket message sequence must be strictly increasing", + )); + }; + if sequence < current_lower_bound { + return Err(Status::invalid_argument( + "WebSocket message sequence must be strictly increasing", + )); + } + *lower_bound = sequence.checked_add(1); + Ok(()) +} + fn find_match_ranges(body: &str, terms: &[String]) -> (Vec>, u32, u32) { let mut ranges = Vec::new(); let mut match_count = 0_u32; @@ -304,6 +472,7 @@ async fn main() -> Result<(), Box> { #[cfg(test)] mod tests { use super::*; + use openshell_core::proto::{WebSocketPreflight, WebSocketSessionEnd, WebSocketSessionStart}; use prost_types::{ListValue, Value}; use std::collections::BTreeMap; @@ -331,6 +500,106 @@ mod tests { Struct { fields } } + fn event(event: web_socket_session_event::Event) -> Result { + Ok(WebSocketSessionEvent { event: Some(event) }) + } + + #[tokio::test] + async fn manifest_advertises_http_and_websocket_bindings() { + let manifest = SupervisorMiddleware::describe(&ContentGuard, Request::new(())) + .await + .expect("describe") + .into_inner(); + + assert_eq!(manifest.bindings.len(), 2); + assert_eq!( + manifest.bindings[0].operation, + SupervisorMiddlewareOperation::HttpRequest as i32 + ); + assert_eq!(manifest.bindings[0].max_payload_bytes, MAX_PAYLOAD_BYTES); + assert_eq!( + manifest.bindings[1].operation, + SupervisorMiddlewareOperation::WebsocketMessage as i32 + ); + assert_eq!(manifest.bindings[1].max_payload_bytes, MAX_PAYLOAD_BYTES); + } + + #[tokio::test] + async fn websocket_stream_redacts_text_messages() { + let events = tokio_stream::iter([ + event(web_socket_session_event::Event::Preflight( + WebSocketPreflight { + phase: PHASE as i32, + config: Some(config("redact", &["prototype-secret"], Some("[FILTERED]"))), + ..Default::default() + }, + )), + event(web_socket_session_event::Event::SessionStart( + WebSocketSessionStart::default(), + )), + event(web_socket_session_event::Event::Message(WebSocketMessage { + sequence: 1, + message_type: WebSocketMessageType::Text as i32, + payload: b"contains prototype-secret".to_vec(), + })), + event(web_socket_session_event::Event::SessionEnd( + WebSocketSessionEnd::default(), + )), + ]); + let mut results = ContentGuard::websocket_stream(events); + + let preflight = results + .next() + .await + .expect("preflight result") + .expect("valid preflight result"); + assert!(matches!( + preflight.result, + Some(web_socket_session_event_result::Result::PreflightDecision( + WebSocketPreflightDecision { action, .. } + )) if action == WebSocketPreflightAction::Inspect as i32 + )); + + let message = results + .next() + .await + .expect("message result") + .expect("valid message result"); + let Some(web_socket_session_event_result::Result::MessageResult(message)) = message.result + else { + panic!("expected message result"); + }; + assert_eq!(message.decision, Decision::Allow as i32); + assert!(message.has_replacement); + assert_eq!( + String::from_utf8(message.replacement).expect("UTF-8 replacement"), + "contains [FILTERED]" + ); + assert_eq!(message.findings[0].count, 1); + assert!(results.next().await.is_none()); + } + + #[test] + fn websocket_deny_preserves_safe_diagnostics() { + let config = GuardConfig::parse(Some(&config("deny", &["prototype-secret"], None))) + .expect("valid config"); + let result = evaluate_websocket_message( + &config, + &WebSocketMessage { + sequence: 7, + message_type: WebSocketMessageType::Text as i32, + payload: b"contains prototype-secret".to_vec(), + }, + ) + .expect("message result"); + + assert_eq!(result.sequence, 7); + assert_eq!(result.decision, Decision::Deny as i32); + assert_eq!(result.reason_code, "content_match"); + assert!(!result.reason.contains("prototype-secret")); + assert!(!result.has_replacement); + } + #[test] fn redact_replaces_every_configured_match() { let config = GuardConfig::parse(Some(&config( diff --git a/proto/sandbox.proto b/proto/sandbox.proto index 16b3ca998d..0aea3c5b69 100644 --- a/proto/sandbox.proto +++ b/proto/sandbox.proto @@ -376,8 +376,9 @@ message SupervisorMiddlewareService { string name = 1; // gRPC endpoint reachable from the sandbox supervisor. string grpc_endpoint = 2; - // Operator-owned body limit applied to every binding exposed by the service. - uint64 max_body_bytes = 3; + // Operator-owned logical payload limit applied to every binding exposed by + // the service. This caps HTTP bodies and complete WebSocket messages. + uint64 max_payload_bytes = 3; // Default RPC timeout for this service. Empty uses the platform default of // 500ms. Values use an integer with an `ms` or `s` suffix and must be // between 10ms and 30s. diff --git a/proto/supervisor_middleware.proto b/proto/supervisor_middleware.proto index dbde411c9f..1460841b1f 100644 --- a/proto/supervisor_middleware.proto +++ b/proto/supervisor_middleware.proto @@ -9,7 +9,8 @@ import "google/protobuf/empty.proto"; import "google/protobuf/struct.proto"; // SupervisorMiddleware lets an operator-run service inspect and transform -// sandbox HTTP egress before OpenShell injects credentials. +// sandbox HTTP requests and client WebSocket text messages before OpenShell +// injects credentials. service SupervisorMiddleware { // Describe returns the service manifest and declared bindings. rpc Describe(google.protobuf.Empty) returns (MiddlewareManifest); @@ -20,6 +21,16 @@ service SupervisorMiddleware { // EvaluateHttpRequest returns an allow, deny, or mutation decision for one // buffered HTTP request. rpc EvaluateHttpRequest(HttpRequestEvaluation) returns (HttpRequestResult); + + // EvaluateWebSocketSession opens one ordered, phase-specific stream for a + // single middleware stage and WebSocket upgrade attempt. The current + // implementation supports client-to-upstream text messages at + // PRE_CREDENTIALS; PRE_RETURN is reserved for upstream-to-client messages. + // A request may go unanswered when the session terminates. For every opened + // stage stream, OpenShell attempts at most one session_end before closing the + // stream when its transport is still writable. + rpc EvaluateWebSocketSession(stream WebSocketSessionEvent) + returns (stream WebSocketSessionEventResult); } // MiddlewareManifest describes one middleware service and the bindings it @@ -38,12 +49,16 @@ message MiddlewareManifest { // MiddlewareBinding declares one operation and phase supported by a service. message MiddlewareBinding { - // Supported operation. V1 supports HTTP_REQUEST. + // Supported operation. SupervisorMiddlewareOperation operation = 1; - // Supported evaluation phase. V1 supports PRE_CREDENTIALS. + // Supported evaluation phase. PR 1 supports PRE_CREDENTIALS. PRE_RETURN is + // reserved for the return-path follow-up and is rejected by current + // manifest validation. SupervisorMiddlewarePhase phase = 2; - // Maximum request or replacement body this binding can process. - uint64 max_body_bytes = 3; + // Maximum logical payload or replacement this binding can process. For + // HTTP_REQUEST this is the request body; for WEBSOCKET_MESSAGE this is one + // complete message. Required for every payload-bearing operation. + uint64 max_payload_bytes = 3; // Optional binding-specific RPC timeout. Empty uses the operator-configured // service timeout, or the 500ms platform default when that is also omitted. // A non-empty value may shorten but cannot extend the operator timeout. @@ -104,12 +119,148 @@ message HttpHeader { enum SupervisorMiddlewareOperation { SUPERVISOR_MIDDLEWARE_OPERATION_UNSPECIFIED = 0; SUPERVISOR_MIDDLEWARE_OPERATION_HTTP_REQUEST = 1; + SUPERVISOR_MIDDLEWARE_OPERATION_WEBSOCKET_MESSAGE = 2; } // Ordered phase within a supervisor operation. enum SupervisorMiddlewarePhase { SUPERVISOR_MIDDLEWARE_PHASE_UNSPECIFIED = 0; SUPERVISOR_MIDDLEWARE_PHASE_PRE_CREDENTIALS = 1; + SUPERVISOR_MIDDLEWARE_PHASE_PRE_RETURN = 2; +} + +// Logical WebSocket message type. Raw frame mechanics are never exposed. +enum WebSocketMessageType { + WEB_SOCKET_MESSAGE_TYPE_UNSPECIFIED = 0; + WEB_SOCKET_MESSAGE_TYPE_TEXT = 1; + WEB_SOCKET_MESSAGE_TYPE_BINARY = 2; +} + +// Why OpenShell is ending a middleware stream. +enum WebSocketSessionEndReason { + WEB_SOCKET_SESSION_END_REASON_UNSPECIFIED = 0; + WEB_SOCKET_SESSION_END_REASON_NORMAL_CLOSE = 1; + WEB_SOCKET_SESSION_END_REASON_PEER_DISCONNECT = 2; + WEB_SOCKET_SESSION_END_REASON_POLICY_RELOAD = 3; + WEB_SOCKET_SESSION_END_REASON_MIDDLEWARE_DENIAL = 4; + WEB_SOCKET_SESSION_END_REASON_MIDDLEWARE_FAILURE = 5; + WEB_SOCKET_SESSION_END_REASON_PROTOCOL_ERROR = 6; + WEB_SOCKET_SESSION_END_REASON_CANCELLATION = 7; + WEB_SOCKET_SESSION_END_REASON_UPSTREAM_REJECTED = 8; + WEB_SOCKET_SESSION_END_REASON_POLICY_DENIAL = 9; +} + +// WebSocketSessionEvent is one ordered event in a stage-local stream. +// Message sequence numbers identify logical messages session-wide. A stage +// receives a strictly increasing subset of those numbers; gaps are valid when +// session messages are not delivered to that stage. +message WebSocketSessionEvent { + oneof event { + WebSocketPreflight preflight = 1; + WebSocketSessionStart session_start = 2; + WebSocketMessage message = 3; + WebSocketSessionEnd session_end = 4; + } +} + +// WebSocketPreflight lets a service decline this upgrade before OpenShell +// contacts upstream. It deliberately excludes query data, arbitrary request +// headers, and message payloads. +message WebSocketPreflight { + string session_id = 1; + SupervisorMiddlewarePhase phase = 2; + RequestContext context = 3; + // Admitted HTTP WebSocket-upgrade target. The method is GET, query is always + // empty, and path never includes a query string. + HttpRequestTarget target = 4; + repeated string requested_subprotocols = 5; + // Built-in middleware name or operator-owned registration name. + string middleware_name = 6; + google.protobuf.Struct config = 7; +} + +// WebSocketSessionStart reports bounded metadata known only after the +// upstream 101 response validates. Empty selected_subprotocol means none. +message WebSocketSessionStart { + string selected_subprotocol = 1; +} + +// WebSocketMessage contains one complete reconstructed logical message. +message WebSocketMessage { + // Session-global sequence starting at 1. Values delivered to one stage must + // strictly increase but need not be contiguous. Reject zero, duplicates, and + // regressions; accept gaps. + uint64 sequence = 1; + WebSocketMessageType message_type = 2; + // Limited to 4 MiB by the platform and the binding-specific cap. + bytes payload = 3; +} + +// WebSocketSessionEnd is OpenShell's best-effort terminal notification for one +// opened stage stream. A stage receives at most one such notification. +message WebSocketSessionEnd { + WebSocketSessionEndReason reason = 1; +} + +// WebSocketPreflightAction is the service's one-time scoping decision. +enum WebSocketPreflightAction { + // Invalid response value handled according to the policy failure mode. + WEB_SOCKET_PREFLIGHT_ACTION_UNSPECIFIED = 0; + // Inspect this session after the upstream accepts the upgrade. + WEB_SOCKET_PREFLIGHT_ACTION_INSPECT = 1; + // Voluntarily decline inspection without denying the upgrade. This is a + // successful decision and does not engage on_error. + WEB_SOCKET_PREFLIGHT_ACTION_SKIP = 2; + // Authoritatively deny the upgrade before upstream contact. This is a + // successful decision and is enforced regardless of on_error. + WEB_SOCKET_PREFLIGHT_ACTION_DENY = 3; +} + +message WebSocketPreflightDecision { + WebSocketPreflightAction action = 1; + // Free-form service diagnostic. OpenShell never exposes this to the + // workload or security logs. Limited to 4 KiB before discarding. + string reason = 2; + // Audit-safe findings produced during preflight. At most 32 findings of at + // most 4 KiB encoded each are accepted. + repeated Finding findings = 3; + // Non-secret service-defined metadata included in diagnostics. At most 64 + // entries and 32 KiB of combined key/value data are accepted. + map metadata = 4; + // Optional stable machine-readable code for a deny decision. Because + // preflight runs before the HTTP upgrade completes, OpenShell may return + // this code to the requester. Codes follow the same format and 64-byte + // maximum as HttpRequestResult.reason_code. + string reason_code = 5; +} + +// WebSocketMessageResult contains the decision and optional replacement for +// one message. A replacement preserves the input message type; text +// replacements must be valid UTF-8. +message WebSocketMessageResult { + // Must exactly match the sequence of the corresponding WebSocketMessage. + uint64 sequence = 1; + Decision decision = 2; + bytes replacement = 3; + // True when replacement should be used, including an empty replacement. + bool has_replacement = 4; + // Free-form service diagnostic. OpenShell never exposes this to the + // workload or security logs. Limited to 4 KiB before discarding. + string reason = 5; + repeated Finding findings = 6; + map metadata = 7; + // Optional stable machine-readable code for OCSF only. Unlike the HTTP + // reason_code, this value is never put in a WebSocket close frame. + string reason_code = 8; +} + +// WebSocketSessionEventResult is an evaluation result for a preflight or message +// event. Session start and end events do not produce results. +message WebSocketSessionEventResult { + oneof result { + WebSocketPreflightDecision preflight_decision = 1; + WebSocketMessageResult message_result = 2; + } } // RequestContext identifies the sandbox request being evaluated. @@ -124,7 +275,7 @@ message RequestContext { // HttpRequestTarget describes the admitted HTTP destination and request target. message HttpRequestTarget { - // Request scheme, such as "http" or "https". + // Request scheme, such as "http", "https", "ws", or "wss". string scheme = 1; // Destination hostname selected by network policy. string host = 2; @@ -148,13 +299,16 @@ message Process { repeated string ancestors = 3; } -// Decision controls whether OpenShell continues processing the request. +// Decision controls whether OpenShell continues processing the current +// evaluation unit. enum Decision { // Invalid response value handled according to the policy failure mode. DECISION_UNSPECIFIED = 0; - // Continue processing the request and apply any returned mutations. + // Continue processing the current request or message and apply any returned + // mutations. DECISION_ALLOW = 1; - // Deny the request before credentials are injected or data is sent upstream. + // Reject the current request or message. The operation-specific result + // defines the enclosing protocol behavior. DECISION_DENY = 2; } diff --git a/rfc/0009-supervisor-middleware/README.md b/rfc/0009-supervisor-middleware/README.md index f43e05ed79..f88c1e91e4 100644 --- a/rfc/0009-supervisor-middleware/README.md +++ b/rfc/0009-supervisor-middleware/README.md @@ -13,6 +13,13 @@ links: # RFC 0009 - Supervisor Middleware +## Revision history + +| Date | References | Change | +|------|------------|--------| +| 2026-07-17 | [#2010](https://github.com/NVIDIA/OpenShell/issues/2010) | Added unary HTTP request middleware with built-in and operator-run services. | +| 2026-07-28 | [#2428](https://github.com/NVIDIA/OpenShell/issues/2428) | Added WebSocket preflight and text-message evaluation, and aligned the middleware API names, limits, and diagnostics. | + ## Summary This RFC proposes the introduction of supervisor middleware: a supervisor-side extension system for hooks that can inspect, transform, block, and annotate supervisor-managed operations at specific operation phases. The first hook family is supervisor egress middleware for outbound sandbox HTTP requests, but the framework is intentionally named and shaped so later supervisor hooks can cover other protocols or supervisor operations without renaming the feature. @@ -134,21 +141,26 @@ This ordering is deliberate: The hook operates on a parsed HTTP request, so it runs wherever OpenShell can parse one. The supervisor proxy TLS-terminates and HTTP-parses every egress connection that is not marked `tls: skip` and is not opaque, non-HTTP traffic, so the hook fires on those requests regardless of whether the endpoint also declares a `protocol`. Declaring a `protocol` additionally subjects the request to L7 Rego policy; an endpoint without one is still terminated and parsed and the middleware hook runs on it. The only traffic the hook cannot inspect is traffic OpenShell never parses: `tls: skip` endpoints and opaque TCP or TLS passthrough. Policy validation rejects any middleware selector whose possible hosts overlap an endpoint configured with `tls: skip`, so selector-based middleware cannot be silently bypassed by an unparsed path. -If matching middleware exists but a request becomes uninspectable at runtime, OpenShell examines the whole matching chain. If any stage is `fail_closed`, the request is denied. If every matching stage is `fail_open`, OpenShell relays the request and emits a bypass `DetectionFinding`. This chain-level rule prevents one permissive stage from overriding a required stage. +> **Update in PR #2477 - WebSocket middleware:** The following text extends the original HTTP-only scope. It adds operation-specific selection, client WebSocket text-message inspection, and explicit coverage for traffic that an attached middleware cannot inspect. + +If a selected operation chain becomes uninspectable at runtime, OpenShell examines that chain. If any selected stage is `fail_closed`, the request is denied. If every selected stage is `fail_open`, OpenShell relays the request and emits a bypass `DetectionFinding`. This chain-level rule prevents one permissive selected stage from overriding a required stage. -WebSocket sits on this boundary. The upgrade request is a normal HTTP/1.1 request that the hook can inspect, allow, or deny, but the frames exchanged after the connection upgrades are out of scope for v1 (a later `message.*` hook may cover them). To keep the v1 boundary unambiguous: +Attachment and operation selection are separate. A destination host selector attaches a policy config, then the implementation manifest decides whether that config participates in `HTTP_REQUEST/PRE_CREDENTIALS`, `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`, or both. The absence of an operation binding is a declared capability boundary rather than a middleware failure, so `on_error` does not apply. OpenShell records informational coverage when an attached config does not join the WebSocket chain. + +WebSocket sits on this boundary. The upgrade request is a normal HTTP/1.1 request that an HTTP binding can inspect, allow, or deny. A separate V1 operation covers complete client-to-upstream text messages after upgrade. Binary messages, control frames, and upstream-to-client messages remain outside that operation. To keep the v1 boundary unambiguous: **In scope for v1:** - Inspectable HTTP/1.x requests that OpenShell terminates and parses, after L4 and SSRF admit them (and L7 policy too, where the endpoint declares a `protocol`). - WebSocket upgrade (handshake) requests - the HTTP request that initiates the upgrade. +- Complete client-to-upstream WebSocket text messages for implementations that advertise `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`. - Bounded request bodies: a `Content-Length` or bounded chunked body OpenShell can buffer within the applicable chain cap. - Safe metadata output for later routing or audit. **Out of scope for v1:** - HTTP/2 and HTTP/3. The proxy's TLS termination pins ALPN to `http/1.1` today, so these are not introspected. -- Post-upgrade WebSocket frames, and response-body scanning. +- Binary and control WebSocket messages, upstream-to-client WebSocket messages, and response-body scanning. - Opaque TCP streams and endpoints with `tls: skip`. - Unbounded streaming uploads or full-duplex request processing. - Multipart or compressed body semantics, unless a selected service's manifest and policy explicitly support them within the size limits. @@ -191,7 +203,7 @@ message MiddlewareBinding { string id = 1; // service-owned stable ID SupervisorMiddlewareOperation operation = 2; SupervisorMiddlewarePhase phase = 3; - uint64 max_body_bytes = 4; + uint64 max_payload_bytes = 4; // one logical request or message payload string timeout = 5; // optional binding-specific RPC timeout } @@ -276,9 +288,15 @@ The evaluation and result are shaped so middleware composes cleanly in a chain. Headers use a repeated representation so duplicate lines and wire order survive evaluation and chaining. Before an external call, OpenShell omits credential-bearing, routing, framing, hop-by-hop, and `Connection`-nominated headers. A result may return ordered writes and removals. Writes support append, overwrite, and skip modes but may target only the `x-openshell-middleware-*` namespace. Removals may target other headers visible to middleware, except credential-bearing, routing, framing, hop-by-hop, and `Connection`-nominated headers. Header values containing control characters are invalid. OpenShell validates and applies a stage's mutations atomically. If any mutation is invalid, none are applied and the stage follows its configured `on_error` behavior. -The interface is gRPC. The protobuf package `openshell.middleware.v1` is the protocol version boundary, so manifests and evaluation messages do not repeat an API-version string. The hot-path v1 RPC is unary: the supervisor buffers the bounded body, sends one `HttpRequestEvaluation`, and receives one `HttpRequestResult`. Streaming is deliberately not baked into `EvaluateHttpRequest`; if OpenShell later needs chunked or incremental processing, it should add a separate operation-specific method rather than changing the v1 method cardinality. Possible extensions are collected in the [protocol-extensions appendix](appendices/protocol-extensions.md). Built-in middleware uses the same logical contract in-process with no network hop. +> **Update in PR #2477 - WebSocket middleware:** The following contract text adds the bidirectional `EvaluateWebSocketSession` RPC, WebSocket preflight, message limits, and the WebSocket binding for the built-in regex middleware. The unary HTTP contract does not change. + +The interface is gRPC. The protobuf package `openshell.middleware.v1` is the protocol version boundary, so manifests and evaluation messages do not repeat an API-version string. HTTP evaluation remains unary: the supervisor buffers the bounded body, sends one `HttpRequestEvaluation`, and receives one `HttpRequestResult`. Complete client-to-upstream WebSocket text messages use the separate bidirectional-streaming `EvaluateWebSocketSession` RPC. The supervisor sends `WebSocketSessionEvent` values; the service returns `WebSocketSessionEventResult` values for preflight and message events, while session start and end are notifications without corresponding results. Streaming is not baked into `EvaluateHttpRequest`; future chunked HTTP transport should add another operation-specific method rather than changing the existing method's cardinality. Possible extensions are collected in the [protocol-extensions appendix](appendices/protocol-extensions.md). Built-in middleware uses the same logical contracts in-process; the `openshell/regex` built-in advertises both V1 operations. + +V1 applies explicit public envelope limits before invoking a service or accepting its result: 64 KiB for encoded config, 4 KiB for request context, 32 KiB for the target, 128 header lines and 64 KiB of encoded headers, 4 MiB for the logical payload, 4 KiB for a reason, 64 header mutations with at most 32 KiB of validated name/value data and 64 KiB encoded, 32 findings per stage with each finding at most 4 KiB encoded, and 64 metadata entries totaling at most 32 KiB. A chain has at most 10 stages and therefore at most 320 findings. Middleware gRPC servers configure request and response message limits to cover the 4 MiB payload plus at least 292 KiB for the remaining envelope. -V1 applies explicit public envelope limits before invoking a service or accepting its result: 64 KiB for encoded config, 4 KiB for request context, 32 KiB for the target, 128 header lines and 64 KiB of encoded headers, 4 MiB for the body, 4 KiB for a reason, 64 header mutations with at most 32 KiB of validated name/value data and 64 KiB encoded, 32 findings per stage with each finding at most 4 KiB encoded, and 64 metadata entries totaling at most 32 KiB. A chain has at most 10 stages and therefore at most 320 findings. Middleware gRPC servers configure request and response message limits to cover the 4 MiB body plus at least 292 KiB for the remaining envelope. +For WebSocket traffic, a service advertises `WEBSOCKET_MESSAGE/PRE_CREDENTIALS` with `max_payload_bytes`, which limits one complete message or replacement rather than the whole session. HTTP bindings use the same field for one request body or replacement. An attached service without that exact binding does not join the chain, does not apply `on_error`, and produces internal `binding_not_selected` coverage. OpenShell opens one phase-specific `EvaluateWebSocketSession` stream per selected stage and upgrade attempt. Future upstream-to-client inspection uses the same RPC with `PRE_RETURN`; a service selected for both phases receives two independent streams for the WebSocket session. A bounded preflight exposes only the admitted destination through `HttpRequestTarget`, with its path separated from query data, requested subprotocols, sandbox context, the attached API middleware name, and validated implementation config. Policy-local config identity remains internal for audit and denial metadata. OpenShell evaluates selected preflights concurrently. Each stage returns `inspect`, voluntary `skip`, or authoritative `deny` before upstream contact, plus optional bounded reason, reason code, findings, and metadata. `deny` is a successful decision enforced independently of `on_error` and takes precedence over concurrent failures; failures alone follow each stage's `on_error`. OpenShell sends the terminal reason to every still-writable stream whose preflight opened successfully, at most once per stage. Inspecting stages that continue receive `session_start`, then complete logical text messages with monotonic sequence numbers. Stages run in global policy order and each sees the prior stage's accepted replacement. Binary logical messages pass through without middleware inspection under both error modes, consume a session-global sequence, and emit `unsupported_message_type` coverage for active stages; a later text RPC may therefore contain a valid sequence gap. `PRE_RETURN` and upstream-to-client inspection are reserved for a later implementation. + +Inspectable WebSocket text input and replacements share the 4 MiB platform cap. The operator's `max_payload_bytes` is the shared HTTP-body and WebSocket-text ceiling, further constrained by each operation binding's capability. It does not bound binary pass-through, which retains a separate raw-frame safety limit. Text replacements must be valid UTF-8 and preserve message type. A complete text message holds one process-wide admission permit for its entire chain; preflight fan-out holds one permit until every stage resolves. Permit waiting is backpressure and does not consume the per-message deadline. Per-stage timeouts are also bounded by a 30-second total chain budget, which applies to HTTP chains too. This bound controls concurrency and peak buffered inspection memory; it is not rate limiting. The `originating_process` is the same identity OpenShell resolves on the egress path - the binary, pid, and ancestor chain it uses for binary-scoped network policy and OCSF audit. It is per-connection rather than strictly per-request and is optional. Middleware must treat missing process data as unavailable rather than as an authorization failure. The initial implementation leaves this field unset until reliable propagation is available. @@ -310,7 +328,7 @@ The protobuf package is the wire-version handshake. `Describe` reports a diagnos ### Registration and delivery -The operator registers available external middleware services in gateway configuration under `openshell.supervisor.middleware`. The namespace identifies the subsystem whose behavior is extended, not the process that reads the configuration. The gateway still loads, validates, and distributes these registrations to supervisors. Each entry has a diagnostic name, gRPC endpoint, maximum body size, optional RPC timeout, and transport settings. The diagnostic name identifies the configured connection in logs but is not a policy key. Policy authors select stable binding IDs returned by `Describe`, so they cannot point traffic at an arbitrary endpoint and do not depend on an operator-local registration name. +The operator registers available external middleware services in gateway configuration under `openshell.supervisor.middleware`. The namespace identifies the subsystem whose behavior is extended, not the process that reads the configuration. The gateway still loads, validates, and distributes these registrations to supervisors. Each entry has a diagnostic name, gRPC endpoint, maximum logical payload size, optional RPC timeout, and transport settings. The diagnostic name identifies the configured connection in logs but is not a policy key. Policy authors select stable binding IDs returned by `Describe`, so they cannot point traffic at an arbitrary endpoint and do not depend on an operator-local registration name. The v1 transport is gRPC over a network endpoint reachable from every supervisor across Docker, Podman, VM, and Kubernetes drivers. In local single-player deployments, a loopback endpoint such as `127.0.0.1:1234` may be translated to `host.openshell.internal` so a supervisor can reach a service running on the local host. That loopback shorthand is not an HA deployment model: Kubernetes and other shared deployments should register a routable service DNS name or address that every supervisor can reach directly. Other deployment shapes are deferred until OpenShell has a universal way to make those endpoints reachable from the relevant supervisor environments. @@ -318,19 +336,19 @@ The v1 transport is gRPC over a network endpoint reachable from every supervisor [[openshell.supervisor.middleware]] name = "anonymizer" grpc_endpoint = "http://127.0.0.1:1234" -max_body_bytes = 4194304 +max_payload_bytes = 4194304 timeout = "500ms" allow_insecure = true [[openshell.supervisor.middleware]] name = "agent-traces-exporter" grpc_endpoint = "https://middleware.example.internal:443" -max_body_bytes = 1048576 +max_payload_bytes = 1048576 ``` The stable transport requirement is confidentiality plus authentication of the intended middleware service. Phase 1 may temporarily accept a plaintext `http://` endpoint only when the same entry explicitly sets `allow_insecure = true`. OpenShell rejects plaintext without that opt-in, warns prominently, and records the insecure registration as auditable configuration state. This escape hatch is limited to trusted local development and isolated research environments. Phase 2 removes plaintext support and the `allow_insecure` field, requiring authenticated encrypted transport. That removal is an intentional research-preview breaking change with no long-term compatibility obligation. The exact phase 2 mechanism, such as mTLS or TLS plus explicit caller authentication, is follow-up protocol work (see [appendices/protocol-extensions.md](appendices/protocol-extensions.md#middleware-authentication)). -For each binding, the operator's `max_body_bytes` must not exceed the binding capability returned by `Describe` or the 4 MiB platform maximum. The gateway rejects an invalid registration rather than silently clamping it. The resulting operator limit applies to every binding exposed by that registration. +For each binding, the operator's `max_payload_bytes` must not exceed the binding capability returned by `Describe` or the 4 MiB platform maximum. The gateway rejects an invalid registration rather than silently clamping it. The resulting operator limit applies to every binding exposed by that registration. RPC timeouts use an integer with an `ms` or `s` suffix, range from 10 ms through 30 s, and default to 500 ms. A binding may advertise its own timeout through `Describe`; that value overrides the service registration timeout. The service timeout applies to `Describe`, while the effective binding timeout applies to `ValidateConfig` and `EvaluateHttpRequest`. @@ -424,8 +442,11 @@ The namespace is the policy-local middleware config map key, not the optional hu A middleware decision is observable sandbox behavior, so it is recorded as an OCSF event, consistent with how the supervisor already logs network and L7 enforcement. This RFC commits to the event categories and the safety rules; exact field mappings are an implementation detail. +> **Update in PR #2477 - WebSocket middleware:** The coverage-boundary event below is new. It distinguishes an unsupported operation or message type from a middleware invocation or failure. + - **Per-invocation decisions** are `HttpActivity` events, since middleware is an L7 enforcement point. Each stage records the policy-local config key, validated binding ID, decision, transformation state, latency, and policy and endpoint context. Allowed requests are `Informational`; denials are `Medium`. - **Enforcement failures and bypasses** also emit `DetectionFinding` events. Required-stage failures, invalid responses, uninspectable traffic with a required stage, and body-aware policy evaluation failures are `High`. A `fail_open` bypass and uninspectable traffic allowed because every matching stage is `fail_open` are still findings so operators can alert on reduced enforcement. +- **Coverage boundaries** emit informational `NetworkActivity` events separately from invocations and failures. `binding_not_selected` records an attached config whose manifest lacks the WebSocket binding. `unsupported_message_type` records binary pass-through for an active stage with its internal config identity, logical sequence, message class, and size. - **Configuration events** are `ConfigStateChange` events: middleware registration validation, registry reload success or failure, and policy validation outcome. These events must never leak the content they describe. The OCSF JSONL may be shipped to external systems, so: @@ -440,7 +461,9 @@ This mirrors the middleware response contract, which already forbids the service Supervisor egress middleware stays opt-in throughout: until a policy declares a matching middleware config, no sandbox invokes one and the proxy hot path is unchanged. The initial usable slice proves built-in and external execution together; the phase boundary is transport hardening, not whether external middleware exists. -**Phase 1 - research-preview contract and execution.** Define `openshell.middleware.v1` with `Describe`, `ValidateConfig`, and `EvaluateHttpRequest`; ship the example `openshell/regex` built-in; and support statically registered operator-run services. Policy uses a top-level selector-based `network_middlewares` map with stable config keys, unique numeric `order`, per-stage `on_error`, bounded bodies, bounded RPC timeouts, atomic header mutations, post-transformation policy re-evaluation, and OCSF observability. Gateway startup validates external manifests, policy writes validate implementation-owned config, effective sandbox config carries only required external registrations, and supervisors install policy plus registry as one last-known-good runtime generation. Phase 1 requires encrypted authenticated transport for normal use but temporarily permits plaintext `http://` only with explicit `allow_insecure = true` for trusted local development or isolated research. OpenShell warns and emits auditable configuration state whenever that exception is used. +> **Update in PR #2477 - WebSocket middleware:** Phase 1 now also includes the forward-text WebSocket operation, bounded WebSocket messages, and the WebSocket binding for the built-in regex middleware. + +**Phase 1 - research-preview contract and execution.** Define `openshell.middleware.v1` with `Describe`, `ValidateConfig`, unary `EvaluateHttpRequest`, and forward-text `EvaluateWebSocketSession`; ship the example `openshell/regex` built-in; and support statically registered operator-run services. Policy uses a top-level selector-based `network_middlewares` map with stable config keys, unique numeric `order`, per-stage `on_error`, bounded bodies and messages, bounded RPC timeouts, atomic header mutations, post-transformation policy re-evaluation, and OCSF observability. Gateway startup validates external manifests, policy writes validate implementation-owned config, effective sandbox config carries only required external registrations, and supervisors install policy plus registry as one last-known-good runtime generation. Phase 1 requires encrypted authenticated transport for normal use but temporarily permits plaintext `http://` only with explicit `allow_insecure = true` for trusted local development or isolated research. OpenShell warns and emits auditable configuration state whenever that exception is used. **Phase 2 - mandatory authenticated encryption.** Remove plaintext middleware transport and remove `allow_insecure`. Every external connection must provide transport confidentiality and authenticate the intended service, with the final mechanism and credential delivery model defined by follow-up protocol work. Because phase 1 is explicitly a research preview, removing its insecure escape hatch is an intentional breaking change and does not create a long-term compatibility obligation. Operator-run service deployment otherwise keeps the same binding, policy, validation, delivery, reload, and invocation model. @@ -456,10 +479,12 @@ The first release is a research preview. The contract, policy surface, and scope Adding a synchronous, content-aware hook to the egress path has real costs. The most significant: +> **Update in PR #2477 - WebSocket middleware:** OpenShell now bounds concurrent middleware work and buffered memory. It does not add request-rate limiting. The updated rate-limit risk below keeps that distinction explicit. + - **Hot-path latency and a new per-request dependency.** Each selected external stage makes a synchronous call and blocks on its reply, so middleware latency becomes request latency and the service becomes a new failure surface on the data plane. This is bounded by opt-in host selectors, per-middleware timeouts, and built-ins running in-process with no network hop, but for matching traffic the tax is unavoidable. - **Fail-closed breaks workloads.** Denying traffic when a required middleware is unavailable, times out, or returns a malformed response is the safe default, but it converts a middleware outage into a sandbox outage. The opposite default leaks the very content the middleware exists to control. There is no choice that is both safe and always available; `on_error` makes the tradeoff explicit per middleware, but operators can still pick a default that surprises them. - **Body buffering and size limits.** Inspecting content means buffering a bounded request body instead of streaming it, which adds memory cost and interacts badly with growing payloads (for example inference requests whose context expands each turn until it exceeds the cap). An over-cap request is treated as an `on_error` event for the middleware that needs the body, so it follows the same `fail_closed` default: it is denied unless the operator has explicitly set `on_error: fail_open` for that middleware. Passing an over-cap request through unprocessed is therefore never the default - it is an opt-in choice made per middleware, and one a security-critical middleware would deliberately leave off so that oversized content is denied rather than silently egressed. -- **No OpenShell-side rate limiting.** OpenShell does not throttle calls to a middleware. A middleware that is slow, overloaded, or unavailable is handled only by its timeout and `on_error`, so operators must size, scale, and protect the service themselves; a struggling middleware degrades every request routed through it. +- **No OpenShell-side rate limiting.** OpenShell bounds concurrent middleware work and buffered memory, but does not throttle fast calls. A middleware that is slow, overloaded, or unavailable is handled by admission backpressure, its timeout, and `on_error`, so operators must still size, scale, and protect the service. - **Trusting an unsandboxed service with raw content.** Middleware receives raw request payloads, and OpenShell does not sandbox it, verify its behavior, or prevent it from mishandling or exfiltrating what it inspects. A buggy or malicious middleware is a direct data-exposure path. Trust in the middleware is the operator's responsibility, the same as trust in a sandbox image, but the blast radius here is in-flight request content. - **A false sense of coverage.** The hook runs only on traffic OpenShell terminates and parses. Opaque TCP or TLS passthrough, encrypted or otherwise opaque bodies, endpoints outside every selector, and content the middleware fails to detect can still leave without effective inspection. Policy validation rejects selector overlap with `tls: skip`, and runtime uninspectability follows the matching chain's failure policy, but detection correctness and traffic outside the selected host set remain inherent limitations. - **Phase 1 plaintext is risky.** The research-preview exception permits plaintext gRPC only with explicit `allow_insecure = true`. Because middleware can allow, deny, or transform egress, an impersonated or eavesdropped service is a policy-enforcement bypass, not just an observability gap. The exception is unsuitable for shared or untrusted networks, produces an explicit warning and audit event, and is removed in phase 2. See [appendices/protocol-extensions.md](appendices/protocol-extensions.md#middleware-authentication). @@ -491,10 +516,12 @@ This section closes the current review themes. ### Decisions +> **Update in PR #2477 - WebSocket middleware:** The operation-scope and failure-behavior decisions below now include WebSocket bindings, client text messages, binary pass-through, and capability coverage. + - **Middleware naming.** Use the feature name "supervisor middleware." The first operation family is egress middleware, but the higher-level feature name stays extensible for future supervisor hooks. The service can inspect, transform, deny, and annotate, so narrower names such as "transformer" or "request processor" describe only part of the contract. - **Middleware binding IDs.** Services own stable binding IDs and policy selects them through the `middleware` field. Gateway registration names are diagnostic only. Binding IDs use `/` for namespaces, `openshell/` is reserved for built-ins, and empty path segments are invalid. - **Operation naming.** Use typed operation and phase enums such as `HTTP_REQUEST/PRE_CREDENTIALS`. The operation describes the middleware API payload, and the phase describes the proxy position. Later protocols can add typed operations such as WebSocket message or TCP connect without renaming the v1 hook. -- **HTTP scope of v1.** `HTTP_REQUEST/PRE_CREDENTIALS` applies to every HTTP/1.x request that OpenShell terminates and parses, whether or not the endpoint declares a `protocol`. WebSocket upgrade requests are included; post-upgrade frames, HTTP/2, HTTP/3, opaque TCP, and `tls: skip` traffic are excluded. +- **Operation scope of v1.** `HTTP_REQUEST/PRE_CREDENTIALS` applies to every HTTP/1.x request that OpenShell terminates and parses, whether or not the endpoint declares a `protocol`; WebSocket upgrade requests are included. `WEBSOCKET_MESSAGE/PRE_CREDENTIALS` applies only to complete client-to-upstream text messages for attachments whose manifest advertises it. Binary and return-path messages, HTTP/2, HTTP/3, opaque TCP, and `tls: skip` traffic are excluded from those operation bindings. - **Route selection and forwarding.** V1 has no `forward_to` decision. Middleware never makes the upstream call. Future route-selection hooks may choose among OpenShell-managed destinations, such as model routes, but must not become arbitrary external endpoint rewrites. - **SigV4/request signing.** AWS SigV4 belongs to a restricted built-in `HttpRequest/post_credentials` hook, not external `HttpRequest/pre_credentials` middleware. The middleware can be configured by policy, but it must run in-process with supervisor host capabilities so it can strip placeholder signatures and sign with real supervisor-resolved credentials without exposing those credentials over the external middleware contract. - **Composability and ordering.** Middleware is chainable and ordered by ascending numeric `order`. Order values must be unique across the policy. A stage receives the previous stage's transformed body and header mutations; `deny` short-circuits the chain; and different config map keys may invoke the same binding as separate stages. @@ -503,7 +530,7 @@ This section closes the current review themes. - **Actor data.** Actor process data is optional and per-connection. Middleware must treat it as context, not a reliable per-request identity or authorization input. - **Metadata namespacing.** Metadata is stored under the policy-local middleware config map key rather than the optional human-readable name. This prevents collisions without a central key registry and lets two configs using the same implementation emit independent metadata. - **Selector-only placement.** V1 uses only config-level `endpoints.include` and `endpoints.exclude` selectors. Policy-level and endpoint-level attachment lists are not part of the schema. Selection is independent of the network rule that admitted the request and therefore remains stable after effective-policy composition. -- **Failure behavior.** Middleware errors, timeouts, malformed responses, and over-cap bodies use `on_error`; `fail_closed` is the default. Passing unprocessed content through is only possible when the operator explicitly sets `on_error: fail_open`. +- **Failure behavior.** Middleware errors, timeouts, malformed responses, and over-cap inspectable payloads use `on_error` after an operation binding is selected; `fail_closed` is the default. An absent operation binding and binary WebSocket messages are capability coverage states, not failures, and pass with informational telemetry under both error modes. - **Limits.** V1 caps policies at 10 middleware configs, selectors at 32 combined patterns per config, bodies at 4 MiB, findings at 32 per stage, and all non-body request and result fields at the public envelope limits in the contract section. - **Delivery and reload.** `GetSandboxConfig` delivers only external registrations required by the effective policy. Built-ins are installed locally. Supervisors prepare candidate policy and registry state off-path, swap them as one generation, reuse connections for policy-only changes, and preserve the complete last-known-good runtime on failure. - **Chunked and compressed bodies.** V1 operates on bounded bytes OpenShell can buffer safely. Known over-cap content length may fail open before consumption when every affected stage permits it. Chunked overflow after consumption is denied because the raw stream cannot be resumed. Compressed bodies remain opaque unless a binding explicitly supports them. diff --git a/rfc/0009-supervisor-middleware/appendices/protocol-extensions.md b/rfc/0009-supervisor-middleware/appendices/protocol-extensions.md index 924c7bd699..bc753382bb 100644 --- a/rfc/0009-supervisor-middleware/appendices/protocol-extensions.md +++ b/rfc/0009-supervisor-middleware/appendices/protocol-extensions.md @@ -2,7 +2,9 @@ > This is an appendix to the [RFC](../README.md). Please familiarize yourself with the RFC before reading this. -The v1 contract is intentionally minimal: one HTTP request hook, buffered unary calls, an `allow`/`deny` decision plus optional transformed content, findings, and metadata. This appendix records extensions the proto should not preclude, so v1 stays small without painting future work into a corner. None of these are committed; they exist to validate that the v1 shape is forward-compatible. +**Update in PR #2477 - WebSocket middleware:** V1 now includes a forward-text WebSocket operation. The updated text below separates this implemented operation from future HTTP streaming and WebSocket return-path operations. + +The v1 contract is intentionally minimal: one buffered unary HTTP request hook and one forward-text WebSocket message hook, each with an `allow`/`deny` decision plus optional transformed content, findings, and metadata. This appendix records extensions the proto should not preclude, so v1 stays small without painting future work into a corner. None of these are committed; they exist to validate that the v1 shape is forward-compatible. ## Streaming @@ -42,16 +44,19 @@ A cleaner phased design using a `oneof` over `context` and `body_chunk`, in the ## Additional operation phases -V1 defines a single typed operation, `HTTP_REQUEST/PRE_CREDENTIALS`, which runs after network and L7 policy admit a request and before credential injection. The same service interface can host more operations, each advertised through the `Describe` manifest and invoked through an operation-specific method such as `EvaluateHttpRequest`. Each operation and phase pair encodes a different position in the proxy flow: +> **Update in PR #2477 - WebSocket middleware:** This section now records `WEBSOCKET_MESSAGE/PRE_CREDENTIALS` as implemented. It keeps `WEBSOCKET_MESSAGE/PRE_RETURN` as a reserved future operation. + +V1 supports `HTTP_REQUEST/PRE_CREDENTIALS` and forward-text `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`. The same service interface can host more operations, each advertised through the `Describe` manifest and invoked through an operation-specific method. Each operation and phase pair encodes a different position in the proxy flow: - `Connection/before_policy` / `HttpRequest/before_policy` - *before* network/L7 policy admits the request, for earlier classification. Riskier, because request content reaches a service before policy has allowed the request. - `HTTP_REQUEST/PRE_CREDENTIALS` (v1) - after policy admits the request, before credential injection. - `HttpRequest/post_credentials` - after credential injection, immediately before the relay writes the request upstream. This hook is credential-visible, so it is built-in-only: OpenShell marks it as a restricted hook and rejects any externally registered middleware that advertises it during manifest validation. The motivating use is request signing that must run after credentials are injected - for example a built-in `openshell/sigv4` that strips placeholder-signed AWS headers and signs the finalized request with supervisor-resolved credentials just before it is sent upstream. - `HttpResponse/completed` - after an upstream request completes, emit metadata such as status, content length, selected route, selected model, and model usage if available. This is notification-only: no body, no transformation, and no allow/deny verdict. It would let reservation-style budget middleware reconcile a pre-dispatch decision without introducing response-body inspection. - `HttpResponse/before_return` - on the return path, after the upstream responds and before the response reaches the sandbox; inspect or redact upstream responses. -- `WebSocketMessage/before_forward` / `WebSocketMessage/before_return` - after a WebSocket or streaming protocol upgrade, on each forwarded or returned message, well past the one-shot request path. +- `WEBSOCKET_MESSAGE/PRE_CREDENTIALS` (v1 forward text) - after a WebSocket upgrade, on each complete client text message before credential placeholder rewriting. Before upstream contact, a concurrent preflight lets each selected stage inspect, voluntarily skip, or authoritatively deny the upgrade. Explicit denial takes precedence over failures and is enforced independently of `on_error`; OpenShell best-effort ends every still-writable opened stage stream with the typed terminal reason. An attached implementation without this binding is not selected and records coverage rather than applying `on_error`. Binary messages pass without inspection, consume a logical sequence, and record unsupported-message coverage for active stages. +- `WEBSOCKET_MESSAGE/PRE_RETURN` - on complete upstream messages before they return to the workload. The enum value is reserved, but manifests advertising it are rejected until return-path inspection is implemented. -Pre-policy phases run earliest, the two request phases bracket credential injection, response notifications and response phases run after the upstream call, and message phases run later, sometimes on a different path entirely. Of these, only `HTTP_REQUEST/PRE_CREDENTIALS` is part of v1. `HttpRequest/post_credentials` is the nearest planned request-path follow-up and is kept built-in-only because it sees injected credentials; `HttpResponse/completed` is a separate future notification hook for metadata-only post-call reconciliation. +Pre-policy phases run earliest, the two request phases bracket credential injection, response notifications and response phases run after the upstream call, and message phases run later on the parsed relay. V1 implements only the two pre-credentials pairs above. `HttpRequest/post_credentials` is the nearest planned request-path follow-up and is kept built-in-only because it sees injected credentials; `HttpResponse/completed` is a separate future notification hook for metadata-only post-call reconciliation. ## Semantic context