diff --git a/.agents/skills/debug-openshell-cluster/SKILL.md b/.agents/skills/debug-openshell-cluster/SKILL.md index 07de687103..3cdc0cc3a2 100644 --- a/.agents/skills/debug-openshell-cluster/SKILL.md +++ b/.agents/skills/debug-openshell-cluster/SKILL.md @@ -102,9 +102,9 @@ 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 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 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. WebSocket message sequences are allocated session-wide; each stage receives a strictly increasing subset, so gaps are valid when messages 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 @@ -400,7 +400,8 @@ 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 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 8da14420c9..3c267f1575 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,12 @@ 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. WebSocket middleware runs for both `ws://` and `wss://`, receives complete client text messages only, and does not inspect binary or upstream-to-client messages. +- 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. - 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. @@ -406,7 +408,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 applies independently of the admitting network rule to every matching HTTP or WebSocket destination. Narrow `endpoints.include` or add exclusions if the stage 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 a36a372df7..adf4071808 100644 --- a/.agents/skills/openshell-cli/SKILL.md +++ b/.agents/skills/openshell-cli/SKILL.md @@ -341,10 +341,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. Binary and upstream-to-client WebSocket 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 31e2104987..64cad51f2b 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/sandbox-limits.md b/architecture/sandbox-limits.md new file mode 100644 index 0000000000..1152e1df4b --- /dev/null +++ b/architecture/sandbox-limits.md @@ -0,0 +1,166 @@ +# 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 because every connection shares +one `MiddlewareRegistry`. + +| 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 | 32 | Additional work is shed when both the active budget and waiter budget are full. | +| 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. | +| 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. + +Parsed WebSocket text assembly acquires shared middleware work before buffering. +The assembly state retains that permit with its buffer, compression and +fragment state, and total deadline. Input progress resets only the idle +deadline. Completion transfers the permit into middleware evaluation; every +timeout and terminal parser error drops it. + +## 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 4f95e1ef69..6443bf7176 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,14 @@ 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 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 +96,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 +203,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 +211,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 +285,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..d49aaa9716 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, WebSocketEvaluationRequest, WebSocketEvaluationResponse, +}; + +/// 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( + &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-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..927c42b52f 100644 --- a/crates/openshell-supervisor-middleware-builtins/src/lib.rs +++ b/crates/openshell-supervisor-middleware-builtins/src/lib.rs @@ -8,17 +8,21 @@ 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, WebSocketDirection, WebSocketEvaluationRequest, + WebSocketEvaluationResponse, WebSocketMessageType, WebSocketPreflightAction, + WebSocketPreflightDecision, web_socket_evaluation_request, web_socket_evaluation_response, }; +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)] } @@ -45,8 +49,149 @@ fn evaluate_http_request(evaluation: &HttpRequestEvaluation) -> Result(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.request { + Some(web_socket_evaluation_request::Request::Preflight(preflight)) + if config.is_none() && !started => + { + if preflight.middleware_name != BUILTIN_REGEX { + Err(Status::invalid_argument( + "unknown built-in WebSocket middleware", + )) + } else if preflight.phase + != SupervisorMiddlewarePhase::PreCredentials as i32 + || preflight.direction != WebSocketDirection::ClientToUpstream as i32 + { + Err(Status::invalid_argument( + "unsupported built-in WebSocket binding", + )) + } else { + let selected_config = preflight.config.unwrap_or_default(); + match regex::validate_config(&selected_config) { + Ok(()) => { + config = Some(selected_config); + Ok(Some(WebSocketEvaluationResponse { + response: Some( + web_socket_evaluation_response::Response::PreflightDecision( + WebSocketPreflightDecision { + action: WebSocketPreflightAction::Inspect as i32, + }, + ), + ), + })) + } + Err(error) => Err(Status::invalid_argument(error.to_string())), + } + } + } + Some(web_socket_evaluation_request::Request::SessionStart(_)) + if config.is_some() && !started => + { + started = true; + Ok(None) + } + Some(web_socket_evaluation_request::Request::Message(message)) if started => { + if let Err(error) = advance_sequence_lower_bound( + &mut sequence_lower_bound, + message.sequence, + ) { + Err(error) + } else if message.direction != WebSocketDirection::ClientToUpstream as i32 + || 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(WebSocketEvaluationResponse { + response: Some( + web_socket_evaluation_response::Response::MessageResult( + result, + ), + ), + })), + Err(error) => Err(Status::invalid_argument(error.to_string())), + } + } + } + Some(web_socket_evaluation_request::Request::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 EvaluateWebSocketStream = WebSocketResponseStream; + async fn describe( &self, _request: Request<()>, @@ -54,7 +199,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(), })) } @@ -86,6 +231,43 @@ impl SupervisorMiddleware for BuiltinMiddlewareService { .map(Response::new) .map_err(|error| Status::invalid_argument(error.to_string())) } + + async fn evaluate_web_socket( + &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( + &self, + receiver: tokio::sync::mpsc::Receiver, + ) -> Result { + Ok(Self::websocket_stream( + tokio_stream::wrappers::ReceiverStream::new(receiver).map(Ok), + )) + } } #[cfg(test)] @@ -109,12 +291,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 @@ -124,6 +305,17 @@ mod tests { SupervisorMiddlewarePhase::PreCredentials as i32 ); assert_eq!(manifest.bindings[0].max_body_bytes, 256 * 1024); + assert_eq!(manifest.bindings[0].max_message_bytes, 0); + 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_body_bytes, 0); + assert_eq!(manifest.bindings[1].max_message_bytes, 256 * 1024); } #[test] @@ -199,4 +391,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..5f41342c0e 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,14 @@ 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_MESSAGE_BYTES: u64 = 256 * 1024; #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Deserialize)] #[serde(default, deny_unknown_fields)] @@ -46,13 +47,23 @@ 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_body_bytes: MAX_BODY_BYTES, + timeout: String::new(), + max_message_bytes: 0, + }, + MiddlewareBinding { + operation: SupervisorMiddlewareOperation::WebsocketMessage as i32, + phase: SupervisorMiddlewarePhase::PreCredentials as i32, + max_body_bytes: 0, + timeout: String::new(), + max_message_bytes: MAX_MESSAGE_BYTES, + }, + ] } struct ReplacementPattern { @@ -85,34 +96,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_MESSAGE_BYTES { + return Err(miette!( + "{NAME} WebSocket text message exceeds {MAX_MESSAGE_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..8dca4014bc 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::{ + 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,14 +31,102 @@ 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( + &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 }) +} + +const MAX_QUEUED_MIDDLEWARE_WORK: usize = MAX_CONCURRENT_MIDDLEWARE_WORK; + +/// 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 + } +} + +/// One slot in the shared persistent middleware session budget. +/// +/// Protocol-specific session runners retain this guard while at least one +/// streaming stage remains active. The generic registry owns admission so +/// future streaming HTTP middleware can share 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. @@ -80,11 +176,13 @@ pub const MIDDLEWARE_GRPC_ENVELOPE_BYTES: usize = pub const MIDDLEWARE_GRPC_MESSAGE_BYTES: usize = MAX_MIDDLEWARE_BODY_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, @@ -144,6 +242,7 @@ pub struct DescribedChainEntry { service: Option>, binding: Option, max_body_bytes: usize, + max_message_bytes: usize, timeout: Duration, } @@ -152,6 +251,10 @@ impl DescribedChainEntry { self.max_body_bytes } + pub fn max_message_bytes(&self) -> usize { + self.max_message_bytes + } + pub fn on_error(&self) -> OnError { self.entry.on_error } @@ -303,7 +406,7 @@ 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, @@ -375,6 +478,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 +490,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 +513,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,12 +540,6 @@ fn validate_registration(registration: &SupervisorMiddlewareService) -> Result MAX_MIDDLEWARE_BODY_BYTES as u64 { return Err(miette!( "middleware registration '{}' max_body_bytes exceeds the platform maximum of {MAX_MIDDLEWARE_BODY_BYTES}", @@ -490,6 +601,50 @@ fn validate_body_limit(source: &str, binding: &MiddlewareBinding) -> Result Result { + if binding.max_message_bytes == 0 { + return Err(miette!("{source} must advertise a non-zero message limit")); + } + if binding.max_message_bytes > MAX_MIDDLEWARE_BODY_BYTES as u64 { + return Err(miette!( + "{source} message limit exceeds the platform maximum of {MAX_MIDDLEWARE_BODY_BYTES}" + )); + } + usize::try_from(binding.max_message_bytes) + .map_err(|_| miette!("{source} reports a message 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, @@ -501,29 +656,47 @@ 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" - )); - } + let kind = 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 = match kind { + SupportedBinding::HttpPreCredentials => { + if binding.max_message_bytes != 0 { + return Err(miette!( + "{source} HTTP_REQUEST binding must omit max_message_bytes" + )); + } + validate_body_limit(source, binding)? + } + SupportedBinding::WebSocketPreCredentials => { + if binding.max_body_bytes != 0 { + return Err(miette!( + "{source} WEBSOCKET_MESSAGE binding must omit max_body_bytes" + )); + } + validate_message_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 kind == SupportedBinding::HttpPreCredentials + && operator_max_body_bytes.is_some_and(|limit| limit > advertised) + { return Err(miette!( "{source} max_body_bytes ({}) exceeds the binding capability ({advertised})", operator_max_body_bytes.expect("operator limit checked above") )); } + if kind == SupportedBinding::HttpPreCredentials && operator_max_body_bytes == Some(0) { + return Err(miette!( + "{source} must configure max_body_bytes for its HTTP_REQUEST binding" + )); + } } Ok(()) } @@ -531,12 +704,12 @@ fn validate_manifest_bindings( fn validate_external_manifest( registration: &SupervisorMiddlewareService, manifest: &MiddlewareManifest, - operator_max_body_bytes: usize, + operator_max_body_bytes: Option, ) -> Result<()> { validate_manifest_bindings( &format!("external middleware registration '{}'", registration.name), manifest, - Some(operator_max_body_bytes), + operator_max_body_bytes, ) } @@ -668,7 +841,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,7 +886,7 @@ 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, @@ -758,17 +931,18 @@ impl MiddlewareRegistry { safe_reason(&error.to_string()) ) })?; - validate_external_manifest(®istration, &manifest, operator_max_body_bytes)?; + validate_external_manifest(®istration, &manifest, Some(operator_max_body_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_body_bytes: (operator_max_body_bytes != 0) + .then_some(operator_max_body_bytes), operator_timeout, })); registered_services.push(RegisteredMiddlewareService { registration }); @@ -778,6 +952,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,12 +1023,16 @@ 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, @@ -859,6 +1040,9 @@ impl ChainRunner { })]), 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 +1053,48 @@ impl ChainRunner { } } + /// 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(MiddlewareWorkAdmission { + _work: permit, + saturated: false, + }) + } else { + let waiter = Arc::clone(&self.registry.work_admission_waiters) + .try_acquire_owned() + .map_err(|_| { + miette!("middleware admission queue is full; refusing additional buffered work") + })?; + let permit = Arc::clone(&self.registry.work_admission) + .acquire_owned() + .await + .map_err(|_| miette!("middleware admission semaphore closed"))?; + drop(waiter); + Ok(MiddlewareWorkAdmission { + _work: permit, + saturated: true, + }) + } + } + + /// 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 +1104,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 +1131,92 @@ 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> { + self.describe_chain_for( + entries, + SupervisorMiddlewareOperation::HttpRequest, + SupervisorMiddlewarePhase::PreCredentials, + ) + .await + } + + pub async fn describe_websocket_chain( + &self, + entries: &[ChainEntry], + ) -> Result> { + self.describe_chain_for( + entries, + SupervisorMiddlewareOperation::WebsocketMessage, + SupervisorMiddlewarePhase::PreCredentials, + ) + .await + } + + 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()); + 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_body_bytes: 0, + max_message_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. + continue; + }; + let timeout = state.timeout_for_binding(&binding)?; + let max_body_bytes = if operation == SupervisorMiddlewareOperation::HttpRequest { + let advertised = validate_body_limit("middleware manifest", &binding)?; + state.operator_max_body_bytes.unwrap_or(advertised) + } else { + 0 + }; + let max_message_bytes = if operation == SupervisorMiddlewareOperation::WebsocketMessage + { + validate_message_limit("middleware manifest", &binding)? + } else { + 0 + }; + described_entries.push(DescribedChainEntry { + entry, + service: Some(Arc::clone(state)), + binding: Some(binding), + max_body_bytes, + max_message_bytes, + timeout, + }); + } + ensure_chain_capacity(described_entries.len())?; + Ok(described_entries) } pub async fn validate_config( @@ -974,19 +1230,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 +1295,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().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 +1326,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 { @@ -1106,11 +1385,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 @@ -1416,15 +1713,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 +1943,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 +1978,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for ScriptedService { + type EvaluateWebSocketStream = WebSocketResponseStream; + + async fn evaluate_web_socket( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -1688,6 +2000,7 @@ mod tests { phase: SupervisorMiddlewarePhase::PreCredentials as i32, max_body_bytes: self.max_body_bytes, timeout: String::new(), + max_message_bytes: 0, }], })) } @@ -1725,6 +2038,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for SlowService { + type EvaluateWebSocketStream = WebSocketResponseStream; + + async fn evaluate_web_socket( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -1737,6 +2060,7 @@ mod tests { phase: SupervisorMiddlewarePhase::PreCredentials as i32, max_body_bytes: 4096, timeout: self.binding_timeout.clone(), + max_message_bytes: 0, }], })) } @@ -1778,6 +2102,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for TwoStageService { + type EvaluateWebSocketStream = WebSocketResponseStream; + + async fn evaluate_web_socket( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -1790,6 +2124,7 @@ mod tests { phase: SupervisorMiddlewarePhase::PreCredentials as i32, max_body_bytes: 256 * 1024, timeout: String::new(), + max_message_bytes: 0, }], })) } @@ -1842,7 +2177,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 +2239,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 +2291,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 +2377,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for RecordingService { + type EvaluateWebSocketStream = WebSocketResponseStream; + + async fn evaluate_web_socket( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -2054,6 +2399,7 @@ mod tests { phase: SupervisorMiddlewarePhase::PreCredentials as i32, max_body_bytes: 4096, timeout: String::new(), + max_message_bytes: 0, }], })) } @@ -2101,6 +2447,16 @@ mod tests { #[tonic::async_trait] impl SupervisorMiddleware for HeaderChainService { + type EvaluateWebSocketStream = WebSocketResponseStream; + + async fn evaluate_web_socket( + &self, + _request: Request>, + ) -> std::result::Result, tonic::Status> + { + Err(tonic::Status::unimplemented("HTTP-only test middleware")) + } + async fn describe( &self, _request: Request<()>, @@ -2113,6 +2469,7 @@ mod tests { phase: SupervisorMiddlewarePhase::PreCredentials as i32, max_body_bytes: 4096, timeout: String::new(), + max_message_bytes: 0, }], })) } @@ -2226,7 +2583,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()) @@ -2285,7 +2642,7 @@ mod tests { } async fn registry_with_external( - service: Arc, + service: Arc, registration: SupervisorMiddlewareService, ) -> MiddlewareRegistry { let builtin_service = services() @@ -2312,7 +2669,7 @@ mod tests { .into_inner(); let operator_max_body_bytes = usize::try_from(registration.max_body_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_body_bytes)) .expect("valid external manifest"); let manifest_cell = OnceCell::new(); manifest_cell.set(manifest).expect("manifest cache"); @@ -2321,7 +2678,7 @@ 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, @@ -2329,7 +2686,7 @@ mod tests { }), 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), @@ -2338,6 +2695,9 @@ mod tests { ]), 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)), } } @@ -2528,9 +2888,10 @@ mod tests { phase: PRE_CREDENTIALS_PHASE as i32, max_body_bytes: 4096, timeout: String::new(), + max_message_bytes: 0, }], }; - 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")); } @@ -2554,9 +2915,10 @@ mod tests { phase: PRE_CREDENTIALS_PHASE as i32, max_body_bytes: u64::MAX, timeout: String::new(), + max_message_bytes: 0, }], }; - let error = validate_external_manifest(®istration, &manifest, 4096) + let error = validate_external_manifest(®istration, &manifest, Some(4096)) .expect_err("extreme advertised body limit must be rejected"); assert!(error.to_string().contains("platform maximum")); } @@ -2569,6 +2931,7 @@ mod tests { phase: PRE_CREDENTIALS_PHASE as i32, max_body_bytes: 4096, timeout: String::new(), + max_message_bytes: 0, }; let manifest = MiddlewareManifest { name: "example/service".into(), @@ -2576,15 +2939,38 @@ 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_body_bytes: 0, + timeout: "500ms".into(), + max_message_bytes: MAX_MIDDLEWARE_BODY_BYTES as u64, + }; + 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_registration_accepts_http_and_https_grpc_endpoints() { for grpc_endpoint in [ @@ -2650,9 +3036,10 @@ mod tests { phase: PRE_CREDENTIALS_PHASE as i32, max_body_bytes: 4096, timeout: timeout.into(), + max_message_bytes: 0, }], }; - 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")); } @@ -3511,4 +3898,829 @@ mod tests { ); assert!(outcome.applied[0].failed); } + + #[derive(Clone, Default)] + struct OpenAiRedactionService { + preflight: Arc>>, + skip: bool, + close_on_first_message: bool, + messages: Arc, + } + + impl OpenAiRedactionService { + fn websocket_stream(&self, mut requests: S) -> WebSocketResponseStream + where + S: Stream< + Item = std::result::Result< + openshell_core::proto::WebSocketEvaluationRequest, + tonic::Status, + >, + > + Send + + Unpin + + 'static, + { + use openshell_core::proto::{ + WebSocketEvaluationResponse, WebSocketMessageResult, WebSocketPreflightAction, + WebSocketPreflightDecision, web_socket_evaluation_request, + web_socket_evaluation_response, + }; + + let preflight = Arc::clone(&self.preflight); + let skip = self.skip; + let close_on_first_message = self.close_on_first_message; + let messages = Arc::clone(&self.messages); + 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.request { + Some(web_socket_evaluation_request::Request::Preflight(value)) => { + *preflight.lock().expect("preflight lock") = Some(value); + Some(WebSocketEvaluationResponse { + response: Some( + web_socket_evaluation_response::Response::PreflightDecision( + WebSocketPreflightDecision { + action: if skip { + WebSocketPreflightAction::Skip as i32 + } else { + WebSocketPreflightAction::Inspect as i32 + }, + }, + ), + ), + }) + } + Some(web_socket_evaluation_request::Request::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(WebSocketEvaluationResponse { + response: Some( + web_socket_evaluation_response::Response::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_evaluation_request::Request::SessionStart(_) + | web_socket_evaluation_request::Request::SessionEnd(_), + ) + | None => 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 EvaluateWebSocketStream = WebSocketResponseStream; + + async fn describe( + &self, + _request: Request<()>, + ) -> std::result::Result, tonic::Status> { + Ok(tonic::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_body_bytes: 0, + timeout: "1s".into(), + max_message_bytes: MAX_MIDDLEWARE_BODY_BYTES as u64, + }], + })) + } + + 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( + &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( + &self, + receiver: tokio::sync::mpsc::Receiver< + openshell_core::proto::WebSocketEvaluationRequest, + >, + ) -> std::result::Result { + Ok( + self.websocket_stream( + tokio_stream::wrappers::ReceiverStream::new(receiver).map(Ok), + ), + ) + } + } + + #[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(0); + 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 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"); + assert_eq!(observed.host, "api.openai.com"); + assert_eq!(observed.path, "/v1/responses"); + 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(0); + 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") + .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() + .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() + .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() + .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_err("work beyond the bounded waiter queue must be shed"); + assert!(overflow.to_string().contains("admission queue is full")); + + drop(active); + for waiter in waiters { + waiter + .await + .expect("waiter task") + .expect("queued admission after capacity is released"); + } + } + + #[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(0); + 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_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..b66e4655c2 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, WebSocketEvaluationRequest, }; 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( + &self, + receiver: tokio::sync::mpsc::Receiver, + ) -> std::result::Result { + let mut client = self.client.clone(); + let responses = client + .evaluate_web_socket(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..ba0fbe1ac1 --- /dev/null +++ b/crates/openshell-supervisor-middleware/src/websocket.rs @@ -0,0 +1,981 @@ +// 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, RequestContext, SupervisorMiddlewarePhase, WebSocketDirection, + WebSocketEvaluationRequest, WebSocketMessage, WebSocketMessageResult, WebSocketMessageType, + WebSocketPreflight, WebSocketPreflightAction, WebSocketSessionEnd, WebSocketSessionEndReason, + WebSocketSessionStart, web_socket_evaluation_request, web_socket_evaluation_response, +}; + +use super::{ + ChainEntry, ChainRunner, DescribedChainEntry, MAX_MIDDLEWARE_BODY_BYTES, + 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_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, +} + +pub struct WebSocketPreflightResult { + pub allowed: bool, + pub reason: String, + pub session: Option, + pub invocations: Vec, + pub saturated: bool, + pub session_capacity_exhausted: bool, +} + +#[derive(Debug)] +pub struct WebSocketSessionStartOutcome { + pub allowed: bool, + 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() + } + + fn disable(&mut self) { + self.transport.take(); + } +} + +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, WebSocketInvocation), + Skip(WebSocketInvocation), + Failed(DescribedChainEntry, String), +} + +impl ChainRunner { + pub async fn preflight_websocket( + &self, + entries: &[ChainEntry], + input: WebSocketPreflightInput, + ) -> miette::Result { + validate_preflight_input(&input)?; + let described = self.describe_websocket_chain(entries).await?; + if described.is_empty() { + return Ok(WebSocketPreflightResult { + allowed: true, + reason: String::new(), + session: None, + invocations: Vec::new(), + saturated: false, + session_capacity_exhausted: false, + }); + } + + // 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().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, 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 fail_closed_reason = None; + for result in opened { + match result { + OpenStage::Inspect(stage, invocation) => { + stages.push(*stage); + invocations.push(invocation); + } + OpenStage::Skip(invocation) => invocations.push(invocation), + 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(reason) = fail_closed_reason { + end_stages(&mut stages, WebSocketSessionEndReason::MiddlewareFailure).await; + return Ok(WebSocketPreflightResult { + allowed: false, + reason, + session: None, + invocations, + saturated, + session_capacity_exhausted: false, + }); + } + + if stages.is_empty() { + drop(session_admission); + return Ok(WebSocketPreflightResult { + allowed: true, + reason: String::new(), + session: None, + invocations, + saturated, + session_capacity_exhausted: false, + }); + } + + Ok(WebSocketPreflightResult { + allowed: true, + reason: String::new(), + session: Some(WebSocketSession { + runner: self.clone(), + stages, + next_sequence: 1, + session_admission: Some(session_admission), + }), + invocations, + saturated, + session_capacity_exhausted: false, + }) + } +} + +fn session_capacity_exhausted( + described: 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, + reason: if fail_closed { + format!("middleware_failed: {reason}") + } else { + String::new() + }, + session: None, + invocations, + 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() + .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, + 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 = WebSocketEvaluationRequest { + request: Some(web_socket_evaluation_request::Request::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(); + 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(), + reason: fail_closed.unwrap_or_default(), + invocations, + } + } + + pub async fn evaluate_text(&mut self, payload: Vec) -> WebSocketMessageOutcome { + if payload.len() > MAX_MIDDLEWARE_BODY_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_BODY_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_message_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(); + 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 = WebSocketEvaluationRequest { + request: Some(web_socket_evaluation_request::Request::Message( + WebSocketMessage { + sequence, + direction: WebSocketDirection::ClientToUpstream as i32, + 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.response { + Some(web_socket_evaluation_response::Response::MessageResult(result)) => result, + Some(web_socket_evaluation_response::Response::PreflightDecision(_)) | None => { + if let Some(outcome) = handle_stage_failure( + stage, + sequence, + original_size, + "unexpected_websocket_response", + ¤t, + &findings, + &metadata, + &mut invocations, + saturated, + ) { + 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, + ) { + 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, + ) { + return outcome; + } + continue; + } + Err(_) => { + if let Some(outcome) = handle_stage_failure( + stage, + sequence, + original_size, + "middleware_timeout", + ¤t, + &findings, + &metadata, + &mut invocations, + saturated, + ) { + return outcome; + } + continue; + } + }; + + let result = + match validate_message_result(result, sequence, stage.entry.max_message_bytes) { + Ok(result) => result, + Err(reason) => { + if let Some(outcome) = handle_stage_failure( + stage, + sequence, + original_size, + reason, + ¤t, + &findings, + &metadata, + &mut invocations, + saturated, + ) { + 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(); + } +} + +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, + } +} + +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, + direction: WebSocketDirection::ClientToUpstream as i32, + context: Some(RequestContext { + request_id: input.request_id, + sandbox_id: input.sandbox_id, + originating_process: None, + }), + scheme: input.scheme, + host: input.host, + port: u32::from(input.port), + path: input.path, + requested_subprotocols: input.requested_subprotocols, + middleware_name: entry.entry.implementation.clone(), + config_name: entry.entry.name.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(WebSocketEvaluationRequest { + request: Some(web_socket_evaluation_request::Request::Preflight(preflight)), + }) + .await + .map_err(|_| tonic::Status::unavailable("request stream closed"))?; + let mut responses = endpoint.open_websocket(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 { + return OpenStage::Failed(entry, "missing_preflight_decision".into()); + }; + let Some(web_socket_evaluation_response::Response::PreflightDecision(decision)) = + response.response + else { + return OpenStage::Failed(entry, "invalid_preflight_decision".into()); + }; + match WebSocketPreflightAction::try_from(decision.action) { + Ok(WebSocketPreflightAction::Inspect) => { + let invocation = success_invocation( + &entry, + WebSocketInvocationOutcome::Inspect, + 0, + 0, + None, + false, + None, + ); + OpenStage::Inspect( + Box::new(WebSocketStage { + entry, + transport: Some(WebSocketStageTransport { sender, responses }), + }), + invocation, + ) + } + Ok(WebSocketPreflightAction::Skip) => { + let invocation = success_invocation( + &entry, + WebSocketInvocationOutcome::Skip, + 0, + 0, + None, + false, + None, + ); + let _ = sender.try_send(session_end_request(WebSocketSessionEndReason::Cancellation)); + OpenStage::Skip(invocation) + } + Ok(WebSocketPreflightAction::Unspecified) | Err(_) => { + OpenStage::Failed(entry, "invalid_preflight_decision".into()) + } + } +} + +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> { + 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_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_BODY_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)] +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(); + 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 { + if let Some(transport) = stage.transport.take() { + let _ = tokio::time::timeout( + Duration::from_millis(10), + transport.sender.send(session_end_request(reason)), + ) + .await; + } + } +} + +fn session_end_request(reason: WebSocketSessionEndReason) -> WebSocketEvaluationRequest { + WebSocketEvaluationRequest { + request: Some(web_socket_evaluation_request::Request::SessionEnd( + WebSocketSessionEnd { + reason: reason as i32, + }, + )), + } +} 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..b555d250d9 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}; @@ -83,6 +84,206 @@ 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); + if outcome.saturated { + emit_websocket_saturation(ctx); + } + if outcome.session_capacity_exhausted { + emit_middleware_session_capacity_exhausted(ctx); + } +} + +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); + } + for finding in &outcome.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); +} + /// 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 @@ -163,6 +364,10 @@ pub async fn apply_middleware_chain_for_scheme, 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..1a3660ead6 100644 --- a/crates/openshell-supervisor-network/src/l7/relay.rs +++ b/crates/openshell-supervisor-network/src/l7/relay.rs @@ -72,12 +72,15 @@ pub(crate) struct UpgradeRelayOptions<'a> { pub(crate) websocket_request: bool, pub(crate) websocket: WebSocketUpgradeBehavior, 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 +473,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 @@ -506,6 +510,32 @@ where 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.allowed { + write_bad_gateway_response(client).await?; + return Ok(()); + } + preflight.session + } else { + None + }; let outcome = crate::l7::rest::relay_http_request_with_options_guarded( &req, client, @@ -525,11 +555,25 @@ where ) .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 +584,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 +685,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 +727,29 @@ 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, { + 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); + } + if !start.allowed { + 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 { @@ -716,8 +806,11 @@ where crate::l7::websocket::RelayOptions { policy_name: &options.policy_name, resolver, + generation_guard: options.generation_guard, inspector, compression, + middleware_session: options.middleware_session.take(), + middleware_context: options.ctx, }, ) .await; @@ -732,6 +825,18 @@ where 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, @@ -764,12 +869,15 @@ pub(crate) fn upgrade_options<'a>( } 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 +1064,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. @@ -990,6 +1099,32 @@ where 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.allowed { + write_bad_gateway_response(client).await?; + return Ok(()); + } + preflight.session + } else { + None + }; let req_with_auth = match crate::l7::token_grant_injection::inject_if_needed(req, ctx).await { Ok(req) => req, @@ -1025,8 +1160,19 @@ where ) .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 +1183,7 @@ where RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + websocket_subprotocol, } => { let mut options = upgrade_options( config, @@ -1047,6 +1194,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 +1224,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 +1241,17 @@ 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 } async fn relay_jsonrpc( @@ -1500,6 +1658,7 @@ where RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + .. } => { let options = UpgradeRelayOptions { websocket: WebSocketUpgradeBehavior { @@ -2192,12 +2351,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) { @@ -2393,6 +2554,293 @@ 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 EvaluateWebSocketStream = 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, + timeout: "2s".into(), + max_message_bytes: + openshell_supervisor_middleware::MAX_MIDDLEWARE_BODY_BYTES as u64, + ..Default::default() + }], + }, + )) + } + + 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( + &self, + request: tonic::Request< + tonic::Streaming, + >, + ) -> std::result::Result, tonic::Status> + { + use openshell_core::proto::{ + WebSocketEvaluationResponse, WebSocketPreflightDecision, + web_socket_evaluation_request, web_socket_evaluation_response, + }; + 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.request, + Some(web_socket_evaluation_request::Request::Preflight(_)) + ) { + let _ = seen.send(()); + release.notified().await; + let _ = responses_tx + .send(Ok(WebSocketEvaluationResponse { + response: Some( + web_socket_evaluation_response::Response::PreflightDecision( + WebSocketPreflightDecision { + action: action as i32, + }, + ), + ), + })) + .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_body_bytes: 0, + 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 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::Unspecified, + 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: @@ -3175,7 +3623,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 +3654,20 @@ network_policies: impl openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware for BodyReplacingService { + type EvaluateWebSocketStream = openshell_supervisor_middleware::WebSocketResponseStream; + + async fn evaluate_web_socket( + &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<()>, @@ -3224,6 +3686,7 @@ network_policies: as i32, max_body_bytes: 8192, timeout: String::new(), + max_message_bytes: 0, }], }, )) @@ -3693,6 +4156,20 @@ network_policies: impl openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware for LimitService { + type EvaluateWebSocketStream = openshell_supervisor_middleware::WebSocketResponseStream; + + async fn evaluate_web_socket( + &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<()>, @@ -3712,6 +4189,7 @@ network_policies: phase: SupervisorMiddlewarePhase::PreCredentials as i32, max_body_bytes: self.max_body_bytes, timeout: String::new(), + max_message_bytes: 0, }], })) } @@ -3764,16 +4242,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(), ) diff --git a/crates/openshell-supervisor-network/src/l7/rest.rs b/crates/openshell-supervisor-network/src/l7/rest.rs index fd9d373f81..a37814b5b0 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() @@ -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"))?; @@ -2760,7 +2766,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 +2785,7 @@ where return Ok(RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + websocket_subprotocol, }); } @@ -2908,9 +2915,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 +2978,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 +2994,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 +3572,7 @@ mod tests { let RelayOutcome::Upgraded { overflow, websocket_permessage_deflate, + .. } = outcome else { panic!("expected upgraded relay outcome"); @@ -6075,7 +6086,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 +6094,7 @@ mod tests { .expect("reordered safe extension params should canonicalize"); assert!(negotiated); + assert_eq!(subprotocol, None); } #[test] @@ -6105,7 +6117,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 +6125,7 @@ mod tests { .expect("preserve mode should not parse or reject raw extension negotiation"); assert!(!negotiated); + assert_eq!(subprotocol, None); } #[test] @@ -6136,7 +6149,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 +6157,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..ae6bbe270d 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,16 @@ use openshell_ocsf::{ ocsf_emit, }; use std::collections::HashMap; +use std::future::Future; +use std::io::{Error as IoError, ErrorKind as IoErrorKind}; +use std::time::Duration as StdDuration; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; -const MAX_TEXT_MESSAGE_BYTES: usize = 1024 * 1024; +const MAX_TEXT_MESSAGE_BYTES: usize = openshell_supervisor_middleware::MAX_MIDDLEWARE_BODY_BYTES; 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 +35,68 @@ const OPCODE_CLOSE: u8 = 0x8; const OPCODE_PING: u8 = 0x9; const OPCODE_PONG: u8 = 0xA; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum WebSocketTerminationCause { + PeerDisconnect, + PolicyReload, + MiddlewareDenial, + MiddlewareFailure, + PolicyDenial, + InvalidUtf8, + ProtocolError, + MessageTooBig, +} + +impl WebSocketTerminationCause { + fn close_code(self) -> Option { + match self { + Self::PeerDisconnect => None, + Self::PolicyReload => Some(1012), + 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::MiddlewareFailure => { + openshell_core::proto::WebSocketSessionEndReason::MiddlewareFailure + } + Self::InvalidUtf8 | Self::ProtocolError | Self::MessageTooBig => { + openshell_core::proto::WebSocketSessionEndReason::ProtocolError + } + } + } +} + +#[derive(Debug)] +struct WebSocketTermination { + cause: WebSocketTerminationCause, + error: miette::Report, +} + +#[derive(Debug)] +enum WebSocketDecompressionError { + MessageTooBig(miette::Report), + Protocol(miette::Report), +} + +type WebSocketRelayResult = std::result::Result; + +fn terminate(cause: WebSocketTerminationCause, error: miette::Report) -> WebSocketTermination { + WebSocketTermination { cause, error } +} + #[derive(Debug)] struct FrameHeader { fin: bool, @@ -43,8 +111,136 @@ struct FrameHeader { #[derive(Debug)] enum FragmentState { None, - Text { payload: Vec, compressed: bool }, - Binary, + Text(TextMessageAssembly), + Binary { fragment_count: usize }, +} + +/// 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, + admission: Option, + total_deadline: tokio::time::Instant, +} + +impl TextMessageAssembly { + fn new( + compressed: bool, + admission: Option, + ) -> Self { + Self { + payload: Vec::new(), + compressed, + fragment_count: 1, + admission, + total_deadline: tokio::time::Instant::now() + TEXT_MESSAGE_ASSEMBLY_TOTAL_TIMEOUT, + } + } + + async fn read_some( + &self, + reader: &mut R, + buffer: &mut [u8], + ) -> std::io::Result { + 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(IoError::new( + IoErrorKind::TimedOut, + "websocket text message assembly total timeout", + )) + } + () = tokio::time::sleep_until(idle_deadline) => { + Err(IoError::new( + IoErrorKind::TimedOut, + "websocket text message assembly idle timeout", + )) + } + result = reader.read(buffer) => result, + } + } + + async fn read_exact( + &self, + reader: &mut R, + buffer: &mut [u8], + ) -> std::io::Result<()> { + let mut filled = 0; + while filled < buffer.len() { + let read = self.read_some(reader, &mut buffer[filled..]).await?; + if read == 0 { + return Err(IoError::new( + IoErrorKind::UnexpectedEof, + "websocket payload ended before declared length", + )); + } + filled += read; + } + Ok(()) + } + + async fn read_payload( + &mut self, + reader: &mut R, + frame: &FrameHeader, + ) -> Result<()> { + let next = read_masked_payload(reader, frame, Some(self)).await?; + append_text_fragment(&mut self.payload, next) + } + + fn ensure_payload_fits(&self, frame: &FrameHeader) -> Result<()> { + let frame_len = usize::try_from(frame.payload_len) + .map_err(|_| miette!("websocket text frame is too large to buffer"))?; + let complete_len = self + .payload + .len() + .checked_add(frame_len) + .ok_or_else(|| miette!("websocket text message length overflow"))?; + if complete_len > MAX_TEXT_MESSAGE_BYTES { + return Err(miette!( + "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" + )); + } + Ok(()) + } + + fn add_fragment(&mut self) -> Result<()> { + self.fragment_count = self.fragment_count.saturating_add(1); + if self.fragment_count > MAX_MESSAGE_FRAGMENTS { + return Err(miette!( + "websocket message exceeds {MAX_MESSAGE_FRAGMENTS} fragment limit" + )); + } + Ok(()) + } + + async fn within_total(&self, future: impl Future>) -> Result { + tokio::time::timeout_at(self.total_deadline, future) + .await + .map_err(|_| miette!("websocket text message assembly total timeout"))? + } + + fn into_parts( + self, + ) -> ( + Vec, + bool, + Option, + ) { + (self.payload, self.compressed, self.admission) + } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -65,8 +261,11 @@ pub(super) struct InspectionOptions<'a> { pub(super) struct RelayOptions<'a> { pub(super) policy_name: &'a str, 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 +276,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 +290,54 @@ 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 + && 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,50 +345,90 @@ 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 + .inspect_err(|e| { + emit_protocol_failure(host, port, options.policy_name, protocol_failure_class(e)); + }) + .map_err(|error| terminate(WebSocketTerminationCause::ProtocolError, error))? 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); + let error = miette!("websocket frame received after close frame"); + emit_protocol_failure( + host, + port, + options.policy_name, + protocol_failure_class(&error), + ); + return Err(terminate(WebSocketTerminationCause::ProtocolError, error)); } 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); + return Err(terminate(WebSocketTerminationCause::ProtocolError, e)); + } + + if matches!( + frame.opcode, + OPCODE_TEXT | OPCODE_BINARY | OPCODE_CONTINUATION + ) && let Some(guard) = options.generation_guard + && let Err(error) = guard.ensure_current() + { + crate::l7::relay::emit_policy_reload(guard, host, port, options.policy_name); + return Err(terminate(WebSocketTerminationCause::PolicyReload, error)); } 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(terminate( + WebSocketTerminationCause::MessageTooBig, + miette!( + "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" + ), + )); + } + 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 + }; let compressed = frame.rsv == 0x40; - if frame.fin { - relay_text_payload( - writer, &frame, payload, false, compressed, host, port, options, - ) + let mut assembly = TextMessageAssembly::new(compressed, admission); + assembly + .read_payload(reader, &frame) .await .inspect_err(|e| { emit_protocol_failure( @@ -167,44 +437,75 @@ where options.policy_name, protocol_failure_class(e), ); + }) + .map_err(|error| terminate(WebSocketTerminationCause::ProtocolError, error))?; + if frame.fin { + let (payload, compressed, admission) = assembly.into_parts(); + relay_text_payload( + writer, &frame, payload, admission, false, compressed, host, port, options, + ) + .await + .inspect_err(|termination| { + observe_termination(host, port, options.policy_name, termination); })?; } else { - fragments = FragmentState::Text { - payload, - compressed, - }; + fragments = FragmentState::Text(assembly); } } - OPCODE_CONTINUATION => match &mut fragments { - FragmentState::Text { - payload, - compressed, - } => { - let next = read_masked_payload(reader, &frame).await.inspect_err(|e| { + OPCODE_CONTINUATION => { + if let FragmentState::Text(assembly) = &mut fragments { + if let Err(error) = assembly.add_fragment() { emit_protocol_failure( host, port, options.policy_name, - protocol_failure_class(e), + protocol_failure_class(&error), ); - })?; - if let Err(e) = append_text_fragment(payload, next) { + return Err(terminate(WebSocketTerminationCause::ProtocolError, error)); + } + if frame.payload_len > MAX_TEXT_MESSAGE_BYTES as u64 { + return Err(terminate( + WebSocketTerminationCause::MessageTooBig, + miette!( + "websocket text message exceeds {MAX_TEXT_MESSAGE_BYTES} byte limit" + ), + )); + } + if let Err(error) = assembly.ensure_payload_fits(&frame) { emit_protocol_failure( host, port, options.policy_name, - protocol_failure_class(&e), + protocol_failure_class(&error), ); - return Err(e); + return Err(terminate(WebSocketTerminationCause::MessageTooBig, error)); } + assembly + .read_payload(reader, &frame) + .await + .inspect_err(|e| { + emit_protocol_failure( + host, + port, + options.policy_name, + protocol_failure_class(e), + ); + }) + .map_err(|error| { + terminate(WebSocketTerminationCause::ProtocolError, error) + })?; 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, admission) = assembly.into_parts(); relay_text_payload( writer, &frame, complete, + admission, true, was_compressed, host, @@ -212,17 +513,20 @@ where options, ) .await - .inspect_err(|e| { - emit_protocol_failure( - host, - port, - options.policy_name, - protocol_failure_class(e), - ); + .inspect_err(|termination| { + observe_termination(host, port, options.policy_name, termination); })?; } - } - FragmentState::Binary => { + } else if let FragmentState::Binary { fragment_count } = &mut fragments { + *fragment_count = fragment_count.saturating_add(1); + if *fragment_count > MAX_MESSAGE_FRAGMENTS { + return Err(terminate( + WebSocketTerminationCause::ProtocolError, + miette!( + "websocket message exceeds {MAX_MESSAGE_FRAGMENTS} fragment limit" + ), + )); + } copy_raw_frame_payload(reader, writer, &frame) .await .inspect_err(|e| { @@ -232,12 +536,14 @@ where options.policy_name, protocol_failure_class(e), ); + }) + .map_err(|error| { + terminate(WebSocketTerminationCause::PeerDisconnect, error) })?; if frame.fin { fragments = FragmentState::None; } - } - FragmentState::None => { + } else { let e = miette!("websocket continuation frame without active fragmented message"); emit_protocol_failure( @@ -246,12 +552,12 @@ where options.policy_name, protocol_failure_class(&e), ); - return Err(e); + return Err(terminate(WebSocketTerminationCause::ProtocolError, e)); } - }, + } OPCODE_BINARY => { if !frame.fin { - fragments = FragmentState::Binary; + fragments = FragmentState::Binary { fragment_count: 1 }; } copy_raw_frame_payload(reader, writer, &frame) .await @@ -262,11 +568,26 @@ where options.policy_name, protocol_failure_class(e), ); - })?; + }) + .map_err(|error| terminate(WebSocketTerminationCause::PeerDisconnect, error))?; } OPCODE_CLOSE | OPCODE_PING | OPCODE_PONG => { - relay_control_frame(reader, writer, &frame) - .await + 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 .inspect_err(|e| { emit_protocol_failure( host, @@ -274,6 +595,14 @@ where options.policy_name, protocol_failure_class(e), ); + }) + .map_err(|error| { + let cause = if is_text_assembly_timeout(&error) { + WebSocketTerminationCause::ProtocolError + } else { + WebSocketTerminationCause::PeerDisconnect + }; + terminate(cause, error) })?; if frame.opcode == OPCODE_CLOSE { close_seen = true; @@ -284,9 +613,29 @@ where } } -async fn read_frame_header(reader: &mut R) -> Result> { - let first = match reader.read_u8().await { - Ok(byte) => byte, +async fn read_exact_for_assembly( + reader: &mut R, + buffer: &mut [u8], + assembly: Option<&TextMessageAssembly>, +) -> std::io::Result<()> { + match assembly { + Some(assembly) => assembly.read_exact(reader, buffer).await, + None => reader.read_exact(buffer).await.map(|_| ()), + } +} + +async fn read_frame_header( + reader: &mut R, + assembly: Option<&TextMessageAssembly>, +) -> Result> { + 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, + }; + let first = match first_read { + Ok(0) => return Ok(None), + Ok(_) => first[0], Err(e) if matches!( e.kind(), @@ -299,10 +648,11 @@ async fn read_frame_header(reader: &mut R) -> Result return Err(miette!("{e}")), }; - let second = reader - .read_u8() + let mut second = [0u8; 1]; + read_exact_for_assembly(reader, &mut second, assembly) .await .map_err(|e| miette!("malformed websocket frame header: {e}"))?; + let second = second[0]; let mut raw_header = vec![first, second]; let len_code = second & 0x7F; @@ -310,8 +660,7 @@ async fn read_frame_header(reader: &mut R) -> Result u64::from(len_code), 126 => { let mut bytes = [0u8; 2]; - reader - .read_exact(&mut bytes) + read_exact_for_assembly(reader, &mut bytes, assembly) .await .map_err(|e| miette!("malformed websocket extended length: {e}"))?; raw_header.extend_from_slice(&bytes); @@ -325,8 +674,7 @@ async fn read_frame_header(reader: &mut R) -> Result { let mut bytes = [0u8; 8]; - reader - .read_exact(&mut bytes) + read_exact_for_assembly(reader, &mut bytes, assembly) .await .map_err(|e| miette!("malformed websocket extended length: {e}"))?; if bytes[0] & 0x80 != 0 { @@ -347,8 +695,7 @@ async fn read_frame_header(reader: &mut R) -> Result MAX_RAW_FRAME_PAYLOAD_BYTES { return Err(miette!( @@ -440,6 +788,7 @@ fn valid_rsv_bits( async fn read_masked_payload( reader: &mut R, frame: &FrameHeader, + assembly: Option<&TextMessageAssembly>, ) -> Result> { let payload_len = usize::try_from(frame.payload_len) .map_err(|_| miette!("websocket text frame is too large to buffer"))?; @@ -449,8 +798,7 @@ async fn read_masked_payload( )); } let mut payload = vec![0u8; payload_len]; - reader - .read_exact(&mut payload) + read_exact_for_assembly(reader, &mut payload, assembly) .await .map_err(|e| miette!("malformed websocket payload: {e}"))?; let mask_key = frame @@ -479,43 +827,126 @@ async fn relay_text_payload( writer: &mut W, frame: &FrameHeader, payload: Vec, + admission: Option, force_reframe: bool, compressed: bool, host: &str, port: u16, - options: &RelayOptions<'_>, -) -> Result<()> { + options: &mut RelayOptions<'_>, +) -> WebSocketRelayResult<()> { let message_payload = if compressed { - decompress_permessage_deflate(&payload)? + decompress_permessage_deflate(&payload).map_err(|error| match error { + WebSocketDecompressionError::MessageTooBig(error) => { + terminate(WebSocketTerminationCause::MessageTooBig, error) + } + WebSocketDecompressionError::Protocol(error) => { + terminate(WebSocketTerminationCause::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(|_| { + terminate( + WebSocketTerminationCause::InvalidUtf8, + 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)?; + } + + 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; + 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(terminate( + WebSocketTerminationCause::MessageTooBig, + miette!("websocket message over middleware platform capacity"), + )); + } + 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)?; + } + 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 }; - 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}"), + ) + })?; 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(|| { + terminate( + WebSocketTerminationCause::ProtocolError, + 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 +954,15 @@ 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| terminate(WebSocketTerminationCause::ProtocolError, error))?; + return write_masked_frame_with_rsv(writer, OPCODE_TEXT, 0x40, &compressed_payload) + .await + .map_err(|error| terminate(WebSocketTerminationCause::PeerDisconnect, error)); } - write_masked_frame(writer, OPCODE_TEXT, text.as_bytes()).await + write_masked_frame(writer, OPCODE_TEXT, text.as_bytes()) + .await + .map_err(|error| terminate(WebSocketTerminationCause::PeerDisconnect, error)) } fn inspect_websocket_text_message( @@ -535,7 +971,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 +983,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 +1000,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 +1014,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 +1054,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 +1074,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,6 +1172,7 @@ async fn relay_control_frame( reader: &mut R, writer: &mut W, frame: &FrameHeader, + assembly: Option<&TextMessageAssembly>, ) -> Result<()> where R: AsyncRead + Unpin, @@ -736,8 +1181,7 @@ where let raw_payload_len = usize::try_from(frame.payload_len) .map_err(|_| miette!("websocket control frame payload length overflow"))?; let mut raw_payload = vec![0u8; raw_payload_len]; - reader - .read_exact(&mut raw_payload) + read_exact_for_assembly(reader, &mut raw_payload, assembly) .await .map_err(|e| miette!("malformed websocket control payload: {e}"))?; @@ -821,6 +1265,28 @@ 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, @@ -855,7 +1321,9 @@ async fn write_masked_frame_with_rsv( Ok(()) } -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 +1336,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 +1369,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) @@ -1050,10 +1530,16 @@ fn graphql_log_summary(info: &crate::l7::graphql::GraphqlRequestInfo) -> String fn protocol_failure_class(error: &miette::Report) -> &'static str { let msg = error.to_string().to_ascii_lowercase(); - if msg.contains("credential") { + if msg.contains("text message assembly idle timeout") { + "message_assembly_idle_timeout" + } else if msg.contains("text message assembly total timeout") { + "message_assembly_total_timeout" + } else if msg.contains("credential") { "credential_resolution_failed" } else if msg.contains("utf-8") { "invalid_utf8" + } else if msg.contains("fragment limit") { + "invalid_fragmentation" } else if msg.contains("close frame") || msg.contains("after close") { "invalid_close_frame" } else if msg.contains("control frame") { @@ -1079,6 +1565,34 @@ fn protocol_failure_class(error: &miette::Report) -> &'static str { } } +fn is_text_assembly_timeout(error: &miette::Report) -> bool { + error + .to_string() + .to_ascii_lowercase() + .contains("text message assembly") +} + +fn observe_termination( + host: &str, + port: u16, + policy_name: &str, + termination: &WebSocketTermination, +) { + if matches!( + termination.cause, + WebSocketTerminationCause::InvalidUtf8 + | WebSocketTerminationCause::ProtocolError + | WebSocketTerminationCause::MessageTooBig + ) { + emit_protocol_failure( + host, + port, + policy_name, + protocol_failure_class(&termination.error), + ); + } +} + fn emit_protocol_failure(host: &str, port: u16, policy_name: &str, failure_class: &str) { let policy_name = if policy_name.is_empty() { "-" @@ -1108,9 +1622,14 @@ 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}; const TEST_POLICY: &str = include_str!("../../data/sandbox-policy.rego"); const GRAPHQL_WS_POLICY: &str = r#" @@ -1137,6 +1656,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 +1775,23 @@ 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 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 +1807,30 @@ network_policies: client_write.write_all(&input).await.unwrap(); drop(client_write); - let options = RelayOptions { + let mut options = RelayOptions { policy_name: "test-policy", 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) } async fn run_client_to_server_with_graphql_policy( @@ -1281,9 +1870,10 @@ network_policies: client_write.write_all(&input).await.unwrap(); drop(client_write); - let options = RelayOptions { + let mut options = RelayOptions { policy_name: "graphql_ws", resolver, + generation_guard: Some(tunnel_engine.generation_guard()), inspector: Some(InspectionOptions { engine: &tunnel_engine, ctx: &ctx, @@ -1293,20 +1883,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 +1911,30 @@ network_policies: client_write.write_all(&input).await.unwrap(); drop(client_write); - let options = RelayOptions { + let mut options = RelayOptions { policy_name: "test-policy", 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,16 +2010,203 @@ network_policies: frame } - #[test] - fn classifies_graphql_transport_ws_subscribe_operation() { - let message = r#"{"type":"subscribe","id":"1","payload":{"query":"subscription NewMessages { messageAdded }"}}"#; + #[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() + .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, Some(admission)); + assembly.read_payload(&mut reader, &frame).await + }); - match classify_graphql_websocket_message(message) { - GraphqlWebSocketMessage::Operation { - message_type, - graphql, - } => { - assert_eq!(message_type, "subscribe"); + 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() + .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() + .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, 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() + .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() + .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, 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() + .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() + .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, 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() + .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() + .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, 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().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 }"}}"#; + + match classify_graphql_websocket_message(message) { + GraphqlWebSocketMessage::Operation { + message_type, + graphql, + } => { + assert_eq!(message_type, "subscribe"); assert!( graphql.error.is_none(), "unexpected error: {:?}", @@ -1557,8 +2343,11 @@ network_policies: RelayOptions { policy_name: "test-policy", resolver: Some(&resolver), + generation_guard: None, inspector: None, compression: WebSocketCompression::None, + middleware_session: None, + middleware_context: None, }, ) .await @@ -1583,6 +2372,847 @@ network_policies: let _ = tokio::time::timeout(std::time::Duration::from_secs(2), relay).await; } + #[derive(Debug)] + enum ObservedWebSocketRequest { + SessionStart, + Message(Vec), + SessionEnd(openshell_core::proto::WebSocketSessionEndReason), + } + + #[derive(Clone, Default)] + struct OpenAiWebSocketRedactor { + observed: Option>, + close_on_first_message: bool, + } + + #[tonic::async_trait] + impl SupervisorMiddleware for OpenAiWebSocketRedactor { + type EvaluateWebSocketStream = 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_body_bytes: 0, + timeout: "1s".into(), + max_message_bytes: openshell_supervisor_middleware::MAX_MIDDLEWARE_BODY_BYTES + as u64, + }], + })) + } + + 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( + &self, + request: Request>, + ) -> std::result::Result, Status> { + use openshell_core::proto::{ + Decision, WebSocketEvaluationResponse, WebSocketMessageResult, + WebSocketPreflightAction, WebSocketPreflightDecision, + web_socket_evaluation_request, web_socket_evaluation_response, + }; + let mut requests = request.into_inner(); + let observed = self.observed.clone(); + let close_on_first_message = self.close_on_first_message; + 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.request { + Some(web_socket_evaluation_request::Request::Preflight(_)) => { + Some(WebSocketEvaluationResponse { + response: Some( + web_socket_evaluation_response::Response::PreflightDecision( + WebSocketPreflightDecision { + action: WebSocketPreflightAction::Inspect as i32, + }, + ), + ), + }) + } + Some(web_socket_evaluation_request::Request::Message(message)) => { + if let Some(observed) = &observed { + let _ = observed.send(ObservedWebSocketRequest::Message( + message.payload.clone(), + )); + } + if close_on_first_message { + break; + } + 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(WebSocketEvaluationResponse { + response: Some( + web_socket_evaluation_response::Response::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_evaluation_request::Request::SessionStart(_)) => { + if let Some(observed) = &observed { + let _ = observed.send(ObservedWebSocketRequest::SessionStart); + } + None + } + Some(web_socket_evaluation_request::Request::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>, + ) { + 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), + ..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_body_bytes: 0, + 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_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 = std::sync::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(1); + 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, + 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) + )); + if fragmented { + client_app + .write_all(&masked_frame(false, OPCODE_TEXT, b"stale-")) + .await + .expect("send initial fragment"); + } + + engine + .replace_middleware_registry(MiddlewareRegistry::default()) + .expect("invalidate generation"); + let mut stale_frame = if fragmented { + masked_frame(true, OPCODE_CONTINUATION, b"message") + } else { + masked_frame(true, OPCODE_TEXT, b"stale-message") + }; + stale_frame.truncate(6); + client_app + .write_all(&stale_frame) + .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 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", + 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_body_bytes: 0, + 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", + 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, + })) + .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_body_bytes: 0, + 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() + .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", + 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", + 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/proxy.rs b/crates/openshell-supervisor-network/src/proxy.rs index 6e9c48220b..79a2072dd4 100644 --- a/crates/openshell-supervisor-network/src/proxy.rs +++ b/crates/openshell-supervisor-network/src/proxy.rs @@ -4915,6 +4915,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( @@ -4952,6 +4953,58 @@ async fn handle_forward_proxy( }; } + 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.allowed { + respond( + client, + &build_json_error_response( + 502, + "Bad Gateway", + "middleware_failed", + "WebSocket middleware preflight denied the upgrade", + ), + ) + .await?; + return Ok(()); + } + preflight.session + } else { + None + }; + forward_request_bytes = match inject_token_grant_for_forward_request( method, &upstream_target, @@ -5077,39 +5130,54 @@ 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.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(()) @@ -5869,7 +5937,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() @@ -6192,6 +6260,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 +6319,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 +6425,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 +6466,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 +6539,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 +6593,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 +8908,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..822bb7c1e3 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 selects the same host-matched chain and opens one ordered `EvaluateWebSocket` stream per WebSocket-capable stage. Each stream receives: + +1. A preflight before the upgrade is sent upstream. The stage chooses `INSPECT` or `SKIP`. +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. + +Allowed replacements are re-framed, re-compressed when required, and forwarded. Binary messages, control frames, and upstream-to-client traffic remain uninspected. OpenShell reserves shared middleware capacity before buffering HTTP bodies or WebSocket text messages; at most 32 evaluations run and 32 additional unbuffered callers wait for capacity. 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. @@ -35,9 +44,9 @@ Middleware receives the request before credential injection. Operator-run servic | 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 @@ -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_body_bytes` | Operator limit applied to HTTP bindings exposed by the service, up to the 4 MiB platform maximum. A WebSocket-only service sets this to `0`; each WebSocket binding declares `max_message_bytes` in its manifest. | | `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,14 +101,14 @@ 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 when middleware is unavailable, rejects its configuration, returns an invalid result, or exceeds a payload limit. | 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`. @@ -130,6 +139,8 @@ The gateway rejects a registration whose operator limit exceeds the service capa 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. +WebSocket bindings use `max_message_bytes` for complete client text messages and replacements. The platform cap is 4 MiB, with at most 4,096 fragments per message. Oversized messages close 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: @@ -171,8 +182,9 @@ 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 to parsed HTTP requests and client text messages on validated RFC 6455 upgrades. 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 pairs are `HTTP_REQUEST/PRE_CREDENTIALS` and `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`. +- WebSocket middleware does not inspect binary messages, control frames, or upstream-to-client messages. - 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..dd18fbaa7d 100644 --- a/docs/observability/logging.mdx +++ b/docs/observability/logging.mdx @@ -137,6 +137,14 @@ 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=- +``` + +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; work beyond the bounded wait queue is rejected before its payload is buffered. + Proxy and SSH servers ready: ```text @@ -167,7 +175,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 75ac6f7382..810d82c081 100644 --- a/docs/reference/gateway-config.mdx +++ b/docs/reference/gateway-config.mdx @@ -187,13 +187,13 @@ max_body_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_body_bytes` is the operator limit for HTTP bindings exposed by the service. It must be greater than zero when the manifest includes an HTTP binding, no larger than each HTTP binding's advertised limit, and no larger than the 4 MiB platform maximum. Set it to `0` for a WebSocket-only service. Each WebSocket binding declares its own `max_message_bytes`, also capped at 4 MiB. OpenShell rejects oversized values instead of silently clamping them. 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. diff --git a/docs/reference/policy-schema.mdx b/docs/reference/policy-schema.mdx index d585517d13..39b7aada81 100644 --- a/docs/reference/policy-schema.mdx +++ b/docs/reference/policy-schema.mdx @@ -490,7 +490,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: @@ -512,10 +512,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 | `fail_closed` denies the HTTP request or closes the WebSocket when the 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. Middleware runs on HTTP requests and validated RFC 6455 upgrades that the supervisor parses; WebSocket bindings inspect only complete client text messages. 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. 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 990f260cea..b02911299f 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 | Sets the OS-level identity for the agent process. `run_as_user` and `run_as_group` default to `sandbox`. Root (`root` or `0`) is 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 failed middleware 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. Middleware applies only to parsed HTTP traffic and validated RFC 6455 upgrades; binary and upstream-to-client WebSocket 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/proto/supervisor_middleware.proto b/proto/supervisor_middleware.proto index dbde411c9f..6975717494 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,13 @@ service SupervisorMiddleware { // EvaluateHttpRequest returns an allow, deny, or mutation decision for one // buffered HTTP request. rpc EvaluateHttpRequest(HttpRequestEvaluation) returns (HttpRequestResult); + + // EvaluateWebSocket opens one ordered stream for a single middleware stage + // and WebSocket upgrade attempt. PR 1 supports client-to-upstream text + // messages at PRE_CREDENTIALS. A request may go unanswered when the session + // terminates; OpenShell then sends session_end on a best-effort basis. + rpc EvaluateWebSocket(stream WebSocketEvaluationRequest) + returns (stream WebSocketEvaluationResponse); } // MiddlewareManifest describes one middleware service and the bindings it @@ -38,11 +46,14 @@ 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. + // Maximum request or replacement HTTP body this binding can process. + // Required for HTTP_REQUEST and omitted for WEBSOCKET_MESSAGE. uint64 max_body_bytes = 3; // Optional binding-specific RPC timeout. Empty uses the operator-configured // service timeout, or the 500ms platform default when that is also omitted. @@ -50,6 +61,9 @@ message MiddlewareBinding { // Values use an integer with an `ms` or `s` suffix and must be between // 10ms and 30s. string timeout = 4; + // Maximum complete WebSocket message or replacement this binding can + // process. Required for WEBSOCKET_MESSAGE and omitted for HTTP_REQUEST. + uint64 max_message_bytes = 5; } // ValidateConfigRequest contains one policy configuration to validate. @@ -104,12 +118,136 @@ 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; +} + +// Direction of one complete logical WebSocket message. +enum WebSocketDirection { + WEB_SOCKET_DIRECTION_UNSPECIFIED = 0; + WEB_SOCKET_DIRECTION_CLIENT_TO_UPSTREAM = 1; + WEB_SOCKET_DIRECTION_UPSTREAM_TO_CLIENT = 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; +} + +// WebSocketEvaluationRequest 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 WebSocketEvaluationRequest { + oneof request { + 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; + WebSocketDirection direction = 3; + RequestContext context = 4; + string scheme = 5; + string host = 6; + uint32 port = 7; + // Raw request path without the query string. + string path = 8; + repeated string requested_subprotocols = 9; + // Built-in middleware name or operator-owned registration name. + string middleware_name = 10; + // Stable policy-local configuration identity. + string config_name = 11; + google.protobuf.Struct config = 12; +} + +// 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; + WebSocketDirection direction = 2; + WebSocketMessageType message_type = 3; + // Limited to 4 MiB by the platform and the binding-specific cap. + bytes payload = 4; +} + +message WebSocketSessionEnd { + WebSocketSessionEndReason reason = 1; +} + +// WebSocketPreflightAction is the service's one-time scoping decision. +enum WebSocketPreflightAction { + WEB_SOCKET_PREFLIGHT_ACTION_UNSPECIFIED = 0; + WEB_SOCKET_PREFLIGHT_ACTION_INSPECT = 1; + WEB_SOCKET_PREFLIGHT_ACTION_SKIP = 2; +} + +message WebSocketPreflightDecision { + WebSocketPreflightAction action = 1; +} + +// 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; +} + +message WebSocketEvaluationResponse { + oneof response { + WebSocketPreflightDecision preflight_decision = 1; + WebSocketMessageResult message_result = 2; + } } // RequestContext identifies the sandbox request being evaluated. diff --git a/rfc/0009-supervisor-middleware/README.md b/rfc/0009-supervisor-middleware/README.md index f43e05ed79..08a72a3da9 100644 --- a/rfc/0009-supervisor-middleware/README.md +++ b/rfc/0009-supervisor-middleware/README.md @@ -276,10 +276,14 @@ 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. +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 `EvaluateWebSocket` RPC. 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 HTTP contract in-process and does not advertise WebSocket support. 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_message_bytes`. OpenShell opens one stream per selected stage and upgrade attempt. A bounded preflight exposes only the admitted destination, path without query data, requested subprotocols, sandbox context, and validated middleware config. The service returns `inspect` or a voluntary `skip` before OpenShell forwards the upgrade upstream. Inspecting stages 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. `PRE_RETURN`, binary messages, and upstream-to-client inspection are reserved for a later implementation. + +WebSocket input and replacements share the 4 MiB platform cap. Text replacements must be valid UTF-8 and preserve message type. A complete 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. ### Relationship to RFC 0010 @@ -440,7 +444,7 @@ 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. +**Phase 1 - research-preview contract and execution.** Define `openshell.middleware.v1` with `Describe`, `ValidateConfig`, unary `EvaluateHttpRequest`, and forward-text `EvaluateWebSocket`; ship the example `openshell/regex` HTTP 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. @@ -459,7 +463,7 @@ Adding a synchronous, content-aware hook to the egress path has real costs. The - **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). diff --git a/rfc/0009-supervisor-middleware/appendices/protocol-extensions.md b/rfc/0009-supervisor-middleware/appendices/protocol-extensions.md index 924c7bd699..6b37a5d1a5 100644 --- a/rfc/0009-supervisor-middleware/appendices/protocol-extensions.md +++ b/rfc/0009-supervisor-middleware/appendices/protocol-extensions.md @@ -2,7 +2,7 @@ > 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. +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 +42,17 @@ 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: +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. A concurrent preflight resolves inspecting stages before upstream contact. +- `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