Bind request body admission to ownership

Author: eldryoth

crates/fluxheim-server/src/native_http1.rs

@@ -261,7 +261,7 @@ where
                     return Ok(());
                 }
             };
-        let body = match read_body(
+        let mut body = match read_body(
             policy,
             request_body_timeout,
             &mut stream,
@@ -320,6 +320,7 @@ where
             }
             Err(error) => return Err(error),
         };
+        body.attach_admission(request_body_reservation);
         let mut request = request;
         request.body = body;
         request_handler.prepare_request_context(&mut request);

crates/fluxheim-server/src/native_http1/request.rs

@@ -3,6 +3,8 @@ use std::ops::Deref;
 
 use fluxheim_protocol::{Http1RequestTarget, Http1Version, http1_request_target};
 
+use crate::request_body_budget::NativeRequestBodyReservation;
+
 #[derive(Debug, Eq, PartialEq)]
 pub struct NativeHttp1Request {
     pub method: String,
@@ -19,48 +21,68 @@ pub struct NativeHttp1Request {
     pub trailers: Vec<(String, String)>,
 }
 
-#[derive(Default, Eq, PartialEq)]
-pub struct NativeHttp1RequestBody(Vec<u8>);
+#[derive(Default)]
+pub struct NativeHttp1RequestBody {
+    bytes: Vec<u8>,
+    admission: Option<NativeRequestBodyReservation>,
+}
 
 impl NativeHttp1RequestBody {
     #[must_use]
     pub const fn empty() -> Self {
-        Self(Vec::new())
+        Self {
+            bytes: Vec::new(),
+            admission: None,
+        }
     }
 
     #[must_use]
     pub const fn from_vec(body: Vec<u8>) -> Self {
-        Self(body)
+        Self {
+            bytes: body,
+            admission: None,
+        }
+    }
+
+    pub(crate) fn attach_admission(&mut self, admission: NativeRequestBodyReservation) {
+        if self.admission.is_some() {
+            log::error!(
+                target: "fluxheim::security",
+                "request body admission was attached more than once"
+            );
+            std::process::abort();
+        }
+        self.admission = Some(admission);
     }
 
     pub(crate) fn extend_from_slice(
         &mut self,
         bytes: &[u8],
     ) -> Result<(), fluxheim_protocol::Http1ParseError> {
         let required = self
-            .0
+            .bytes
             .len()
             .checked_add(bytes.len())
             .ok_or(fluxheim_protocol::Http1ParseError::BodyTooLarge)?;
         self.reserve_capacity(required)?;
-        self.0.extend_from_slice(bytes);
+        self.bytes.extend_from_slice(bytes);
         Ok(())
     }
 
     pub(crate) fn reserve_capacity(
         &mut self,
         required: usize,
     ) -> Result<(), fluxheim_protocol::Http1ParseError> {
-        if required <= self.0.capacity() {
+        if required <= self.bytes.capacity() {
             return Ok(());
         }
         let mut replacement = Vec::new();
         replacement
             .try_reserve_exact(required)
             .map_err(|_| fluxheim_protocol::Http1ParseError::BodyTooLarge)?;
-        replacement.extend_from_slice(&self.0);
-        sanitization::unsafe_wipe::volatile_sanitize_vec(&mut self.0);
-        self.0 = replacement;
+        replacement.extend_from_slice(&self.bytes);
+        sanitization::unsafe_wipe::volatile_sanitize_vec(&mut self.bytes);
+        self.bytes = replacement;
         Ok(())
     }
 
@@ -72,22 +94,30 @@ impl NativeHttp1RequestBody {
         replacement
             .try_reserve_exact(required)
             .map_err(|_| fluxheim_protocol::Http1ParseError::BodyTooLarge)?;
-        replacement.extend_from_slice(&self.0);
-        sanitization::unsafe_wipe::volatile_sanitize_vec(&mut self.0);
-        self.0 = replacement;
+        replacement.extend_from_slice(&self.bytes);
+        sanitization::unsafe_wipe::volatile_sanitize_vec(&mut self.bytes);
+        self.bytes = replacement;
         Ok(())
     }
 
     pub(crate) fn capacity(&self) -> usize {
-        self.0.capacity()
+        self.bytes.capacity()
+    }
+}
+
+impl PartialEq for NativeHttp1RequestBody {
+    fn eq(&self, other: &Self) -> bool {
+        self.bytes == other.bytes
     }
 }
 
+impl Eq for NativeHttp1RequestBody {}
+
 impl std::fmt::Debug for NativeHttp1RequestBody {
     fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         formatter
             .debug_struct("NativeHttp1RequestBody")
-            .field("len", &self.0.len())
+            .field("len", &self.bytes.len())
             .field("contents", &"<redacted>")
             .finish()
     }
@@ -97,19 +127,19 @@ impl Deref for NativeHttp1RequestBody {
     type Target = [u8];
 
     fn deref(&self) -> &Self::Target {
-        &self.0
+        &self.bytes
     }
 }
 
 impl AsRef<[u8]> for NativeHttp1RequestBody {
     fn as_ref(&self) -> &[u8] {
-        &self.0
+        &self.bytes
     }
 }
 
 impl fluxheim_protocol::Http1ChunkSink for NativeHttp1RequestBody {
     fn len(&self) -> usize {
-        self.0.len()
+        self.bytes.len()
     }
 
     fn append(&mut self, bytes: &[u8]) -> Result<(), fluxheim_protocol::Http1ParseError> {
@@ -119,7 +149,7 @@ impl fluxheim_protocol::Http1ChunkSink for NativeHttp1RequestBody {
 
 impl Drop for NativeHttp1RequestBody {
     fn drop(&mut self) {
-        sanitization::unsafe_wipe::volatile_sanitize_vec(&mut self.0);
+        sanitization::unsafe_wipe::volatile_sanitize_vec(&mut self.bytes);
     }
 }
 

crates/fluxheim-server/src/native_http1_tests.rs

@@ -22,6 +22,11 @@ struct PinningBodyBudgetHandler {
     pinned: Arc<HoldingBodyBudgetHandler>,
 }
 
+struct RetainingBodyBudgetHandler {
+    budget: NativeRequestBodyBudget,
+    retained: std::sync::Mutex<Option<crate::NativeHttp1RequestBody>>,
+}
+
 impl NativeHttp1Handler for PinningBodyBudgetHandler {
     fn pin_request_handler(&self) -> Option<Arc<dyn NativeHttp1Handler>> {
         Some(self.pinned.clone())
@@ -42,16 +47,33 @@ impl NativeHttp1Handler for HoldingBodyBudgetHandler {
 
     fn handle<'a>(
         &'a self,
-        _request: NativeHttp1Request,
+        request: NativeHttp1Request,
     ) -> std::pin::Pin<Box<dyn std::future::Future<Output = NativeHttp1Response> + Send + 'a>> {
         Box::pin(async move {
             self.entered.notify_one();
             self.release.notified().await;
+            drop(request);
             NativeHttp1Response::new(200, "OK", b"ok\n").close_connection()
         })
     }
 }
 
+impl NativeHttp1Handler for RetainingBodyBudgetHandler {
+    fn request_body_budget(&self) -> NativeRequestBodyBudget {
+        self.budget.clone()
+    }
+
+    fn handle<'a>(
+        &'a self,
+        request: NativeHttp1Request,
+    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = NativeHttp1Response> + Send + 'a>> {
+        Box::pin(async move {
+            *self.retained.lock().unwrap() = Some(request.body);
+            NativeHttp1Response::new(200, "OK", b"retained\n").close_connection()
+        })
+    }
+}
+
 pub(crate) async fn spawn_server(
     handler: impl Fn(NativeHttp1Request) -> NativeHttp1Response + Send + Sync + 'static,
 ) -> std::net::SocketAddr {
@@ -572,3 +594,44 @@ async fn native_http1_rejects_aggregate_request_body_overcommit() {
     );
     server.await.unwrap();
 }
+
+#[tokio::test]
+async fn native_http1_body_retains_admission_after_handler_returns() {
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let address = listener.local_addr().unwrap();
+    let budget = NativeRequestBodyBudget::new(64 * 1024);
+    let handler = Arc::new(RetainingBodyBudgetHandler {
+        budget: budget.clone(),
+        retained: std::sync::Mutex::new(None),
+    });
+    let server_handler = handler.clone();
+    let server = tokio::spawn(async move {
+        let (stream, peer) = listener.accept().await.unwrap();
+        serve_native_http1_connection(
+            stream,
+            Some(peer),
+            DownstreamHttp1Policy::default(),
+            server_handler,
+        )
+        .await
+        .unwrap();
+    });
+    let mut stream = TcpStream::connect(address).await.unwrap();
+
+    stream
+        .write_all(
+            b"POST /retain HTTP/1.1\r\nHost: local.test\r\nContent-Length: 1\r\nConnection: close\r\n\r\na",
+        )
+        .await
+        .unwrap();
+    assert!(
+        read_response(&mut stream)
+            .await
+            .starts_with("HTTP/1.1 200 OK\r\n")
+    );
+    server.await.unwrap();
+
+    assert!(budget.reserve(1).await.is_err());
+    handler.retained.lock().unwrap().take();
+    assert!(budget.reserve(1).await.is_ok());
+}

crates/fluxheim-server/src/native_http2_client.rs

@@ -452,14 +452,20 @@ async fn drain_response_body(
 mod owned_body_tests {
     use super::*;
 
-    #[test]
-    fn h2_bytes_retains_sanitizing_request_body_allocation() {
-        let body = NativeHttp1RequestBody::from_vec(b"owned request body".to_vec());
+    #[tokio::test]
+    async fn h2_bytes_retains_body_allocation_and_admission() {
+        let budget = crate::NativeRequestBodyBudget::new(64 * 1024);
+        let admission = budget.reserve(1).await.unwrap();
+        let mut body = NativeHttp1RequestBody::from_vec(b"owned request body".to_vec());
+        body.attach_admission(admission);
         let original_pointer = body.as_ref().as_ptr();
 
         let bytes = owned_request_body_bytes(body);
 
         assert_eq!(bytes.as_ptr(), original_pointer);
         assert_eq!(bytes.as_ref(), b"owned request body");
+        assert!(budget.reserve(1).await.is_err());
+        drop(bytes);
+        assert!(budget.reserve(1).await.is_ok());
     }
 }

crates/fluxheim-server/src/native_http2_route_adapter.rs

@@ -79,7 +79,7 @@ fn native_http2_request_to_http1(
         method,
         uri,
         headers,
-        mut body,
+        body,
         trailers,
     } = request;
     let authority = uri.authority()?.as_str().to_owned();
@@ -107,7 +107,7 @@ fn native_http2_request_to_http1(
         target,
         version: fluxheim_protocol::Http1Version::Http11,
         headers: owned_headers,
-        body: crate::NativeHttp1RequestBody::from_vec(std::mem::take(&mut *body)),
+        body,
         trailers: native_http2_headers_to_native(trailers.as_ref())?,
     })
 }
@@ -160,7 +160,7 @@ mod tests {
             method: http::Method::POST,
             uri: "https://native.test/upload?x=1".parse().unwrap(),
             headers,
-            body: zeroize::Zeroizing::new(b"body".to_vec()),
+            body: crate::NativeHttp1RequestBody::from_vec(b"body".to_vec()),
             trailers: Some({
                 let mut trailers = http::HeaderMap::new();
                 trailers.insert("grpc-status", http::HeaderValue::from_static("0"));
@@ -224,7 +224,7 @@ mod tests {
             method: http::Method::GET,
             uri: "https://public.test/resource".parse().unwrap(),
             headers,
-            body: zeroize::Zeroizing::new(Vec::new()),
+            body: crate::NativeHttp1RequestBody::empty(),
             trailers: None,
         };
 
@@ -247,7 +247,7 @@ mod tests {
             method: http::Method::GET,
             uri: "/resource".parse().unwrap(),
             headers: http::HeaderMap::new(),
-            body: zeroize::Zeroizing::new(Vec::new()),
+            body: crate::NativeHttp1RequestBody::empty(),
             trailers: None,
         };