fix(middleware): buffer request bodies at the chain's largest stage limit

Buffering at the smallest stage limit forced the whole chain onto the
unbuffered over-capacity path whenever the body exceeded any single
stage, losing per-stage on_error behavior: a small fail-open stage
could deny a request that a larger fail-closed stage was able to
inspect. Buffer for the most capable stage instead, so stages whose own
limit is smaller fail individually with request_body_over_capacity
through their own on_error while the rest of the chain still runs. A
body over every stage limit keeps the existing aggregate handling,
which now matches the per-stage semantics by construction.

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

Author: pimlock

Cargo.lock

@@ -1322,6 +1322,89 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn undersized_stage_fails_open_while_later_stage_runs() {
+        // A body over one stage's limit must fail only that stage through its
+        // own `on_error`, not the whole chain: the 1 KiB fail-open guard is
+        // skipped while the 256 KiB fail-closed redactor still runs.
+        let external = Arc::new(ScriptedService {
+            binding_id: "example/content-guard".into(),
+            max_body_bytes: 4096,
+            result: allow_result(),
+        });
+        let registry = registry_with_external(external, external_registration(1024)).await;
+        let runner = ChainRunner::from_registry(registry);
+        let guard_entry = ChainEntry {
+            name: "guard".into(),
+            implementation: "example/content-guard".into(),
+            order: 0,
+            config: prost_types::Struct::default(),
+            on_error: OnError::FailOpen,
+        };
+        let mut redact_entry = entry("redact", OnError::FailClosed);
+        redact_entry.order = 10;
+        let entries = [guard_entry, redact_entry];
+
+        let body = format!("{}password=\"top-secret\"", "x".repeat(1500));
+        let outcome = runner
+            .evaluate(&entries, input(&body))
+            .await
+            .expect("evaluate mixed-limit chain");
+
+        assert!(outcome.allowed);
+        assert_eq!(outcome.applied.len(), 2);
+        assert!(
+            outcome.applied[0].failed,
+            "undersized guard must be skipped"
+        );
+        assert_eq!(outcome.applied[0].decision, Decision::Allow);
+        assert!(!outcome.applied[1].failed);
+        assert!(outcome.applied[1].transformed);
+        let body = String::from_utf8(outcome.body).expect("utf8");
+        assert!(body.contains("[REDACTED]"));
+        assert!(!body.contains("top-secret"));
+    }
+
+    #[tokio::test]
+    async fn transformed_body_still_over_later_stage_capacity_honors_on_error() {
+        // Per-stage capacity applies to the current body: the redactor's
+        // replacement is still over the 1 KiB guard limit, so the fail-closed
+        // guard denies through its own `on_error` after the redactor ran.
+        let external = Arc::new(ScriptedService {
+            binding_id: "example/content-guard".into(),
+            max_body_bytes: 4096,
+            result: allow_result(),
+        });
+        let registry = registry_with_external(external, external_registration(1024)).await;
+        let runner = ChainRunner::from_registry(registry);
+        let guard_entry = ChainEntry {
+            name: "guard".into(),
+            implementation: "example/content-guard".into(),
+            order: 10,
+            config: prost_types::Struct::default(),
+            on_error: OnError::FailClosed,
+        };
+        let entries = [entry("redact", OnError::FailClosed), guard_entry];
+
+        let body = format!("{}password=\"top-secret\"", "x".repeat(1500));
+        let outcome = runner
+            .evaluate(&entries, input(&body))
+            .await
+            .expect("evaluate mixed-limit chain");
+
+        assert!(!outcome.allowed);
+        assert_eq!(
+            outcome.reason,
+            "middleware_failed: request_body_over_capacity"
+        );
+        assert_eq!(outcome.applied.len(), 2);
+        assert!(
+            outcome.applied[0].transformed,
+            "redactor ran before the deny"
+        );
+        assert!(outcome.applied[1].failed);
+    }
+
     #[test]
     fn external_manifest_rejects_operator_limit_above_capability() {
         let registration = external_registration(4097);

crates/openshell-supervisor-middleware/src/lib.rs

@@ -52,6 +52,7 @@ webpki-roots = { workspace = true }
 [dev-dependencies]
 openshell-core = { path = "../openshell-core", features = ["test-helpers"] }
 tempfile = "3"
+tonic = { workspace = true }
 temp-env = "0.3"
 tokio-tungstenite = { workspace = true }
 futures = { workspace = true }

crates/openshell-supervisor-network/Cargo.toml

@@ -19,11 +19,14 @@ pub enum MiddlewareApplyResult {
     Denied(String),
 }
 
-/// Smallest body-buffering limit across the entries that actually resolved to a
-/// registered binding. Unresolved entries (`is_resolved() == false`) report a
-/// zero limit and are excluded here: they are handled by their `on_error` policy
-/// in `evaluate_described` without inspecting the body, so letting a zero drag
-/// the chain limit to zero would spuriously fail the whole chain over capacity.
+/// 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
+/// individually with `request_body_over_capacity` through their `on_error`
+/// policy in `evaluate_described`, instead of one undersized stage forcing the
+/// whole chain onto the unbuffered path. Unresolved entries
+/// (`is_resolved() == false`) report a zero limit and are excluded here: they
+/// are handled by their `on_error` policy without inspecting the body.
 /// Returns `None` when no entry resolved, so the caller can skip buffering.
 pub(super) fn middleware_chain_body_limit(
     chain: &[openshell_supervisor_middleware::DescribedChainEntry],
@@ -32,7 +35,7 @@ pub(super) fn middleware_chain_body_limit(
         .iter()
         .filter(|entry| entry.is_resolved())
         .map(openshell_supervisor_middleware::DescribedChainEntry::max_body_bytes)
-        .min()
+        .max()
 }
 
 pub async fn apply_middleware_chain<C: AsyncRead + AsyncWrite + Unpin + Send>(
@@ -147,10 +150,11 @@ pub(super) fn raw_query_from_request_headers(headers: &[u8]) -> Result<String> {
         .map_or_else(String::new, |(_, query)| query.to_string()))
 }
 
-/// Apply the chain's `on_error` policy when the request body cannot be buffered
-/// for inspection because it exceeds the size cap. The RFC treats an unbufferable
-/// body as an `on_error` event: it is denied unless every attached middleware is
-/// `fail_open`, and passing it through is only safe when no bytes were consumed.
+/// Apply the chain's `on_error` policy when the request body exceeds every
+/// stage's buffering limit. No stage can inspect such a body, so each stage
+/// would individually fail with `request_body_over_capacity`; the aggregate is
+/// a deny unless every attached middleware is `fail_open`, and passing the
+/// body through is only safe when no bytes were consumed.
 pub(super) fn resolve_unbuffered_body(
     ctx: &L7EvalContext,
     req: crate::l7::provider::L7Request,

crates/openshell-supervisor-network/src/l7/middleware.rs

@@ -2875,6 +2875,146 @@ network_policies:
         assert_eq!(middleware_chain_body_limit(&none), None);
     }
 
+    /// A middleware service advertising two bindings with different body
+    /// limits, for exercising mixed-limit chain buffering at the relay level.
+    /// The redactor binding replaces the body; the guard binding allows as-is.
+    struct TwoLimitService;
+
+    #[tonic::async_trait]
+    impl openshell_core::proto::middleware::v1::supervisor_middleware_server::SupervisorMiddleware
+        for TwoLimitService
+    {
+        async fn describe(
+            &self,
+            _request: tonic::Request<()>,
+        ) -> std::result::Result<
+            tonic::Response<openshell_core::proto::MiddlewareManifest>,
+            tonic::Status,
+        > {
+            use openshell_core::proto::{
+                MiddlewareBinding, MiddlewareManifest, SupervisorMiddlewareOperation,
+                SupervisorMiddlewarePhase,
+            };
+            let binding = |id: &str, max_body_bytes: u64| MiddlewareBinding {
+                id: id.into(),
+                operation: SupervisorMiddlewareOperation::HttpRequest as i32,
+                phase: SupervisorMiddlewarePhase::PreCredentials as i32,
+                max_body_bytes,
+            };
+            Ok(tonic::Response::new(MiddlewareManifest {
+                name: "test/two-limits".into(),
+                service_version: "test".into(),
+                bindings: vec![binding("test/redactor", 8192), binding("test/guard", 16)],
+            }))
+        }
+
+        async fn validate_config(
+            &self,
+            _request: tonic::Request<openshell_core::proto::ValidateConfigRequest>,
+        ) -> std::result::Result<
+            tonic::Response<openshell_core::proto::ValidateConfigResponse>,
+            tonic::Status,
+        > {
+            Ok(tonic::Response::new(
+                openshell_core::proto::ValidateConfigResponse {
+                    valid: true,
+                    reason: String::new(),
+                },
+            ))
+        }
+
+        async fn evaluate_http_request(
+            &self,
+            request: tonic::Request<openshell_core::proto::HttpRequestEvaluation>,
+        ) -> std::result::Result<
+            tonic::Response<openshell_core::proto::HttpRequestResult>,
+            tonic::Status,
+        > {
+            let evaluation = request.into_inner();
+            let mut result = openshell_core::proto::HttpRequestResult {
+                decision: openshell_core::proto::Decision::Allow as i32,
+                ..Default::default()
+            };
+            if evaluation.binding_id == "test/redactor" {
+                result.body = b"[SCRUBBED BY TEST REDACTOR]".to_vec();
+                result.has_body = true;
+            }
+            Ok(tonic::Response::new(result))
+        }
+    }
+
+    #[tokio::test]
+    async fn body_over_smallest_stage_limit_is_buffered_and_evaluated() {
+        use openshell_supervisor_middleware::{ChainEntry, ChainRunner, OnError};
+
+        // A 64-byte body exceeds the 16-byte guard limit but fits the 8 KiB
+        // redactor. The chain must buffer for its largest stage so the
+        // redactor runs and replaces the body, while the undersized fail-open
+        // guard is skipped through its own on_error, instead of the whole
+        // chain taking the unbuffered over-capacity path.
+        let (_config, tunnel_engine, ctx) =
+            middleware_relay_context("openshell/secrets", "fail_closed");
+        let runner = ChainRunner::new(Arc::new(TwoLimitService));
+        let chain = vec![
+            ChainEntry {
+                name: "redact".into(),
+                implementation: "test/redactor".into(),
+                order: 0,
+                config: prost_types::Struct::default(),
+                on_error: OnError::FailClosed,
+            },
+            ChainEntry {
+                name: "guard".into(),
+                implementation: "test/guard".into(),
+                order: 10,
+                config: prost_types::Struct::default(),
+                on_error: OnError::FailOpen,
+            },
+        ];
+        let described = runner.describe_chain(&chain).await.expect("describe chain");
+        assert_eq!(middleware_chain_body_limit(&described), Some(8192));
+
+        let body = [b'a'; 64];
+        let raw_header = format!(
+            "POST /v1/messages HTTP/1.1\r\nHost: api.example.test\r\nContent-Length: {}\r\n\r\n",
+            body.len()
+        );
+        let req = crate::l7::provider::L7Request {
+            action: "POST".into(),
+            target: "/v1/messages".into(),
+            query_params: std::collections::HashMap::new(),
+            raw_header: raw_header.into_bytes(),
+            body_length: crate::l7::provider::BodyLength::ContentLength(body.len() as u64),
+        };
+        let (mut app, mut relay_client) = tokio::io::duplex(8192);
+        app.write_all(&body).await.unwrap();
+
+        let result = crate::l7::middleware::apply_middleware_chain_for_scheme(
+            req,
+            &mut relay_client,
+            &ctx,
+            "https",
+            chain,
+            &runner,
+            tunnel_engine.generation_guard(),
+        )
+        .await
+        .expect("apply middleware chain");
+
+        match result {
+            MiddlewareApplyResult::Allowed(rebuilt) => {
+                let raw = String::from_utf8(rebuilt.raw_header).expect("utf8 request");
+                assert!(
+                    raw.ends_with("[SCRUBBED BY TEST REDACTOR]"),
+                    "redactor must replace the body: {raw}"
+                );
+            }
+            MiddlewareApplyResult::Denied(reason) => {
+                panic!("body within the largest stage limit must not fail the chain: {reason}")
+            }
+        }
+    }
+
     #[tokio::test]
     async fn all_unresolved_fail_open_forwards_body_unbuffered() {
         // A chain whose only entry is an unregistered binding has no resolvable

crates/openshell-supervisor-network/src/l7/relay.rs

@@ -17,7 +17,7 @@ For each inspected HTTP request, the supervisor:
 
 1. Evaluates network and L7 policy.
 2. Selects middleware whose host selectors match the admitted destination.
-3. Buffers the request body using the smallest body limit in the selected chain.
+3. Buffers the request body using the largest body limit in the selected chain.
 4. Runs matching middleware by ascending `order`, using the policy-local name to break ties.
 5. Applies allowed transformations, injects provider credentials, and forwards the request.
 
@@ -99,10 +99,10 @@ Every middleware binding declares the largest request or replacement body it sup
 
 - Built-in middleware uses its OpenShell-defined limit.
 - Each operator-run registration sets `max_body_bytes` no higher than the service capability.
-- A selected chain buffers using its smallest stage limit.
+- A selected chain buffers using its largest stage limit, so every stage that can process the body receives it.
 - The same per-stage limit applies to request bodies and replacement bodies.
 
-The gateway rejects a registration whose operator limit exceeds the service capability instead of silently clamping it. At request time, exceeding a selected stage's limit is a middleware failure and follows that config's `on_error` behavior.
+The gateway rejects a registration whose operator limit exceeds the service capability instead of silently clamping it. 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. A body larger than every stage limit cannot be inspected at all and is denied unless every selected stage is `fail_open`.
 
 ## Add Request Headers