feat(middleware): extend content guard response example

Signed-off-by: Piotr Mlocek <pmlocek@nvidia.com>
This commit is contained in:
Piotr Mlocek
2026-08-31 23:07:44 -07:00
parent 8d6e9c11a0
commit 3f8a40efcd
8 changed files with 760 additions and 14 deletions
+2
View File
@@ -154,12 +154,14 @@ jobs:
cargo fmt --all -- --check
cargo fmt --manifest-path e2e/rust/Cargo.toml --all -- --check
cargo fmt --manifest-path examples/governance-interceptor/Cargo.toml --all -- --check
cargo fmt --manifest-path examples/supervisor-middleware-content-guard/Cargo.toml --all -- --check
- name: Lint
run: |
cargo clippy --workspace --all-targets -- -D warnings
cargo clippy --manifest-path e2e/rust/Cargo.toml --all-targets -- -D warnings
cargo check --manifest-path examples/governance-interceptor/Cargo.toml --all-targets
cargo check --manifest-path examples/supervisor-middleware-content-guard/Cargo.toml --all-targets
- name: Test
env:
+264
View File
@@ -206,6 +206,17 @@ version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527"
[[package]]
name = "chacha20"
version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06"
dependencies = [
"cfg-if",
"cpufeatures",
"rand_core",
]
[[package]]
name = "clap"
version = "4.6.1"
@@ -252,6 +263,16 @@ version = "1.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570"
[[package]]
name = "combine"
version = "4.6.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e"
dependencies = [
"bytes",
"memchr",
]
[[package]]
name = "core-foundation"
version = "0.10.1"
@@ -268,6 +289,27 @@ version = "0.8.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
[[package]]
name = "cpufeatures"
version = "0.3.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5ca28b0ae3115b884660db4118d803791fd6756b6e88f39c0f3f7859060d7566"
dependencies = [
"libc",
]
[[package]]
name = "critical-section"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b"
[[package]]
name = "data-encoding"
version = "2.11.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06"
[[package]]
name = "displaydoc"
version = "0.2.6"
@@ -399,6 +441,7 @@ dependencies = [
"cfg-if",
"libc",
"r-efi",
"rand_core",
]
[[package]]
@@ -453,6 +496,25 @@ version = "0.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
[[package]]
name = "hickory-proto"
version = "0.26.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0bab31817bfb44672a252e97fe81cd0c18d1b2cf892108922f6818820df8c643"
dependencies = [
"data-encoding",
"idna",
"ipnet",
"jni",
"once_cell",
"rand",
"ring",
"thiserror",
"tinyvec",
"tracing",
"url",
]
[[package]]
name = "http"
version = "1.4.2"
@@ -699,12 +761,71 @@ version = "1.0.18"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682"
[[package]]
name = "jni"
version = "0.22.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498"
dependencies = [
"cfg-if",
"combine",
"jni-macros",
"jni-sys",
"log",
"simd_cesu8",
"thiserror",
"walkdir",
"windows-link",
]
[[package]]
name = "jni-macros"
version = "0.22.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3"
dependencies = [
"proc-macro2",
"quote",
"rustc_version",
"simd_cesu8",
"syn",
]
[[package]]
name = "jni-sys"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2"
dependencies = [
"jni-sys-macros",
]
[[package]]
name = "jni-sys-macros"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264"
dependencies = [
"quote",
"syn",
]
[[package]]
name = "libc"
version = "0.2.186"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66"
[[package]]
name = "libyml"
version = "0.0.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3302702afa434ffa30847a83305f0a69d6abd74293b6554c18ec85c7ef30c980"
dependencies = [
"anyhow",
"version_check",
]
[[package]]
name = "linux-raw-sys"
version = "0.12.1"
@@ -832,6 +953,10 @@ name = "once_cell"
version = "1.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
dependencies = [
"critical-section",
"portable-atomic",
]
[[package]]
name = "once_cell_polyfill"
@@ -853,6 +978,7 @@ dependencies = [
"prost",
"prost-types",
"protoc-bin-vendored",
"rustix",
"serde",
"serde_json",
"thiserror",
@@ -877,12 +1003,26 @@ dependencies = [
"tower",
]
[[package]]
name = "openshell-policy"
version = "0.0.0"
dependencies = [
"hickory-proto",
"miette",
"openshell-core",
"prost-types",
"serde",
"serde_json",
"serde_yml",
]
[[package]]
name = "openshell-supervisor-middleware-content-guard"
version = "0.0.0"
dependencies = [
"clap",
"openshell-core",
"openshell-policy",
"prost-types",
"tokio",
"tokio-stream",
@@ -967,6 +1107,12 @@ version = "0.2.17"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
[[package]]
name = "portable-atomic"
version = "1.15.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
[[package]]
name = "potential_utf"
version = "0.1.5"
@@ -1147,6 +1293,23 @@ version = "6.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
[[package]]
name = "rand"
version = "0.10.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80"
dependencies = [
"chacha20",
"getrandom 0.4.3",
"rand_core",
]
[[package]]
name = "rand_core"
version = "0.10.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69"
[[package]]
name = "redox_syscall"
version = "0.5.18"
@@ -1205,6 +1368,15 @@ version = "0.1.27"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b50b8869d9fc858ce7266cce0194bd74df58b9d0e3f6df3a9fc8eb470d95c09d"
[[package]]
name = "rustc_version"
version = "0.4.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92"
dependencies = [
"semver",
]
[[package]]
name = "rustix"
version = "1.1.4"
@@ -1265,6 +1437,21 @@ dependencies = [
"untrusted",
]
[[package]]
name = "ryu"
version = "1.0.23"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f"
[[package]]
name = "same-file"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502"
dependencies = [
"winapi-util",
]
[[package]]
name = "schannel"
version = "0.1.29"
@@ -1303,6 +1490,12 @@ dependencies = [
"libc",
]
[[package]]
name = "semver"
version = "1.0.28"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd"
[[package]]
name = "serde"
version = "1.0.228"
@@ -1346,6 +1539,21 @@ dependencies = [
"zmij",
]
[[package]]
name = "serde_yml"
version = "0.0.12"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "59e2dd588bf1597a252c3b920e0143eb99b0f76e4e082f4c92ce34fbc9e71ddd"
dependencies = [
"indexmap",
"itoa",
"libyml",
"memchr",
"ryu",
"serde",
"version_check",
]
[[package]]
name = "shlex"
version = "2.0.1"
@@ -1362,6 +1570,22 @@ dependencies = [
"libc",
]
[[package]]
name = "simd_cesu8"
version = "1.2.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "11031e251abf8611c80f460e19dbdeb54a66db918e49c65a7065b46ac7aec520"
dependencies = [
"rustc_version",
"simdutf8",
]
[[package]]
name = "simdutf8"
version = "0.1.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e"
[[package]]
name = "slab"
version = "0.4.12"
@@ -1514,6 +1738,21 @@ dependencies = [
"zerovec",
]
[[package]]
name = "tinyvec"
version = "1.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f"
dependencies = [
"tinyvec_macros",
]
[[package]]
name = "tinyvec_macros"
version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
[[package]]
name = "tokio"
version = "1.52.3"
@@ -1774,6 +2013,22 @@ version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821"
[[package]]
name = "version_check"
version = "0.9.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
[[package]]
name = "walkdir"
version = "2.5.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b"
dependencies = [
"same-file",
"winapi-util",
]
[[package]]
name = "want"
version = "0.3.1"
@@ -1789,6 +2044,15 @@ version = "0.11.1+wasi-snapshot-preview1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b"
[[package]]
name = "winapi-util"
version = "0.1.11"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22"
dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "windows-link"
version = "0.2.1"
@@ -20,6 +20,9 @@ tokio = { version = "1.43", features = ["macros", "rt-multi-thread"] }
tokio-stream = "0.1"
tonic = { version = "0.14", features = ["transport"] }
[dev-dependencies]
openshell-policy = { path = "../../crates/openshell-policy" }
[[bin]]
name = "supervisor-middleware-content-guard"
path = "src/main.rs"
@@ -8,14 +8,14 @@ SPDX-License-Identifier: Apache-2.0
> [!WARNING]
> Supervisor middleware is a research preview. Its policy and service contracts may change without compatibility guarantees. Use it only to prototype and evaluate middleware integrations.
This example implements an operator-run supervisor middleware service. It scans UTF-8 HTTP request bodies and complete client-to-upstream WebSocket text messages for configured literal strings, then either replaces every match or denies the request or message. Findings report only aggregate counts and never include configured terms or inspected content.
This example implements one operator-run service with HTTP request, HTTP response, and client WebSocket bindings. Request bodies and WebSocket text messages use the content-guard behavior. The response binding demonstrates headers-only, whole-body, and normalized streaming processing.
> [!WARNING]
> This intentionally simple implementation demonstrates the supervisor middleware service contract. It is not a complete or reliable content guard and must not be used as a security control. It handles only UTF-8 HTTP request bodies and WebSocket text messages with case-sensitive literal terms, merges overlapping literal match ranges before redaction, and does not address encodings, transformations, normalization, binary WebSocket messages, upstream-to-client messages, or adversarial inputs that a production content guard must handle.
> This intentionally simple implementation demonstrates the supervisor middleware service contract. It is not a complete or reliable content guard and must not be used as a security control. Request and WebSocket handling supports only case-sensitive literal matches in UTF-8 text. The response paths demonstrate framing and body modes with fixed transformations; they do not implement content detection. The example does not address content encodings, application-level normalization, binary WebSocket messages, or adversarial inputs.
## Prerequisites
Install `cargo`, `curl`, `jq`, and `openssl` on the host before running the smoke script.
Install `cargo`, `curl`, `jq`, `openssl`, and Python 3 on the host before running the example.
## Run the smoke example
@@ -54,6 +54,7 @@ Add the service registration to your local gateway TOML:
[[openshell.supervisor.middleware]]
name = "content-guard-example"
grpc_endpoint = "http://host.openshell.internal:50051"
allow_insecure_transport = true
max_payload_bytes = 262144
timeout = "500ms"
```
@@ -84,6 +85,32 @@ curl -sS https://httpbin.org/anything \
The echoed JSON body contains `[FILTERED]` instead of the configured term.
## HTTP response behavior
The same service advertises `HTTP_RESPONSE/PRE_RETURN`. Start the included raw HTTP upstream in another terminal:
```shell
python3 examples/supervisor-middleware-content-guard/upstream.py
```
The example policy allows these requests from the sandbox:
```shell
curl -i http://host.openshell.internal:18081/headers-only
curl -i http://host.openshell.internal:18081/whole-body
curl -i --raw http://host.openshell.internal:18081/stream
```
The paths exercise each response mode:
| Path | Mode | Result |
| --- | --- | --- |
| `/headers-only` | `HEADERS_ONLY` | Adds `x-example-response-mode` and passes the body through. |
| `/whole-body` | `WHOLE_BODY_BYTES` | Prefixes the normalized complete body with `[whole]` plus a space. |
| `/stream` | `STREAM_BYTES` | Uppercases each normalized unit and writes the declared `x-example-body-bytes` trailer. |
The upstream uses content-length, chunked, and close-delimited responses. OpenShell removes those transport boundaries before invoking middleware and repairs downstream framing after a transformation. Paths outside this table return `SKIP` with `path_not_selected`.
## WebSocket behavior
For a selected WebSocket upgrade, the service accepts preflight, waits for the session-start notification, and evaluates each complete client-to-upstream text message. Redact mode returns a replacement message, while deny mode returns `content_match` and OpenShell closes the session according to middleware policy. Session-start and session-end events are notifications and do not produce results.
@@ -107,4 +134,4 @@ config:
- prototype-secret
```
The implementation supports `HTTP_REQUEST/PRE_CREDENTIALS` and `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`, advertises a 256 KiB limit for each operation, and inherits the service-wide RPC timeout. The gateway registration's `max_payload_bytes` may set a smaller shared limit. A binding can advertise a shorter timeout, but it cannot extend the operator-configured timeout.
The implementation supports `HTTP_REQUEST/PRE_CREDENTIALS`, `HTTP_RESPONSE/PRE_RETURN`, and `WEBSOCKET_MESSAGE/PRE_CREDENTIALS`. It advertises a 256 KiB limit for each operation and inherits the service-wide RPC timeout. The gateway registration's `max_payload_bytes` may set a smaller shared limit. A binding can advertise a shorter timeout, but it cannot extend the operator-configured timeout.
@@ -18,6 +18,7 @@ network_middlewares:
endpoints:
include:
- httpbin.org
- host.openshell.internal
network_policies:
httpbin:
@@ -44,3 +45,21 @@ network_policies:
path: /anything
binaries:
- path: /usr/bin/curl
response-framing-demo:
name: Response framing demo
endpoints:
- host: host.openshell.internal
port: 18081
protocol: rest
rules:
- allow:
method: GET
path: /headers-only
- allow:
method: GET
path: /whole-body
- allow:
method: GET
path: /stream
binaries:
- path: /usr/bin/curl
@@ -6,20 +6,31 @@ use std::net::SocketAddr;
use std::ops::Range;
use clap::Parser;
use openshell_core::middleware::WebSocketResponseStream;
use openshell_core::middleware::{HttpResponseResultStream, WebSocketResponseStream};
use openshell_core::proto::middleware::v1::http_response_pre_return_server::{
HttpResponsePreReturn, HttpResponsePreReturnServer,
};
use openshell_core::proto::middleware::v1::supervisor_middleware_server::{
SupervisorMiddleware, SupervisorMiddlewareServer,
};
use openshell_core::proto::{
Decision, Finding, HttpRequestEvaluation, HttpRequestResult, MiddlewareBinding,
MiddlewareManifest, SupervisorMiddlewareOperation, SupervisorMiddlewarePhase,
ValidateConfigRequest, ValidateConfigResponse, WebSocketMessage, WebSocketMessageResult,
WebSocketPreflightAction, WebSocketPreflightDecision, WebSocketSessionEvent,
WebSocketSessionEventResult, web_socket_message, web_socket_message_result,
web_socket_session_event, web_socket_session_event_result,
Decision, ExistingHeaderAction, Finding, HeaderMutation, HttpRequestEvaluation,
HttpRequestResult, HttpResponseBodyMode, HttpResponseBodyResult, HttpResponseBodyTransform,
HttpResponseEvent, HttpResponseEventResult, HttpResponsePreflightDecision,
HttpResponsePreflightInspect, HttpResponsePreflightSkip, HttpResponseTrailersResult,
MiddlewareBinding, MiddlewareManifest, SupervisorMiddlewareOperation,
SupervisorMiddlewarePhase, ValidateConfigRequest, ValidateConfigResponse, WebSocketMessage,
WebSocketMessageResult, WebSocketPreflightAction, WebSocketPreflightDecision,
WebSocketSessionEvent, WebSocketSessionEventResult, WriteHeader, header_mutation,
http_response_body_result, http_response_body_transform, http_response_body_unit,
http_response_event, http_response_event_result, http_response_preflight_decision,
web_socket_message, web_socket_message_result, web_socket_session_event,
web_socket_session_event_result,
};
use prost_types::Struct;
use prost_types::value::Kind;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tokio_stream::{Stream, StreamExt};
use tonic::transport::Server;
use tonic::{Request, Response, Status};
@@ -238,6 +249,12 @@ impl SupervisorMiddleware for ContentGuard {
max_payload_bytes: MAX_PAYLOAD_BYTES,
timeout: String::new(),
},
MiddlewareBinding {
operation: SupervisorMiddlewareOperation::HttpResponse as i32,
phase: SupervisorMiddlewarePhase::PreReturn as i32,
max_payload_bytes: MAX_PAYLOAD_BYTES,
timeout: String::new(),
},
],
expected_audience: String::new(),
}))
@@ -282,6 +299,240 @@ impl SupervisorMiddleware for ContentGuard {
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum ResponseMode {
HeadersOnly,
WholeBody,
Stream,
}
#[derive(Debug, Default)]
struct ResponseSessionState {
selected: Option<ResponseMode>,
next_sequence: u64,
input_bytes: u64,
}
impl ResponseSessionState {
fn preflight(
&mut self,
preflight: openshell_core::proto::HttpResponsePreflight,
) -> Result<HttpResponseEventResult, Status> {
if self.selected.is_some() {
return Err(Status::failed_precondition("duplicate response preflight"));
}
GuardConfig::parse(preflight.config.as_ref()).map_err(Status::invalid_argument)?;
let path = preflight
.target
.as_ref()
.map(|target| target.path.as_str())
.unwrap_or_default();
let Some(selected) = response_mode_for_path(path) else {
return Ok(response_preflight_skip());
};
self.selected = Some(selected);
self.next_sequence = 1;
Ok(response_preflight_inspect(selected))
}
fn body(
&mut self,
body: openshell_core::proto::HttpResponseBodyUnit,
) -> Result<HttpResponseEventResult, Status> {
let selected = self
.selected
.ok_or_else(|| Status::failed_precondition("body arrived before preflight"))?;
if selected == ResponseMode::HeadersOnly {
return Err(Status::failed_precondition(
"headers-only sessions do not receive body events",
));
}
if body.sequence != self.next_sequence {
return Err(Status::invalid_argument(format!(
"expected body sequence {}, received {}",
self.next_sequence, body.sequence
)));
}
self.next_sequence = self.next_sequence.saturating_add(1);
let Some(http_response_body_unit::Payload::Data(data)) = body.payload else {
return Err(Status::invalid_argument("body data is required"));
};
self.input_bytes = self.input_bytes.saturating_add(data.len() as u64);
let replacement = match selected {
ResponseMode::WholeBody => [b"[whole] ".as_slice(), &data].concat(),
ResponseMode::Stream => data.to_ascii_uppercase(),
ResponseMode::HeadersOnly => unreachable!(),
};
Ok(HttpResponseEventResult {
result: Some(http_response_event_result::Result::BodyResult(
HttpResponseBodyResult {
sequence: body.sequence,
decision: Some(http_response_body_result::Decision::Transform(
HttpResponseBodyTransform {
replacement: Some(http_response_body_transform::Replacement::Data(
replacement,
)),
},
)),
..Default::default()
},
)),
})
}
fn body_end(&self, final_sequence: u64) -> Result<(), Status> {
let expected = self.next_sequence.saturating_sub(1);
if final_sequence != expected {
return Err(Status::invalid_argument(format!(
"expected final body sequence {expected}, received {final_sequence}"
)));
}
Ok(())
}
fn trailers(&self) -> Result<HttpResponseEventResult, Status> {
let selected = self
.selected
.ok_or_else(|| Status::failed_precondition("trailers arrived before preflight"))?;
if selected == ResponseMode::HeadersOnly {
return Err(Status::failed_precondition(
"headers-only sessions do not receive trailers",
));
}
let trailer_mutations = if selected == ResponseMode::Stream {
vec![write_header(
"x-example-body-bytes",
&self.input_bytes.to_string(),
)]
} else {
Vec::new()
};
Ok(HttpResponseEventResult {
result: Some(http_response_event_result::Result::TrailersResult(
HttpResponseTrailersResult {
trailer_mutations,
..Default::default()
},
)),
})
}
}
fn response_mode_for_path(path: &str) -> Option<ResponseMode> {
match path {
"/headers-only" => Some(ResponseMode::HeadersOnly),
"/whole-body" => Some(ResponseMode::WholeBody),
"/stream" => Some(ResponseMode::Stream),
_ => None,
}
}
fn response_preflight_skip() -> HttpResponseEventResult {
HttpResponseEventResult {
result: Some(http_response_event_result::Result::PreflightDecision(
HttpResponsePreflightDecision {
decision: Some(http_response_preflight_decision::Decision::Skip(
HttpResponsePreflightSkip {
reason: "path is outside the response example".into(),
reason_code: "path_not_selected".into(),
..Default::default()
},
)),
},
)),
}
}
fn response_preflight_inspect(selected: ResponseMode) -> HttpResponseEventResult {
let (body_mode, declared_trailer_names) = match selected {
ResponseMode::HeadersOnly => (HttpResponseBodyMode::HeadersOnly, Vec::new()),
ResponseMode::WholeBody => (HttpResponseBodyMode::WholeBodyBytes, Vec::new()),
ResponseMode::Stream => (
HttpResponseBodyMode::StreamBytes,
vec!["x-example-body-bytes".into()],
),
};
let mode_name = match selected {
ResponseMode::HeadersOnly => "headers-only",
ResponseMode::WholeBody => "whole-body",
ResponseMode::Stream => "stream",
};
HttpResponseEventResult {
result: Some(http_response_event_result::Result::PreflightDecision(
HttpResponsePreflightDecision {
decision: Some(http_response_preflight_decision::Decision::Inspect(
HttpResponsePreflightInspect {
body_mode: body_mode as i32,
header_mutations: vec![write_header("x-example-response-mode", mode_name)],
declared_trailer_names,
..Default::default()
},
)),
},
)),
}
}
fn write_header(name: &str, value: &str) -> HeaderMutation {
HeaderMutation {
operation: Some(header_mutation::Operation::Write(WriteHeader {
name: name.into(),
value: value.into(),
on_existing: ExistingHeaderAction::Overwrite as i32,
})),
}
}
#[tonic::async_trait]
impl HttpResponsePreReturn for ContentGuard {
type EvaluateStream = HttpResponseResultStream;
async fn evaluate(
&self,
request: Request<tonic::Streaming<HttpResponseEvent>>,
) -> Result<Response<Self::EvaluateStream>, Status> {
let mut events = request.into_inner();
let (sender, receiver) = mpsc::channel(4);
tokio::spawn(async move {
let mut state = ResponseSessionState::default();
while let Some(event) = events.next().await {
let event = match event {
Ok(event) => event,
Err(error) => {
let _ = sender.send(Err(error)).await;
break;
}
};
let result = match event.event {
Some(http_response_event::Event::Preflight(preflight)) => {
state.preflight(preflight).map(Some)
}
Some(http_response_event::Event::Body(body)) => state.body(body).map(Some),
Some(http_response_event::Event::BodyEnd(body_end)) => {
state.body_end(body_end.final_sequence).map(|()| None)
}
Some(http_response_event::Event::Trailers(_)) => state.trailers().map(Some),
Some(http_response_event::Event::SessionEnd(_)) => break,
None => Err(Status::invalid_argument("response event is required")),
};
match result {
Ok(Some(result)) => {
if sender.send(Ok(result)).await.is_err() {
break;
}
}
Ok(None) => {}
Err(error) => {
let _ = sender.send(Err(error)).await;
break;
}
}
}
});
Ok(Response::new(Box::pin(ReceiverStream::new(receiver))))
}
}
fn validate_phase(phase: i32) -> Result<(), String> {
if phase != PHASE as i32 {
return Err(format!("unsupported phase '{phase}'"));
@@ -470,6 +721,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("serving {MANIFEST_NAME} on http://{}", cli.bind);
Server::builder()
.add_service(SupervisorMiddlewareServer::new(ContentGuard))
.add_service(HttpResponsePreReturnServer::new(ContentGuard))
.serve(cli.bind)
.await?;
Ok(())
@@ -478,7 +730,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
#[cfg(test)]
mod tests {
use super::*;
use openshell_core::proto::{WebSocketPreflight, WebSocketSessionEnd, WebSocketSessionStart};
use openshell_core::proto::{
HttpRequestTarget, HttpResponseBodyUnit, HttpResponsePreflight, WebSocketPreflight,
WebSocketSessionEnd, WebSocketSessionStart,
};
use prost_types::{ListValue, Value};
use std::collections::BTreeMap;
@@ -511,13 +766,13 @@ mod tests {
}
#[tokio::test]
async fn manifest_advertises_http_and_websocket_bindings() {
async fn manifest_advertises_request_response_and_websocket_bindings() {
let manifest = SupervisorMiddleware::describe(&ContentGuard, Request::new(()))
.await
.expect("describe")
.into_inner();
assert_eq!(manifest.bindings.len(), 2);
assert_eq!(manifest.bindings.len(), 3);
assert_eq!(
manifest.bindings[0].operation,
SupervisorMiddlewareOperation::HttpRequest as i32
@@ -528,6 +783,121 @@ mod tests {
SupervisorMiddlewareOperation::WebsocketMessage as i32
);
assert_eq!(manifest.bindings[1].max_payload_bytes, MAX_PAYLOAD_BYTES);
assert_eq!(
manifest.bindings[2].operation,
SupervisorMiddlewareOperation::HttpResponse as i32
);
assert_eq!(
manifest.bindings[2].phase,
SupervisorMiddlewarePhase::PreReturn as i32
);
assert_eq!(manifest.bindings[2].max_payload_bytes, MAX_PAYLOAD_BYTES);
}
fn response_preflight(path: &str) -> HttpResponsePreflight {
HttpResponsePreflight {
target: Some(HttpRequestTarget {
path: path.into(),
..Default::default()
}),
config: Some(config("redact", &["prototype-secret"], Some("[FILTERED]"))),
..Default::default()
}
}
#[test]
fn paths_select_all_three_response_modes() {
for (path, expected) in [
("/headers-only", HttpResponseBodyMode::HeadersOnly),
("/whole-body", HttpResponseBodyMode::WholeBodyBytes),
("/stream", HttpResponseBodyMode::StreamBytes),
] {
let mut state = ResponseSessionState::default();
let result = state.preflight(response_preflight(path)).unwrap();
let Some(http_response_event_result::Result::PreflightDecision(decision)) =
result.result
else {
panic!("expected preflight decision");
};
let Some(http_response_preflight_decision::Decision::Inspect(inspect)) =
decision.decision
else {
panic!("expected inspect decision");
};
assert_eq!(inspect.body_mode, expected as i32);
}
}
#[test]
fn whole_body_and_stream_transform_differently() {
for (path, expected) in [
("/whole-body", b"[whole] hello".as_slice()),
("/stream", b"HELLO".as_slice()),
] {
let mut state = ResponseSessionState::default();
state.preflight(response_preflight(path)).unwrap();
let result = state
.body(HttpResponseBodyUnit {
sequence: 1,
payload: Some(http_response_body_unit::Payload::Data(b"hello".to_vec())),
})
.unwrap();
let Some(http_response_event_result::Result::BodyResult(body)) = result.result else {
panic!("expected body result");
};
let Some(http_response_body_result::Decision::Transform(transform)) = body.decision
else {
panic!("expected body transform");
};
let Some(http_response_body_transform::Replacement::Data(data)) = transform.replacement
else {
panic!("expected replacement data");
};
assert_eq!(data, expected);
}
}
#[test]
fn stream_declares_and_writes_byte_count_trailer() {
let mut state = ResponseSessionState::default();
let preflight = state.preflight(response_preflight("/stream")).unwrap();
let Some(http_response_event_result::Result::PreflightDecision(decision)) =
preflight.result
else {
panic!("expected preflight decision");
};
let Some(http_response_preflight_decision::Decision::Inspect(inspect)) = decision.decision
else {
panic!("expected inspect decision");
};
assert_eq!(inspect.declared_trailer_names, ["x-example-body-bytes"]);
state
.body(HttpResponseBodyUnit {
sequence: 1,
payload: Some(http_response_body_unit::Payload::Data(b"hello".to_vec())),
})
.unwrap();
state.body_end(1).unwrap();
let trailers = state.trailers().unwrap();
let Some(http_response_event_result::Result::TrailersResult(trailers)) = trailers.result
else {
panic!("expected trailer result");
};
assert_eq!(trailers.trailer_mutations.len(), 1);
}
#[test]
fn paths_outside_the_response_demo_are_skipped() {
let mut state = ResponseSessionState::default();
let result = state.preflight(response_preflight("/outside")).unwrap();
let Some(http_response_event_result::Result::PreflightDecision(decision)) = result.result
else {
panic!("expected preflight decision");
};
let Some(http_response_preflight_decision::Decision::Skip(skip)) = decision.decision else {
panic!("expected skip decision");
};
assert_eq!(skip.reason_code, "path_not_selected");
}
#[tokio::test]
@@ -737,4 +1107,11 @@ mod tests {
assert_eq!(parsed.mode, Mode::Redact);
assert_eq!(parsed.replacement, DEFAULT_REPLACEMENT);
}
#[test]
fn example_policy_is_valid() {
let policy = openshell_policy::parse_sandbox_policy(include_str!("../policy.yaml"))
.expect("example policy must parse");
openshell_policy::validate_sandbox_policy(&policy).expect("example policy must be valid");
}
}
@@ -0,0 +1,51 @@
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import socketserver
class Handler(socketserver.BaseRequestHandler):
def handle(self):
request = b""
while b"\r\n\r\n" not in request:
block = self.request.recv(4096)
if not block:
return
request += block
path = request.split(b" ", 2)[1]
if path == b"/headers-only":
response = (
b"HTTP/1.1 200 OK\r\n"
b"Content-Type: text/plain\r\n"
b"Content-Length: 12\r\n\r\n"
b"headers-only"
)
elif path == b"/whole-body":
response = (
b"HTTP/1.1 200 OK\r\n"
b"Content-Type: text/plain\r\n"
b"Transfer-Encoding: chunked\r\n\r\n"
b"6\r\nwhole \r\n4\r\nbody\r\n0\r\n\r\n"
)
elif path == b"/stream":
response = (
b"HTTP/1.1 200 OK\r\n"
b"Content-Type: text/plain\r\n"
b"Connection: close\r\n\r\n"
b"stream body"
)
else:
response = b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n"
self.request.sendall(response)
class DemoServer(socketserver.ThreadingTCPServer):
allow_reuse_address = True
with DemoServer(("0.0.0.0", 18081), Handler) as server:
print("response framing demo upstream listening on 0.0.0.0:18081", flush=True)
try:
server.serve_forever()
except KeyboardInterrupt:
pass
+3
View File
@@ -15,6 +15,7 @@ run = [
"cargo clippy --workspace --all-targets -- -D warnings",
"cargo clippy --manifest-path e2e/rust/Cargo.toml --all-targets -- -D warnings",
"cargo check --manifest-path examples/governance-interceptor/Cargo.toml --all-targets",
"cargo check --manifest-path examples/supervisor-middleware-content-guard/Cargo.toml --all-targets",
]
run_windows = "powershell -NoProfile -ExecutionPolicy Bypass -File tasks/scripts/windows-msvc.ps1 lint native"
hide = true
@@ -25,6 +26,7 @@ run = [
"cargo fmt --all",
"cargo fmt --manifest-path e2e/rust/Cargo.toml --all",
"cargo fmt --manifest-path examples/governance-interceptor/Cargo.toml --all",
"cargo fmt --manifest-path examples/supervisor-middleware-content-guard/Cargo.toml --all",
]
hide = true
@@ -34,6 +36,7 @@ run = [
"cargo fmt --all -- --check",
"cargo fmt --manifest-path e2e/rust/Cargo.toml --all -- --check",
"cargo fmt --manifest-path examples/governance-interceptor/Cargo.toml --all -- --check",
"cargo fmt --manifest-path examples/supervisor-middleware-content-guard/Cargo.toml --all -- --check",
]
hide = true