Resolve PR 624 round-6 review findings

Blocking fixes:
- Restore drop(streaming_body) on error path so client sees EOF/truncated
  response instead of a silently completed partial response (cd4c621 regression)
- Add HandlerOutcome::AuthChallenge variant to distinguish our enforce_basic_auth
  401s from origin-forwarded 401s; gate geo lookup on AuthChallenge only so
  origin-forwarded 401s still carry geo headers

Non-blocking fixes:
- Fix stream_publisher_body doc: remove incorrect async claim; function is sync
- Clarify to_fastly_request_ref doc: non-empty body is a caller error
- Add enforce_max_body_size helper in http_util; replace 6 duplicated body-size
  cap blocks across proxy, auction, and request_signing endpoints
- Add test for to_fastly_response_skeleton covering status + header round-trip
- Restore "bytes" unit in enforce_max_body_size error message for log clarity
- Widen enforce_max_body_size what param from &'static str to &str

Author: prk-Jr

crates/trusted-server-adapter-fastly/src/main.rs

@@ -48,8 +48,13 @@ use crate::platform::{build_runtime_services, open_kv_store, UnavailableKvStore}
 ///
 /// The streaming arm keeps the publisher body out of WASM heap until it is written directly
 /// to the client via [`fastly::Response::stream_to_client`].  All other routes are buffered.
+///
+/// [`AuthChallenge`](HandlerOutcome::AuthChallenge) marks responses produced by this server's
+/// own `enforce_basic_auth` so the geo-lookup gate can distinguish them from origin-forwarded
+/// 401s, which should still carry geo headers.
 enum HandlerOutcome {
     Buffered(HttpResponse),
+    AuthChallenge(HttpResponse),
     Streaming {
         response: HttpResponse,
         body: EdgeBody,
@@ -58,9 +63,10 @@ enum HandlerOutcome {
 }
 
 impl HandlerOutcome {
+    #[cfg(test)]
     fn status(&self) -> edgezero_core::http::StatusCode {
         match self {
-            HandlerOutcome::Buffered(resp) => resp.status(),
+            HandlerOutcome::Buffered(resp) | HandlerOutcome::AuthChallenge(resp) => resp.status(),
             HandlerOutcome::Streaming { response, .. } => response.status(),
         }
     }
@@ -142,8 +148,10 @@ fn main() {
     ))
     .unwrap_or_else(|e| HandlerOutcome::Buffered(http_error_response(&e)));
 
-    // Skip geo lookup for 401s: avoids exposing geo headers to unauthenticated callers.
-    let geo_info = if outcome.status() == edgezero_core::http::StatusCode::UNAUTHORIZED {
+    // Skip geo lookup for our own auth challenges: avoids exposing geo headers to
+    // unauthenticated callers.  Origin-forwarded 401s are not AuthChallenge and
+    // do receive geo headers — the client already reached the origin anyway.
+    let geo_info = if matches!(outcome, HandlerOutcome::AuthChallenge(_)) {
         None
     } else {
         runtime_services
@@ -156,7 +164,7 @@ fn main() {
     };
 
     match outcome {
-        HandlerOutcome::Buffered(mut response) => {
+        HandlerOutcome::Buffered(mut response) | HandlerOutcome::AuthChallenge(mut response) => {
             finalize_response(&settings, geo_info.as_ref(), &mut response);
             compat::to_fastly_response(response).send_to_client();
         }
@@ -182,9 +190,9 @@ fn main() {
                 }
                 Err(e) => {
                     log::error!("streaming processing failed: {e:?}");
-                    if let Err(finish_err) = streaming_body.finish() {
-                        log::error!("failed to finish streaming body after error: {finish_err}");
-                    }
+                    // Headers already committed. Drop the body so the client sees a
+                    // truncated response (EOF mid-stream) — standard proxy behavior.
+                    drop(streaming_body);
                 }
             }
         }
@@ -244,7 +252,7 @@ async fn route_request(
     // Keep this fallback so manually-constructed or otherwise unprepared
     // settings still become an error response instead of panicking.
     match enforce_basic_auth(settings, &req) {
-        Ok(Some(response)) => return Ok(HandlerOutcome::Buffered(response)),
+        Ok(Some(response)) => return Ok(HandlerOutcome::AuthChallenge(response)),
         Ok(None) => {}
         Err(e) => return Err(e),
     }

crates/trusted-server-core/src/auction/endpoints.rs

@@ -10,6 +10,7 @@ use crate::consent;
 use crate::cookies::handle_request_cookies;
 use crate::edge_cookie::get_or_generate_ec_id_from_http_request;
 use crate::error::TrustedServerError;
+use crate::http_util::enforce_max_body_size;
 use crate::platform::RuntimeServices;
 use crate::settings::Settings;
 
@@ -42,15 +43,7 @@ pub async fn handle_auction(
 
     // Parse request body
     let body_bytes = body.into_bytes();
-    if body_bytes.len() > AUCTION_MAX_BODY_BYTES {
-        return Err(Report::new(TrustedServerError::RequestTooLarge {
-            message: format!(
-                "auction payload {} exceeds limit of {}",
-                body_bytes.len(),
-                AUCTION_MAX_BODY_BYTES,
-            ),
-        }));
-    }
+    enforce_max_body_size(&body_bytes, AUCTION_MAX_BODY_BYTES, "auction")?;
     let body: AdRequest =
         serde_json::from_slice(&body_bytes).change_context(TrustedServerError::Auction {
             message: "Failed to parse auction request body".to_string(),

crates/trusted-server-core/src/compat.rs

@@ -96,7 +96,10 @@ pub fn to_fastly_request(req: http::Request<EdgeBody>) -> fastly::Request {
 
 /// Convert a borrowed `http::Request<EdgeBody>` into a `fastly::Request`.
 ///
-/// Headers, method, and URI are copied; the body is empty.
+/// Headers, method, and URI are copied; the body is always empty. This function
+/// requires that the caller has already consumed or discarded the body — passing
+/// a request with a non-empty body is a caller error: the body bytes will be
+/// silently lost with no warning or panic.
 ///
 /// # PR 15 removal target
 pub fn to_fastly_request_ref(req: &http::Request<EdgeBody>) -> fastly::Request {
@@ -670,6 +673,38 @@ mod tests {
         );
     }
 
+    #[test]
+    fn to_fastly_response_skeleton_copies_status_and_headers_discards_body() {
+        let http_resp = http::Response::builder()
+            .status(206)
+            .header("content-type", "text/html; charset=utf-8")
+            .header("x-custom", "value")
+            .body(EdgeBody::from(b"some body bytes".as_ref()))
+            .expect("should build response");
+
+        let fastly_resp = to_fastly_response_skeleton(http_resp);
+
+        assert_eq!(
+            fastly_resp.get_status().as_u16(),
+            206,
+            "should copy status code"
+        );
+        assert_eq!(
+            fastly_resp
+                .get_header("content-type")
+                .and_then(|v| v.to_str().ok()),
+            Some("text/html; charset=utf-8"),
+            "should copy content-type header"
+        );
+        assert_eq!(
+            fastly_resp
+                .get_header("x-custom")
+                .and_then(|v| v.to_str().ok()),
+            Some("value"),
+            "should copy custom header"
+        );
+    }
+
     #[test]
     fn to_fastly_request_with_streaming_body_produces_empty_body() {
         // Stream bodies cannot cross the compat boundary: the Fastly SDK has no

crates/trusted-server-core/src/http_util.rs

@@ -1,11 +1,13 @@
 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
 use chacha20poly1305::{aead::Aead, aead::KeyInit, XChaCha20Poly1305, XNonce};
 use edgezero_core::body::Body as EdgeBody;
+use error_stack::Report;
 use http::{header, Request, Response, StatusCode};
 use sha2::{Digest, Sha256};
 use subtle::ConstantTimeEq as _;
 
 use crate::constants::INTERNAL_HEADERS;
+use crate::error::TrustedServerError;
 use crate::platform::ClientInfo;
 use crate::settings::Settings;
 
@@ -399,6 +401,27 @@ pub fn compute_encrypted_sha256_token(settings: &Settings, full_url: &str) -> St
     URL_SAFE_NO_PAD.encode(digest)
 }
 
+/// Return an error if `bytes` exceeds `limit`.
+///
+/// # Errors
+///
+/// Returns [`TrustedServerError::RequestTooLarge`] when `bytes.len() > limit`.
+pub fn enforce_max_body_size(
+    bytes: &[u8],
+    limit: usize,
+    what: &str,
+) -> Result<(), Report<TrustedServerError>> {
+    if bytes.len() > limit {
+        return Err(Report::new(TrustedServerError::RequestTooLarge {
+            message: format!(
+                "{what} payload {} exceeds limit of {limit} bytes",
+                bytes.len()
+            ),
+        }));
+    }
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/trusted-server-core/src/proxy.rs

@@ -1,5 +1,5 @@
 use crate::backend::DEFAULT_FIRST_BYTE_TIMEOUT;
-use crate::http_util::{compute_encrypted_sha256_token, ct_str_eq};
+use crate::http_util::{compute_encrypted_sha256_token, ct_str_eq, enforce_max_body_size};
 use edgezero_core::body::Body as EdgeBody;
 use edgezero_core::http::{request_builder as edge_request_builder, Uri as EdgeUri};
 use error_stack::{Report, ResultExt};
@@ -883,15 +883,7 @@ pub async fn handle_first_party_proxy_sign(
 
     let payload = if method == Method::POST {
         let body_bytes = req.into_body().into_bytes();
-        if body_bytes.len() > SIGN_MAX_BODY_BYTES {
-            return Err(Report::new(TrustedServerError::RequestTooLarge {
-                message: format!(
-                    "payload size {} exceeds limit of {} bytes",
-                    body_bytes.len(),
-                    SIGN_MAX_BODY_BYTES,
-                ),
-            }));
-        }
+        enforce_max_body_size(&body_bytes, SIGN_MAX_BODY_BYTES, "first-party sign")?;
         let body =
             std::str::from_utf8(&body_bytes).change_context(TrustedServerError::InvalidUtf8 {
                 message: "first-party sign request body should be valid UTF-8".to_string(),
@@ -1006,15 +998,7 @@ pub async fn handle_first_party_proxy_rebuild(
     let req_url = req.uri().to_string();
     let payload = if method == Method::POST {
         let body_bytes = req.into_body().into_bytes();
-        if body_bytes.len() > REBUILD_MAX_BODY_BYTES {
-            return Err(Report::new(TrustedServerError::RequestTooLarge {
-                message: format!(
-                    "payload size {} exceeds limit of {} bytes",
-                    body_bytes.len(),
-                    REBUILD_MAX_BODY_BYTES,
-                ),
-            }));
-        }
+        enforce_max_body_size(&body_bytes, REBUILD_MAX_BODY_BYTES, "first-party rebuild")?;
         let body =
             std::str::from_utf8(&body_bytes).change_context(TrustedServerError::InvalidUtf8 {
                 message: "first-party rebuild request body should be valid UTF-8".to_string(),

crates/trusted-server-core/src/publisher.rs

@@ -396,13 +396,10 @@ pub struct OwnedProcessResponseParams {
 
 /// Stream the publisher response body through the processing pipeline.
 ///
-/// Called by the adapter after `stream_to_client()` has committed the
-/// response headers. Writes processed chunks directly to `output`.
-///
-/// This is `async` because it uses `services.http_client().send(...).await` rather
-/// than the synchronous Fastly SDK `req.send()`. The only caller wraps the entire
-/// route handler in `block_on`, so behavior is equivalent — the change reflects the
-/// migration to the platform-agnostic HTTP client.
+/// Called by the adapter after `stream_to_client()` has committed the response
+/// headers. Runs synchronously against an already-materialised body; the async
+/// I/O happened upstream in [`handle_publisher_request`]. Writes processed
+/// chunks directly to `output`.
 ///
 /// # Errors
 ///

crates/trusted-server-core/src/request_signing/endpoints.rs

@@ -9,6 +9,7 @@ use http::{header, Request, Response, StatusCode};
 use serde::{Deserialize, Serialize};
 
 use crate::error::{IntoHttpResponse, TrustedServerError};
+use crate::http_util::enforce_max_body_size;
 use crate::platform::RuntimeServices;
 use crate::request_signing::discovery::TrustedServerDiscovery;
 use crate::request_signing::rotation::KeyRotationManager;
@@ -100,15 +101,7 @@ pub fn handle_verify_signature(
     req: Request<EdgeBody>,
 ) -> Result<Response<EdgeBody>, Report<TrustedServerError>> {
     let body = req.into_body().into_bytes();
-    if body.len() > VERIFY_MAX_BODY_BYTES {
-        return Err(Report::new(TrustedServerError::RequestTooLarge {
-            message: format!(
-                "verify-signature payload {} exceeds limit of {}",
-                body.len(),
-                VERIFY_MAX_BODY_BYTES,
-            ),
-        }));
-    }
+    enforce_max_body_size(&body, VERIFY_MAX_BODY_BYTES, "verify-signature")?;
     let verify_req: VerifySignatureRequest =
         serde_json::from_slice(&body).change_context(TrustedServerError::Configuration {
             message: "invalid JSON request body".into(),
@@ -251,15 +244,7 @@ pub fn handle_rotate_key(
     } = signing_store_ids(settings)?;
 
     let body = req.into_body().into_bytes();
-    if body.len() > ADMIN_MAX_BODY_BYTES {
-        return Err(Report::new(TrustedServerError::RequestTooLarge {
-            message: format!(
-                "rotate-key payload {} exceeds limit of {}",
-                body.len(),
-                ADMIN_MAX_BODY_BYTES,
-            ),
-        }));
-    }
+    enforce_max_body_size(&body, ADMIN_MAX_BODY_BYTES, "rotate-key")?;
     let rotate_req: RotateKeyRequest = if body.is_empty() {
         RotateKeyRequest { kid: None }
     } else {
@@ -378,15 +363,7 @@ pub fn handle_deactivate_key(
     } = signing_store_ids(settings)?;
 
     let body = req.into_body().into_bytes();
-    if body.len() > ADMIN_MAX_BODY_BYTES {
-        return Err(Report::new(TrustedServerError::RequestTooLarge {
-            message: format!(
-                "deactivate-key payload {} exceeds limit of {}",
-                body.len(),
-                ADMIN_MAX_BODY_BYTES,
-            ),
-        }));
-    }
+    enforce_max_body_size(&body, ADMIN_MAX_BODY_BYTES, "deactivate-key")?;
     let deactivate_req: DeactivateKeyRequest =
         serde_json::from_slice(&body).change_context(TrustedServerError::Configuration {
             message: "invalid JSON request body".into(),