From 35e872bf621962c814d13866fac26b5c75f4d44c Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Wed, 8 Jul 2026 20:56:59 +0530 Subject: [PATCH 01/15] Stream publisher origin bodies end-to-end on Fastly Publisher pages were fully buffered before the first byte reached the client: the platform client materialized the origin body (10 MiB cap), the rewrite pipeline ran over an in-memory cursor, and the EdgeZero finalize buffered the assembled response while awaiting auction collection. TTFB therefore tracked full origin transfer plus the auction instead of origin first byte. - Add supports_streaming_responses() to PlatformHttpClient (default false, Fastly true) and request with_stream_response() on the publisher origin fetch only where honored - Teach the pipeline to consume Body::Stream asynchronously: BodyChunkSource (cumulative raw-byte cap via publisher.max_buffered_body_bytes), push-style BodyStreamDecoder/BodyStreamEncoder in streaming_processor - Replace the Fastly buffered finalize with publisher_response_into_streaming_response: a lazy Body::Stream that commits headers at origin first byte, streams rewritten chunks, and holds only the tail for auction collection; bids still inject before body close - Share one hold implementation (hold_step_decoded_chunk / hold_finish_segments) between the lazy body and the writer-driven loop so the paths cannot drift; collect_non_html_auction dedupes the collect-before-stream path - Finalize brotli decode with close() so truncated origin streams error instead of silently truncating; decode failures emit stream_decode_error telemetry - Guard bodiless (HEAD/204/304) responses and log wasted auction dispatch, matching the buffered finalizer Local A/B on a 183 KB gzip publisher page with a live 3-slot auction (release builds, 20 interleaved rounds): TTFB median 741 ms buffered vs 161 ms streamed (-78%); guest wall time and wasm heap unchanged. --- Cargo.lock | 1 + Cargo.toml | 1 + .../trusted-server-adapter-fastly/src/app.rs | 43 +- .../trusted-server-adapter-fastly/src/main.rs | 11 +- .../src/platform.rs | 4 + crates/trusted-server-core/Cargo.toml | 1 + .../trusted-server-core/src/platform/http.rs | 11 + .../src/platform/test_support.rs | 15 + crates/trusted-server-core/src/proxy.rs | 9 +- crates/trusted-server-core/src/publisher.rs | 1915 +++++++++++++++-- crates/trusted-server-core/src/settings.rs | 15 +- .../src/streaming_processor.rs | 191 ++ ...2026-07-08-true-origin-streaming-fastly.md | 1039 +++++++++ 13 files changed, 3078 insertions(+), 178 deletions(-) create mode 100644 docs/superpowers/plans/2026-07-08-true-origin-streaming-fastly.md diff --git a/Cargo.lock b/Cargo.lock index a19be7abb..cd2227882 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5277,6 +5277,7 @@ dependencies = [ name = "trusted-server-core" version = "0.1.0" dependencies = [ + "async-stream", "async-trait", "base64", "brotli", diff --git a/Cargo.toml b/Cargo.toml index 27411acbd..0512371e7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ debug = 1 [workspace.dependencies] anyhow = "1" +async-stream = "0.3" async-trait = "0.1" axum = "0.8" base64 = "0.22" diff --git a/crates/trusted-server-adapter-fastly/src/app.rs b/crates/trusted-server-adapter-fastly/src/app.rs index e56498b10..d5e37f91a 100644 --- a/crates/trusted-server-adapter-fastly/src/app.rs +++ b/crates/trusted-server-adapter-fastly/src/app.rs @@ -65,10 +65,10 @@ //! run on these responses. Legacy ran EC finalization on its own auth //! challenges. Like the 401 geo-skip, this is privacy-conservative: no EC //! cookies are issued to unauthenticated callers. -//! - **Publisher responses** are buffered (bounded by -//! `publisher.max_buffered_body_bytes`) instead of streamed to the client. -//! Asset responses are streamed straight to the client (see -//! [`dispatch_asset_fallback`]), matching legacy. +//! - **Publisher responses** keep Fastly origin bodies streaming through the +//! `EdgeZero` response body when the body is processable or pass-through. +//! Adapters without streaming-body support still use the bounded buffered +//! finalizer. //! - **Router-level 405s** (unregistered verbs) skip EC finalization along //! with the middleware chain; the entry point still adds TS headers. //! @@ -116,8 +116,8 @@ use trusted_server_core::proxy::{ handle_first_party_proxy_rebuild, handle_first_party_proxy_sign, AssetProxyCachePolicy, }; use trusted_server_core::publisher::{ - buffer_publisher_response_async, handle_page_bids, handle_publisher_request, - handle_tsjs_dynamic, page_bids_preflight_denied, AuctionDispatch, + handle_page_bids, handle_publisher_request, handle_tsjs_dynamic, page_bids_preflight_denied, + publisher_response_into_streaming_response, AuctionDispatch, }; use trusted_server_core::request_signing::{ handle_deactivate_key, handle_rotate_key, handle_trusted_server_discovery, @@ -721,10 +721,9 @@ async fn dispatch_fallback( let result = if uses_dynamic_tsjs_fallback(&method, &path) { handle_tsjs_dynamic(&req, &state.registry) } else if state.registry.has_route(&method, &path) { - // Integration-proxy responses are not bounded by publisher.max_buffered_body_bytes. - // Only the handle_publisher_request branch below routes through - // buffer_publisher_response_async. Integration responses are small in practice - // and the EdgeZero flag is off by default; extend the cap here if that changes. + // Integration-proxy responses are not bounded by + // publisher.max_buffered_body_bytes. Publisher fallback below uses the + // publisher-specific streaming finalizer instead. state .registry .handle_proxy(ProxyDispatchInput { @@ -773,9 +772,8 @@ async fn dispatch_fallback( match runtime_services_for_consent_route(&state.settings, services) { Ok(publisher_services) => { // Run the server-side auction with the configured creative- - // opportunity slots and collect the dispatched bids in the - // buffered finalize (`buffer_publisher_response_async`), matching - // the legacy streaming path. `handle_publisher_request` matches the + // opportunity slots and collect dispatched bids from the lazy + // publisher body stream. `handle_publisher_request` matches the // slots against the request path. The partner registry plus the // EC identity-graph KV (`ec.kv_graph`) enrich the bid request with // server-side EIDs, same as the legacy auction. @@ -797,17 +795,14 @@ async fn dispatch_fallback( ) .await { - Ok(pub_response) => { - buffer_publisher_response_async( - pub_response, - &method, - &state.settings, - &state.registry, - &state.orchestrator, - &publisher_services, - ) - .await - } + Ok(pub_response) => publisher_response_into_streaming_response( + pub_response, + &method, + Arc::clone(&state.settings), + state.registry.as_ref(), + Arc::clone(&state.orchestrator), + publisher_services.clone(), + ), Err(e) => Err(e), } } diff --git a/crates/trusted-server-adapter-fastly/src/main.rs b/crates/trusted-server-adapter-fastly/src/main.rs index d20de533d..963686cb8 100644 --- a/crates/trusted-server-adapter-fastly/src/main.rs +++ b/crates/trusted-server-adapter-fastly/src/main.rs @@ -321,10 +321,9 @@ fn run_edgezero_pull_sync_after_send( /// Sends a finalized `EdgeZero` response to the client. /// -/// Asset streams commit headers first, then pipe the origin body chunk by chunk -/// so large responses do not materialize in the Wasm heap. Publisher responses -/// are buffered by the server-side auction path so bids can be injected into the -/// document, and are sent in one shot along with all other responses. +/// Streaming `EdgeZero` bodies commit headers first, then pipe chunks to Fastly's +/// client stream so large asset and publisher-origin responses do not +/// materialize in the Wasm heap. fn send_edgezero_response( mut response: HttpResponse, request_filter_effects: Option<&RequestFilterEffects>, @@ -350,11 +349,11 @@ fn send_edgezero_response( match futures::executor::block_on(stream_asset_body(body, &mut streaming_body)) { Ok(()) => { if let Err(e) = streaming_body.finish() { - log::error!("failed to finish EdgeZero asset streaming body: {e}"); + log::error!("failed to finish EdgeZero streaming body: {e}"); } } Err(e) => { - log::error!("EdgeZero asset streaming failed: {e:?}"); + log::error!("EdgeZero streaming failed: {e:?}"); drop(streaming_body); } } diff --git a/crates/trusted-server-adapter-fastly/src/platform.rs b/crates/trusted-server-adapter-fastly/src/platform.rs index 106ced787..65d0f4b0d 100644 --- a/crates/trusted-server-adapter-fastly/src/platform.rs +++ b/crates/trusted-server-adapter-fastly/src/platform.rs @@ -426,6 +426,10 @@ pub struct FastlyPlatformHttpClient; #[async_trait::async_trait(?Send)] impl PlatformHttpClient for FastlyPlatformHttpClient { + fn supports_streaming_responses(&self) -> bool { + true + } + async fn send( &self, request: PlatformHttpRequest, diff --git a/crates/trusted-server-core/Cargo.toml b/crates/trusted-server-core/Cargo.toml index ba88f361f..bedefc327 100644 --- a/crates/trusted-server-core/Cargo.toml +++ b/crates/trusted-server-core/Cargo.toml @@ -13,6 +13,7 @@ workspace = true [dependencies] async-trait = { workspace = true } +async-stream = { workspace = true } base64 = { workspace = true } brotli = { workspace = true } bytes = { workspace = true } diff --git a/crates/trusted-server-core/src/platform/http.rs b/crates/trusted-server-core/src/platform/http.rs index 80bb23121..9e9337edd 100644 --- a/crates/trusted-server-core/src/platform/http.rs +++ b/crates/trusted-server-core/src/platform/http.rs @@ -276,6 +276,17 @@ pub trait PlatformHttpClient: Send + Sync { true } + /// Whether [`send`](Self::send) can preserve upstream response bodies as + /// [`Body::Stream`](edgezero_core::body::Body::Stream) when requested via + /// [`PlatformHttpRequest::with_stream_response`]. + /// + /// Adapters that cannot preserve streaming response bodies must keep the + /// default `false` so callers do not request a contract the adapter will + /// reject or silently buffer. + fn supports_streaming_responses(&self) -> bool { + false + } + /// Wait for one of the in-flight requests to complete. /// /// # Errors diff --git a/crates/trusted-server-core/src/platform/test_support.rs b/crates/trusted-server-core/src/platform/test_support.rs index ee7201fb8..0c86b9594 100644 --- a/crates/trusted-server-core/src/platform/test_support.rs +++ b/crates/trusted-server-core/src/platform/test_support.rs @@ -224,6 +224,9 @@ pub(crate) struct StubHttpClient { // Reported by supports_concurrent_fanout(); set false to emulate // platforms whose send_async executes eagerly (e.g. Cloudflare Workers). concurrent_fanout: std::sync::atomic::AtomicBool, + // Reported by supports_streaming_responses(); set true to emulate Fastly's + // streaming response support. + streaming_responses_supported: std::sync::atomic::AtomicBool, image_optimizer_options: Mutex>>, stream_response_flags: Mutex>, request_methods: Mutex>, @@ -246,6 +249,7 @@ impl StubHttpClient { request_headers: Mutex::new(Vec::new()), select_errors: Mutex::new(VecDeque::new()), concurrent_fanout: std::sync::atomic::AtomicBool::new(true), + streaming_responses_supported: std::sync::atomic::AtomicBool::new(false), image_optimizer_options: Mutex::new(Vec::new()), stream_response_flags: Mutex::new(Vec::new()), request_methods: Mutex::new(Vec::new()), @@ -260,6 +264,12 @@ impl StubHttpClient { .store(supported, std::sync::atomic::Ordering::Relaxed); } + /// Make `supports_streaming_responses()` report the given value. + pub fn set_streaming_responses_supported(&self, supported: bool) { + self.streaming_responses_supported + .store(supported, std::sync::atomic::Ordering::Relaxed); + } + /// Queue a canned response by status code and body bytes. pub fn push_response(&self, status: u16, body: Vec) { self.push_response_with_headers(status, body, Vec::<(String, String)>::new()); @@ -363,6 +373,11 @@ impl PlatformHttpClient for StubHttpClient { .load(std::sync::atomic::Ordering::Relaxed) } + fn supports_streaming_responses(&self) -> bool { + self.streaming_responses_supported + .load(std::sync::atomic::Ordering::Relaxed) + } + async fn send( &self, request: PlatformHttpRequest, diff --git a/crates/trusted-server-core/src/proxy.rs b/crates/trusted-server-core/src/proxy.rs index 19c10a80d..00444b384 100644 --- a/crates/trusted-server-core/src/proxy.rs +++ b/crates/trusted-server-core/src/proxy.rs @@ -246,7 +246,10 @@ fn platform_response_to_fastly_asset(platform_resp: PlatformResponse) -> AssetPr } } -/// Stream an asset response body directly to a writable client stream. +/// Stream a platform response body directly to a writable client stream. +/// +/// Asset routes and Fastly `EdgeZero` publisher fallback both use this bridge +/// after headers have been committed through `stream_to_client()`. /// /// # Errors /// @@ -261,7 +264,7 @@ pub async fn stream_asset_body( output .write_all(bytes.as_ref()) .change_context(TrustedServerError::Proxy { - message: "failed to write buffered asset response body".to_string(), + message: "failed to write buffered platform response body".to_string(), })?; } EdgeBody::Stream(mut stream) => { @@ -274,7 +277,7 @@ pub async fn stream_asset_body( output .write_all(chunk.as_ref()) .change_context(TrustedServerError::Proxy { - message: "failed to write streaming asset response body".to_string(), + message: "failed to write streaming platform response body".to_string(), })?; } } diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index 6b0ea3a5d..c10ac0b39 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -22,9 +22,15 @@ use std::io::Write; use std::sync::{Arc, Mutex}; use std::time::Duration; +use brotli::enc::writer::CompressorWriter; +use brotli::enc::BrotliEncoderParams; +use brotli::Decompressor; use cookie::CookieJar; use edgezero_core::body::Body as EdgeBody; use error_stack::{Report, ResultExt}; +use flate2::read::{GzDecoder, ZlibDecoder}; +use flate2::write::{GzEncoder, ZlibEncoder}; +use futures::StreamExt as _; use http::{header, HeaderValue, Method, Request, Response, StatusCode, Uri}; use crate::auction::endpoints::{ @@ -53,19 +59,134 @@ use crate::platform::{GeoInfo, PlatformBackendSpec, PlatformHttpRequest, Runtime use crate::price_bucket::{price_bucket, PriceGranularity}; use crate::rsc_flight::RscFlightUrlRewriter; use crate::settings::Settings; -use crate::streaming_processor::{Compression, PipelineConfig, StreamProcessor, StreamingPipeline}; +use crate::streaming_processor::{ + BodyStreamDecoder, BodyStreamEncoder, Compression, PipelineConfig, StreamProcessor, + StreamingPipeline, STREAM_CHUNK_SIZE, +}; use crate::streaming_replacer::create_url_replacer; const SUPPORTED_ENCODING_VALUES: [&str; 3] = ["gzip", "deflate", "br"]; const DEFAULT_PUBLISHER_FIRST_BYTE_TIMEOUT: Duration = Duration::from_secs(15); -/// Read buffer size for streaming body processing and brotli internal buffers. -/// Both the `Decompressor` and `CompressorWriter` use this value so all -/// brotli I/O layers operate on consistently-sized chunks. -const STREAM_CHUNK_SIZE: usize = 8192; +fn body_as_reader( + body: EdgeBody, +) -> Result, Report> { + let bytes = body.into_bytes().ok_or_else(|| { + Report::new(TrustedServerError::Proxy { + message: "streaming body cannot be processed by sync publisher pipeline".to_string(), + }) + })?; + Ok(std::io::Cursor::new(bytes)) +} + +struct BodyChunkSource { + body: Option, + chunk_size: usize, + max_bytes: usize, + bytes_seen: usize, + once_offset: usize, +} + +impl BodyChunkSource { + fn new(body: EdgeBody, chunk_size: usize) -> Self { + Self { + body: Some(body), + chunk_size, + max_bytes: usize::MAX, + bytes_seen: 0, + once_offset: 0, + } + } + + fn with_max_bytes(mut self, max_bytes: usize) -> Self { + self.max_bytes = max_bytes; + self + } + + async fn next_chunk(&mut self) -> Result, Report> { + let Some(body) = self.body.take() else { + return Ok(None); + }; + + let chunk = match body { + EdgeBody::Once(bytes) => { + if self.once_offset >= bytes.len() { + None + } else { + let end = (self.once_offset + self.chunk_size).min(bytes.len()); + let chunk = bytes.slice(self.once_offset..end); + self.once_offset = end; + if self.once_offset < bytes.len() { + self.body = Some(EdgeBody::Once(bytes)); + } + Some(chunk) + } + } + EdgeBody::Stream(mut stream) => match stream.next().await { + Some(Ok(chunk)) => { + self.body = Some(EdgeBody::Stream(stream)); + Some(chunk) + } + Some(Err(err)) => { + return Err(Report::new(TrustedServerError::Proxy { + message: format!("Failed to read publisher origin body stream: {err}"), + })); + } + None => None, + }, + }; + + let Some(chunk) = chunk else { + return Ok(None); + }; + + self.bytes_seen = self.bytes_seen.checked_add(chunk.len()).ok_or_else(|| { + Report::new(TrustedServerError::Proxy { + message: "publisher origin body byte count overflowed".to_string(), + }) + })?; + if self.bytes_seen > self.max_bytes { + return Err(Report::new(TrustedServerError::Proxy { + message: format!( + "publisher origin body exceeded {}-byte streaming limit", + self.max_bytes + ), + })); + } + + Ok(Some(chunk)) + } +} + +fn process_and_encode_chunk( + processor: &mut P, + encoder: &mut BodyStreamEncoder, + chunk: &[u8], + is_last: bool, + process_error: &str, +) -> Result, Report> { + let processed = + processor + .process_chunk(chunk, is_last) + .change_context(TrustedServerError::Proxy { + message: process_error.to_string(), + })?; + if processed.is_empty() { + return Ok(None); + } + let encoded = encoder.encode_chunk(&processed)?; + if encoded.is_empty() { + return Ok(None); + } + Ok(Some(bytes::Bytes::from(encoded))) +} -fn body_as_reader(body: EdgeBody) -> std::io::Cursor { - std::io::Cursor::new(body.into_bytes().unwrap_or_default()) +fn publisher_stream_error(err: Report) -> std::io::Error { + let message = format!("{err:?}"); + // Consume the report so clippy's needless_pass_by_value accepts the + // by-value signature that `map_err(publisher_stream_error)` requires. + drop(err); + std::io::Error::other(message) } fn not_found_response() -> Response { @@ -233,6 +354,55 @@ struct ProcessResponseParams<'a> { ad_bids_state: &'a Arc>>, } +struct PublisherBodyProcessor { + inner: Box, +} + +impl PublisherBodyProcessor { + fn new( + params: &OwnedProcessResponseParams, + settings: &Settings, + integration_registry: &IntegrationRegistry, + ) -> Result> { + let is_html = is_html_content_type(¶ms.content_type); + let is_rsc_flight = + content_type_contains_ascii_case_insensitive(¶ms.content_type, "text/x-component"); + let inner: Box = if is_html { + Box::new(create_html_stream_processor( + ¶ms.origin_host, + ¶ms.request_host, + ¶ms.request_scheme, + settings, + integration_registry, + params.ad_slots_script.as_deref().map(str::to_string), + Arc::clone(¶ms.ad_bids_state), + )?) + } else if is_rsc_flight { + Box::new(RscFlightUrlRewriter::new( + ¶ms.origin_host, + ¶ms.origin_url, + ¶ms.request_host, + ¶ms.request_scheme, + )) + } else { + Box::new(create_url_replacer( + ¶ms.origin_host, + ¶ms.origin_url, + ¶ms.request_host, + ¶ms.request_scheme, + )) + }; + + Ok(Self { inner }) + } +} + +impl StreamProcessor for PublisherBodyProcessor { + fn process_chunk(&mut self, chunk: &[u8], is_last: bool) -> Result, std::io::Error> { + self.inner.process_chunk(chunk, is_last) + } +} + /// Process response body through the streaming pipeline. /// /// Selects the appropriate processor based on content type (HTML rewriter, @@ -276,7 +446,7 @@ fn process_response_streaming( params.ad_slots_script.map(str::to_string), params.ad_bids_state.clone(), )?; - StreamingPipeline::new(config, processor).process(body_as_reader(body), output)?; + StreamingPipeline::new(config, processor).process(body_as_reader(body)?, output)?; } else if is_rsc_flight { // RSC Flight responses are length-prefixed (T rows). A naive string replacement will // corrupt the stream by changing byte lengths without updating the prefixes. @@ -286,7 +456,7 @@ fn process_response_streaming( params.request_host, params.request_scheme, ); - StreamingPipeline::new(config, processor).process(body_as_reader(body), output)?; + StreamingPipeline::new(config, processor).process(body_as_reader(body)?, output)?; } else { let replacer = create_url_replacer( params.origin_host, @@ -294,12 +464,352 @@ fn process_response_streaming( params.request_host, params.request_scheme, ); - StreamingPipeline::new(config, replacer).process(body_as_reader(body), output)?; + StreamingPipeline::new(config, replacer).process(body_as_reader(body)?, output)?; } Ok(()) } +async fn process_response_streaming_async( + body: EdgeBody, + output: &mut W, + params: &ProcessResponseParams<'_>, + max_raw_body_bytes: usize, +) -> Result<(), Report> { + let is_html = is_html_content_type(params.content_type); + let is_rsc_flight = + content_type_contains_ascii_case_insensitive(params.content_type, "text/x-component"); + log::debug!( + "process_response_streaming_async: content_type={}, content_encoding={}, is_html={}, is_rsc_flight={}", + params.content_type, + params.content_encoding, + is_html, + is_rsc_flight + ); + + let compression = Compression::from_content_encoding(params.content_encoding); + + if is_html { + let mut processor = create_html_stream_processor( + params.origin_host, + params.request_host, + params.request_scheme, + params.settings, + params.integration_registry, + params.ad_slots_script.map(str::to_string), + params.ad_bids_state.clone(), + )?; + process_body_chunks_async( + body, + output, + &mut processor, + compression, + max_raw_body_bytes, + ) + .await + } else if is_rsc_flight { + let mut processor = RscFlightUrlRewriter::new( + params.origin_host, + params.origin_url, + params.request_host, + params.request_scheme, + ); + process_body_chunks_async( + body, + output, + &mut processor, + compression, + max_raw_body_bytes, + ) + .await + } else { + let mut replacer = create_url_replacer( + params.origin_host, + params.origin_url, + params.request_host, + params.request_scheme, + ); + process_body_chunks_async(body, output, &mut replacer, compression, max_raw_body_bytes) + .await + } +} + +async fn process_body_chunks_async( + body: EdgeBody, + writer: &mut W, + processor: &mut P, + compression: Compression, + max_raw_body_bytes: usize, +) -> Result<(), Report> { + let mut decoder = BodyStreamDecoder::new(compression); + let mut encoder = BodyStreamEncoder::new(compression); + let mut source = + BodyChunkSource::new(body, STREAM_CHUNK_SIZE).with_max_bytes(max_raw_body_bytes); + + while let Some(chunk) = source.next_chunk().await? { + let decoded = decoder.decode_chunk(&chunk)?; + if decoded.is_empty() { + continue; + } + if let Some(encoded) = process_and_encode_chunk( + processor, + &mut encoder, + &decoded, + false, + "Failed to process chunk", + )? { + write_encoded_segment(writer, &encoded)?; + } + } + + for encoded in passthrough_finish_segments(processor, &mut decoder, &mut encoder)? { + write_encoded_segment(writer, &encoded)?; + } + writer.flush().change_context(TrustedServerError::Proxy { + message: "Failed to flush output".to_string(), + })?; + + Ok(()) +} + +/// Write one encoded output segment produced by the chunk pipeline. +fn write_encoded_segment( + writer: &mut W, + encoded: &[u8], +) -> Result<(), Report> { + writer + .write_all(encoded) + .change_context(TrustedServerError::Proxy { + message: "Failed to write encoded chunk".to_string(), + }) +} + +/// Finalize a no-hold chunk pipeline: drain the decoder tail through the +/// processor, signal end-of-stream to the processor, and emit the encoder +/// trailer. Returns the encoded segments for the caller to emit. +fn passthrough_finish_segments( + processor: &mut P, + decoder: &mut BodyStreamDecoder, + encoder: &mut BodyStreamEncoder, +) -> Result, Report> { + let mut segments = Vec::new(); + let decoded_tail = decoder.finish()?; + if !decoded_tail.is_empty() { + if let Some(encoded) = process_and_encode_chunk( + processor, + encoder, + &decoded_tail, + false, + "Failed to process decoded tail", + )? { + segments.push(encoded); + } + } + if let Some(encoded) = process_and_encode_chunk( + processor, + encoder, + &[], + true, + "Failed to finalize processor", + )? { + segments.push(encoded); + } + let trailer = encoder.finish()?; + if !trailer.is_empty() { + segments.push(bytes::Bytes::from(trailer)); + } + Ok(segments) +} + +/// Mutable auction-hold state threaded through the streaming hold pipeline. +struct AuctionHoldState { + hold: Option, + dispatched: Option, + telemetry: AuctionTelemetryCarry, +} + +impl AuctionHoldState { + fn new(dispatched: DispatchedAuction, telemetry: AuctionTelemetryCarry) -> Self { + Self { + hold: Some(BodyCloseHoldBuffer::new()), + dispatched: Some(dispatched), + telemetry, + } + } +} + +/// Abandon the in-flight auction (if still pending) with the given telemetry +/// reason. No-op once the auction has been collected or already abandoned. +async fn abandon_hold_auction( + state: &mut AuctionHoldState, + services: &RuntimeServices, + reason: &'static str, +) { + if let Some(dispatched) = state.dispatched.take() { + emit_abandoned_auction( + services, + state.telemetry.observation.take(), + dispatched, + reason, + ) + .await; + } +} + +/// Feed one decoded chunk through the close-body hold and processor. +/// +/// Returns the encoded output segments for the caller to emit — written to a +/// client stream by [`body_close_hold_loop_stream`], yielded from the lazy +/// body by [`publisher_response_into_streaming_response`]. Both async hold +/// paths share this function so their behavior cannot drift apart. +/// +/// When the raw `( + processor: &mut P, + encoder: &mut BodyStreamEncoder, + chunk: &[u8], + state: &mut AuctionHoldState, + collect_refs: &AuctionHoldCollectRefs<'_>, +) -> Result, Report> { + let mut segments = Vec::new(); + if let Some(hold_buffer) = state.hold.as_mut() { + let ready = hold_buffer.push(chunk); + match process_and_encode_chunk(processor, encoder, &ready, false, "Failed to process chunk") + { + Ok(Some(encoded)) => segments.push(encoded), + Ok(None) => {} + Err(err) => { + abandon_hold_auction(state, collect_refs.services, "stream_process_error").await; + return Err(err); + } + } + + if state + .hold + .as_ref() + .is_some_and(BodyCloseHoldBuffer::found_close) + { + let dispatched = state + .dispatched + .take() + .expect("should have dispatched auction to collect"); + collect_stream_auction( + dispatched, + state.telemetry.take(), + collect_refs.price_granularity, + collect_refs.ad_bids_state, + collect_refs.orchestrator, + collect_refs.services, + collect_refs.settings, + ) + .await; + + let held = state + .hold + .take() + .expect("should have close-body hold buffer") + .finish(); + if let Some(encoded) = process_and_encode_chunk( + processor, + encoder, + &held, + false, + "Failed to process held body close", + )? { + segments.push(encoded); + } + } + } else { + match process_and_encode_chunk(processor, encoder, chunk, false, "Failed to process chunk") + { + Ok(Some(encoded)) => segments.push(encoded), + Ok(None) => {} + Err(err) => { + abandon_hold_auction(state, collect_refs.services, "stream_process_error").await; + return Err(err); + } + } + } + Ok(segments) +} + +/// Finalize the close-body hold pipeline at end of the origin stream. +/// +/// Drains the decoder tail through the hold (or straight through when the +/// hold was already released mid-stream), collects the auction if the +/// close-body tag never streamed, processes the held tail plus the +/// processor's final chunk, and emits the encoder trailer. Returns the +/// encoded segments for the caller to emit. On decoder failure the pending +/// auction is abandoned before the error is returned. +async fn hold_finish_segments( + processor: &mut P, + decoder: &mut BodyStreamDecoder, + encoder: &mut BodyStreamEncoder, + state: &mut AuctionHoldState, + collect_refs: &AuctionHoldCollectRefs<'_>, +) -> Result, Report> { + let mut segments = Vec::new(); + + let decoded_tail = match decoder.finish() { + Ok(decoded_tail) => decoded_tail, + Err(err) => { + abandon_hold_auction(state, collect_refs.services, "stream_decode_error").await; + return Err(err); + } + }; + if !decoded_tail.is_empty() { + segments.extend( + hold_step_decoded_chunk(processor, encoder, &decoded_tail, state, collect_refs).await?, + ); + } + + if let Some(hold) = state.hold.take() { + let dispatched = state + .dispatched + .take() + .expect("should have dispatched auction to collect"); + collect_stream_auction( + dispatched, + state.telemetry.take(), + collect_refs.price_granularity, + collect_refs.ad_bids_state, + collect_refs.orchestrator, + collect_refs.services, + collect_refs.settings, + ) + .await; + + let held = hold.finish(); + if let Some(encoded) = process_and_encode_chunk( + processor, + encoder, + &held, + false, + "Failed to process held body close", + )? { + segments.push(encoded); + } + } + + if let Some(encoded) = process_and_encode_chunk( + processor, + encoder, + &[], + true, + "Failed to finalize processor", + )? { + segments.push(encoded); + } + let trailer = encoder.finish()?; + if !trailer.is_empty() { + segments.push(bytes::Bytes::from(trailer)); + } + Ok(segments) +} + /// Create a unified HTML stream processor. /// /// Builds the config via [`HtmlProcessorConfig::from_settings`] and then @@ -339,16 +849,13 @@ pub enum PublisherResponse { /// content on any status (2xx or non-2xx — e.g., branded 404/500 HTML and /// error JSON still get URL rewriting) where the encoding is supported. /// Post-processors run inside the streaming processor, so processable HTML - /// is streamed regardless of whether any are registered. The caller must: - /// 1. Call `finalize_response()` on the response - /// 2. Call `response.stream_to_client()` to get a `StreamingBody` - /// 3. Call `stream_publisher_body()` with the body and streaming writer - /// 4. Call `StreamingBody::finish()` + /// is streamed regardless of whether any are registered. /// - /// **Interim (PR 15):** `body` has already been fully materialised into - /// WASM heap by the platform HTTP client. `stream_publisher_body` reads - /// from an in-memory buffer, not a live origin stream. The origin-side - /// peak is bounded by `MAX_PLATFORM_RESPONSE_BODY_BYTES`. + /// Adapters with platform streaming support preserve `body` as + /// [`EdgeBody::Stream`] and attach a lazy processed stream via + /// [`publisher_response_into_streaming_response`]. Buffered adapters use + /// [`buffer_publisher_response_async`] and are bounded by + /// `settings.publisher.max_buffered_body_bytes`. Stream { /// Response with all headers set (EC ID, cookies, etc.) /// but body not yet written. `Content-Length` already removed. @@ -363,12 +870,9 @@ pub enum PublisherResponse { /// `finalize_response()` and `send_to_client()` are applied at the outer /// response-dispatch level, not in this arm. /// - /// `Content-Length` is preserved — the body is unmodified. - /// - /// **Interim (PR 15):** `body` has been fully materialised into WASM heap. - /// Previously, binary assets streamed lazily from origin with no WASM - /// buffering. This path is now bounded by `MAX_PLATFORM_RESPONSE_BODY_BYTES`; - /// assets exceeding that limit return an error instead of exhausting heap. + /// `Content-Length` is preserved — the body is unmodified. Streaming + /// adapters reattach the origin body directly so non-processable 2xx bodies + /// can pass through without materializing in WASM memory. PassThrough { /// Response with all headers set but body not yet written. response: Response, @@ -465,7 +969,7 @@ pub struct OwnedProcessResponseParams { /// statuses (204, 304) carry no body but may advertise the `GET` representation's /// length, so they skip the buffer and length rewrite. /// -/// Every adapter (Axum, Cloudflare, Spin, and the Fastly `EdgeZero` path) calls +/// Buffered adapters (Axum, Cloudflare, Spin, and non-streaming fallbacks) call /// this: it drives /// [`stream_publisher_body_async`], which awaits /// [`AuctionOrchestrator::collect_dispatched_auction`], writes the winning bids @@ -530,48 +1034,225 @@ pub async fn buffer_publisher_response_async( } } -/// Returns `true` when a buffered publisher response should carry a body and a -/// recomputed `Content-Length`. +/// Convert a [`PublisherResponse`] into a response that preserves streaming +/// bodies where possible. /// -/// `HEAD` responses and bodiless statuses (204, 304) carry no body; rewriting -/// their `Content-Length` to the (empty) buffered length would mislead clients -/// and caches, so the origin metadata is preserved instead. -fn response_carries_body(method: &Method, status: StatusCode) -> bool { - *method != Method::HEAD - && status != StatusCode::NO_CONTENT - && status != StatusCode::NOT_MODIFIED -} - -/// A [`Write`] sink that buffers into a `Vec` but fails once the configured -/// byte limit would be exceeded. +/// Buffered adapters should keep using [`buffer_publisher_response_async`]. +/// Fastly uses this helper before the entry point commits headers, allowing the +/// response body to be pulled lazily by `stream_to_client()`. /// -/// Used to bound in-WASM-heap buffering of decoded/re-written publisher bodies. -/// A highly-compressible origin response can sit under the platform raw-body cap -/// yet expand past a safe heap size after decode and post-processing; this writer -/// turns that into a recoverable error instead of an out-of-memory abort. -pub struct BoundedWriter { - inner: Vec, - limit: usize, -} - -impl BoundedWriter { - /// Creates a writer that accepts at most `limit` bytes before erroring. - #[must_use] - pub fn new(limit: usize) -> Self { - Self { - inner: Vec::new(), - limit, +/// # Errors +/// +/// Returns an error if processor construction fails before the streaming body is +/// created. +pub fn publisher_response_into_streaming_response( + publisher_response: PublisherResponse, + method: &Method, + settings: Arc, + integration_registry: &IntegrationRegistry, + orchestrator: Arc, + services: RuntimeServices, +) -> Result, Report> { + match publisher_response { + PublisherResponse::Buffered(response) => Ok(response), + PublisherResponse::PassThrough { mut response, body } => { + if response_carries_body(method, response.status()) { + *response.body_mut() = body; + } + Ok(response) } - } - - /// Consumes the writer and returns the buffered bytes. - #[must_use] - pub fn into_inner(self) -> Vec { - self.inner - } -} - -impl Write for BoundedWriter { + PublisherResponse::Stream { + mut response, + body, + params, + } => { + if !response_carries_body(method, response.status()) { + if params.dispatched_auction.is_some() { + // A bodiless response (HEAD navigation, 204/304) has no + // `` to inject bids into, so the dispatched SSP + // requests are wasted — surface it for quota observability, + // matching the buffered finalizer. + log::warn!( + "Server-side auction dispatched but response is bodiless (method: {}, status: {}); in-flight SSP bid requests will not be collected", + method, + response.status(), + ); + } + return Ok(response); + } + + response.headers_mut().remove(header::CONTENT_LENGTH); + let mut params = *params; + let mut processor = + PublisherBodyProcessor::new(¶ms, &settings, integration_registry)?; + let stream = async_stream::try_stream! { + let compression = Compression::from_content_encoding(¶ms.content_encoding); + let mut decoder = BodyStreamDecoder::new(compression); + let mut encoder = BodyStreamEncoder::new(compression); + let mut source = BodyChunkSource::new(body, STREAM_CHUNK_SIZE) + .with_max_bytes(settings.publisher.max_buffered_body_bytes); + + // HTML rides the close-body hold so bids land before ``; + // non-HTML has no injection point, so its auction is collected + // before any byte streams (matching the buffered finalizer). + let mut hold_auction = None; + if let Some(dispatched) = params.dispatched_auction.take() { + let telemetry = AuctionTelemetryCarry { + observation: params.auction_observation.take(), + auction_request: params.auction_request.take(), + }; + if is_html_content_type(¶ms.content_type) { + hold_auction = Some((dispatched, telemetry)); + } else { + collect_non_html_auction( + dispatched, + telemetry, + ¶ms, + &orchestrator, + &services, + &settings, + ) + .await; + } + } + + if let Some((dispatched, telemetry)) = hold_auction { + let mut state = AuctionHoldState::new(dispatched, telemetry); + let collect_refs = AuctionHoldCollectRefs { + price_granularity: params.price_granularity, + ad_bids_state: ¶ms.ad_bids_state, + orchestrator: &orchestrator, + services: &services, + settings: &settings, + }; + + loop { + let raw_chunk = match source.next_chunk().await { + Ok(Some(chunk)) => chunk, + Ok(None) => break, + Err(err) => { + abandon_hold_auction(&mut state, &services, "stream_read_error") + .await; + Err(publisher_stream_error(err))?; + unreachable!("error should have returned"); + } + }; + let decoded = match decoder.decode_chunk(&raw_chunk) { + Ok(decoded) => decoded, + Err(err) => { + abandon_hold_auction(&mut state, &services, "stream_decode_error") + .await; + Err(publisher_stream_error(err))?; + unreachable!("error should have returned"); + } + }; + if decoded.is_empty() { + continue; + } + for encoded in hold_step_decoded_chunk( + &mut processor, + &mut encoder, + &decoded, + &mut state, + &collect_refs, + ) + .await + .map_err(publisher_stream_error)? + { + yield encoded; + } + } + + for encoded in hold_finish_segments( + &mut processor, + &mut decoder, + &mut encoder, + &mut state, + &collect_refs, + ) + .await + .map_err(publisher_stream_error)? + { + yield encoded; + } + } else { + while let Some(raw_chunk) = + source.next_chunk().await.map_err(publisher_stream_error)? + { + let decoded = decoder + .decode_chunk(&raw_chunk) + .map_err(publisher_stream_error)?; + if decoded.is_empty() { + continue; + } + if let Some(encoded) = process_and_encode_chunk( + &mut processor, + &mut encoder, + &decoded, + false, + "Failed to process chunk", + ) + .map_err(publisher_stream_error)? + { + yield encoded; + } + } + for encoded in + passthrough_finish_segments(&mut processor, &mut decoder, &mut encoder) + .map_err(publisher_stream_error)? + { + yield encoded; + } + } + }; + *response.body_mut() = EdgeBody::from_stream::<_, std::io::Error>(stream); + Ok(response) + } + } +} + +/// Returns `true` when a buffered publisher response should carry a body and a +/// recomputed `Content-Length`. +/// +/// `HEAD` responses and bodiless statuses (204, 304) carry no body; rewriting +/// their `Content-Length` to the (empty) buffered length would mislead clients +/// and caches, so the origin metadata is preserved instead. +fn response_carries_body(method: &Method, status: StatusCode) -> bool { + *method != Method::HEAD + && status != StatusCode::NO_CONTENT + && status != StatusCode::NOT_MODIFIED +} + +/// A [`Write`] sink that buffers into a `Vec` but fails once the configured +/// byte limit would be exceeded. +/// +/// Used to bound in-WASM-heap buffering of decoded/re-written publisher bodies. +/// A highly-compressible origin response can sit under the platform raw-body cap +/// yet expand past a safe heap size after decode and post-processing; this writer +/// turns that into a recoverable error instead of an out-of-memory abort. +pub struct BoundedWriter { + inner: Vec, + limit: usize, +} + +impl BoundedWriter { + /// Creates a writer that accepts at most `limit` bytes before erroring. + #[must_use] + pub fn new(limit: usize) -> Self { + Self { + inner: Vec::new(), + limit, + } + } + + /// Consumes the writer and returns the buffered bytes. + #[must_use] + pub fn into_inner(self) -> Vec { + self.inner + } +} + +impl Write for BoundedWriter { fn write(&mut self, buf: &[u8]) -> std::io::Result { if self.inner.len() + buf.len() > self.limit { return Err(std::io::Error::other( @@ -652,7 +1333,29 @@ pub async fn stream_publisher_body_async( services: &RuntimeServices, ) -> Result<(), Report> { let Some(dispatched) = params.dispatched_auction.take() else { - // No auction — use the existing sync pipeline unchanged. + if body.is_stream() { + let borrowed = ProcessResponseParams { + content_encoding: ¶ms.content_encoding, + origin_host: ¶ms.origin_host, + origin_url: ¶ms.origin_url, + request_host: ¶ms.request_host, + request_scheme: ¶ms.request_scheme, + settings, + content_type: ¶ms.content_type, + integration_registry, + ad_slots_script: params.ad_slots_script.as_deref(), + ad_bids_state: ¶ms.ad_bids_state, + }; + return process_response_streaming_async( + body, + output, + &borrowed, + settings.publisher.max_buffered_body_bytes, + ) + .await; + } + + // No auction and already-buffered body — keep the existing sync pipeline. return stream_publisher_body(body, output, params, settings, integration_registry); }; let telemetry = AuctionTelemetryCarry { @@ -665,35 +1368,36 @@ pub async fn stream_publisher_body_async( if !is_html { // Non-HTML: collect auction first, then stream. There is no // to hold, so delaying the entire body until collection is acceptable. - let placeholder = mediator_placeholder_request(); - let result = orchestrator - .collect_dispatched_auction( - dispatched, - services, - &make_collect_context(settings, services, &placeholder), + collect_non_html_auction( + dispatched, + telemetry, + params, + orchestrator, + services, + settings, + ) + .await; + if body.is_stream() { + let borrowed = ProcessResponseParams { + content_encoding: ¶ms.content_encoding, + origin_host: ¶ms.origin_host, + origin_url: ¶ms.origin_url, + request_host: ¶ms.request_host, + request_scheme: ¶ms.request_scheme, + settings, + content_type: ¶ms.content_type, + integration_registry, + ad_slots_script: params.ad_slots_script.as_deref(), + ad_bids_state: ¶ms.ad_bids_state, + }; + return process_response_streaming_async( + body, + output, + &borrowed, + settings.publisher.max_buffered_body_bytes, ) .await; - if let (Some(observation), Some(auction_request)) = - (telemetry.observation, telemetry.auction_request.as_ref()) - { - emit_auction_events_best_effort_lazy(services, || { - build_auction_events( - observation, - AuctionTerminalOutcome::Completed { - request: auction_request, - result: &result, - }, - ) - }) - .await; } - - write_bids_to_state( - &result.winning_bids, - params.price_granularity, - ¶ms.ad_bids_state, - settings.debug.inject_adm_for_testing, - ); return stream_publisher_body(body, output, params, settings, integration_registry); } @@ -917,6 +1621,14 @@ struct AuctionCollectCtx<'a> { settings: &'a Settings, } +struct AuctionHoldCollectRefs<'a> { + price_granularity: PriceGranularity, + ad_bids_state: &'a Arc>>, + orchestrator: &'a AuctionOrchestrator, + services: &'a RuntimeServices, + settings: &'a Settings, +} + /// Run the close-body hold loop for HTML bodies, collecting the auction before /// the raw `( @@ -926,13 +1638,20 @@ async fn stream_html_with_auction_hold( compression: Compression, ctx: AuctionCollectCtx<'_>, ) -> Result<(), Report> { - use brotli::enc::writer::CompressorWriter; - use brotli::enc::BrotliEncoderParams; - use brotli::Decompressor; - use flate2::read::{GzDecoder, ZlibDecoder}; - use flate2::write::{GzEncoder, ZlibEncoder}; + if body.is_stream() { + let max_raw_body_bytes = ctx.settings.publisher.max_buffered_body_bytes; + return body_close_hold_loop_stream( + body, + output, + processor, + compression, + ctx, + max_raw_body_bytes, + ) + .await; + } - let body = body_as_reader(body); + let body = body_as_reader(body)?; match compression { Compression::None => body_close_hold_loop(body, output, processor, ctx).await, Compression::Gzip => { @@ -969,6 +1688,85 @@ async fn stream_html_with_auction_hold( } } +/// Async-pull variant of [`body_close_hold_loop`] for live origin streams. +/// +/// Shares [`hold_step_decoded_chunk`] and [`hold_finish_segments`] with the +/// lazy streaming body built by [`publisher_response_into_streaming_response`], +/// so the two async hold paths cannot drift apart. +async fn body_close_hold_loop_stream( + body: EdgeBody, + writer: &mut W, + processor: &mut P, + compression: Compression, + ctx: AuctionCollectCtx<'_>, + max_raw_body_bytes: usize, +) -> Result<(), Report> { + let AuctionCollectCtx { + dispatched, + telemetry, + price_granularity, + ad_bids_state, + orchestrator, + services, + settings, + } = ctx; + let mut decoder = BodyStreamDecoder::new(compression); + let mut encoder = BodyStreamEncoder::new(compression); + let mut source = + BodyChunkSource::new(body, STREAM_CHUNK_SIZE).with_max_bytes(max_raw_body_bytes); + let mut state = AuctionHoldState::new(dispatched, telemetry); + let collect_refs = AuctionHoldCollectRefs { + price_granularity, + ad_bids_state, + orchestrator, + services, + settings, + }; + + loop { + let raw_chunk = match source.next_chunk().await { + Ok(Some(chunk)) => chunk, + Ok(None) => break, + Err(err) => { + abandon_hold_auction(&mut state, services, "stream_read_error").await; + return Err(err); + } + }; + let decoded = match decoder.decode_chunk(&raw_chunk) { + Ok(decoded) => decoded, + Err(err) => { + abandon_hold_auction(&mut state, services, "stream_decode_error").await; + return Err(err); + } + }; + if decoded.is_empty() { + continue; + } + for encoded in + hold_step_decoded_chunk(processor, &mut encoder, &decoded, &mut state, &collect_refs) + .await? + { + write_encoded_segment(writer, &encoded)?; + } + } + + for encoded in hold_finish_segments( + processor, + &mut decoder, + &mut encoder, + &mut state, + &collect_refs, + ) + .await? + { + write_encoded_segment(writer, &encoded)?; + } + writer.flush().change_context(TrustedServerError::Proxy { + message: "Failed to flush output".to_string(), + })?; + Ok(()) +} + const BODY_CLOSE_PREFIX: &[u8] = b"` to inject into, so bids are written to state up front and the +/// auction telemetry completes immediately. +async fn collect_non_html_auction( + dispatched: DispatchedAuction, + telemetry: AuctionTelemetryCarry, + params: &OwnedProcessResponseParams, + orchestrator: &AuctionOrchestrator, + services: &RuntimeServices, + settings: &Settings, +) { + let placeholder = mediator_placeholder_request(); + let result = orchestrator + .collect_dispatched_auction( + dispatched, + services, + &make_collect_context(settings, services, &placeholder), + ) + .await; + if let (Some(observation), Some(auction_request)) = + (telemetry.observation, telemetry.auction_request.as_ref()) + { + emit_auction_events_best_effort_lazy(services, || { + build_auction_events( + observation, + AuctionTerminalOutcome::Completed { + request: auction_request, + result: &result, + }, + ) + }) + .await; + } + write_bids_to_state( + &result.winning_bids, + params.price_granularity, + ¶ms.ad_bids_state, + settings.debug.inject_adm_for_testing, + ); +} + async fn collect_stream_auction( dispatched: DispatchedAuction, telemetry: AuctionTelemetryCarry, @@ -1600,11 +2439,12 @@ pub async fn handle_publisher_request( // SSP requests are already racing through the platform HTTP client, so // origin TTFB tracks origin latency rather than the auction timeout. - let mut response = match services - .http_client() - .send(PlatformHttpRequest::new(req, backend_name)) - .await - { + let mut platform_request = PlatformHttpRequest::new(req, backend_name); + if services.http_client().supports_streaming_responses() { + platform_request = platform_request.with_stream_response(); + } + + let mut response = match services.http_client().send(platform_request).await { Ok(platform_response) => platform_response.response, Err(err) => { if let Some(dispatched) = dispatched_auction.take() { @@ -2505,6 +3345,24 @@ mod tests { output } + fn deflate_encode(input: &[u8]) -> Vec { + let mut encoder = + flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(input) + .expect("should write deflate test input"); + encoder.finish().expect("should finish deflate encoding") + } + + fn deflate_decode(input: &[u8]) -> Vec { + let mut decoder = flate2::read::ZlibDecoder::new(input); + let mut output = Vec::new(); + decoder + .read_to_end(&mut output) + .expect("should decode deflate test output"); + output + } + fn brotli_encode(input: &[u8]) -> Vec { let mut encoder = CompressorWriter::new(Vec::new(), 4096, 5, 22); encoder @@ -2734,50 +3592,107 @@ mod tests { } #[tokio::test] - async fn handle_publisher_request_does_not_self_generate_ec() { - // EC generation is the adapter's real-browser-gated responsibility. This - // handler must never mint an EC ID on its own: for a navigation from a - // client the adapter did not pre-generate for (e.g. a non-real browser), - // `ec_value` must stay `None` so no IP-derived identifier reaches the - // auction. Consent allows EC creation and a client IP is present here — - // exactly the conditions under which the old inline call would have - // generated one. + async fn publisher_origin_fetch_leaves_stream_response_disabled_when_unsupported() { let settings = create_test_settings(); let stub = Arc::new(StubHttpClient::new()); - stub.push_response(200, b"ok".to_vec()); + stub.push_response_with_headers( + 200, + b"origin".to_vec(), + vec![("content-type", "text/html; charset=utf-8")], + ); let services = build_services_with_http_client( Arc::clone(&stub) as Arc ); - - let consent = crate::consent::ConsentContext { - jurisdiction: crate::consent::jurisdiction::Jurisdiction::NonRegulated, - ..Default::default() - }; - let mut ec_context = - EcContext::new_for_test_with_ip(None, consent, Some("203.0.113.7".to_string())); - assert!( - ec_context.ec_allowed(), - "test precondition: consent must allow EC creation" - ); - - let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); let req = HttpRequest::builder() .method(Method::GET) - .uri("https://publisher.example/article") + .uri("https://publisher.example/page") .header(header::HOST, "publisher.example") - .header("sec-fetch-dest", "document") .body(EdgeBody::empty()) .expect("should build request"); - let _ = handle_publisher_request( - &settings, - &services, - None, - &mut ec_context, - AuctionDispatch { - orchestrator: &orchestrator, - slots: &[], - registry: None, + let _ = run_publisher_proxy(&settings, &services, req).await; + + assert_eq!( + stub.recorded_stream_response_flags(), + vec![false], + "publisher origin fetch must not request streams when the platform does not support them" + ); + } + + #[tokio::test] + async fn publisher_origin_fetch_sets_stream_response_when_supported() { + let settings = create_test_settings(); + let stub = Arc::new(StubHttpClient::new()); + stub.set_streaming_responses_supported(true); + stub.push_response_with_headers( + 200, + b"origin".to_vec(), + vec![("content-type", "text/html; charset=utf-8")], + ); + let services = build_services_with_http_client( + Arc::clone(&stub) as Arc + ); + let req = HttpRequest::builder() + .method(Method::GET) + .uri("https://publisher.example/page") + .header(header::HOST, "publisher.example") + .body(EdgeBody::empty()) + .expect("should build request"); + + let _ = run_publisher_proxy(&settings, &services, req).await; + + assert_eq!( + stub.recorded_stream_response_flags(), + vec![true], + "publisher origin fetch should request streams when the platform supports them" + ); + } + + #[tokio::test] + async fn handle_publisher_request_does_not_self_generate_ec() { + // EC generation is the adapter's real-browser-gated responsibility. This + // handler must never mint an EC ID on its own: for a navigation from a + // client the adapter did not pre-generate for (e.g. a non-real browser), + // `ec_value` must stay `None` so no IP-derived identifier reaches the + // auction. Consent allows EC creation and a client IP is present here — + // exactly the conditions under which the old inline call would have + // generated one. + let settings = create_test_settings(); + let stub = Arc::new(StubHttpClient::new()); + stub.push_response(200, b"ok".to_vec()); + let services = build_services_with_http_client( + Arc::clone(&stub) as Arc + ); + + let consent = crate::consent::ConsentContext { + jurisdiction: crate::consent::jurisdiction::Jurisdiction::NonRegulated, + ..Default::default() + }; + let mut ec_context = + EcContext::new_for_test_with_ip(None, consent, Some("203.0.113.7".to_string())); + assert!( + ec_context.ec_allowed(), + "test precondition: consent must allow EC creation" + ); + + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let req = HttpRequest::builder() + .method(Method::GET) + .uri("https://publisher.example/article") + .header(header::HOST, "publisher.example") + .header("sec-fetch-dest", "document") + .body(EdgeBody::empty()) + .expect("should build request"); + + let _ = handle_publisher_request( + &settings, + &services, + None, + &mut ec_context, + AuctionDispatch { + orchestrator: &orchestrator, + slots: &[], + registry: None, }, req, ) @@ -3659,6 +4574,734 @@ mod tests { ); } + #[test] + fn stream_publisher_body_rejects_stream_body_in_sync_path() { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let params = OwnedProcessResponseParams { + content_encoding: String::new(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/html; charset=utf-8".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: None, + dispatched_auction: None, + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let body = EdgeBody::from_stream(futures::stream::iter(vec![Ok::<_, io::Error>( + bytes::Bytes::from_static(b"live"), + )])); + let mut output = Vec::new(); + + let err = stream_publisher_body(body, &mut output, ¶ms, &settings, ®istry) + .expect_err("should reject stream body in sync path"); + + assert!( + format!("{err:?}").contains("streaming body"), + "should explain that Body::Stream is not supported by the sync path: {err:?}" + ); + } + + #[test] + fn body_chunk_source_yields_once_body_in_chunks() { + futures::executor::block_on(async { + let body = EdgeBody::from_bytes(bytes::Bytes::from_static(b"abcdef")); + let mut source = BodyChunkSource::new(body, 3).with_max_bytes(16); + + assert_eq!( + source.next_chunk().await.expect("should read").as_deref(), + Some(&b"abc"[..]), + "should yield the first chunk" + ); + assert_eq!( + source.next_chunk().await.expect("should read").as_deref(), + Some(&b"def"[..]), + "should yield the second chunk" + ); + assert!( + source.next_chunk().await.expect("should read").is_none(), + "should end after buffered bytes are exhausted" + ); + }); + } + + #[test] + fn body_chunk_source_preserves_stream_chunks() { + futures::executor::block_on(async { + let body = EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::from_static(b"first"), + bytes::Bytes::from_static(b"second"), + ])); + let mut source = BodyChunkSource::new(body, 3).with_max_bytes(16); + + assert_eq!( + source.next_chunk().await.expect("should read").as_deref(), + Some(&b"first"[..]), + "stream chunks should pass through without re-chunking" + ); + assert_eq!( + source.next_chunk().await.expect("should read").as_deref(), + Some(&b"second"[..]), + "stream chunks should preserve upstream boundaries" + ); + assert!( + source.next_chunk().await.expect("should read").is_none(), + "should end after stream is exhausted" + ); + }); + } + + #[test] + fn body_chunk_source_enforces_cumulative_raw_cap() { + futures::executor::block_on(async { + let body = EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::from_static(b"1234"), + bytes::Bytes::from_static(b"5678"), + ])); + let mut source = BodyChunkSource::new(body, STREAM_CHUNK_SIZE).with_max_bytes(6); + + assert!( + source + .next_chunk() + .await + .expect("first chunk should pass") + .is_some(), + "first chunk should stay under cap" + ); + let err = source + .next_chunk() + .await + .expect_err("second chunk should exceed cap"); + + assert!( + format!("{err:?}").contains("publisher origin body exceeded"), + "should report cumulative cap: {err:?}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_processes_stream_without_auction() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = OwnedProcessResponseParams { + content_encoding: String::new(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/css".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: None, + dispatched_auction: None, + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let body = EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::from_static(b"body{background:url('https://origin.example.com/"), + bytes::Bytes::from_static(b"asset.png')}"), + ])); + let mut output = Vec::new(); + + stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect("stream body should process on async path"); + + let css = String::from_utf8(output).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "should rewrite origin host while streaming. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "should not leave origin host after rewrite. Got: {css}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_processes_gzip_stream_without_auction() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = OwnedProcessResponseParams { + content_encoding: "gzip".to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/css".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: None, + dispatched_auction: None, + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let compressed = + gzip_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + let split_at = compressed.len() / 2; + let body = EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::copy_from_slice(&compressed[..split_at]), + bytes::Bytes::copy_from_slice(&compressed[split_at..]), + ])); + let mut output = Vec::new(); + + stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect("gzip stream body should process on async path"); + + let css = String::from_utf8(gzip_decode(&output)).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "should rewrite origin host while streaming gzip. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "should not leave origin host after gzip rewrite. Got: {css}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_processes_deflate_stream_without_auction() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = OwnedProcessResponseParams { + content_encoding: "deflate".to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/css".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: None, + dispatched_auction: None, + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let compressed = + deflate_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + let split_at = compressed.len() / 2; + let body = EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::copy_from_slice(&compressed[..split_at]), + bytes::Bytes::copy_from_slice(&compressed[split_at..]), + ])); + let mut output = Vec::new(); + + stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect("deflate stream body should process on async path"); + + let css = String::from_utf8(deflate_decode(&output)).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "should rewrite origin host while streaming deflate. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "should not leave origin host after deflate rewrite. Got: {css}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_processes_brotli_stream_without_auction() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = OwnedProcessResponseParams { + content_encoding: "br".to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/css".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: None, + dispatched_auction: None, + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let compressed = + brotli_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + let split_at = compressed.len() / 2; + let body = EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::copy_from_slice(&compressed[..split_at]), + bytes::Bytes::copy_from_slice(&compressed[split_at..]), + ])); + let mut output = Vec::new(); + + stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect("brotli stream body should process on async path"); + + let css = String::from_utf8(brotli_decode(&output)).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "should rewrite origin host while streaming brotli. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "should not leave origin host after brotli rewrite. Got: {css}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_rejects_truncated_brotli_stream() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = OwnedProcessResponseParams { + content_encoding: "br".to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/css".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: None, + dispatched_auction: None, + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let compressed = + brotli_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + let truncated = &compressed[..compressed.len() - 3]; + let body = + EdgeBody::stream(futures::stream::iter(vec![bytes::Bytes::copy_from_slice( + truncated, + )])); + let mut output = Vec::new(); + + let err = stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect_err("truncated brotli stream must fail instead of truncating silently"); + + assert!( + format!("{err:?}").contains("brotli"), + "should surface the brotli finalization failure: {err:?}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_processes_stream_with_auction_hold() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let state = Arc::new(Mutex::new(None)); + let mut params = OwnedProcessResponseParams { + content_encoding: String::new(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/html; charset=utf-8".to_string(), + ad_slots_script: Some( + r#""# + .to_string(), + ), + ad_bids_state: state, + auction_observation: None, + auction_request: Some(test_auction_request()), + dispatched_auction: Some(DispatchedAuction::empty_for_test( + test_auction_request(), + 10, + )), + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let body = EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::from_static(b"hello"), + bytes::Bytes::from_static(b""), + ])); + let mut output = Vec::new(); + + stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect("stream body with auction should process on async path"); + + let html = String::from_utf8(output).expect("should be valid UTF-8"); + assert!( + html.contains("hello"), + "should preserve streamed HTML content. Got: {html}" + ); + assert!( + html.contains(".adSlots=JSON.parse"), + "should still inject ad slots. Got: {html}" + ); + assert!( + html.contains(".bids=JSON.parse"), + "should collect auction and inject bids before body close. Got: {html}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_processes_non_html_stream_after_auction_collect() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = OwnedProcessResponseParams { + content_encoding: String::new(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/css".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: Some(test_auction_request()), + dispatched_auction: Some(DispatchedAuction::empty_for_test( + test_auction_request(), + 10, + )), + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let body = EdgeBody::stream(futures::stream::iter(vec![bytes::Bytes::from_static( + b"body{background:url('https://origin.example.com/asset.png')}", + )])); + let mut output = Vec::new(); + + stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect("non-html stream body should process after auction collection"); + + let css = String::from_utf8(output).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "should rewrite non-html stream after auction collection. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "should not leave origin host after rewrite. Got: {css}" + ); + }); + } + + fn drain_streaming_finalize_body(content_encoding: &str, body: EdgeBody) -> Vec { + let settings = Arc::new(create_test_settings()); + let registry = Arc::new( + IntegrationRegistry::new(&settings).expect("should create integration registry"), + ); + let orchestrator = Arc::new(AuctionOrchestrator::new(settings.auction.clone())); + let services = noop_services(); + let response = Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "text/css") + .body(EdgeBody::empty()) + .expect("should build response"); + let params = OwnedProcessResponseParams { + content_encoding: content_encoding.to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/css".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: None, + dispatched_auction: None, + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let publisher_response = PublisherResponse::Stream { + response, + body, + params: Box::new(params), + }; + + let response = publisher_response_into_streaming_response( + publisher_response, + &Method::GET, + Arc::clone(&settings), + registry.as_ref(), + orchestrator, + services, + ) + .expect("should build streaming response"); + + assert!( + matches!(response.body(), EdgeBody::Stream(_)), + "streaming finalize should keep a lazy Body::Stream" + ); + + futures::executor::block_on( + response + .into_body() + .into_bytes_bounded(settings.publisher.max_buffered_body_bytes), + ) + .expect("streaming body should drain") + .to_vec() + } + + #[test] + fn publisher_response_streaming_finalize_keeps_stream_body_lazy() { + let body_bytes = drain_streaming_finalize_body( + "", + EdgeBody::stream(futures::stream::iter(vec![bytes::Bytes::from_static( + b"body{background:url('https://origin.example.com/asset.png')}", + )])), + ); + let css = String::from_utf8(body_bytes).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "streaming response body should still run publisher rewriting. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "streaming response body should not leave origin URLs unrewritten. Got: {css}" + ); + } + + #[test] + fn publisher_response_streaming_finalize_processes_gzip_stream() { + let compressed = + gzip_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + let split_at = compressed.len() / 2; + let output = drain_streaming_finalize_body( + "gzip", + EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::copy_from_slice(&compressed[..split_at]), + bytes::Bytes::copy_from_slice(&compressed[split_at..]), + ])), + ); + + let css = String::from_utf8(gzip_decode(&output)).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "streaming response finalize should rewrite gzip body. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "streaming response finalize should not leave gzip origin URLs. Got: {css}" + ); + } + + #[test] + fn publisher_response_streaming_finalize_processes_deflate_stream() { + let compressed = + deflate_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + let split_at = compressed.len() / 2; + let output = drain_streaming_finalize_body( + "deflate", + EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::copy_from_slice(&compressed[..split_at]), + bytes::Bytes::copy_from_slice(&compressed[split_at..]), + ])), + ); + + let css = String::from_utf8(deflate_decode(&output)).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "streaming response finalize should rewrite deflate body. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "streaming response finalize should not leave deflate origin URLs. Got: {css}" + ); + } + + #[test] + fn publisher_response_streaming_finalize_processes_brotli_stream() { + let compressed = + brotli_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + let split_at = compressed.len() / 2; + let output = drain_streaming_finalize_body( + "br", + EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::copy_from_slice(&compressed[..split_at]), + bytes::Bytes::copy_from_slice(&compressed[split_at..]), + ])), + ); + + let css = String::from_utf8(brotli_decode(&output)).expect("should be valid UTF-8"); + assert!( + css.contains("proxy.example.com"), + "streaming response finalize should rewrite brotli body. Got: {css}" + ); + assert!( + !css.contains("origin.example.com"), + "streaming response finalize should not leave brotli origin URLs. Got: {css}" + ); + } + + #[test] + fn publisher_response_streaming_finalize_holds_auction_and_keeps_gzip_tail() { + let settings = Arc::new(create_test_settings()); + let registry = Arc::new( + IntegrationRegistry::new(&settings).expect("should create integration registry"), + ); + let orchestrator = Arc::new(AuctionOrchestrator::new(settings.auction.clone())); + let services = noop_services(); + let response = Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "text/html; charset=utf-8") + .body(EdgeBody::empty()) + .expect("should build response"); + // The trailing content after `` must exceed the flate2 write + // decoder's 32 KiB internal output buffer: the close-body tag then + // surfaces (and releases the auction hold) mid-stream, while the + // trailing markup only surfaces at decoder finalization. This guards + // against the EOF decoded tail being dropped once the hold is gone. + let trailing_comment = format!("", "trailing-content ".repeat(3 * 1024)); + let page = format!("hello{trailing_comment}"); + let compressed = gzip_encode(page.as_bytes()); + let chunks: Vec = compressed + .chunks(STREAM_CHUNK_SIZE) + .map(bytes::Bytes::copy_from_slice) + .collect(); + let params = OwnedProcessResponseParams { + content_encoding: "gzip".to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/html; charset=utf-8".to_string(), + ad_slots_script: Some( + r#""# + .to_string(), + ), + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: Some(test_auction_request()), + dispatched_auction: Some(DispatchedAuction::empty_for_test( + test_auction_request(), + 10, + )), + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + let publisher_response = PublisherResponse::Stream { + response, + body: EdgeBody::stream(futures::stream::iter(chunks)), + params: Box::new(params), + }; + + let response = publisher_response_into_streaming_response( + publisher_response, + &Method::GET, + Arc::clone(&settings), + registry.as_ref(), + orchestrator, + services, + ) + .expect("should build streaming response"); + + let output = futures::executor::block_on( + response + .into_body() + .into_bytes_bounded(settings.publisher.max_buffered_body_bytes), + ) + .expect("streaming body should drain") + .to_vec(); + + let html = String::from_utf8(gzip_decode(&output)).expect("should be valid UTF-8"); + assert!( + html.contains(".bids=JSON.parse"), + "should collect the held auction and inject bids. Got tail: {}", + &html[html.len().saturating_sub(200)..] + ); + assert!( + html.contains("trailing-content"), + "should preserve content after the close-body tag" + ); + assert!( + html.trim_end().ends_with(""), + "should not drop the decoded tail once the auction hold is released. Got tail: {}", + &html[html.len().saturating_sub(200)..] + ); + } + #[test] fn stream_publisher_body_treats_mixed_case_html_as_html() { let settings = create_test_settings(); diff --git a/crates/trusted-server-core/src/settings.rs b/crates/trusted-server-core/src/settings.rs index 9cbb2a546..008b4a2d5 100644 --- a/crates/trusted-server-core/src/settings.rs +++ b/crates/trusted-server-core/src/settings.rs @@ -50,15 +50,12 @@ pub struct Publisher { /// exceeding it fails the response rather than allocating past the cap. /// Defaults to 16 MiB — a conservative cap that prevents Wasm-heap OOM. /// - /// On Fastly the *effective* ceiling for a publisher page is lower: the - /// platform HTTP client rejects any origin response whose raw (still - /// compressed) body exceeds 10 MiB before this buffer is ever filled, so - /// raising this value only helps highly compressible pages whose decoded - /// size exceeds the 16 MiB default while their compressed origin body stays - /// under 10 MiB. Raising it above ~10 MiB does not lift the platform cap for - /// uncompressed pages. That platform limit is removed once true streaming - /// lands (tracked for PR 15, issue #495), after which this setting becomes - /// the sole ceiling. + /// Fastly origin bodies are preserved as streams on the publisher path, so + /// this setting is also the cumulative raw-byte cap while the streaming + /// processor decodes and rewrites chunks. Buffered adapters keep using it + /// as the post-rewrite output buffer cap. On the streaming path headers + /// are already committed when the cap trips, so the response is truncated + /// mid-body (with the error logged) rather than replaced with a 5xx. /// /// Must be at least 1: a zero-byte cap fails every non-empty buffered /// publisher response at request time, so it is rejected at config diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index 5692118a8..ef1a0bc55 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -349,6 +349,197 @@ impl StreamProcessor for StreamingReplacer { } } +/// Read buffer size for streaming body processing and brotli internal buffers. +/// Both the `Decompressor` and `CompressorWriter` use this value so all +/// brotli I/O layers operate on consistently-sized chunks. +pub(crate) const STREAM_CHUNK_SIZE: usize = 8192; + +/// Incremental push-style decompressor for the async chunk pipeline. +/// +/// Compressed bytes go in via [`Self::decode_chunk`]; decoded bytes drain +/// out of the internal buffer after every push. Write-based decoders are +/// used because the async publisher path cannot wrap a blocking `Read`. +pub(crate) enum BodyStreamDecoder { + None, + Gzip(flate2::write::GzDecoder>), + Deflate(flate2::write::ZlibDecoder>), + Brotli(Box>>), +} + +impl BodyStreamDecoder { + pub(crate) fn new(compression: Compression) -> Self { + match compression { + Compression::None => Self::None, + Compression::Gzip => Self::Gzip(flate2::write::GzDecoder::new(Vec::new())), + Compression::Deflate => Self::Deflate(flate2::write::ZlibDecoder::new(Vec::new())), + Compression::Brotli => Self::Brotli(Box::new(brotli::DecompressorWriter::new( + Vec::new(), + STREAM_CHUNK_SIZE, + ))), + } + } + + pub(crate) fn decode_chunk( + &mut self, + chunk: &[u8], + ) -> Result, Report> { + match self { + Self::None => Ok(chunk.to_vec()), + Self::Gzip(decoder) => { + decoder + .write_all(chunk) + .change_context(TrustedServerError::Proxy { + message: "Failed to decode gzip publisher body chunk".to_string(), + })?; + Ok(std::mem::take(decoder.get_mut())) + } + Self::Deflate(decoder) => { + decoder + .write_all(chunk) + .change_context(TrustedServerError::Proxy { + message: "Failed to decode deflate publisher body chunk".to_string(), + })?; + Ok(std::mem::take(decoder.get_mut())) + } + Self::Brotli(decoder) => { + decoder + .write_all(chunk) + .change_context(TrustedServerError::Proxy { + message: "Failed to decode brotli publisher body chunk".to_string(), + })?; + Ok(std::mem::take(decoder.get_mut())) + } + } + } + + pub(crate) fn finish(&mut self) -> Result, Report> { + match self { + Self::None => Ok(Vec::new()), + Self::Gzip(decoder) => { + decoder + .try_finish() + .change_context(TrustedServerError::Proxy { + message: "Failed to finalize gzip publisher body decoder".to_string(), + })?; + Ok(std::mem::take(decoder.get_mut())) + } + Self::Deflate(decoder) => { + decoder + .try_finish() + .change_context(TrustedServerError::Proxy { + message: "Failed to finalize deflate publisher body decoder".to_string(), + })?; + Ok(std::mem::take(decoder.get_mut())) + } + Self::Brotli(decoder) => { + // `close()` (not `flush()`): flush accepts a truncated brotli + // stream silently, while close validates end-of-stream and + // errors on incomplete input, matching the gzip/deflate arms. + decoder.close().change_context(TrustedServerError::Proxy { + message: "Failed to finalize brotli publisher body decoder".to_string(), + })?; + Ok(std::mem::take(decoder.get_mut())) + } + } + } +} + +/// Incremental push-style compressor mirroring [`BodyStreamDecoder`]. +/// +/// Processed bytes go in via [`Self::encode_chunk`]; encoded bytes drain out +/// after every push, and [`Self::finish`] emits the stream trailer. +pub(crate) enum BodyStreamEncoder { + None, + Gzip(flate2::write::GzEncoder>), + Deflate(flate2::write::ZlibEncoder>), + Brotli(Box>>), +} + +fn new_brotli_vec_encoder() -> brotli::enc::writer::CompressorWriter> { + let params = brotli::enc::BrotliEncoderParams { + quality: 4, + lgwin: 22, + ..Default::default() + }; + brotli::enc::writer::CompressorWriter::with_params(Vec::new(), STREAM_CHUNK_SIZE, ¶ms) +} + +impl BodyStreamEncoder { + pub(crate) fn new(compression: Compression) -> Self { + match compression { + Compression::None => Self::None, + Compression::Gzip => Self::Gzip(flate2::write::GzEncoder::new( + Vec::new(), + flate2::Compression::default(), + )), + Compression::Deflate => Self::Deflate(flate2::write::ZlibEncoder::new( + Vec::new(), + flate2::Compression::default(), + )), + Compression::Brotli => Self::Brotli(Box::new(new_brotli_vec_encoder())), + } + } + + pub(crate) fn encode_chunk( + &mut self, + chunk: &[u8], + ) -> Result, Report> { + match self { + Self::None => Ok(chunk.to_vec()), + Self::Gzip(encoder) => { + encoder + .write_all(chunk) + .change_context(TrustedServerError::Proxy { + message: "Failed to encode gzip publisher body chunk".to_string(), + })?; + Ok(std::mem::take(encoder.get_mut())) + } + Self::Deflate(encoder) => { + encoder + .write_all(chunk) + .change_context(TrustedServerError::Proxy { + message: "Failed to encode deflate publisher body chunk".to_string(), + })?; + Ok(std::mem::take(encoder.get_mut())) + } + Self::Brotli(encoder) => { + encoder + .write_all(chunk) + .change_context(TrustedServerError::Proxy { + message: "Failed to encode brotli publisher body chunk".to_string(), + })?; + Ok(std::mem::take(encoder.get_mut())) + } + } + } + + pub(crate) fn finish(&mut self) -> Result, Report> { + match self { + Self::None => Ok(Vec::new()), + Self::Gzip(encoder) => { + encoder + .try_finish() + .change_context(TrustedServerError::Proxy { + message: "Failed to finalize gzip publisher body encoder".to_string(), + })?; + Ok(std::mem::take(encoder.get_mut())) + } + Self::Deflate(encoder) => { + encoder + .try_finish() + .change_context(TrustedServerError::Proxy { + message: "Failed to finalize deflate publisher body encoder".to_string(), + })?; + Ok(std::mem::take(encoder.get_mut())) + } + Self::Brotli(encoder) => { + let encoder = std::mem::replace(encoder, Box::new(new_brotli_vec_encoder())); + Ok((*encoder).into_inner()) + } + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/docs/superpowers/plans/2026-07-08-true-origin-streaming-fastly.md b/docs/superpowers/plans/2026-07-08-true-origin-streaming-fastly.md new file mode 100644 index 000000000..bc1983a97 --- /dev/null +++ b/docs/superpowers/plans/2026-07-08-true-origin-streaming-fastly.md @@ -0,0 +1,1039 @@ +# True Origin Streaming Fastly Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Fix issue #849 for the production Fastly path so publisher HTML origin bodies stream through the rewrite pipeline to the client, with auction collection outside TTFB except for the held `` tail. + +**Architecture:** Keep one cohesive Fastly PR because the core pipeline, Fastly origin fetch, and Fastly finalize path are incomplete in isolation. Convert publisher body processing from sync `Read` over buffered bodies to async chunk-pull over `edgezero_core::body::Body`, then enable Fastly `with_stream_response()` and return a lazy streaming body for publisher responses. Leave Cloudflare, Spin, and Axum streaming as follow-up work. + +**Tech Stack:** Rust 2024, `edgezero_core::body::Body`, `futures::StreamExt`, `error-stack`, `flate2`, `brotli`, Fastly Compute, Viceroy tests. + +--- + +## Scope + +In scope: + +- Publisher HTML and processable publisher responses on the Fastly adapter. +- Core publisher pipeline support for `Body::Stream`. +- Fastly platform capability signaling for streaming origin responses. +- Fastly EdgeZero response delivery that streams publisher bodies to clients. +- Tests proving stream-vs-buffer parity, bodiless handling, stream caps, and Fastly routing behavior. + +Out of scope: + +- Cloudflare origin streaming. Current adapter rejects `PlatformHttpRequest::stream_response`. +- Spin streaming. Current adapter and upstream EdgeZero Spin conversion are buffered/blocking issues. +- Axum client streaming. Axum is dev-only and has `LocalBoxStream`/`Send` constraints. +- Parser-context `` scan fix from issue #850. +- Origin template caching and transformed HTML caching from issue #852. + +## Current Failure Points + +- `crates/trusted-server-adapter-fastly/src/platform.rs`: `fastly_response_to_platform(..., stream_response: false)` uses `take_body_bytes()` and the 10 MiB platform cap for publisher origin responses. +- `crates/trusted-server-core/src/publisher.rs`: `body_as_reader()` calls `body.into_bytes().unwrap_or_default()`, so `Body::Stream` becomes an empty body. +- `crates/trusted-server-core/src/publisher.rs`: `stream_html_with_auction_hold()` and `body_close_hold_loop()` are sync-`Read` based. +- `crates/trusted-server-adapter-fastly/src/app.rs`: publisher route calls `buffer_publisher_response_async()`, buffering all processed output and awaiting auction before any client bytes are sent. +- `crates/trusted-server-adapter-fastly/src/main.rs`: `send_edgezero_response()` already streams `EdgeBody::Stream`, but currently only asset responses reach that arm. + +## File Structure + +- Modify `crates/trusted-server-core/src/platform/http.rs` + - Add `PlatformHttpClient::supports_streaming_responses()` with default `false`. +- Modify adapter platform implementations: + - `crates/trusted-server-adapter-fastly/src/platform.rs`: return `true` for `supports_streaming_responses()`. + - `crates/trusted-server-adapter-cloudflare/src/platform.rs`: inherit default `false`. + - `crates/trusted-server-adapter-spin/src/platform.rs`: inherit default `false`. + - `crates/trusted-server-adapter-axum/src/platform.rs`: inherit default `false`. + - `crates/trusted-server-core/src/platform/test_support.rs`: configurable test support if needed. +- Modify `crates/trusted-server-core/src/streaming_processor.rs` + - Add small push decoder/encoder helpers only if keeping them here reduces duplication. + - Keep the existing `StreamingPipeline::process(Read, Write)` API for existing call sites. +- Modify `crates/trusted-server-core/src/publisher.rs` + - Replace publisher async processing internals with async chunk-pull. + - Keep public `buffer_publisher_response_async()` for buffered adapters. + - Add a streaming response constructor/helper for Fastly to use. + - Make `body_as_reader()` reject `Body::Stream` loudly or remove its use from any stream-capable path. +- Modify `crates/trusted-server-adapter-fastly/src/app.rs` + - Replace publisher `buffer_publisher_response_async()` call with streaming finalize for streamable publisher responses. + - Preserve buffered behavior for `PublisherResponse::Buffered`, pass-through/bodiless responses, and error paths. +- Modify `crates/trusted-server-adapter-fastly/src/main.rs` + - Reuse existing `EdgeBody::Stream` delivery. + - If publisher streaming needs a different log message from asset streaming, split the helper name/log text without changing behavior. +- Tests: + - `crates/trusted-server-core/src/publisher.rs` unit tests. + - `crates/trusted-server-core/src/streaming_processor.rs` unit tests if push codec helpers are introduced there. + - `crates/trusted-server-core/src/platform/test_support.rs` tests for capability behavior. + - `crates/trusted-server-adapter-fastly/src/app.rs` route tests for publisher streaming response shape. + +## Design Decisions + +- Use one PR for the full Fastly production fix. Intermediate merged PRs would create incomplete behavior and review confusion. +- Use existing `publisher.max_buffered_body_bytes` as the publisher body ceiling after streaming. `settings.rs` already documents that this becomes the sole ceiling after true streaming removes the 10 MiB Fastly materialization cap. +- Keep `Content-Length` removed for rewritten stream responses. Streaming output can change size due to URL rewriting and bid injection. +- Preserve bodiless behavior for `HEAD`, `204`, and `304`: do not attach or drive a body stream, and log abandoned/wasted auctions as current code does. +- Do not build a sync `Read` bridge over `Body::Stream`; nested `block_on` can panic on Fastly because the router already runs under `futures::executor::block_on`. +- Avoid adding `async-stream` initially. Use `futures::stream::unfold` or a custom stream type so the PR does not add a dependency unless the implementation becomes materially clearer. + +## Task 1: Baseline Tests for Stream Input Safety + +**Files:** + +- Modify: `crates/trusted-server-core/src/publisher.rs` + +- [ ] **Step 1: Add a failing test for `Body::Stream` not becoming empty** + +Add a test near the existing `stream_publisher_body` tests: + +```rust +#[test] +fn stream_publisher_body_rejects_stream_body_in_sync_path() { + let settings = create_test_settings(); + let registry = IntegrationRegistry::new(&settings).expect("should build registry"); + let body = EdgeBody::from_stream(futures::stream::iter(vec![Ok(Bytes::from_static( + b"live", + ))])); + let params = test_process_params("text/html", ""); + let mut output = Vec::new(); + + let err = stream_publisher_body(body, &mut output, ¶ms, &settings, ®istry) + .expect_err("should reject stream body in sync path"); + + assert!( + format!("{err:?}").contains("streaming body"), + "should explain that Body::Stream is not supported by the sync path: {err:?}" + ); +} +``` + +- [ ] **Step 2: Run the targeted test and verify it fails** + +Run: + +```bash +cargo test-axum stream_publisher_body_rejects_stream_body_in_sync_path +``` + +Expected: FAIL because current `body_as_reader()` silently returns empty bytes. + +- [ ] **Step 3: Replace `body_as_reader()` with a fallible helper** + +Change `body_as_reader(body: EdgeBody) -> Cursor` to return `Result, Report>` and return a proxy error for `Body::Stream`. + +Minimal shape: + +```rust +fn body_as_reader(body: EdgeBody) -> Result, Report> { + let bytes = body.into_bytes().ok_or_else(|| { + Report::new(TrustedServerError::Proxy { + message: "streaming body cannot be processed by sync publisher pipeline".to_owned(), + }) + })?; + Ok(std::io::Cursor::new(bytes)) +} +``` + +Update existing sync call sites to use `body_as_reader(body)?`. + +- [ ] **Step 4: Run the targeted test and existing publisher sync tests** + +Run: + +```bash +cargo test-axum stream_publisher_body +``` + +Expected: PASS. + +- [ ] **Step 5: Commit** + +```bash +git add crates/trusted-server-core/src/publisher.rs +git commit -m "Reject publisher stream bodies on sync path" +``` + +## Task 2: Async Chunk Source and Cumulative Cap + +**Files:** + +- Modify: `crates/trusted-server-core/src/publisher.rs` + +- [ ] **Step 1: Add tests for async chunk pulling** + +Add focused tests for a private helper that will pull chunks from both body variants: + +```rust +#[test] +fn body_chunk_source_yields_once_body_in_chunks() { + futures::executor::block_on(async { + let body = EdgeBody::from(Bytes::from_static(b"abcdef")); + let mut source = BodyChunkSource::new(body, 3); + + assert_eq!(source.next_chunk().await.expect("should read").as_deref(), Some(&b"abc"[..])); + assert_eq!(source.next_chunk().await.expect("should read").as_deref(), Some(&b"def"[..])); + assert!(source.next_chunk().await.expect("should read").is_none()); + }); +} +``` + +Add a separate test for `Body::Stream` preserving chunk boundaries and surfacing stream errors. + +- [ ] **Step 2: Add a failing test for the cumulative cap** + +Use a stream with two chunks whose total exceeds a small cap: + +```rust +#[test] +fn body_chunk_source_enforces_cumulative_raw_cap() { + futures::executor::block_on(async { + let body = EdgeBody::from_stream(futures::stream::iter(vec![ + Ok(Bytes::from_static(b"1234")), + Ok(Bytes::from_static(b"5678")), + ])); + let mut source = BodyChunkSource::new(body, STREAM_CHUNK_SIZE).with_max_bytes(6); + + assert!(source.next_chunk().await.expect("first chunk should pass").is_some()); + let err = source.next_chunk().await.expect_err("second chunk should exceed cap"); + assert!( + format!("{err:?}").contains("publisher origin body exceeded"), + "should report cumulative cap: {err:?}" + ); + }); +} +``` + +- [ ] **Step 3: Run the new tests and verify they fail** + +Run: + +```bash +cargo test-axum body_chunk_source +``` + +Expected: FAIL because helper does not exist. + +- [ ] **Step 4: Implement `BodyChunkSource`** + +Implement a private helper near `STREAM_CHUNK_SIZE`: + +- Owns `EdgeBody`. +- For `Body::Once`, yields `Bytes` slices up to `chunk_size` without copying more than necessary. +- For `Body::Stream`, awaits `stream.next()`. +- Tracks cumulative raw bytes and errors when total exceeds `max_bytes`. +- Maps stream errors to `TrustedServerError::Proxy`. + +Do not use `block_on` inside the helper. + +- [ ] **Step 5: Run helper tests** + +Run: + +```bash +cargo test-axum body_chunk_source +``` + +Expected: PASS. + +- [ ] **Step 6: Commit** + +```bash +git add crates/trusted-server-core/src/publisher.rs +git commit -m "Add async publisher body chunk source" +``` + +## Task 3: Push Compression Helpers + +**Files:** + +- Modify: `crates/trusted-server-core/src/streaming_processor.rs` +- Or modify: `crates/trusted-server-core/src/publisher.rs` if helpers are publisher-only + +- [ ] **Step 1: Add parity tests for compressed chunk processing** + +For each compression mode used by publisher HTML (`gzip`, `deflate`, `br`), add a test that feeds compressed HTML in multiple raw chunks through the future async path and verifies the decompressed/processed/recompressed output decodes to expected HTML. + +Start with gzip: + +```rust +#[test] +fn async_publisher_pipeline_preserves_gzip_html_across_stream_chunks() { + futures::executor::block_on(async { + let compressed = gzip_bytes(b"Hello"); + let body = EdgeBody::from_stream(bytes_to_two_chunk_stream(compressed)); + let output = process_test_body_async(body, "text/html", "gzip") + .await + .expect("should process gzip stream"); + + assert_eq!( + gunzip_bytes(&output), + b"Hello" + ); + }); +} +``` + +Use existing HTML processor expectations rather than inventing new behavior. + +- [ ] **Step 2: Run the gzip test and verify it fails** + +Run: + +```bash +cargo test-axum async_publisher_pipeline_preserves_gzip_html_across_stream_chunks +``` + +Expected: FAIL because async compressed processing does not exist. + +- [ ] **Step 3: Implement write-based push decoders** + +Use write-based APIs: + +- `flate2::write::GzDecoder` +- `flate2::write::ZlibDecoder` +- `brotli::DecompressorWriter` + +The helper should: + +- Accept raw compressed chunks. +- Write decoded bytes into an internal `Vec` sink. +- Return newly decoded bytes after each input chunk. +- Finalize at EOF and return any decoder tail bytes. +- Surface decoder errors as `TrustedServerError::Proxy`. + +Keep this helper private unless tests or other modules need it. + +- [ ] **Step 4: Implement output encoding wrapper** + +Continue to use existing write-based encoders: + +- `flate2::write::GzEncoder` +- `flate2::write::ZlibEncoder` +- `brotli::enc::writer::CompressorWriter` + +The async loop should write processed decoded chunks into the encoder and finalize once. + +- [ ] **Step 5: Add deflate and brotli tests** + +Run: + +```bash +cargo test-axum async_publisher_pipeline_preserves_ +``` + +Expected: gzip, deflate, and brotli async parity tests PASS. + +- [ ] **Step 6: Commit** + +```bash +git add crates/trusted-server-core/src/streaming_processor.rs crates/trusted-server-core/src/publisher.rs +git commit -m "Add push compression support for publisher streams" +``` + +## Task 4: Async Publisher Pipeline Without Auction + +**Files:** + +- Modify: `crates/trusted-server-core/src/publisher.rs` + +- [ ] **Step 1: Add stream-vs-once parity tests for no-auction paths** + +Cover: + +- HTML rewrite. +- RSC flight rewrite. +- Generic URL replacement. +- Unsupported stream body cannot reach sync path. + +Example: + +```rust +#[test] +fn stream_publisher_body_async_matches_buffered_html_without_auction() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = IntegrationRegistry::new(&settings).expect("should build registry"); + let html = Bytes::from_static(b"x"); + + let mut once_params = test_process_params("text/html", ""); + let mut once_output = Vec::new(); + stream_publisher_body_async( + EdgeBody::from(html.clone()), + &mut once_output, + &mut once_params, + &settings, + ®istry, + &AuctionOrchestrator::new(settings.auction.clone()), + &noop_services(), + ) + .await + .expect("once body should process"); + + let mut stream_params = test_process_params("text/html", ""); + let mut stream_output = Vec::new(); + stream_publisher_body_async( + EdgeBody::from_stream(futures::stream::iter(vec![Ok(html)])), + &mut stream_output, + &mut stream_params, + &settings, + ®istry, + &AuctionOrchestrator::new(settings.auction.clone()), + &noop_services(), + ) + .await + .expect("stream body should process"); + + assert_eq!(stream_output, once_output); + }); +} +``` + +- [ ] **Step 2: Run parity tests and verify they fail** + +Run: + +```bash +cargo test-axum stream_publisher_body_async_matches_buffered +``` + +Expected: FAIL because no-auction async path still delegates to sync processing. + +- [ ] **Step 3: Refactor `process_response_streaming` into reusable processor construction** + +Extract the shared routing logic into a helper such as: + +```rust +enum PublisherProcessor { + Html(HtmlRewriterAdapter), + Rsc(RscFlightUrlRewriter), + Url(StreamingReplacer), +} +``` + +Or use a generic closure/helper if that fits existing patterns better. The goal is to avoid duplicating content-type routing between sync and async paths. + +- [ ] **Step 4: Drive all `stream_publisher_body_async()` calls through async chunk-pull** + +Even when `params.dispatched_auction` is `None`, build the same processor and use `BodyChunkSource`. This prevents stream bodies from falling into the sync path. + +- [ ] **Step 5: Keep `stream_publisher_body()` for compatibility** + +The sync function should remain for old tests and any current non-stream callers, but it must not be used by the async path once this task is complete. + +- [ ] **Step 6: Run targeted tests** + +Run: + +```bash +cargo test-axum stream_publisher_body_async_matches_buffered +``` + +Expected: PASS. + +- [ ] **Step 7: Commit** + +```bash +git add crates/trusted-server-core/src/publisher.rs +git commit -m "Drive publisher async processing from body chunks" +``` + +## Task 5: Async Auction Hold Loop + +**Files:** + +- Modify: `crates/trusted-server-core/src/publisher.rs` + +- [ ] **Step 1: Replace reader-based hold-loop test with stream-based test** + +Update or add a test based on `body_close_hold_loop_processes_close_tail_before_reading_post_body_chunks()`: + +- Feed pre-`` chunk. +- Feed held `` chunk. +- Feed post-body chunk. +- Assert the loop collects auction immediately when `` test** + +Verify that auction collection happens at EOF and finalization still calls `processor.process_chunk(&[], true)`. + +- [ ] **Step 3: Add stream error abandonment test** + +Feed a stream error after dispatch and assert telemetry abandonment uses `stream_read_error` or the current expected reason. + +- [ ] **Step 4: Run tests and verify failures** + +Run: + +```bash +cargo test-axum body_close_hold_loop +``` + +Expected: FAIL until the loop consumes `BodyChunkSource`. + +- [ ] **Step 5: Change `body_close_hold_loop` to async chunk-pull** + +Replace: + +```rust +async fn body_close_hold_loop(...) +``` + +with a shape that accepts decoded chunks from an async driver, or accepts `BodyChunkSource` plus codec state. Keep the control flow: + +- Push decoded chunks into `BodyCloseHoldBuffer`. +- Write ready bytes immediately. +- On first ` bool { + false +} +``` + +In `FastlyPlatformHttpClient`: + +```rust +fn supports_streaming_responses(&self) -> bool { + true +} +``` + +In `StubHttpClient`, add a configurable flag if tests need both states. + +- [ ] **Step 4: Enable publisher origin streaming behind the gate** + +At the publisher origin fetch: + +```rust +let mut platform_request = PlatformHttpRequest::new(req, backend_name); +if services.http_client().supports_streaming_responses() { + platform_request = platform_request.with_stream_response(); +} +let mut response = services.http_client().send(platform_request).await?; +``` + +- [ ] **Step 5: Run capability tests** + +Run: + +```bash +cargo test-axum publisher_origin_fetch_sets_stream_response_when_supported +cargo test-axum publisher_origin_fetch_leaves_stream_response_disabled_when_unsupported +``` + +Expected: PASS. + +- [ ] **Step 6: Commit** + +```bash +git add crates/trusted-server-core/src/platform/http.rs crates/trusted-server-adapter-fastly/src/platform.rs crates/trusted-server-core/src/platform/test_support.rs crates/trusted-server-core/src/publisher.rs +git commit -m "Gate publisher origin streaming by platform capability" +``` + +## Task 8: Fastly Publisher Streaming Finalize + +**Files:** + +- Modify: `crates/trusted-server-adapter-fastly/src/app.rs` +- Modify: `crates/trusted-server-core/src/publisher.rs` if a helper is needed +- Possibly modify: `crates/trusted-server-adapter-fastly/src/main.rs` for logging/helper naming + +- [ ] **Step 1: Add Fastly route test for publisher response body shape** + +In Fastly app tests, configure a publisher HTML origin response and assert the router returns `Body::Stream` for processable publisher responses on `GET`. + +Expected assertion: + +```rust +assert!( + matches!(response.body(), Body::Stream(_)), + "processable publisher response should remain streaming on Fastly" +); +``` + +- [ ] **Step 2: Add bodiless route tests** + +Assert `HEAD`, `204`, and `304` publisher responses do not carry a stream body, preserving existing metadata. + +- [ ] **Step 3: Run tests and verify failure** + +Run: + +```bash +cargo test-fastly publisher_response_streams +``` + +Expected: FAIL because app still calls `buffer_publisher_response_async()`. + +- [ ] **Step 4: Add a core helper to convert `PublisherResponse` to streaming response** + +Preferred shape in `publisher.rs`: + +```rust +pub fn publisher_response_into_streaming_body( + publisher_response: PublisherResponse, + method: &Method, + settings: Arc, + integration_registry: Arc, + orchestrator: Arc, + services: RuntimeServices, +) -> Result, Report> +``` + +The helper should: + +- Return `PublisherResponse::Buffered` unchanged. +- Return `PublisherResponse::PassThrough` with body attached, except bodiless responses. +- For `PublisherResponse::Stream`, build `EdgeBody::from_stream(futures::stream::unfold(...))` or equivalent. +- Move `OwnedProcessResponseParams`, origin body, settings, registry, orchestrator, and services into the stream state. +- Yield processed chunks as they become available. +- On mid-stream processing error, log and end the stream. The client sees a truncated body, matching existing mid-stream error behavior. + +If borrowing/lifetime pressure is high, keep the helper in Fastly `app.rs` and call core `stream_publisher_body_async()` from inside the stream. Prefer core if it avoids Fastly-specific body processing logic. + +- [ ] **Step 5: Replace Fastly buffered finalize for publisher route** + +In `handle_publisher_route`, replace the `buffer_publisher_response_async()` call for Fastly with the streaming helper. Keep non-Fastly adapters on `buffer_publisher_response_async()`. + +- [ ] **Step 6: Preserve entry-point finalization** + +Verify the returned `Response` still carries extensions needed by `main.rs`: + +- `EcFinalizeState` +- `RequestFilterEffects` +- Final cache privacy guard + +Headers must be finalized before `send_edgezero_response()` splits the response and commits headers. + +- [ ] **Step 7: Run Fastly route tests** + +Run: + +```bash +cargo test-fastly publisher_response +``` + +Expected: PASS. + +- [ ] **Step 8: Commit** + +```bash +git add crates/trusted-server-adapter-fastly/src/app.rs crates/trusted-server-core/src/publisher.rs crates/trusted-server-adapter-fastly/src/main.rs +git commit -m "Stream Fastly publisher responses to clients" +``` + +## Task 9: Pass-Through Publisher Bodies + +**Files:** + +- Modify: `crates/trusted-server-core/src/publisher.rs` +- Modify: `crates/trusted-server-adapter-fastly/src/app.rs` + +- [ ] **Step 1: Add pass-through large body test** + +Verify a non-processable successful publisher response (`image/png`, font, video) is returned as `Body::Stream` when origin streaming is supported and body is not bodiless. + +- [ ] **Step 2: Add pass-through bodiless test** + +Verify `HEAD`, `204`, and `304` pass-through arms preserve headers but do not attach/drain the stream body. + +- [ ] **Step 3: Run targeted tests** + +Run: + +```bash +cargo test-fastly publisher_pass_through +``` + +Expected: FAIL if pass-through still buffers or attaches a body for bodiless responses. + +- [ ] **Step 4: Make pass-through use the same body-carrying guard** + +Mirror `asset_response_carries_body()` semantics in publisher finalize. If `response_carries_body(method, status)` is false, drop the body and return headers only. + +- [ ] **Step 5: Run targeted tests** + +Run: + +```bash +cargo test-fastly publisher_pass_through +``` + +Expected: PASS. + +- [ ] **Step 6: Commit** + +```bash +git add crates/trusted-server-core/src/publisher.rs crates/trusted-server-adapter-fastly/src/app.rs +git commit -m "Preserve publisher pass-through streaming semantics" +``` + +## Task 10: Headers, Length, and Error Semantics + +**Files:** + +- Modify: `crates/trusted-server-core/src/publisher.rs` +- Modify: `crates/trusted-server-adapter-fastly/src/app.rs` +- Modify: `crates/trusted-server-adapter-fastly/src/main.rs` if logs are misleading + +- [ ] **Step 1: Add header tests** + +Assert for streamed processed publisher responses: + +- `Content-Length` is absent. +- `Transfer-Encoding` is not manually set. +- `Content-Encoding` is preserved when recompression is used. +- Cache/privacy headers still downgrade when `Set-Cookie` is present. + +- [ ] **Step 2: Add mid-stream cap/error test** + +Use a body stream that exceeds `publisher.max_buffered_body_bytes` after headers would be committed. Assert the stream returns an error/truncates consistently with existing mid-stream asset behavior and logs enough context. + +- [ ] **Step 3: Run tests and verify failures** + +Run: + +```bash +cargo test-fastly publisher_stream +``` + +Expected: FAIL until header/error cleanup is complete. + +- [ ] **Step 4: Clean header handling** + +Ensure the existing `response.headers_mut().remove(header::CONTENT_LENGTH)` remains on `PublisherResponse::Stream`. Do not re-add content length for streaming finalize. + +- [ ] **Step 5: Improve log wording** + +If `main.rs` still logs "asset streaming" for all `EdgeBody::Stream` responses, rename log messages to "EdgeZero streaming body" or split publisher/asset helpers. + +- [ ] **Step 6: Run tests** + +Run: + +```bash +cargo test-fastly publisher_stream +``` + +Expected: PASS. + +- [ ] **Step 7: Commit** + +```bash +git add crates/trusted-server-core/src/publisher.rs crates/trusted-server-adapter-fastly/src/app.rs crates/trusted-server-adapter-fastly/src/main.rs +git commit -m "Tighten publisher streaming headers and errors" +``` + +## Task 11: End-to-End Regression Coverage + +**Files:** + +- Modify: `crates/trusted-server-adapter-fastly/src/app.rs` +- Modify: existing integration/parity tests if appropriate + +- [ ] **Step 1: Add a slow-origin behavior test if feasible in existing harness** + +Preferred test shape: + +- Origin response body is a stream with first chunk available immediately and second chunk delayed or instrumented. +- Router returns `Body::Stream` without collecting the whole body. +- Pulling the first output chunk does not require pulling the entire origin stream. + +If the harness cannot model time cleanly, use an instrumented stream that panics if polled past the first chunk before the returned response body is consumed. + +- [ ] **Step 2: Add auction timing test** + +Verify response construction does not await `collect_dispatched_auction`; collection happens when the body stream is pulled and reaches `` or EOF. + +- [ ] **Step 3: Run targeted tests** + +Run: + +```bash +cargo test-fastly publisher_streaming_does_not_buffer_origin_before_response +``` + +Expected: PASS after Fastly finalize is lazy. + +- [ ] **Step 4: Commit** + +```bash +git add crates/trusted-server-adapter-fastly/src/app.rs crates/trusted-server-core/src/publisher.rs +git commit -m "Cover lazy publisher streaming behavior" +``` + +## Task 12: Documentation Cleanup + +**Files:** + +- Modify: `crates/trusted-server-core/src/publisher.rs` +- Modify: `crates/trusted-server-core/src/settings.rs` +- Modify: `crates/trusted-server-adapter-fastly/src/app.rs` +- Optionally modify: issue/PR description only, not repo docs + +- [ ] **Step 1: Update stale interim comments** + +Remove or rewrite comments saying publisher stream bodies are already materialized into WASM heap. + +Targets: + +- `PublisherResponse::Stream` doc. +- `PublisherResponse::PassThrough` doc if Fastly now preserves stream bodies. +- `settings.rs` comments that reference future true streaming. +- Fastly app module comment that says publisher responses are buffered by `publisher.max_buffered_body_bytes`. + +- [ ] **Step 2: Run doc-related checks locally** + +Run: + +```bash +cargo fmt --all -- --check +``` + +Expected: PASS. + +- [ ] **Step 3: Commit** + +```bash +git add crates/trusted-server-core/src/publisher.rs crates/trusted-server-core/src/settings.rs crates/trusted-server-adapter-fastly/src/app.rs +git commit -m "Update publisher streaming documentation" +``` + +## Task 13: Full Verification + +**Files:** + +- No source edits unless failures reveal issues. + +- [ ] **Step 1: Run formatting** + +Run: + +```bash +cargo fmt --all -- --check +``` + +Expected: PASS. + +- [ ] **Step 2: Run target checks** + +Run: + +```bash +cargo check-fastly +cargo check-axum +cargo check-cloudflare +``` + +Expected: PASS. + +- [ ] **Step 3: Run target tests** + +Run: + +```bash +cargo test-fastly +cargo test-axum +cargo test-cloudflare +``` + +Expected: PASS. + +- [ ] **Step 4: Run Spin if touched by shared trait changes** + +Run: + +```bash +cargo test-spin +``` + +Expected: PASS. + +- [ ] **Step 5: Run clippy gates** + +Run: + +```bash +cargo clippy-fastly +cargo clippy-axum +cargo clippy-cloudflare +cargo clippy-cloudflare-wasm +cargo clippy-spin-native +cargo clippy-spin-wasm +``` + +Expected: PASS. + +- [ ] **Step 6: Run parity suite if available locally** + +Run: + +```bash +cargo test --manifest-path crates/trusted-server-integration-tests/Cargo.toml --test parity +``` + +Expected: PASS. + +- [ ] **Step 7: Optional local TTFB smoke** + +Run Fastly local serve against an artificially slow publisher origin: + +- Origin sends headers and first HTML chunk immediately. +- Origin delays later body chunks. +- Verify browser/curl receives response headers and first chunk before full origin drain. +- Verify bids still inject before `` when auction completes. + +Expected: TTFB tracks origin first byte, not full origin transfer or auction collection. + +## Review Checklist + +- [ ] No `block_on` inside stream body processing or `Read::read` equivalents. +- [ ] `Body::Stream` never falls through `into_bytes().unwrap_or_default()`. +- [ ] Fastly publisher origin fetch sets `with_stream_response()` only through capability gate. +- [ ] Cloudflare, Spin, and Axum do not start receiving stream-response requests. +- [ ] `HEAD`, `204`, and `304` do not drive or attach response bodies. +- [ ] `Content-Length` is absent on processed streaming responses. +- [ ] Existing buffered adapters still work through `buffer_publisher_response_async()`. +- [ ] Auction telemetry handles completed and abandoned stream cases. +- [ ] Mid-stream errors do not panic; they log and truncate consistently with current streaming behavior. +- [ ] Comments no longer describe publisher streaming as interim/in-memory cursor based. + +## PR Description Skeleton + +```markdown +## Summary + +- convert publisher response processing to async chunk-pull over `Body::Stream` +- enable Fastly publisher origin streaming behind a platform capability gate +- stream Fastly publisher responses to clients instead of buffering and awaiting auction before send + +## Scope + +Fastly production path for issue #849. Cloudflare, Spin, and Axum streaming remain follow-up work. + +## Tests + +- [ ] cargo fmt --all -- --check +- [ ] cargo check-fastly +- [ ] cargo check-axum +- [ ] cargo check-cloudflare +- [ ] cargo test-fastly +- [ ] cargo test-axum +- [ ] cargo test-cloudflare +- [ ] cargo test-spin +- [ ] cargo clippy-fastly +- [ ] cargo clippy-axum +- [ ] cargo clippy-cloudflare +- [ ] cargo clippy-cloudflare-wasm +- [ ] cargo clippy-spin-native +- [ ] cargo clippy-spin-wasm +``` + +## Known Follow-Ups + +- Cloudflare origin streaming once Worker `ReadableStream` is wrapped into `Body::Stream` and response header/set-cookie behavior is verified. +- Spin streaming after upstream EdgeZero Spin response conversion supports incremental body writes. +- Axum streaming only if the dev server needs it enough to justify a `Send` bridge. +- Issue #850 parser-context `` detection. From affb46c9781239ae78695ad5157826ea232e143e Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Wed, 8 Jul 2026 22:39:26 +0530 Subject: [PATCH 02/15] Harden the streaming publisher pipeline after review Address the deep-review findings on the streaming cutover: - Cap cumulative decoded bytes in BodyStreamDecoder against publisher.max_buffered_body_bytes: the chunk source only bounds raw compressed bytes, so a decompression bomb could expand ~1000x past it and push unbounded decoded volume through the rewrite pipeline - Detect truncated deflate streams: write::ZlibDecoder::try_finish accepts truncated input silently, so the deflate arm now drives flate2::Decompress directly and requires Status::StreamEnd at finalization; trailing bytes after the end marker stay ignored. Add truncated-gzip and truncated-deflate regression tests - Make BodyChunkSource::next_chunk cancellation-safe by polling the body in place instead of moving it out across an await; a cancelled pull no longer turns into a silent EOF - Log dispatched auctions dropped uncollected (client disconnect mid-stream or never-polled body) via DispatchedAuctionGuard; the guard is created before the lazy stream so unpolled drops log too - Share the pull+decode step between the lazy publisher body and the write-sink drivers (hold_step_next_chunk / passthrough_step), removing the unreachable!() error plumbing and the triplicated processor selection; document body_close_hold_loop_stream as groundwork for the buffered adapters' streaming cutover - Pass identity-encoded chunks through zero-copy and finish encoders by consuming them instead of allocating a throwaway replacement - Add a Fastly dispatch test asserting the publisher fallback returns Body::Stream without a stale Content-Length, plus a comment on why the publisher fetch gates streaming on capability while the asset path does not Behavior note: gzip bodies with trailing garbage after the trailer now error mid-stream; the old read-path decoder ignored them. --- .../trusted-server-adapter-fastly/src/app.rs | 34 + crates/trusted-server-core/src/publisher.rs | 639 ++++++++++++------ crates/trusted-server-core/src/settings.rs | 13 +- .../src/streaming_processor.rs | 297 ++++++-- 4 files changed, 685 insertions(+), 298 deletions(-) diff --git a/crates/trusted-server-adapter-fastly/src/app.rs b/crates/trusted-server-adapter-fastly/src/app.rs index d5e37f91a..ae9a2749b 100644 --- a/crates/trusted-server-adapter-fastly/src/app.rs +++ b/crates/trusted-server-adapter-fastly/src/app.rs @@ -2150,6 +2150,10 @@ mod tests { #[async_trait::async_trait(?Send)] impl PlatformHttpClient for StreamingHttpClient { + fn supports_streaming_responses(&self) -> bool { + true + } + async fn send( &self, request: PlatformHttpRequest, @@ -2260,6 +2264,36 @@ mod tests { ); } + #[test] + fn dispatch_fallback_streams_publisher_body_without_buffering() { + // Regression guard for the publisher streaming cutover (#849): a + // successful publisher origin fetch must hand `edgezero_main` a lazy + // streaming body (`Body::Stream`) so headers commit at origin first + // byte, rather than draining the processed page into a buffered + // `Body::Once`. Core tests cover the rewrite pipeline itself; this + // guards the adapter wiring that could silently re-buffer. + let settings = test_settings(); + let state = build_state_from_settings(settings).expect("should build state"); + let services = streaming_runtime_services(); + let req = empty_request(Method::GET, "/article"); + + let response = block_on(super::dispatch_fallback(&state, &services, req)); + + assert_eq!( + response.status(), + StatusCode::OK, + "publisher proxy should succeed against the streaming origin stub" + ); + assert!( + matches!(response.body(), Body::Stream(_)), + "EdgeZero publisher dispatch must attach the lazy streaming body, not buffer it" + ); + assert!( + !response.headers().contains_key(header::CONTENT_LENGTH), + "processed streaming publisher responses must not carry a stale Content-Length" + ); + } + #[test] fn dispatch_runs_request_filter_and_threads_response_effects() { // Regression guard for the EdgeZero request-filter bypass: the publisher diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index c10ac0b39..fe6467e12 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -104,40 +104,40 @@ impl BodyChunkSource { } async fn next_chunk(&mut self) -> Result, Report> { - let Some(body) = self.body.take() else { - return Ok(None); - }; - - let chunk = match body { - EdgeBody::Once(bytes) => { - if self.once_offset >= bytes.len() { - None + // The body is polled in place (never moved out across an await) so a + // cancelled `next_chunk` future leaves the source resumable instead of + // silently reporting end-of-stream on the next call. + let pulled = match &mut self.body { + None => Ok(None), + Some(EdgeBody::Once(bytes)) => { + let end = (self.once_offset + self.chunk_size).min(bytes.len()); + if self.once_offset >= end { + Ok(None) } else { - let end = (self.once_offset + self.chunk_size).min(bytes.len()); let chunk = bytes.slice(self.once_offset..end); self.once_offset = end; - if self.once_offset < bytes.len() { - self.body = Some(EdgeBody::Once(bytes)); - } - Some(chunk) + Ok(Some(chunk)) } } - EdgeBody::Stream(mut stream) => match stream.next().await { - Some(Ok(chunk)) => { - self.body = Some(EdgeBody::Stream(stream)); - Some(chunk) - } - Some(Err(err)) => { - return Err(Report::new(TrustedServerError::Proxy { - message: format!("Failed to read publisher origin body stream: {err}"), - })); - } - None => None, + Some(EdgeBody::Stream(stream)) => match stream.next().await { + Some(Ok(chunk)) => Ok(Some(chunk)), + Some(Err(err)) => Err(Report::new(TrustedServerError::Proxy { + message: format!("Failed to read publisher origin body stream: {err}"), + })), + None => Ok(None), }, }; - let Some(chunk) = chunk else { - return Ok(None); + let chunk = match pulled { + Ok(Some(chunk)) => chunk, + Ok(None) => { + self.body = None; + return Ok(None); + } + Err(err) => { + self.body = None; + return Err(err); + } }; self.bytes_seen = self.bytes_seen.checked_add(chunk.len()).ok_or_else(|| { @@ -174,19 +174,17 @@ fn process_and_encode_chunk( if processed.is_empty() { return Ok(None); } - let encoded = encoder.encode_chunk(&processed)?; + let encoded = encoder.encode_chunk(processed)?; if encoded.is_empty() { return Ok(None); } Ok(Some(bytes::Bytes::from(encoded))) } +// By-value signature so `map_err(publisher_stream_error)` works directly. +#[allow(clippy::needless_pass_by_value)] fn publisher_stream_error(err: Report) -> std::io::Error { - let message = format!("{err:?}"); - // Consume the report so clippy's needless_pass_by_value accepts the - // by-value signature that `map_err(publisher_stream_error)` requires. - drop(err); - std::io::Error::other(message) + std::io::Error::other(format!("{err:?}")) } fn not_found_response() -> Response { @@ -473,65 +471,58 @@ fn process_response_streaming( async fn process_response_streaming_async( body: EdgeBody, output: &mut W, - params: &ProcessResponseParams<'_>, - max_raw_body_bytes: usize, + params: &OwnedProcessResponseParams, + settings: &Settings, + integration_registry: &IntegrationRegistry, ) -> Result<(), Report> { - let is_html = is_html_content_type(params.content_type); - let is_rsc_flight = - content_type_contains_ascii_case_insensitive(params.content_type, "text/x-component"); log::debug!( - "process_response_streaming_async: content_type={}, content_encoding={}, is_html={}, is_rsc_flight={}", + "process_response_streaming_async: content_type={}, content_encoding={}", params.content_type, - params.content_encoding, - is_html, - is_rsc_flight + params.content_encoding ); - let compression = Compression::from_content_encoding(params.content_encoding); + let compression = Compression::from_content_encoding(¶ms.content_encoding); + let mut processor = PublisherBodyProcessor::new(params, settings, integration_registry)?; + process_body_chunks_async( + body, + output, + &mut processor, + compression, + settings.publisher.max_buffered_body_bytes, + ) + .await +} - if is_html { - let mut processor = create_html_stream_processor( - params.origin_host, - params.request_host, - params.request_scheme, - params.settings, - params.integration_registry, - params.ad_slots_script.map(str::to_string), - params.ad_bids_state.clone(), - )?; - process_body_chunks_async( - body, - output, - &mut processor, - compression, - max_raw_body_bytes, - ) - .await - } else if is_rsc_flight { - let mut processor = RscFlightUrlRewriter::new( - params.origin_host, - params.origin_url, - params.request_host, - params.request_scheme, - ); - process_body_chunks_async( - body, - output, - &mut processor, - compression, - max_raw_body_bytes, - ) - .await - } else { - let mut replacer = create_url_replacer( - params.origin_host, - params.origin_url, - params.request_host, - params.request_scheme, - ); - process_body_chunks_async(body, output, &mut replacer, compression, max_raw_body_bytes) - .await +/// Pull, decode, process, and encode the next chunk of a no-hold pipeline. +/// +/// Returns `Ok(None)` when the source is exhausted; the caller must then emit +/// [`passthrough_finish_segments`]. Shared by the write-sink driver +/// ([`process_body_chunks_async`]) and the lazy publisher body stream so the +/// two no-hold paths cannot drift apart. +async fn passthrough_step( + source: &mut BodyChunkSource, + decoder: &mut BodyStreamDecoder, + encoder: &mut BodyStreamEncoder, + processor: &mut P, +) -> Result>, Report> { + let Some(raw_chunk) = source.next_chunk().await? else { + return Ok(None); + }; + let decoded = decoder.decode_chunk(raw_chunk)?; + if decoded.is_empty() { + return Ok(Some(Vec::new())); } + let mut segments = Vec::new(); + if let Some(encoded) = process_and_encode_chunk( + processor, + encoder, + &decoded, + false, + "Failed to process chunk", + )? { + segments.push(encoded); + } + Ok(Some(segments)) } async fn process_body_chunks_async( @@ -539,25 +530,16 @@ async fn process_body_chunks_async( writer: &mut W, processor: &mut P, compression: Compression, - max_raw_body_bytes: usize, + max_body_bytes: usize, ) -> Result<(), Report> { - let mut decoder = BodyStreamDecoder::new(compression); + let mut decoder = BodyStreamDecoder::new(compression, max_body_bytes); let mut encoder = BodyStreamEncoder::new(compression); - let mut source = - BodyChunkSource::new(body, STREAM_CHUNK_SIZE).with_max_bytes(max_raw_body_bytes); + let mut source = BodyChunkSource::new(body, STREAM_CHUNK_SIZE).with_max_bytes(max_body_bytes); - while let Some(chunk) = source.next_chunk().await? { - let decoded = decoder.decode_chunk(&chunk)?; - if decoded.is_empty() { - continue; - } - if let Some(encoded) = process_and_encode_chunk( - processor, - &mut encoder, - &decoded, - false, - "Failed to process chunk", - )? { + while let Some(segments) = + passthrough_step(&mut source, &mut decoder, &mut encoder, processor).await? + { + for encoded in segments { write_encoded_segment(writer, &encoded)?; } } @@ -621,18 +603,51 @@ fn passthrough_finish_segments( Ok(segments) } +/// Owns a [`DispatchedAuction`] and logs if it is dropped uncollected. +/// +/// The lazy publisher body stream can be dropped at any await point — a +/// client disconnect aborts the transfer mid-body, or the response may never +/// be polled at all. Async telemetry cannot run in `Drop`, so the loss is +/// surfaced in logs; the abandoned-auction telemetry event is only emitted on +/// error paths that can still await (see [`abandon_hold_auction`]). +struct DispatchedAuctionGuard { + dispatched: Option, +} + +impl DispatchedAuctionGuard { + fn new(dispatched: DispatchedAuction) -> Self { + Self { + dispatched: Some(dispatched), + } + } + + fn take(&mut self) -> Option { + self.dispatched.take() + } +} + +impl Drop for DispatchedAuctionGuard { + fn drop(&mut self) { + if self.dispatched.is_some() { + log::warn!( + "Dispatched server-side auction dropped without collection; SSP bid responses discarded (publisher body stream aborted or never polled)" + ); + } + } +} + /// Mutable auction-hold state threaded through the streaming hold pipeline. struct AuctionHoldState { hold: Option, - dispatched: Option, + dispatched: DispatchedAuctionGuard, telemetry: AuctionTelemetryCarry, } impl AuctionHoldState { - fn new(dispatched: DispatchedAuction, telemetry: AuctionTelemetryCarry) -> Self { + fn new(dispatched: DispatchedAuctionGuard, telemetry: AuctionTelemetryCarry) -> Self { Self { hold: Some(BodyCloseHoldBuffer::new()), - dispatched: Some(dispatched), + dispatched, telemetry, } } @@ -736,6 +751,45 @@ async fn hold_step_decoded_chunk( Ok(segments) } +/// Pull and decode the next chunk of the close-body hold pipeline, feeding it +/// through [`hold_step_decoded_chunk`]. +/// +/// Returns `Ok(None)` when the source is exhausted; the caller must then emit +/// [`hold_finish_segments`]. On read or decode failure the pending auction is +/// abandoned before the error is returned. Shared by the write-sink driver +/// ([`body_close_hold_loop_stream`]) and the lazy publisher body stream so +/// the two hold paths cannot drift apart. +async fn hold_step_next_chunk( + source: &mut BodyChunkSource, + decoder: &mut BodyStreamDecoder, + encoder: &mut BodyStreamEncoder, + processor: &mut P, + state: &mut AuctionHoldState, + collect_refs: &AuctionHoldCollectRefs<'_>, +) -> Result>, Report> { + let raw_chunk = match source.next_chunk().await { + Ok(Some(chunk)) => chunk, + Ok(None) => return Ok(None), + Err(err) => { + abandon_hold_auction(state, collect_refs.services, "stream_read_error").await; + return Err(err); + } + }; + let decoded = match decoder.decode_chunk(raw_chunk) { + Ok(decoded) => decoded, + Err(err) => { + abandon_hold_auction(state, collect_refs.services, "stream_decode_error").await; + return Err(err); + } + }; + if decoded.is_empty() { + return Ok(Some(Vec::new())); + } + hold_step_decoded_chunk(processor, encoder, &decoded, state, collect_refs) + .await + .map(Some) +} + /// Finalize the close-body hold pipeline at end of the origin stream. /// /// Drains the decoder tail through the hold (or straight through when the @@ -843,7 +897,11 @@ fn create_html_stream_processor( /// Result of publisher request handling, indicating whether the response body /// should be streamed or has already been buffered. pub enum PublisherResponse { - /// Response is fully buffered and ready to send via `send_to_client()`. + /// Response returned unmodified, ready to send via `send_to_client()`. + /// + /// On streaming adapters the unmodified body may still be a live + /// [`EdgeBody::Stream`] (the origin fetch requested streaming before the + /// response was classified); it passes through to the client untouched. Buffered(Response), /// Response headers are ready for a streaming response. Covers processable /// content on any status (2xx or non-2xx — e.g., branded 404/500 HTML and @@ -1085,25 +1143,31 @@ pub fn publisher_response_into_streaming_response( let mut params = *params; let mut processor = PublisherBodyProcessor::new(¶ms, &settings, integration_registry)?; + // The guard is created before the lazy stream so an auction whose + // response body is dropped unpolled still logs the loss. + let dispatched_auction = params.dispatched_auction.take().map(|dispatched| { + let telemetry = AuctionTelemetryCarry { + observation: params.auction_observation.take(), + auction_request: params.auction_request.take(), + }; + (DispatchedAuctionGuard::new(dispatched), telemetry) + }); let stream = async_stream::try_stream! { let compression = Compression::from_content_encoding(¶ms.content_encoding); - let mut decoder = BodyStreamDecoder::new(compression); + let max_body_bytes = settings.publisher.max_buffered_body_bytes; + let mut decoder = BodyStreamDecoder::new(compression, max_body_bytes); let mut encoder = BodyStreamEncoder::new(compression); let mut source = BodyChunkSource::new(body, STREAM_CHUNK_SIZE) - .with_max_bytes(settings.publisher.max_buffered_body_bytes); + .with_max_bytes(max_body_bytes); // HTML rides the close-body hold so bids land before ``; // non-HTML has no injection point, so its auction is collected // before any byte streams (matching the buffered finalizer). let mut hold_auction = None; - if let Some(dispatched) = params.dispatched_auction.take() { - let telemetry = AuctionTelemetryCarry { - observation: params.auction_observation.take(), - auction_request: params.auction_request.take(), - }; + if let Some((mut guard, telemetry)) = dispatched_auction { if is_html_content_type(¶ms.content_type) { - hold_auction = Some((dispatched, telemetry)); - } else { + hold_auction = Some((guard, telemetry)); + } else if let Some(dispatched) = guard.take() { collect_non_html_auction( dispatched, telemetry, @@ -1116,8 +1180,8 @@ pub fn publisher_response_into_streaming_response( } } - if let Some((dispatched, telemetry)) = hold_auction { - let mut state = AuctionHoldState::new(dispatched, telemetry); + if let Some((guard, telemetry)) = hold_auction { + let mut state = AuctionHoldState::new(guard, telemetry); let collect_refs = AuctionHoldCollectRefs { price_granularity: params.price_granularity, ad_bids_state: ¶ms.ad_bids_state, @@ -1126,39 +1190,18 @@ pub fn publisher_response_into_streaming_response( settings: &settings, }; - loop { - let raw_chunk = match source.next_chunk().await { - Ok(Some(chunk)) => chunk, - Ok(None) => break, - Err(err) => { - abandon_hold_auction(&mut state, &services, "stream_read_error") - .await; - Err(publisher_stream_error(err))?; - unreachable!("error should have returned"); - } - }; - let decoded = match decoder.decode_chunk(&raw_chunk) { - Ok(decoded) => decoded, - Err(err) => { - abandon_hold_auction(&mut state, &services, "stream_decode_error") - .await; - Err(publisher_stream_error(err))?; - unreachable!("error should have returned"); - } - }; - if decoded.is_empty() { - continue; - } - for encoded in hold_step_decoded_chunk( - &mut processor, - &mut encoder, - &decoded, - &mut state, - &collect_refs, - ) - .await - .map_err(publisher_stream_error)? - { + while let Some(segments) = hold_step_next_chunk( + &mut source, + &mut decoder, + &mut encoder, + &mut processor, + &mut state, + &collect_refs, + ) + .await + .map_err(publisher_stream_error)? + { + for encoded in segments { yield encoded; } } @@ -1176,24 +1219,16 @@ pub fn publisher_response_into_streaming_response( yield encoded; } } else { - while let Some(raw_chunk) = - source.next_chunk().await.map_err(publisher_stream_error)? + while let Some(segments) = passthrough_step( + &mut source, + &mut decoder, + &mut encoder, + &mut processor, + ) + .await + .map_err(publisher_stream_error)? { - let decoded = decoder - .decode_chunk(&raw_chunk) - .map_err(publisher_stream_error)?; - if decoded.is_empty() { - continue; - } - if let Some(encoded) = process_and_encode_chunk( - &mut processor, - &mut encoder, - &decoded, - false, - "Failed to process chunk", - ) - .map_err(publisher_stream_error)? - { + for encoded in segments { yield encoded; } } @@ -1334,23 +1369,12 @@ pub async fn stream_publisher_body_async( ) -> Result<(), Report> { let Some(dispatched) = params.dispatched_auction.take() else { if body.is_stream() { - let borrowed = ProcessResponseParams { - content_encoding: ¶ms.content_encoding, - origin_host: ¶ms.origin_host, - origin_url: ¶ms.origin_url, - request_host: ¶ms.request_host, - request_scheme: ¶ms.request_scheme, - settings, - content_type: ¶ms.content_type, - integration_registry, - ad_slots_script: params.ad_slots_script.as_deref(), - ad_bids_state: ¶ms.ad_bids_state, - }; return process_response_streaming_async( body, output, - &borrowed, - settings.publisher.max_buffered_body_bytes, + params, + settings, + integration_registry, ) .await; } @@ -1378,23 +1402,12 @@ pub async fn stream_publisher_body_async( ) .await; if body.is_stream() { - let borrowed = ProcessResponseParams { - content_encoding: ¶ms.content_encoding, - origin_host: ¶ms.origin_host, - origin_url: ¶ms.origin_url, - request_host: ¶ms.request_host, - request_scheme: ¶ms.request_scheme, - settings, - content_type: ¶ms.content_type, - integration_registry, - ad_slots_script: params.ad_slots_script.as_deref(), - ad_bids_state: ¶ms.ad_bids_state, - }; return process_response_streaming_async( body, output, - &borrowed, - settings.publisher.max_buffered_body_bytes, + params, + settings, + integration_registry, ) .await; } @@ -1639,14 +1652,14 @@ async fn stream_html_with_auction_hold( ctx: AuctionCollectCtx<'_>, ) -> Result<(), Report> { if body.is_stream() { - let max_raw_body_bytes = ctx.settings.publisher.max_buffered_body_bytes; + let max_body_bytes = ctx.settings.publisher.max_buffered_body_bytes; return body_close_hold_loop_stream( body, output, processor, compression, ctx, - max_raw_body_bytes, + max_body_bytes, ) .await; } @@ -1690,16 +1703,22 @@ async fn stream_html_with_auction_hold( /// Async-pull variant of [`body_close_hold_loop`] for live origin streams. /// -/// Shares [`hold_step_decoded_chunk`] and [`hold_finish_segments`] with the +/// Shares [`hold_step_next_chunk`] and [`hold_finish_segments`] with the /// lazy streaming body built by [`publisher_response_into_streaming_response`], /// so the two async hold paths cannot drift apart. +/// +/// No production caller reaches this today: it is only entered through +/// [`buffer_publisher_response_async`], and the buffered adapters (Axum, +/// Cloudflare, Spin) never produce `Body::Stream` because the publisher fetch +/// is gated on `supports_streaming_responses()`. It is groundwork for those +/// adapters' streaming cutover; Fastly uses the lazy stream instead. async fn body_close_hold_loop_stream( body: EdgeBody, writer: &mut W, processor: &mut P, compression: Compression, ctx: AuctionCollectCtx<'_>, - max_raw_body_bytes: usize, + max_body_bytes: usize, ) -> Result<(), Report> { let AuctionCollectCtx { dispatched, @@ -1710,11 +1729,10 @@ async fn body_close_hold_loop_stream( services, settings, } = ctx; - let mut decoder = BodyStreamDecoder::new(compression); + let mut decoder = BodyStreamDecoder::new(compression, max_body_bytes); let mut encoder = BodyStreamEncoder::new(compression); - let mut source = - BodyChunkSource::new(body, STREAM_CHUNK_SIZE).with_max_bytes(max_raw_body_bytes); - let mut state = AuctionHoldState::new(dispatched, telemetry); + let mut source = BodyChunkSource::new(body, STREAM_CHUNK_SIZE).with_max_bytes(max_body_bytes); + let mut state = AuctionHoldState::new(DispatchedAuctionGuard::new(dispatched), telemetry); let collect_refs = AuctionHoldCollectRefs { price_granularity, ad_bids_state, @@ -1723,29 +1741,17 @@ async fn body_close_hold_loop_stream( settings, }; - loop { - let raw_chunk = match source.next_chunk().await { - Ok(Some(chunk)) => chunk, - Ok(None) => break, - Err(err) => { - abandon_hold_auction(&mut state, services, "stream_read_error").await; - return Err(err); - } - }; - let decoded = match decoder.decode_chunk(&raw_chunk) { - Ok(decoded) => decoded, - Err(err) => { - abandon_hold_auction(&mut state, services, "stream_decode_error").await; - return Err(err); - } - }; - if decoded.is_empty() { - continue; - } - for encoded in - hold_step_decoded_chunk(processor, &mut encoder, &decoded, &mut state, &collect_refs) - .await? - { + while let Some(segments) = hold_step_next_chunk( + &mut source, + &mut decoder, + &mut encoder, + processor, + &mut state, + &collect_refs, + ) + .await? + { + for encoded in segments { write_encoded_segment(writer, &encoded)?; } } @@ -2439,6 +2445,11 @@ pub async fn handle_publisher_request( // SSP requests are already racing through the platform HTTP client, so // origin TTFB tracks origin latency rather than the auction timeout. + // + // Streaming is gated on the capability (unlike the asset-proxy path, which + // sets the flag unconditionally and tolerates buffered fallback): adapters + // without streaming support may reject the flag outright rather than + // silently buffering, which would fail every publisher fetch. let mut platform_request = PlatformHttpRequest::new(req, backend_name); if services.http_client().supports_streaming_responses() { platform_request = platform_request.with_stream_response(); @@ -3268,6 +3279,7 @@ pub async fn handle_page_bids( #[cfg(test)] mod tests { + use std::future::Future as _; use std::io::{self, Read as _, Write as _}; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -4952,6 +4964,185 @@ mod tests { }); } + fn non_html_stream_params(content_encoding: &str) -> OwnedProcessResponseParams { + OwnedProcessResponseParams { + content_encoding: content_encoding.to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/css".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: None, + dispatched_auction: None, + price_granularity: crate::price_bucket::PriceGranularity::default(), + } + } + + #[test] + fn stream_publisher_body_async_rejects_truncated_gzip_stream() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = non_html_stream_params("gzip"); + let compressed = + gzip_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + let truncated = &compressed[..compressed.len() - 3]; + let body = + EdgeBody::stream(futures::stream::iter(vec![bytes::Bytes::copy_from_slice( + truncated, + )])); + let mut output = Vec::new(); + + let err = stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect_err("truncated gzip stream must fail instead of truncating silently"); + + assert!( + format!("{err:?}").contains("gzip"), + "should surface the gzip finalization failure: {err:?}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_rejects_truncated_deflate_stream() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = non_html_stream_params("deflate"); + let compressed = + deflate_encode(b"body{background:url('https://origin.example.com/asset.png')}"); + // Cut into the deflate data itself, not just the adler32 trailer. + let truncated = &compressed[..compressed.len() / 2]; + let body = + EdgeBody::stream(futures::stream::iter(vec![bytes::Bytes::copy_from_slice( + truncated, + )])); + let mut output = Vec::new(); + + let err = stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect_err("truncated deflate stream must fail instead of truncating silently"); + + assert!( + format!("{err:?}").contains("deflate"), + "should surface the deflate finalization failure: {err:?}" + ); + }); + } + + #[test] + fn stream_publisher_body_async_enforces_decoded_byte_cap() { + futures::executor::block_on(async { + let mut settings = create_test_settings(); + // Raw compressed input stays tiny (well under the cap); only the + // decoded expansion exceeds it — the decompression-bomb case the + // raw-byte cap alone cannot catch. + settings.publisher.max_buffered_body_bytes = 1024; + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let mut params = non_html_stream_params("gzip"); + let compressed = gzip_encode(&vec![b'a'; 64 * 1024]); + assert!( + compressed.len() < 1024, + "test precondition: compressed input must stay under the raw cap" + ); + let body = + EdgeBody::stream(futures::stream::iter(vec![bytes::Bytes::from(compressed)])); + let mut output = Vec::new(); + + let err = stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect_err("decoded expansion past the cap must fail"); + + assert!( + format!("{err:?}").contains("decoded size exceeded"), + "should report the cumulative decoded cap: {err:?}" + ); + }); + } + + #[test] + fn body_chunk_source_resumes_after_cancelled_poll() { + futures::executor::block_on(async { + let mut pending_once = true; + let mut yielded = false; + let stream = futures::stream::poll_fn(move |cx| { + if pending_once { + pending_once = false; + cx.waker().wake_by_ref(); + return std::task::Poll::Pending; + } + if yielded { + return std::task::Poll::Ready(None); + } + yielded = true; + std::task::Poll::Ready(Some(Ok::<_, io::Error>(bytes::Bytes::from_static( + b"chunk", + )))) + }); + let body = EdgeBody::from_stream(stream); + let mut source = BodyChunkSource::new(body, STREAM_CHUNK_SIZE); + + { + // Poll the pull future once (Pending), then drop it — + // simulating a cancelled await (select/timeout wrapper). + let mut pull = Box::pin(source.next_chunk()); + let waker = futures::task::noop_waker(); + let mut context = std::task::Context::from_waker(&waker); + assert!( + pull.as_mut().poll(&mut context).is_pending(), + "first poll should be pending" + ); + } + + let chunk = source + .next_chunk() + .await + .expect("should read after cancelled poll"); + assert_eq!( + chunk.as_deref(), + Some(&b"chunk"[..]), + "cancelled pull must not lose the origin stream" + ); + }); + } + #[test] fn stream_publisher_body_async_processes_stream_with_auction_hold() { futures::executor::block_on(async { diff --git a/crates/trusted-server-core/src/settings.rs b/crates/trusted-server-core/src/settings.rs index 008b4a2d5..0a6178a07 100644 --- a/crates/trusted-server-core/src/settings.rs +++ b/crates/trusted-server-core/src/settings.rs @@ -51,11 +51,14 @@ pub struct Publisher { /// Defaults to 16 MiB — a conservative cap that prevents Wasm-heap OOM. /// /// Fastly origin bodies are preserved as streams on the publisher path, so - /// this setting is also the cumulative raw-byte cap while the streaming - /// processor decodes and rewrites chunks. Buffered adapters keep using it - /// as the post-rewrite output buffer cap. On the streaming path headers - /// are already committed when the cap trips, so the response is truncated - /// mid-body (with the error logged) rather than replaced with a 5xx. + /// this setting also caps the streaming pipeline twice over: cumulative + /// raw (still compressed) bytes pulled from origin, and cumulative decoded + /// bytes emitted by the decompressor — the latter so a decompression bomb + /// cannot push an unbounded decoded volume through the rewrite pipeline. + /// Buffered adapters keep using it as the post-rewrite output buffer cap. + /// On the streaming path headers are already committed when either cap + /// trips, so the response is truncated mid-body (with the error logged) + /// rather than replaced with a 5xx. /// /// Must be at least 1: a zero-byte cap fails every non-empty buffered /// publisher response at request time, so it is rejected at config diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index ef1a0bc55..963c69daa 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -359,88 +359,177 @@ pub(crate) const STREAM_CHUNK_SIZE: usize = 8192; /// Compressed bytes go in via [`Self::decode_chunk`]; decoded bytes drain /// out of the internal buffer after every push. Write-based decoders are /// used because the async publisher path cannot wrap a blocking `Read`. -pub(crate) enum BodyStreamDecoder { +/// +/// Decoded output is capped cumulatively: the chunk source only bounds raw +/// (still compressed) bytes, and a decompression bomb can expand ~1000x past +/// that, so the decoder enforces its own ceiling on the total bytes it emits. +/// +/// Every codec validates end-of-stream at [`Self::finish`] so a truncated +/// origin body errors instead of silently truncating the page: gzip via its +/// trailer checksum, brotli via `close()`, and deflate via an explicit +/// [`flate2::Status::StreamEnd`] check (`write::ZlibDecoder` accepts +/// truncated input silently, so the deflate arm drives [`flate2::Decompress`] +/// directly). +pub(crate) struct BodyStreamDecoder { + codec: BodyStreamDecoderCodec, + decoded_bytes: usize, + max_decoded_bytes: usize, +} + +enum BodyStreamDecoderCodec { None, Gzip(flate2::write::GzDecoder>), - Deflate(flate2::write::ZlibDecoder>), + Deflate(DeflateStreamDecoder), Brotli(Box>>), } +/// Streaming zlib decoder that tracks whether the stream reached its end +/// marker, so truncated deflate bodies fail at finalization. +struct DeflateStreamDecoder { + decompress: flate2::Decompress, + stream_ended: bool, +} + +impl DeflateStreamDecoder { + fn new() -> Self { + Self { + decompress: flate2::Decompress::new(true), + stream_ended: false, + } + } + + fn decode(&mut self, chunk: &[u8]) -> Result, Report> { + let mut output = Vec::with_capacity(STREAM_CHUNK_SIZE); + let mut offset = 0usize; + // Trailing bytes after the zlib end marker are ignored, matching the + // read-based decoder used by the buffered pipeline. + while offset < chunk.len() && !self.stream_ended { + if output.len() == output.capacity() { + output.reserve(STREAM_CHUNK_SIZE); + } + let before_in = self.decompress.total_in(); + let before_out = self.decompress.total_out(); + let status = self + .decompress + .decompress_vec(&chunk[offset..], &mut output, flate2::FlushDecompress::None) + .change_context(TrustedServerError::Proxy { + message: "Failed to decode deflate publisher body chunk".to_string(), + })?; + let consumed = (self.decompress.total_in() - before_in) as usize; + let produced = (self.decompress.total_out() - before_out) as usize; + offset += consumed; + match status { + flate2::Status::StreamEnd => self.stream_ended = true, + flate2::Status::Ok | flate2::Status::BufError => { + if consumed == 0 && produced == 0 && output.len() < output.capacity() { + return Err(Report::new(TrustedServerError::Proxy { + message: "deflate publisher body decoder made no progress".to_string(), + })); + } + } + } + } + Ok(output) + } +} + impl BodyStreamDecoder { - pub(crate) fn new(compression: Compression) -> Self { - match compression { - Compression::None => Self::None, - Compression::Gzip => Self::Gzip(flate2::write::GzDecoder::new(Vec::new())), - Compression::Deflate => Self::Deflate(flate2::write::ZlibDecoder::new(Vec::new())), - Compression::Brotli => Self::Brotli(Box::new(brotli::DecompressorWriter::new( - Vec::new(), - STREAM_CHUNK_SIZE, - ))), + pub(crate) fn new(compression: Compression, max_decoded_bytes: usize) -> Self { + let codec = match compression { + Compression::None => BodyStreamDecoderCodec::None, + Compression::Gzip => { + BodyStreamDecoderCodec::Gzip(flate2::write::GzDecoder::new(Vec::new())) + } + Compression::Deflate => BodyStreamDecoderCodec::Deflate(DeflateStreamDecoder::new()), + Compression::Brotli => BodyStreamDecoderCodec::Brotli(Box::new( + brotli::DecompressorWriter::new(Vec::new(), STREAM_CHUNK_SIZE), + )), + }; + Self { + codec, + decoded_bytes: 0, + max_decoded_bytes, } } pub(crate) fn decode_chunk( &mut self, - chunk: &[u8], - ) -> Result, Report> { - match self { - Self::None => Ok(chunk.to_vec()), - Self::Gzip(decoder) => { + chunk: bytes::Bytes, + ) -> Result> { + let decoded = match &mut self.codec { + BodyStreamDecoderCodec::None => chunk, + BodyStreamDecoderCodec::Gzip(decoder) => { decoder - .write_all(chunk) + .write_all(&chunk) .change_context(TrustedServerError::Proxy { message: "Failed to decode gzip publisher body chunk".to_string(), })?; - Ok(std::mem::take(decoder.get_mut())) + bytes::Bytes::from(std::mem::take(decoder.get_mut())) } - Self::Deflate(decoder) => { + BodyStreamDecoderCodec::Deflate(decoder) => bytes::Bytes::from(decoder.decode(&chunk)?), + BodyStreamDecoderCodec::Brotli(decoder) => { decoder - .write_all(chunk) - .change_context(TrustedServerError::Proxy { - message: "Failed to decode deflate publisher body chunk".to_string(), - })?; - Ok(std::mem::take(decoder.get_mut())) - } - Self::Brotli(decoder) => { - decoder - .write_all(chunk) + .write_all(&chunk) .change_context(TrustedServerError::Proxy { message: "Failed to decode brotli publisher body chunk".to_string(), })?; - Ok(std::mem::take(decoder.get_mut())) + bytes::Bytes::from(std::mem::take(decoder.get_mut())) } - } + }; + self.track_decoded(decoded.len())?; + Ok(decoded) } pub(crate) fn finish(&mut self) -> Result, Report> { - match self { - Self::None => Ok(Vec::new()), - Self::Gzip(decoder) => { + let tail = match &mut self.codec { + BodyStreamDecoderCodec::None => Vec::new(), + BodyStreamDecoderCodec::Gzip(decoder) => { decoder .try_finish() .change_context(TrustedServerError::Proxy { message: "Failed to finalize gzip publisher body decoder".to_string(), })?; - Ok(std::mem::take(decoder.get_mut())) + std::mem::take(decoder.get_mut()) } - Self::Deflate(decoder) => { - decoder - .try_finish() - .change_context(TrustedServerError::Proxy { - message: "Failed to finalize deflate publisher body decoder".to_string(), - })?; - Ok(std::mem::take(decoder.get_mut())) + BodyStreamDecoderCodec::Deflate(decoder) => { + if !decoder.stream_ended { + return Err(Report::new(TrustedServerError::Proxy { + message: + "Failed to finalize deflate publisher body decoder: truncated stream" + .to_string(), + })); + } + Vec::new() } - Self::Brotli(decoder) => { + BodyStreamDecoderCodec::Brotli(decoder) => { // `close()` (not `flush()`): flush accepts a truncated brotli // stream silently, while close validates end-of-stream and // errors on incomplete input, matching the gzip/deflate arms. decoder.close().change_context(TrustedServerError::Proxy { message: "Failed to finalize brotli publisher body decoder".to_string(), })?; - Ok(std::mem::take(decoder.get_mut())) + std::mem::take(decoder.get_mut()) } + }; + self.track_decoded(tail.len())?; + Ok(tail) + } + + fn track_decoded(&mut self, len: usize) -> Result<(), Report> { + self.decoded_bytes = self.decoded_bytes.checked_add(len).ok_or_else(|| { + Report::new(TrustedServerError::Proxy { + message: "publisher origin body decoded byte count overflowed".to_string(), + }) + })?; + if self.decoded_bytes > self.max_decoded_bytes { + return Err(Report::new(TrustedServerError::Proxy { + message: format!( + "publisher origin body decoded size exceeded {}-byte streaming limit", + self.max_decoded_bytes + ), + })); } + Ok(()) } } @@ -482,13 +571,14 @@ impl BodyStreamEncoder { pub(crate) fn encode_chunk( &mut self, - chunk: &[u8], + chunk: Vec, ) -> Result, Report> { match self { - Self::None => Ok(chunk.to_vec()), + // Identity encoding passes the processed chunk through untouched. + Self::None => Ok(chunk), Self::Gzip(encoder) => { encoder - .write_all(chunk) + .write_all(&chunk) .change_context(TrustedServerError::Proxy { message: "Failed to encode gzip publisher body chunk".to_string(), })?; @@ -496,7 +586,7 @@ impl BodyStreamEncoder { } Self::Deflate(encoder) => { encoder - .write_all(chunk) + .write_all(&chunk) .change_context(TrustedServerError::Proxy { message: "Failed to encode deflate publisher body chunk".to_string(), })?; @@ -504,7 +594,7 @@ impl BodyStreamEncoder { } Self::Brotli(encoder) => { encoder - .write_all(chunk) + .write_all(&chunk) .change_context(TrustedServerError::Proxy { message: "Failed to encode brotli publisher body chunk".to_string(), })?; @@ -513,29 +603,18 @@ impl BodyStreamEncoder { } } + /// Emits the encoder trailer. Consumes the codec state (the encoder + /// becomes identity afterwards); terminal — call once at end of stream. pub(crate) fn finish(&mut self) -> Result, Report> { - match self { + match std::mem::replace(self, Self::None) { Self::None => Ok(Vec::new()), - Self::Gzip(encoder) => { - encoder - .try_finish() - .change_context(TrustedServerError::Proxy { - message: "Failed to finalize gzip publisher body encoder".to_string(), - })?; - Ok(std::mem::take(encoder.get_mut())) - } - Self::Deflate(encoder) => { - encoder - .try_finish() - .change_context(TrustedServerError::Proxy { - message: "Failed to finalize deflate publisher body encoder".to_string(), - })?; - Ok(std::mem::take(encoder.get_mut())) - } - Self::Brotli(encoder) => { - let encoder = std::mem::replace(encoder, Box::new(new_brotli_vec_encoder())); - Ok((*encoder).into_inner()) - } + Self::Gzip(encoder) => encoder.finish().change_context(TrustedServerError::Proxy { + message: "Failed to finalize gzip publisher body encoder".to_string(), + }), + Self::Deflate(encoder) => encoder.finish().change_context(TrustedServerError::Proxy { + message: "Failed to finalize deflate publisher body encoder".to_string(), + }), + Self::Brotli(encoder) => Ok((*encoder).into_inner()), } } } @@ -545,6 +624,86 @@ mod tests { use super::*; use crate::streaming_replacer::{Replacement, StreamingReplacer}; + #[test] + fn body_stream_decoder_enforces_cumulative_decoded_cap() { + let compressed = { + let mut encoder = + flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(&vec![b'a'; 64 * 1024]) + .expect("should write gzip test input"); + encoder.finish().expect("should finish gzip encoding") + }; + assert!( + compressed.len() < 1024, + "test precondition: compressed input must stay small" + ); + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, 1024); + + let err = decoder + .decode_chunk(bytes::Bytes::from(compressed)) + .expect_err("decoded expansion past the cap must fail"); + + assert!( + format!("{err:?}").contains("decoded size exceeded"), + "should report the cumulative decoded cap: {err:?}" + ); + } + + #[test] + fn body_stream_decoder_rejects_truncated_deflate_stream() { + let compressed = { + let mut encoder = + flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(b"deflate payload that spans more than one deflate block boundary") + .expect("should write deflate test input"); + encoder.finish().expect("should finish deflate encoding") + }; + let truncated = &compressed[..compressed.len() / 2]; + let mut decoder = BodyStreamDecoder::new(Compression::Deflate, usize::MAX); + decoder + .decode_chunk(bytes::Bytes::copy_from_slice(truncated)) + .expect("partial deflate input should decode incrementally"); + + let err = decoder + .finish() + .expect_err("truncated deflate stream must fail at finalization"); + + assert!( + format!("{err:?}").contains("truncated stream"), + "should report the missing deflate end marker: {err:?}" + ); + } + + #[test] + fn body_stream_decoder_ignores_deflate_trailing_bytes() { + let compressed = { + let mut encoder = + flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(b"deflate payload") + .expect("should write deflate test input"); + encoder.finish().expect("should finish deflate encoding") + }; + let mut with_trailing = compressed; + with_trailing.extend_from_slice(b"junk"); + let mut decoder = BodyStreamDecoder::new(Compression::Deflate, usize::MAX); + + let decoded = decoder + .decode_chunk(bytes::Bytes::from(with_trailing)) + .expect("complete deflate stream should decode"); + decoder + .finish() + .expect("trailing bytes after the end marker should be ignored"); + + assert_eq!( + decoded.as_ref(), + b"deflate payload", + "should decode the payload and drop trailing junk" + ); + } + /// Verify that `lol_html` fragments text nodes when input chunks split /// mid-text-node. Script rewriters must be fragment-safe — they accumulate /// text fragments internally until `is_last_in_text_node` is true. From 2c64f3c613e4079ac741fab75762376c97573ca1 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Wed, 8 Jul 2026 23:02:08 +0530 Subject: [PATCH 03/15] Emit processor_init_error telemetry on streaming finalize failure The buffered finalizer abandons a dispatched auction with processor_init_error telemetry when HTML processor construction fails; the streaming finalizer dropped the in-flight SSP responses silently. Make publisher_response_into_streaming_response async and emit the same abandonment before returning the construction error. --- .../trusted-server-adapter-fastly/src/app.rs | 19 +++++----- crates/trusted-server-core/src/publisher.rs | 35 ++++++++++++++----- 2 files changed, 38 insertions(+), 16 deletions(-) diff --git a/crates/trusted-server-adapter-fastly/src/app.rs b/crates/trusted-server-adapter-fastly/src/app.rs index ae9a2749b..b68897077 100644 --- a/crates/trusted-server-adapter-fastly/src/app.rs +++ b/crates/trusted-server-adapter-fastly/src/app.rs @@ -795,14 +795,17 @@ async fn dispatch_fallback( ) .await { - Ok(pub_response) => publisher_response_into_streaming_response( - pub_response, - &method, - Arc::clone(&state.settings), - state.registry.as_ref(), - Arc::clone(&state.orchestrator), - publisher_services.clone(), - ), + Ok(pub_response) => { + publisher_response_into_streaming_response( + pub_response, + &method, + Arc::clone(&state.settings), + state.registry.as_ref(), + Arc::clone(&state.orchestrator), + publisher_services.clone(), + ) + .await + } Err(e) => Err(e), } } diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index fe6467e12..4d7e38cf9 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -1101,9 +1101,10 @@ pub async fn buffer_publisher_response_async( /// /// # Errors /// -/// Returns an error if processor construction fails before the streaming body is -/// created. -pub fn publisher_response_into_streaming_response( +/// Returns an error if processor construction fails before the streaming body +/// is created; a dispatched auction is abandoned with `processor_init_error` +/// telemetry first, matching the buffered finalizer. +pub async fn publisher_response_into_streaming_response( publisher_response: PublisherResponse, method: &Method, settings: Arc, @@ -1142,7 +1143,25 @@ pub fn publisher_response_into_streaming_response( response.headers_mut().remove(header::CONTENT_LENGTH); let mut params = *params; let mut processor = - PublisherBodyProcessor::new(¶ms, &settings, integration_registry)?; + match PublisherBodyProcessor::new(¶ms, &settings, integration_registry) { + Ok(processor) => processor, + Err(err) => { + // Parity with the buffered finalizer: a processor + // construction failure abandons the dispatched auction + // with telemetry instead of dropping the in-flight SSP + // responses silently. + if let Some(dispatched) = params.dispatched_auction.take() { + emit_abandoned_auction( + &services, + params.auction_observation.take(), + dispatched, + "processor_init_error", + ) + .await; + } + return Err(err); + } + }; // The guard is created before the lazy stream so an auction whose // response body is dropped unpolled still logs the loss. let dispatched_auction = params.dispatched_auction.take().map(|dispatched| { @@ -5292,14 +5311,14 @@ mod tests { params: Box::new(params), }; - let response = publisher_response_into_streaming_response( + let response = futures::executor::block_on(publisher_response_into_streaming_response( publisher_response, &Method::GET, Arc::clone(&settings), registry.as_ref(), orchestrator, services, - ) + )) .expect("should build streaming response"); assert!( @@ -5458,14 +5477,14 @@ mod tests { params: Box::new(params), }; - let response = publisher_response_into_streaming_response( + let response = futures::executor::block_on(publisher_response_into_streaming_response( publisher_response, &Method::GET, Arc::clone(&settings), registry.as_ref(), orchestrator, services, - ) + )) .expect("should build streaming response"); let output = futures::executor::block_on( From c8ae53f78014458241753f53db630420ff119132 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Tue, 14 Jul 2026 15:49:42 +0530 Subject: [PATCH 04/15] Harden Fastly publisher streaming against review findings Resolve five correctness and resource-safety findings from the PR #867 review of the end-to-end Fastly publisher streaming path. - Drive the deflate decoder to StreamEnd at finalization so a valid stream that exactly fills the internal output buffer is no longer rejected as truncated; the inflater is also drained after all input is consumed within a chunk. - Decode concatenated (multi-member) gzip bodies via MultiGzDecoder on both the streaming decoder and the buffered read pipeline so adapters agree. - Enforce the decoded-body cap during decompression through a bounded sink shared by the gzip and brotli codecs, so a compression bomb errors before its expanded bytes are buffered instead of after a full chunk expands; the deflate codec charges each produced block as it is emitted. - Drop the body of bodiless responses (HEAD, 204, 205, 304) in both the streaming and buffered finalizer Buffered arms, and add RESET_CONTENT to response_carries_body, so a buffered-unmodified stream body is never streamed to the client for a response that must be bodiless. - Keep the dispatched-auction guard armed across the collection await and disarm it only once collection reaches a terminal result, so a body dropped while collection is pending still logs the discarded SSP work. Add regression tests for the deflate output-buffer boundary, multi-member gzip, bodiless buffered stream bodies, and the auction guard sentinel. --- crates/trusted-server-core/src/publisher.rs | 166 +++++++- .../src/streaming_processor.rs | 391 +++++++++++++++--- 2 files changed, 484 insertions(+), 73 deletions(-) diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index 4d7e38cf9..3f9160201 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -612,23 +612,40 @@ fn passthrough_finish_segments( /// error paths that can still await (see [`abandon_hold_auction`]). struct DispatchedAuctionGuard { dispatched: Option, + /// Stays `true` from dispatch until collection (or telemetry-emitting + /// abandonment) reaches a terminal result. [`Self::take`] removes the + /// dispatched auction to hand it to the async collector but deliberately + /// leaves the guard armed, so a drop *while collection is still pending* — + /// a client disconnect at the collection await point — still logs the + /// loss. [`Self::disarm`] clears it only once collection has completed. + armed: bool, } impl DispatchedAuctionGuard { fn new(dispatched: DispatchedAuction) -> Self { Self { dispatched: Some(dispatched), + armed: true, } } + /// Remove the dispatched auction to begin collection. The guard stays armed + /// until [`Self::disarm`] is called, so a drop before collection reaches a + /// terminal result is still reported. fn take(&mut self) -> Option { self.dispatched.take() } + + /// Disarm the drop warning once collection (or telemetry-emitting + /// abandonment) has reached a terminal result. + fn disarm(&mut self) { + self.armed = false; + } } impl Drop for DispatchedAuctionGuard { fn drop(&mut self) { - if self.dispatched.is_some() { + if self.armed { log::warn!( "Dispatched server-side auction dropped without collection; SSP bid responses discarded (publisher body stream aborted or never polled)" ); @@ -668,6 +685,10 @@ async fn abandon_hold_auction( reason, ) .await; + // Abandonment with telemetry is a terminal result, so the drop warning + // is no longer warranted. (A drop *during* the emit above still fires + // it, since the guard stays armed until here.) + state.dispatched.disarm(); } } @@ -721,6 +742,9 @@ async fn hold_step_decoded_chunk( collect_refs.settings, ) .await; + // Collection reached a terminal result; disarm only now so a drop + // while the collect await above was still pending is reported. + state.dispatched.disarm(); let held = state .hold @@ -835,6 +859,9 @@ async fn hold_finish_segments( collect_refs.settings, ) .await; + // Collection reached a terminal result; disarm only now so a drop while + // the collect await above was still pending is reported. + state.dispatched.disarm(); let held = hold.finish(); if let Some(encoded) = process_and_encode_chunk( @@ -1046,7 +1073,17 @@ pub async fn buffer_publisher_response_async( services: &RuntimeServices, ) -> Result, Report> { match publisher_response { - PublisherResponse::Buffered(response) => Ok(response), + PublisherResponse::Buffered(mut response) => { + // A buffered-unmodified response can carry an origin body (a stream + // on streaming-capable adapters). A bodiless response (HEAD, 204, + // 205, 304) must stay bodiless, so drop the body while preserving + // metadata such as `Content-Length`, matching the streaming + // finalizer. + if !response_carries_body(method, response.status()) { + *response.body_mut() = EdgeBody::empty(); + } + Ok(response) + } PublisherResponse::Stream { mut response, body, @@ -1113,7 +1150,18 @@ pub async fn publisher_response_into_streaming_response( services: RuntimeServices, ) -> Result, Report> { match publisher_response { - PublisherResponse::Buffered(response) => Ok(response), + PublisherResponse::Buffered(mut response) => { + // Fastly requests the origin body as a stream before the response is + // classified, so a buffered-unmodified response can still hold an + // `EdgeBody::Stream`. A bodiless response (HEAD, 204, 205, 304) must + // stay bodiless — `send_edgezero_response` streams any + // `EdgeBody::Stream` to the client — so drop the body while + // preserving metadata such as `Content-Length`. + if !response_carries_body(method, response.status()) { + *response.body_mut() = EdgeBody::empty(); + } + Ok(response) + } PublisherResponse::PassThrough { mut response, body } => { if response_carries_body(method, response.status()) { *response.body_mut() = body; @@ -1196,6 +1244,10 @@ pub async fn publisher_response_into_streaming_response( &settings, ) .await; + // Collection reached a terminal result; disarm only now + // so a drop while the collect await above was still + // pending is reported. + guard.disarm(); } } @@ -1268,12 +1320,15 @@ pub async fn publisher_response_into_streaming_response( /// Returns `true` when a buffered publisher response should carry a body and a /// recomputed `Content-Length`. /// -/// `HEAD` responses and bodiless statuses (204, 304) carry no body; rewriting -/// their `Content-Length` to the (empty) buffered length would mislead clients -/// and caches, so the origin metadata is preserved instead. +/// `HEAD` responses and bodiless statuses (204, 205, 304) carry no body; +/// rewriting their `Content-Length` to the (empty) buffered length — or +/// streaming an origin body for them at all — would mislead clients and caches +/// and violate HTTP framing, so the origin metadata is preserved and the body +/// is dropped instead. fn response_carries_body(method: &Method, status: StatusCode) -> bool { *method != Method::HEAD && status != StatusCode::NO_CONTENT + && status != StatusCode::RESET_CONTENT && status != StatusCode::NOT_MODIFIED } @@ -3755,12 +3810,44 @@ mod tests { !super::response_carries_body(&Method::GET, StatusCode::NO_CONTENT), "204 responses must not get a recomputed Content-Length" ); + assert!( + !super::response_carries_body(&Method::GET, StatusCode::RESET_CONTENT), + "205 responses must not get a recomputed Content-Length" + ); assert!( !super::response_carries_body(&Method::GET, StatusCode::NOT_MODIFIED), "304 responses must not get a recomputed Content-Length" ); } + #[test] + fn dispatched_auction_guard_stays_armed_until_collection_completes() { + // `take()` hands the dispatched auction to the async collector, but the + // guard must stay armed across the collection await so a drop while + // collection is still pending (a client disconnect at the await point) + // still logs the loss. Only `disarm()` — called once collection reaches + // a terminal result — clears the warning. + let mut guard = DispatchedAuctionGuard::new(DispatchedAuction::empty_for_test( + test_auction_request(), + 10, + )); + assert!(guard.armed, "a freshly dispatched guard should be armed"); + + let _dispatched = guard + .take() + .expect("guard should yield the dispatched auction for collection"); + assert!( + guard.armed, + "guard must stay armed across the collection await so a drop mid-collection is reported" + ); + + guard.disarm(); + assert!( + !guard.armed, + "guard must disarm once collection reaches a terminal result" + ); + } + fn response_body_string(response: http::Response) -> String { String::from_utf8( response @@ -5354,6 +5441,73 @@ mod tests { ); } + #[test] + fn publisher_response_streaming_finalize_drops_bodiless_buffered_stream_body() { + // Fastly requests the origin body as a stream before classification, so + // a buffered-unmodified response can hold an `EdgeBody::Stream`. The + // adapter streams any `EdgeBody::Stream` to the client, so bodiless + // responses must be normalized to carry no body while keeping metadata. + let settings = Arc::new(create_test_settings()); + let registry = Arc::new( + IntegrationRegistry::new(&settings).expect("should create integration registry"), + ); + let orchestrator = Arc::new(AuctionOrchestrator::new(settings.auction.clone())); + + let cases = [ + (Method::HEAD, StatusCode::OK), + (Method::GET, StatusCode::NO_CONTENT), + (Method::GET, StatusCode::RESET_CONTENT), + (Method::GET, StatusCode::NOT_MODIFIED), + ]; + + for (method, status) in cases { + let response = Response::builder() + .status(status) + .header(header::CONTENT_LENGTH, "42") + .body(EdgeBody::stream(futures::stream::iter(vec![ + bytes::Bytes::from_static(b"origin body bytes that must not reach the client"), + ]))) + .expect("should build response"); + let publisher_response = PublisherResponse::Buffered(response); + + let response = futures::executor::block_on(publisher_response_into_streaming_response( + publisher_response, + &method, + Arc::clone(&settings), + registry.as_ref(), + Arc::clone(&orchestrator), + noop_services(), + )) + .expect("should finalize buffered response"); + + assert!( + !matches!(response.body(), EdgeBody::Stream(_)), + "bodiless {method} {status} must not carry a streaming body" + ); + assert_eq!( + response + .headers() + .get(header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()), + Some("42"), + "bodiless {method} {status} must preserve the origin Content-Length" + ); + + let drained = futures::executor::block_on( + response + .into_body() + .into_bytes_bounded(settings.publisher.max_buffered_body_bytes), + ) + .expect("body should drain") + .to_vec(); + assert!( + drained.is_empty(), + "bodiless {method} {status} must deliver zero body bytes, got {} bytes", + drained.len() + ); + } + } + #[test] fn publisher_response_streaming_finalize_processes_gzip_stream() { let compressed = diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index 963c69daa..e7e9a92bb 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -19,7 +19,7 @@ //! streaming interface. See `crate::platform` module doc for the //! authoritative note. -use std::cell::RefCell; +use std::cell::{Cell, RefCell}; use std::io::{self, Read, Write}; use std::rc::Rc; @@ -27,7 +27,7 @@ use brotli::enc::writer::CompressorWriter; use brotli::enc::BrotliEncoderParams; use brotli::Decompressor; use error_stack::{Report, ResultExt as _}; -use flate2::read::{GzDecoder, ZlibDecoder}; +use flate2::read::{MultiGzDecoder, ZlibDecoder}; use flate2::write::{GzEncoder, ZlibEncoder}; use crate::error::TrustedServerError; @@ -144,7 +144,10 @@ impl StreamingPipeline

{ ) { (Compression::None, Compression::None) => self.process_chunks(input, output), (Compression::Gzip, Compression::Gzip) => { - let decoder = GzDecoder::new(input); + // Multi-member decoder: RFC 1952 permits concatenated gzip + // members, so a single-member reader would stop after the first. + // Matches the streaming `BodyStreamDecoder` gzip codec. + let decoder = MultiGzDecoder::new(input); let mut encoder = GzEncoder::new(output, flate2::Compression::default()); self.process_chunks(decoder, &mut encoder)?; encoder.finish().change_context(TrustedServerError::Proxy { @@ -153,7 +156,7 @@ impl StreamingPipeline

{ Ok(()) } (Compression::Gzip, Compression::None) => { - self.process_chunks(GzDecoder::new(input), output) + self.process_chunks(MultiGzDecoder::new(input), output) } (Compression::Deflate, Compression::Deflate) => { let decoder = ZlibDecoder::new(input); @@ -360,27 +363,106 @@ pub(crate) const STREAM_CHUNK_SIZE: usize = 8192; /// out of the internal buffer after every push. Write-based decoders are /// used because the async publisher path cannot wrap a blocking `Read`. /// -/// Decoded output is capped cumulatively: the chunk source only bounds raw -/// (still compressed) bytes, and a decompression bomb can expand ~1000x past -/// that, so the decoder enforces its own ceiling on the total bytes it emits. +/// Decoded output is capped cumulatively and the cap is enforced *during* +/// decompression, not after: the chunk source only bounds raw (still +/// compressed) bytes, and a decompression bomb can expand ~1000x past that, so +/// a small compressed chunk must not be allowed to fully expand before the +/// ceiling is checked. The gzip and brotli codecs decode into a +/// [`BoundedDecodeSink`] that errors the moment a write would exceed the limit; +/// the deflate codec charges each produced output block as it is emitted. /// /// Every codec validates end-of-stream at [`Self::finish`] so a truncated /// origin body errors instead of silently truncating the page: gzip via its -/// trailer checksum, brotli via `close()`, and deflate via an explicit -/// [`flate2::Status::StreamEnd`] check (`write::ZlibDecoder` accepts -/// truncated input silently, so the deflate arm drives [`flate2::Decompress`] -/// directly). +/// trailer checksum, brotli via `close()`, and deflate by driving +/// [`flate2::Decompress`] to its [`flate2::Status::StreamEnd`] marker (the +/// `write`-based zlib decoder accepts truncated input silently, so the deflate +/// arm drives [`flate2::Decompress`] directly). Concatenated gzip members +/// (RFC 1952) are decoded via [`flate2::write::MultiGzDecoder`]. pub(crate) struct BodyStreamDecoder { codec: BodyStreamDecoderCodec, - decoded_bytes: usize, + /// Cumulative decoded byte count, shared with the codec sinks so the cap is + /// enforced from inside the decompressor writes rather than after them. + decoded_bytes: Rc>, max_decoded_bytes: usize, } enum BodyStreamDecoderCodec { None, - Gzip(flate2::write::GzDecoder>), + Gzip(flate2::write::MultiGzDecoder), Deflate(DeflateStreamDecoder), - Brotli(Box>>), + Brotli(Box>), +} + +/// A [`Write`] sink that buffers decoded bytes while enforcing a shared +/// cumulative decode budget. +/// +/// The gzip and brotli decoders write their decompressed output here as they +/// process input. Rejecting the write as soon as it would push the cumulative +/// decoded total past `max_decoded_bytes` makes the cap a hard ceiling on +/// Wasm-heap growth: a decompression bomb errors before its expanded bytes are +/// buffered, rather than after a full chunk has already expanded. +struct BoundedDecodeSink { + buffer: Vec, + decoded_bytes: Rc>, + max_decoded_bytes: usize, +} + +impl BoundedDecodeSink { + fn new(decoded_bytes: Rc>, max_decoded_bytes: usize) -> Self { + Self { + buffer: Vec::new(), + decoded_bytes, + max_decoded_bytes, + } + } +} + +impl Write for BoundedDecodeSink { + fn write(&mut self, data: &[u8]) -> io::Result { + let next = self + .decoded_bytes + .get() + .checked_add(data.len()) + .ok_or_else(|| { + io::Error::other("publisher origin body decoded byte count overflowed") + })?; + if next > self.max_decoded_bytes { + return Err(io::Error::other(format!( + "publisher origin body decoded size exceeded {}-byte streaming limit", + self.max_decoded_bytes + ))); + } + self.decoded_bytes.set(next); + self.buffer.extend_from_slice(data); + Ok(data.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +/// Charge `len` decoded bytes against `decoded_bytes`, erroring if the +/// cumulative total would exceed `max_decoded_bytes`. +fn charge_decoded( + decoded_bytes: &Cell, + max_decoded_bytes: usize, + len: usize, +) -> Result<(), Report> { + let next = decoded_bytes.get().checked_add(len).ok_or_else(|| { + Report::new(TrustedServerError::Proxy { + message: "publisher origin body decoded byte count overflowed".to_string(), + }) + })?; + if next > max_decoded_bytes { + return Err(Report::new(TrustedServerError::Proxy { + message: format!( + "publisher origin body decoded size exceeded {max_decoded_bytes}-byte streaming limit" + ), + })); + } + decoded_bytes.set(next); + Ok(()) } /// Streaming zlib decoder that tracks whether the stream reached its end @@ -388,22 +470,40 @@ enum BodyStreamDecoderCodec { struct DeflateStreamDecoder { decompress: flate2::Decompress, stream_ended: bool, + decoded_bytes: Rc>, + max_decoded_bytes: usize, } impl DeflateStreamDecoder { - fn new() -> Self { + fn new(decoded_bytes: Rc>, max_decoded_bytes: usize) -> Self { Self { decompress: flate2::Decompress::new(true), stream_ended: false, + decoded_bytes, + max_decoded_bytes, } } + /// Charge `len` decoded bytes against the shared budget. + fn charge(&self, len: usize) -> Result<(), Report> { + charge_decoded(&self.decoded_bytes, self.max_decoded_bytes, len) + } + + /// Decode as much of `chunk` as possible, draining any output the inflater + /// can still produce once all input is consumed. + /// + /// flate2 fills the output buffer up to its capacity, so a chunk that + /// exactly fills the buffer leaves decoded bytes (and possibly the + /// end-of-stream marker) pending with all input already consumed. The loop + /// keeps driving the inflater — reserving more output space — until it + /// makes no further progress, so those pending bytes are never stranded and + /// a valid stream is not mistaken for a truncated one at `finish`. fn decode(&mut self, chunk: &[u8]) -> Result, Report> { let mut output = Vec::with_capacity(STREAM_CHUNK_SIZE); let mut offset = 0usize; // Trailing bytes after the zlib end marker are ignored, matching the // read-based decoder used by the buffered pipeline. - while offset < chunk.len() && !self.stream_ended { + while !self.stream_ended { if output.len() == output.capacity() { output.reserve(STREAM_CHUNK_SIZE); } @@ -418,36 +518,85 @@ impl DeflateStreamDecoder { let consumed = (self.decompress.total_in() - before_in) as usize; let produced = (self.decompress.total_out() - before_out) as usize; offset += consumed; + self.charge(produced)?; match status { flate2::Status::StreamEnd => self.stream_ended = true, flate2::Status::Ok | flate2::Status::BufError => { + // Stop only when the inflater is starved for input: it made + // no progress and there is still spare output capacity, so + // the stall is missing input (arriving in a later chunk, or + // resolved at `finish`), not an exhausted output buffer. if consumed == 0 && produced == 0 && output.len() < output.capacity() { - return Err(Report::new(TrustedServerError::Proxy { - message: "deflate publisher body decoder made no progress".to_string(), - })); + break; } } } } Ok(output) } + + /// Drive the inflater to completion at end of input, draining the final + /// decoded bytes and validating the end-of-stream marker. + /// + /// A valid stream whose last decoded byte exactly filled the previous + /// output buffer still has its end marker pending here; a genuinely + /// truncated stream makes no further progress and errors. + fn finish(&mut self) -> Result, Report> { + let mut output = Vec::new(); + while !self.stream_ended { + if output.len() == output.capacity() { + output.reserve(STREAM_CHUNK_SIZE); + } + let before_out = self.decompress.total_out(); + let status = self + .decompress + .decompress_vec(&[], &mut output, flate2::FlushDecompress::Finish) + .change_context(TrustedServerError::Proxy { + message: "Failed to finalize deflate publisher body decoder".to_string(), + })?; + let produced = (self.decompress.total_out() - before_out) as usize; + self.charge(produced)?; + match status { + flate2::Status::StreamEnd => self.stream_ended = true, + flate2::Status::Ok | flate2::Status::BufError => { + if produced == 0 { + break; + } + } + } + } + if !self.stream_ended { + return Err(Report::new(TrustedServerError::Proxy { + message: "Failed to finalize deflate publisher body decoder: truncated stream" + .to_string(), + })); + } + Ok(output) + } } impl BodyStreamDecoder { pub(crate) fn new(compression: Compression, max_decoded_bytes: usize) -> Self { + let decoded_bytes = Rc::new(Cell::new(0usize)); let codec = match compression { Compression::None => BodyStreamDecoderCodec::None, - Compression::Gzip => { - BodyStreamDecoderCodec::Gzip(flate2::write::GzDecoder::new(Vec::new())) - } - Compression::Deflate => BodyStreamDecoderCodec::Deflate(DeflateStreamDecoder::new()), - Compression::Brotli => BodyStreamDecoderCodec::Brotli(Box::new( - brotli::DecompressorWriter::new(Vec::new(), STREAM_CHUNK_SIZE), + Compression::Gzip => BodyStreamDecoderCodec::Gzip(flate2::write::MultiGzDecoder::new( + BoundedDecodeSink::new(Rc::clone(&decoded_bytes), max_decoded_bytes), + )), + Compression::Deflate => BodyStreamDecoderCodec::Deflate(DeflateStreamDecoder::new( + Rc::clone(&decoded_bytes), + max_decoded_bytes, )), + Compression::Brotli => { + BodyStreamDecoderCodec::Brotli(Box::new(brotli::DecompressorWriter::new( + BoundedDecodeSink::new(Rc::clone(&decoded_bytes), max_decoded_bytes), + STREAM_CHUNK_SIZE, + ))) + } }; Self { codec, - decoded_bytes: 0, + decoded_bytes, max_decoded_bytes, } } @@ -456,51 +605,52 @@ impl BodyStreamDecoder { &mut self, chunk: bytes::Bytes, ) -> Result> { - let decoded = match &mut self.codec { - BodyStreamDecoderCodec::None => chunk, + match &mut self.codec { + BodyStreamDecoderCodec::None => { + // No sink guards the pass-through path, so charge the raw chunk + // directly against the shared budget. + charge_decoded(&self.decoded_bytes, self.max_decoded_bytes, chunk.len())?; + Ok(chunk) + } BodyStreamDecoderCodec::Gzip(decoder) => { decoder .write_all(&chunk) .change_context(TrustedServerError::Proxy { message: "Failed to decode gzip publisher body chunk".to_string(), })?; - bytes::Bytes::from(std::mem::take(decoder.get_mut())) + // The sink charged the decoded bytes during `write_all`. + Ok(bytes::Bytes::from(std::mem::take( + &mut decoder.get_mut().buffer, + ))) + } + BodyStreamDecoderCodec::Deflate(decoder) => { + Ok(bytes::Bytes::from(decoder.decode(&chunk)?)) } - BodyStreamDecoderCodec::Deflate(decoder) => bytes::Bytes::from(decoder.decode(&chunk)?), BodyStreamDecoderCodec::Brotli(decoder) => { decoder .write_all(&chunk) .change_context(TrustedServerError::Proxy { message: "Failed to decode brotli publisher body chunk".to_string(), })?; - bytes::Bytes::from(std::mem::take(decoder.get_mut())) + Ok(bytes::Bytes::from(std::mem::take( + &mut decoder.get_mut().buffer, + ))) } - }; - self.track_decoded(decoded.len())?; - Ok(decoded) + } } pub(crate) fn finish(&mut self) -> Result, Report> { - let tail = match &mut self.codec { - BodyStreamDecoderCodec::None => Vec::new(), + match &mut self.codec { + BodyStreamDecoderCodec::None => Ok(Vec::new()), BodyStreamDecoderCodec::Gzip(decoder) => { decoder .try_finish() .change_context(TrustedServerError::Proxy { message: "Failed to finalize gzip publisher body decoder".to_string(), })?; - std::mem::take(decoder.get_mut()) - } - BodyStreamDecoderCodec::Deflate(decoder) => { - if !decoder.stream_ended { - return Err(Report::new(TrustedServerError::Proxy { - message: - "Failed to finalize deflate publisher body decoder: truncated stream" - .to_string(), - })); - } - Vec::new() + Ok(std::mem::take(&mut decoder.get_mut().buffer)) } + BodyStreamDecoderCodec::Deflate(decoder) => decoder.finish(), BodyStreamDecoderCodec::Brotli(decoder) => { // `close()` (not `flush()`): flush accepts a truncated brotli // stream silently, while close validates end-of-stream and @@ -508,28 +658,9 @@ impl BodyStreamDecoder { decoder.close().change_context(TrustedServerError::Proxy { message: "Failed to finalize brotli publisher body decoder".to_string(), })?; - std::mem::take(decoder.get_mut()) + Ok(std::mem::take(&mut decoder.get_mut().buffer)) } - }; - self.track_decoded(tail.len())?; - Ok(tail) - } - - fn track_decoded(&mut self, len: usize) -> Result<(), Report> { - self.decoded_bytes = self.decoded_bytes.checked_add(len).ok_or_else(|| { - Report::new(TrustedServerError::Proxy { - message: "publisher origin body decoded byte count overflowed".to_string(), - }) - })?; - if self.decoded_bytes > self.max_decoded_bytes { - return Err(Report::new(TrustedServerError::Proxy { - message: format!( - "publisher origin body decoded size exceeded {}-byte streaming limit", - self.max_decoded_bytes - ), - })); } - Ok(()) } } @@ -704,6 +835,132 @@ mod tests { ); } + #[test] + fn body_stream_decoder_decodes_deflate_filling_output_buffer_exactly() { + // A decoded length one byte past the decoder's internal output buffer + // (`STREAM_CHUNK_SIZE`) hits the boundary where flate2 consumes all + // input while exactly filling the output buffer and returns + // `Status::Ok` with the stream-end marker still pending. The decoder + // must drive the inflater to completion instead of reporting a + // truncated stream. + let payload = vec![b'a'; STREAM_CHUNK_SIZE + 1]; + let compressed = { + let mut encoder = + flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(&payload) + .expect("should write deflate test input"); + encoder.finish().expect("should finish deflate encoding") + }; + let mut decoder = BodyStreamDecoder::new(Compression::Deflate, usize::MAX); + + let mut decoded = decoder + .decode_chunk(bytes::Bytes::from(compressed)) + .expect("complete deflate stream should decode") + .to_vec(); + decoded.extend( + decoder + .finish() + .expect("a complete deflate stream must not report truncation"), + ); + + assert_eq!( + decoded, payload, + "should decode the full payload across the output-buffer boundary" + ); + } + + #[test] + fn body_stream_decoder_decodes_deflate_split_across_many_chunks() { + let payload = vec![b'x'; STREAM_CHUNK_SIZE * 3 + 7]; + let compressed = { + let mut encoder = + flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(&payload) + .expect("should write deflate test input"); + encoder.finish().expect("should finish deflate encoding") + }; + let mut decoder = BodyStreamDecoder::new(Compression::Deflate, usize::MAX); + + let mut decoded = Vec::new(); + // Feed the compressed stream a few bytes at a time to exercise many + // input split points, including splits inside the end-of-stream marker. + for piece in compressed.chunks(3) { + decoded.extend( + decoder + .decode_chunk(bytes::Bytes::copy_from_slice(piece)) + .expect("partial deflate input should decode incrementally"), + ); + } + decoded.extend( + decoder + .finish() + .expect("a complete deflate stream must finalize"), + ); + + assert_eq!( + decoded, payload, + "should decode the full payload regardless of input split points" + ); + } + + fn gzip_member(data: &[u8]) -> Vec { + let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); + encoder + .write_all(data) + .expect("should write gzip test input"); + encoder.finish().expect("should finish gzip encoding") + } + + #[test] + fn body_stream_decoder_decodes_multi_member_gzip_single_chunk() { + let mut compressed = gzip_member(b"first member "); + compressed.extend(gzip_member(b"second member")); + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, usize::MAX); + + let mut decoded = decoder + .decode_chunk(bytes::Bytes::from(compressed)) + .expect("a multi-member gzip body must decode all members") + .to_vec(); + decoded.extend( + decoder + .finish() + .expect("a multi-member gzip body must finalize"), + ); + + assert_eq!( + decoded, b"first member second member", + "should concatenate the decoded output of every gzip member" + ); + } + + #[test] + fn body_stream_decoder_decodes_multi_member_gzip_split_across_chunks() { + let mut compressed = gzip_member(b"alpha"); + compressed.extend(gzip_member(b"omega")); + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, usize::MAX); + + let mut decoded = Vec::new(); + for piece in compressed.chunks(4) { + decoded.extend( + decoder + .decode_chunk(bytes::Bytes::copy_from_slice(piece)) + .expect("multi-member gzip should decode across chunk boundaries"), + ); + } + decoded.extend( + decoder + .finish() + .expect("a multi-member gzip body must finalize"), + ); + + assert_eq!( + decoded, b"alphaomega", + "should decode both gzip members split across chunk boundaries" + ); + } + /// Verify that `lol_html` fragments text nodes when input chunks split /// mid-text-node. Script rewriters must be fragment-safe — they accumulate /// text fragments internally until `is_last_in_text_node` is true. From e80cbb767c27e48c8f827a2d5e0d8cab0c92c73b Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Thu, 16 Jul 2026 16:31:30 +0530 Subject: [PATCH 05/15] Flush streaming body encoder per chunk so output decodes before finish MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BodyStreamEncoder::encode_chunk wrote each processed chunk into the gzip/deflate/brotli codec and drained the inner buffer, but never flushed the codec. Compressors buffer internally: write_all alone left gzip emitting at most its 10-byte header and deflate/brotli emitting nothing until finish(). Fastly commits the response headers early, so the browser saw an open stream with no decodable HTML until the entire origin transfer completed — the FCP regression tracked in #849. Sync-flush the flate2 encoders (Flush::Sync — byte-aligned, no trailer) and emit a brotli flush marker (BROTLI_OPERATION_FLUSH) after each write. finish() still writes the terminating trailer, so the full stream stays valid. Add a regression test that decodes the first chunk with a flushed-but-unfinished decoder for every codec. --- .../src/streaming_processor.rs | 77 +++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index 91db43a8d..32532edbf 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -711,6 +711,14 @@ impl BodyStreamEncoder { .change_context(TrustedServerError::Proxy { message: "Failed to encode gzip publisher body chunk".to_string(), })?; + // Sync-flush so the bytes written so far are byte-aligned and + // decodable now; `write_all` alone leaves them buffered inside + // the codec, withholding all output until `finish()` and + // defeating progressive rendering. `flush()` does not emit the + // gzip trailer — `finish()` still does. + encoder.flush().change_context(TrustedServerError::Proxy { + message: "Failed to flush gzip publisher body chunk".to_string(), + })?; Ok(std::mem::take(encoder.get_mut())) } Self::Deflate(encoder) => { @@ -719,6 +727,11 @@ impl BodyStreamEncoder { .change_context(TrustedServerError::Proxy { message: "Failed to encode deflate publisher body chunk".to_string(), })?; + // Sync-flush the deflate codec for the same reason as gzip: make + // the chunk decodable now without writing the stream trailer. + encoder.flush().change_context(TrustedServerError::Proxy { + message: "Failed to flush deflate publisher body chunk".to_string(), + })?; Ok(std::mem::take(encoder.get_mut())) } Self::Brotli(encoder) => { @@ -727,6 +740,13 @@ impl BodyStreamEncoder { .change_context(TrustedServerError::Proxy { message: "Failed to encode brotli publisher body chunk".to_string(), })?; + // `flush()` emits a brotli flush marker (`BROTLI_OPERATION_FLUSH`) + // so the chunk decodes now; the stream is not finished, so + // `finish()` (`into_inner`, `BROTLI_OPERATION_FINISH`) still + // writes the terminating metadata block afterwards. + encoder.flush().change_context(TrustedServerError::Proxy { + message: "Failed to flush brotli publisher body chunk".to_string(), + })?; Ok(std::mem::take(encoder.get_mut())) } } @@ -753,6 +773,63 @@ mod tests { use super::*; use crate::streaming_replacer::{Replacement, StreamingReplacer}; + /// Decode `encoded` with a streaming decoder that is flushed but never + /// finished, mirroring how a browser decodes bytes received so far while + /// the response body is still open. + fn decode_without_finish(compression: Compression, encoded: &[u8]) -> Vec { + match compression { + Compression::None => encoded.to_vec(), + Compression::Gzip => { + let mut decoder = flate2::write::MultiGzDecoder::new(Vec::new()); + decoder.write_all(encoded).expect("should write gzip bytes"); + decoder.flush().expect("should flush gzip decoder"); + decoder.get_ref().clone() + } + Compression::Deflate => { + let mut decoder = flate2::write::ZlibDecoder::new(Vec::new()); + decoder.write_all(encoded).expect("should write deflate bytes"); + decoder.flush().expect("should flush deflate decoder"); + decoder.get_ref().clone() + } + Compression::Brotli => { + let mut decoder = + brotli::DecompressorWriter::new(Vec::new(), STREAM_CHUNK_SIZE); + decoder.write_all(encoded).expect("should write brotli bytes"); + decoder.flush().expect("should flush brotli decoder"); + decoder.get_ref().clone() + } + } + } + + #[test] + fn body_stream_encoder_emits_decodable_chunk_before_finish() { + // A single origin chunk arrives and the origin then stays pending + // (no EOF, so `finish()` is never called). Every compressed codec must + // already emit browser-decodable output for that chunk, otherwise the + // client stalls on an empty stream (or a bare gzip header) until the + // whole origin transfer completes — the FCP regression from #849. + let prefix = b"Example

hello

"; + for compression in [Compression::Gzip, Compression::Deflate, Compression::Brotli] { + let mut encoder = BodyStreamEncoder::new(compression); + let encoded = encoder + .encode_chunk(prefix.to_vec()) + .expect("should encode the first chunk"); + + assert!( + !encoded.is_empty(), + "{compression:?}: first chunk must be emitted, not withheld until finish()" + ); + + // The pre-finish output must already decode to the document prefix. + let decoded = decode_without_finish(compression, &encoded); + assert_eq!( + decoded.as_slice(), + prefix.as_slice(), + "{compression:?}: flushed chunk must decode to the prefix before finish()" + ); + } + } + #[test] fn body_stream_decoder_enforces_cumulative_decoded_cap() { let compressed = { From 5cdaf3de50dd1f5a27f85af1598b1ab89b6c3474 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Thu, 16 Jul 2026 16:51:08 +0530 Subject: [PATCH 06/15] Stream the ready prefix before collecting the held auction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The close-body hold step processed a chunk into ready segments and the held closing tail together, but awaited collect_stream_auction() before returning any of them. When landed in the first source chunk (common for small pages), the client received no HTML until the auction finished; for larger pages the whole prefix from the close-containing chunk was delayed — contradicting the contract that only the tail is held while the auction rides alongside transfer. Split hold_step_decoded_chunk so it returns the ready prefix plus a close_found flag without collecting, and move collection plus held-tail processing into a new hold_collect_close_tail. Both the lazy Fastly stream and the write-sink driver now emit the ready prefix first, then await collection, then emit the tail. Add a regression test asserting the prefix is ready with close_found set while the auction is still uncollected (ad_bids_state stays None until the tail step). --- crates/trusted-server-core/src/publisher.rs | 313 +++++++++++++------- 1 file changed, 207 insertions(+), 106 deletions(-) diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index e271af431..f1cef320c 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -692,85 +692,105 @@ async fn abandon_hold_auction( } } +/// Output of a single close-body hold step, split at the auction-collection +/// barrier. +/// +/// `ready` is the prefix the caller must emit *before* collecting the auction, +/// so a small page whose `` lands in the first source chunk still +/// streams its document prefix immediately instead of stalling behind the +/// auction. `close_found` signals that `, + close_found: bool, +} + /// Feed one decoded chunk through the close-body hold and processor. /// -/// Returns the encoded output segments for the caller to emit — written to a -/// client stream by [`body_close_hold_loop_stream`], yielded from the lazy -/// body by [`publisher_response_into_streaming_response`]. Both async hold -/// paths share this function so their behavior cannot drift apart. +/// Returns the ready prefix for the caller to emit — written to a client stream +/// by [`body_close_hold_loop_stream`], yielded from the lazy body by +/// [`publisher_response_into_streaming_response`]. Both async hold paths share +/// this function so their behavior cannot drift apart. /// -/// When the raw `( processor: &mut P, encoder: &mut BodyStreamEncoder, chunk: &[u8], state: &mut AuctionHoldState, collect_refs: &AuctionHoldCollectRefs<'_>, -) -> Result, Report> { - let mut segments = Vec::new(); - if let Some(hold_buffer) = state.hold.as_mut() { - let ready = hold_buffer.push(chunk); - match process_and_encode_chunk(processor, encoder, &ready, false, "Failed to process chunk") - { - Ok(Some(encoded)) => segments.push(encoded), - Ok(None) => {} - Err(err) => { - abandon_hold_auction(state, collect_refs.services, "stream_process_error").await; - return Err(err); - } +) -> Result> { + let mut ready = Vec::new(); + let bytes = match state.hold.as_mut() { + // Once the hold has been released the chunk streams straight through. + None => chunk.to_vec(), + Some(hold_buffer) => hold_buffer.push(chunk), + }; + match process_and_encode_chunk(processor, encoder, &bytes, false, "Failed to process chunk") { + Ok(Some(encoded)) => ready.push(encoded), + Ok(None) => {} + Err(err) => { + abandon_hold_auction(state, collect_refs.services, "stream_process_error").await; + return Err(err); } + } + let close_found = state + .hold + .as_ref() + .is_some_and(BodyCloseHoldBuffer::found_close); + Ok(HoldStepSegments { ready, close_found }) +} - if state - .hold - .as_ref() - .is_some_and(BodyCloseHoldBuffer::found_close) - { - let dispatched = state - .dispatched - .take() - .expect("should have dispatched auction to collect"); - collect_stream_auction( - dispatched, - state.telemetry.take(), - collect_refs.price_granularity, - collect_refs.ad_bids_state, - collect_refs.orchestrator, - collect_refs.services, - collect_refs.settings, - ) - .await; - // Collection reached a terminal result; disarm only now so a drop - // while the collect await above was still pending is reported. - state.dispatched.disarm(); - - let held = state - .hold - .take() - .expect("should have close-body hold buffer") - .finish(); - if let Some(encoded) = process_and_encode_chunk( - processor, - encoder, - &held, - false, - "Failed to process held body close", - )? { - segments.push(encoded); - } - } - } else { - match process_and_encode_chunk(processor, encoder, chunk, false, "Failed to process chunk") - { - Ok(Some(encoded)) => segments.push(encoded), - Ok(None) => {} - Err(err) => { - abandon_hold_auction(state, collect_refs.services, "stream_process_error").await; - return Err(err); - } - } +/// Collect the dispatched auction and process the held `` tail. +/// +/// Call only after [`hold_step_decoded_chunk`] (or [`hold_finish_segments`]) +/// reports `close_found` and the ready prefix has already been emitted: +/// collecting here — after the prefix streams — is what keeps the auction +/// riding alongside transfer instead of blocking it. Collection runs before the +/// tail is processed so `lol_html` sees live bids at the injection point. +async fn hold_collect_close_tail( + processor: &mut P, + encoder: &mut BodyStreamEncoder, + state: &mut AuctionHoldState, + collect_refs: &AuctionHoldCollectRefs<'_>, +) -> Result, Report> { + let mut segments = Vec::new(); + let dispatched = state + .dispatched + .take() + .expect("should have dispatched auction to collect"); + collect_stream_auction( + dispatched, + state.telemetry.take(), + collect_refs.price_granularity, + collect_refs.ad_bids_state, + collect_refs.orchestrator, + collect_refs.services, + collect_refs.settings, + ) + .await; + // Collection reached a terminal result; disarm only now so a drop while the + // collect await above was still pending is reported. + state.dispatched.disarm(); + + let held = state + .hold + .take() + .expect("should have close-body hold buffer") + .finish(); + if let Some(encoded) = process_and_encode_chunk( + processor, + encoder, + &held, + false, + "Failed to process held body close", + )? { + segments.push(encoded); } Ok(segments) } @@ -790,7 +810,7 @@ async fn hold_step_next_chunk( processor: &mut P, state: &mut AuctionHoldState, collect_refs: &AuctionHoldCollectRefs<'_>, -) -> Result>, Report> { +) -> Result, Report> { let raw_chunk = match source.next_chunk().await { Ok(Some(chunk)) => chunk, Ok(None) => return Ok(None), @@ -807,7 +827,10 @@ async fn hold_step_next_chunk( } }; if decoded.is_empty() { - return Ok(Some(Vec::new())); + return Ok(Some(HoldStepSegments { + ready: Vec::new(), + close_found: false, + })); } hold_step_decoded_chunk(processor, encoder, &decoded, state, collect_refs) .await @@ -839,40 +862,16 @@ async fn hold_finish_segments( } }; if !decoded_tail.is_empty() { - segments.extend( - hold_step_decoded_chunk(processor, encoder, &decoded_tail, state, collect_refs).await?, - ); + let step = + hold_step_decoded_chunk(processor, encoder, &decoded_tail, state, collect_refs).await?; + segments.extend(step.ready); } - if let Some(hold) = state.hold.take() { - let dispatched = state - .dispatched - .take() - .expect("should have dispatched auction to collect"); - collect_stream_auction( - dispatched, - state.telemetry.take(), - collect_refs.price_granularity, - collect_refs.ad_bids_state, - collect_refs.orchestrator, - collect_refs.services, - collect_refs.settings, - ) - .await; - // Collection reached a terminal result; disarm only now so a drop while - // the collect await above was still pending is reported. - state.dispatched.disarm(); - - let held = hold.finish(); - if let Some(encoded) = process_and_encode_chunk( - processor, - encoder, - &held, - false, - "Failed to process held body close", - )? { - segments.push(encoded); - } + // If the hold is still armed the auction was never collected mid-stream: + // `` arrived only in this tail, or the document had none at all. + // Collect now and flush the held remainder before finalizing. + if state.hold.is_some() { + segments.extend(hold_collect_close_tail(processor, encoder, state, collect_refs).await?); } if let Some(encoded) = process_and_encode_chunk( @@ -1266,7 +1265,7 @@ pub async fn publisher_response_into_streaming_response( settings: &settings, }; - while let Some(segments) = hold_step_next_chunk( + while let Some(step) = hold_step_next_chunk( &mut source, &mut decoder, &mut encoder, @@ -1277,9 +1276,25 @@ pub async fn publisher_response_into_streaming_response( .await .map_err(publisher_stream_error)? { - for encoded in segments { + // Emit the ready prefix before collecting the auction so + // the client receives the document up to `` while + // the auction rides alongside transfer. + for encoded in step.ready { yield encoded; } + if step.close_found { + for encoded in hold_collect_close_tail( + &mut processor, + &mut encoder, + &mut state, + &collect_refs, + ) + .await + .map_err(publisher_stream_error)? + { + yield encoded; + } + } } for encoded in hold_finish_segments( @@ -1820,7 +1835,7 @@ async fn body_close_hold_loop_stream( settings, }; - while let Some(segments) = hold_step_next_chunk( + while let Some(step) = hold_step_next_chunk( &mut source, &mut decoder, &mut encoder, @@ -1830,9 +1845,18 @@ async fn body_close_hold_loop_stream( ) .await? { - for encoded in segments { + // Write the ready prefix before collecting the auction, matching the + // lazy Fastly stream: only the held `` tail waits on collection. + for encoded in step.ready { write_encoded_segment(writer, &encoded)?; } + if step.close_found { + for encoded in + hold_collect_close_tail(processor, &mut encoder, &mut state, &collect_refs).await? + { + write_encoded_segment(writer, &encoded)?; + } + } } for encoded in hold_finish_segments( @@ -4010,6 +4034,83 @@ mod tests { ); } + #[tokio::test] + async fn hold_step_yields_ready_prefix_before_collecting_auction() { + // A small page whose `` lands in the first source chunk must + // still stream its document prefix immediately. `hold_step_decoded_chunk` + // reports the ready prefix and `close_found` without collecting; only + // `hold_collect_close_tail` awaits collection. Regression guard for the + // #849 FCP objective: the prefix must become ready while collection + // remains pending. + let settings = create_test_settings(); + let services = noop_services(); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let ad_bids_state = Arc::new(Mutex::new(None)); + let mut state = AuctionHoldState::new( + DispatchedAuctionGuard::new(DispatchedAuction::empty_for_test( + test_auction_request(), + 500, + )), + AuctionTelemetryCarry { + observation: None, + auction_request: None, + }, + ); + let collect_refs = AuctionHoldCollectRefs { + price_granularity: PriceGranularity::default(), + ad_bids_state: &ad_bids_state, + orchestrator: &orchestrator, + services: &services, + settings: &settings, + }; + // Passthrough processor: the ordering contract is about collection, not + // HTML rewriting, so keep the emitted bytes verbatim. + let mut processor = RecordingProcessor { + read_count: Arc::new(AtomicUsize::new(0)), + body_close_processed_at: Arc::new(AtomicUsize::new(0)), + }; + let mut encoder = BodyStreamEncoder::new(Compression::None); + + let step = hold_step_decoded_chunk( + &mut processor, + &mut encoder, + b"painted", + &mut state, + &collect_refs, + ) + .await + .expect("hold step should succeed"); + + assert!( + step.close_found, + " in the first chunk must be detected" + ); + let ready: Vec = step.ready.iter().flat_map(|b| b.to_vec()).collect(); + assert_eq!( + std::str::from_utf8(&ready).expect("ready prefix should be utf8"), + "painted", + "the prefix up to must be ready before collection" + ); + assert!( + ad_bids_state.lock().expect("should lock bid state").is_none(), + "auction must not be collected while the ready prefix is emitted" + ); + + let tail = hold_collect_close_tail(&mut processor, &mut encoder, &mut state, &collect_refs) + .await + .expect("collect should succeed"); + let tail_bytes: Vec = tail.iter().flat_map(|b| b.to_vec()).collect(); + assert_eq!( + std::str::from_utf8(&tail_bytes).expect("held tail should be utf8"), + "", + "the held close tail must be emitted after collection" + ); + assert!( + ad_bids_state.lock().expect("should lock bid state").is_some(), + "collection must run when the held tail is emitted" + ); + } + #[test] fn body_close_hold_buffer_holds_close_body_tail_in_single_chunk() { let mut hold = BodyCloseHoldBuffer::new(); From 43499f8c095f3c949ed383c9cc5305b05aa7fb70 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Thu, 16 Jul 2026 16:56:30 +0530 Subject: [PATCH 07/15] Correct Content-Length framing for bodiless 204/205 responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both publisher finalizers dropped the body for bodiless responses but kept the origin Content-Length verbatim, including a possible nonzero value. RFC 9110 §8.6 forbids Content-Length on 204, and a nonzero length on a now-empty 205 (§15.4.6) is invalid; inconsistent framing risks interoperability failures and response-splitting across intermediaries. Add make_response_bodiless: it empties the body and removes Content-Length for 204, normalizes it to 0 for 205, and preserves it for HEAD and 304 (which legitimately advertise the GET representation length). Call it from the buffered-unmodified arm of both buffer_publisher_response_async and publisher_response_into_streaming_response. Split the combined bodiless regression test into per-status expectations and add the matching coverage for the buffered finalizer. --- crates/trusted-server-core/src/publisher.rs | 119 ++++++++++++++++---- 1 file changed, 100 insertions(+), 19 deletions(-) diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index f1cef320c..380e95a27 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -1080,11 +1080,11 @@ pub async fn buffer_publisher_response_async( PublisherResponse::Buffered(mut response) => { // A buffered-unmodified response can carry an origin body (a stream // on streaming-capable adapters). A bodiless response (HEAD, 204, - // 205, 304) must stay bodiless, so drop the body while preserving - // metadata such as `Content-Length`, matching the streaming - // finalizer. + // 205, 304) must stay bodiless, so drop the body and correct its + // framing (204/205 Content-Length; HEAD/304 preserved), matching + // the streaming finalizer. if !response_carries_body(method, response.status()) { - *response.body_mut() = EdgeBody::empty(); + make_response_bodiless(&mut response); } Ok(response) } @@ -1159,10 +1159,10 @@ pub async fn publisher_response_into_streaming_response( // classified, so a buffered-unmodified response can still hold an // `EdgeBody::Stream`. A bodiless response (HEAD, 204, 205, 304) must // stay bodiless — `send_edgezero_response` streams any - // `EdgeBody::Stream` to the client — so drop the body while - // preserving metadata such as `Content-Length`. + // `EdgeBody::Stream` to the client — so drop the body and correct + // its framing (204/205 Content-Length; HEAD/304 preserved). if !response_carries_body(method, response.status()) { - *response.body_mut() = EdgeBody::empty(); + make_response_bodiless(&mut response); } Ok(response) } @@ -1352,6 +1352,30 @@ fn response_carries_body(method: &Method, status: StatusCode) -> bool { && status != StatusCode::NOT_MODIFIED } +/// Drop a bodiless response's body and correct its framing headers. +/// +/// The response keeps no body, and its `Content-Length` is corrected where the +/// origin's value would be invalid for the now-empty message: +/// - **204 No Content**: RFC 9110 §8.6 forbids `Content-Length`; remove it. +/// - **205 Reset Content**: the reset response carries nothing, so a nonzero +/// origin length is wrong (RFC 9110 §15.4.6); normalize it to `0`. +/// - **HEAD and 304 Not Modified**: `Content-Length` legitimately advertises +/// the `GET` representation length, so it is preserved untouched. +fn make_response_bodiless(response: &mut Response) { + *response.body_mut() = EdgeBody::empty(); + match response.status() { + StatusCode::NO_CONTENT => { + response.headers_mut().remove(header::CONTENT_LENGTH); + } + StatusCode::RESET_CONTENT => { + response + .headers_mut() + .insert(header::CONTENT_LENGTH, HeaderValue::from(0_u64)); + } + _ => {} + } +} + /// A [`Write`] sink that buffers into a `Vec` but fails once the configured /// byte limit would be exceeded. /// @@ -5547,26 +5571,29 @@ mod tests { ); } + // (method, status, expected Content-Length after bodiless normalization). + // 204 forbids Content-Length (removed); 205 must advertise a zero-length + // body; HEAD and 304 legitimately advertise the GET representation length. + const BODILESS_FRAMING_CASES: [(Method, StatusCode, Option<&str>); 4] = [ + (Method::HEAD, StatusCode::OK, Some("42")), + (Method::GET, StatusCode::NO_CONTENT, None), + (Method::GET, StatusCode::RESET_CONTENT, Some("0")), + (Method::GET, StatusCode::NOT_MODIFIED, Some("42")), + ]; + #[test] - fn publisher_response_streaming_finalize_drops_bodiless_buffered_stream_body() { + fn publisher_response_streaming_finalize_normalizes_bodiless_framing() { // Fastly requests the origin body as a stream before classification, so // a buffered-unmodified response can hold an `EdgeBody::Stream`. The // adapter streams any `EdgeBody::Stream` to the client, so bodiless - // responses must be normalized to carry no body while keeping metadata. + // responses must carry no body and correct framing per status. let settings = Arc::new(create_test_settings()); let registry = Arc::new( IntegrationRegistry::new(&settings).expect("should create integration registry"), ); let orchestrator = Arc::new(AuctionOrchestrator::new(settings.auction.clone())); - let cases = [ - (Method::HEAD, StatusCode::OK), - (Method::GET, StatusCode::NO_CONTENT), - (Method::GET, StatusCode::RESET_CONTENT), - (Method::GET, StatusCode::NOT_MODIFIED), - ]; - - for (method, status) in cases { + for (method, status, expected_length) in BODILESS_FRAMING_CASES { let response = Response::builder() .status(status) .header(header::CONTENT_LENGTH, "42") @@ -5595,8 +5622,62 @@ mod tests { .headers() .get(header::CONTENT_LENGTH) .and_then(|v| v.to_str().ok()), - Some("42"), - "bodiless {method} {status} must preserve the origin Content-Length" + expected_length, + "bodiless {method} {status} must carry the corrected Content-Length" + ); + + let drained = futures::executor::block_on( + response + .into_body() + .into_bytes_bounded(settings.publisher.max_buffered_body_bytes), + ) + .expect("body should drain") + .to_vec(); + assert!( + drained.is_empty(), + "bodiless {method} {status} must deliver zero body bytes, got {} bytes", + drained.len() + ); + } + } + + #[test] + fn buffer_publisher_response_normalizes_bodiless_framing() { + // The buffered finalizer (Axum/Cloudflare/Spin) must correct bodiless + // framing identically to the streaming finalizer. + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + + for (method, status, expected_length) in BODILESS_FRAMING_CASES { + let response = Response::builder() + .status(status) + .header(header::CONTENT_LENGTH, "42") + .body(EdgeBody::from( + b"origin body bytes that must not reach the client".to_vec(), + )) + .expect("should build response"); + let publisher_response = PublisherResponse::Buffered(response); + + let response = futures::executor::block_on(buffer_publisher_response_async( + publisher_response, + &method, + &settings, + ®istry, + &orchestrator, + &services, + )) + .expect("should finalize buffered response"); + + assert_eq!( + response + .headers() + .get(header::CONTENT_LENGTH) + .and_then(|v| v.to_str().ok()), + expected_length, + "bodiless {method} {status} must carry the corrected Content-Length" ); let drained = futures::executor::block_on( From 0976cfb7a567a0b9131f69f8e90e475f4df0ffb8 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Thu, 16 Jul 2026 17:03:53 +0530 Subject: [PATCH 08/15] Emit abandoned-auction telemetry for bodiless dispatched responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A conditional navigation can dispatch a server-side auction and receive a processable HTML 304 that classification routes to Stream. Both finalizers' bodiless branches logged a warning and returned, dropping params.dispatched_auction without emit_abandoned_auction — so the SSP work and quota consumption had no terminal auction event, leaving a gap in auction observability on a legitimate conditional-request path. Take the dispatched auction in the bodiless branch of both buffer_publisher_response_async and publisher_response_into_streaming_response and await emit_abandoned_auction(..., "bodiless_response") with the retained observation, matching the processor-init-error path. Add a recording-telemetry-sink test (via a new test_support helper) asserting both finalizers emit the bodiless_response abandonment for a 304 with a dispatched auction. --- .../src/platform/test_support.rs | 18 ++ crates/trusted-server-core/src/publisher.rs | 156 +++++++++++++++++- 2 files changed, 165 insertions(+), 9 deletions(-) diff --git a/crates/trusted-server-core/src/platform/test_support.rs b/crates/trusted-server-core/src/platform/test_support.rs index ab31dd40b..e9a731a64 100644 --- a/crates/trusted-server-core/src/platform/test_support.rs +++ b/crates/trusted-server-core/src/platform/test_support.rs @@ -636,6 +636,24 @@ pub(crate) fn noop_services() -> RuntimeServices { build_services_with_config(NoopConfigStore) } +/// Build a [`RuntimeServices`] whose auction telemetry sink is the supplied +/// recording (or otherwise custom) sink, so tests can assert which terminal +/// auction events were emitted. +pub(crate) fn noop_services_with_telemetry_sink( + auction_telemetry_sink: Arc, +) -> RuntimeServices { + RuntimeServices::builder() + .config_store(Arc::new(NoopConfigStore)) + .secret_store(Arc::new(NoopSecretStore)) + .kv_store(Arc::new(edgezero_core::key_value_store::NoopKvStore)) + .backend(Arc::new(NoopBackend)) + .http_client(Arc::new(NoopHttpClient)) + .geo(Arc::new(NoopGeo)) + .auction_telemetry_sink(auction_telemetry_sink) + .client_info(ClientInfo::default()) + .build() +} + /// Build a [`RuntimeServices`] with a caller-supplied HTTP client and a [`StubBackend`]. /// /// Uses [`StubBackend`] (always returns `Ok("stub-backend")`) rather than diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index 380e95a27..7046346f6 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -1094,16 +1094,24 @@ pub async fn buffer_publisher_response_async( mut params, } => { if !response_carries_body(method, response.status()) { - if params.dispatched_auction.is_some() { - // A bodiless response (HEAD navigation, 204/304) has no + if let Some(dispatched) = params.dispatched_auction.take() { + // A bodiless response (HEAD navigation, 304) has no // `` to inject bids into, so the dispatched SSP - // requests are wasted — surface it for quota observability, - // matching the pass-through / buffered-unmodified arms. + // requests are wasted. Emit a terminal abandonment event so + // the SSP work and quota consumption stay observable instead + // of vanishing, matching the streaming finalizer. log::warn!( "Server-side auction dispatched but response is bodiless (method: {}, status: {}); in-flight SSP bid requests will not be collected", method, response.status(), ); + emit_abandoned_auction( + services, + params.auction_observation.take(), + dispatched, + "bodiless_response", + ) + .await; } return Ok(response); } @@ -1175,19 +1183,28 @@ pub async fn publisher_response_into_streaming_response( PublisherResponse::Stream { mut response, body, - params, + mut params, } => { if !response_carries_body(method, response.status()) { - if params.dispatched_auction.is_some() { - // A bodiless response (HEAD navigation, 204/304) has no + if let Some(dispatched) = params.dispatched_auction.take() { + // A bodiless response (HEAD navigation, 304) has no // `` to inject bids into, so the dispatched SSP - // requests are wasted — surface it for quota observability, - // matching the buffered finalizer. + // requests are wasted. Emit a terminal abandonment event so + // the SSP work and quota consumption stay observable instead + // of vanishing, matching the processor-init-error path and + // the buffered finalizer. log::warn!( "Server-side auction dispatched but response is bodiless (method: {}, status: {}); in-flight SSP bid requests will not be collected", method, response.status(), ); + emit_abandoned_auction( + &services, + params.auction_observation.take(), + dispatched, + "bodiless_response", + ) + .await; } return Ok(response); } @@ -3420,6 +3437,7 @@ mod tests { use crate::integrations::IntegrationRegistry; use crate::platform::test_support::{ StubHttpClient, build_services_with_http_client, noop_services, + noop_services_with_telemetry_sink, }; use crate::test_support::tests::create_test_settings; use edgezero_core::body::Body as EdgeBody; @@ -5695,6 +5713,126 @@ mod tests { } } + #[derive(Default)] + struct RecordingTelemetrySink { + batches: Mutex>, + } + + #[async_trait::async_trait(?Send)] + impl crate::auction::telemetry::AuctionTelemetrySink for RecordingTelemetrySink { + async fn emit_auction_events( + &self, + _services: &RuntimeServices, + batch: crate::auction::telemetry::AuctionEventBatch, + ) -> Result<(), Report> { + self.batches + .lock() + .expect("should lock telemetry batches") + .push(batch); + Ok(()) + } + } + + #[test] + fn finalizers_emit_abandoned_auction_for_bodiless_dispatched_response() { + // A conditional navigation can dispatch an auction and receive a + // processable HTML 304 that classification routes to Stream. The + // bodiless response has no `` to inject into, so the dispatched + // auction is never collected — but both finalizers must still emit a + // terminal abandonment event so the SSP work and quota consumption stay + // observable instead of vanishing silently. + let settings = Arc::new(create_test_settings()); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = Arc::new(AuctionOrchestrator::new(settings.auction.clone())); + + let make_params = || { + let ec_context = + EcContext::new_for_test(None, crate::consent::types::ConsentContext::default()); + OwnedProcessResponseParams { + content_encoding: String::new(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/html; charset=utf-8".to_string(), + ad_slots_script: None, + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: Some(AuctionObservationContext::from_parts( + AuctionSource::SpaNavigation, + "proxy.example.com", + "/article", + 1, + &ec_context, + )), + auction_request: Some(test_auction_request()), + dispatched_auction: Some(DispatchedAuction::empty_for_test( + test_auction_request(), + 10, + )), + price_granularity: PriceGranularity::default(), + } + }; + let make_stream_response = || PublisherResponse::Stream { + response: Response::builder() + .status(StatusCode::NOT_MODIFIED) + .header(header::CONTENT_TYPE, "text/html; charset=utf-8") + .body(EdgeBody::empty()) + .expect("should build 304 response"), + body: EdgeBody::stream(futures::stream::iter(vec![bytes::Bytes::from_static( + b"", + )])), + params: Box::new(make_params()), + }; + + fn assert_bodiless_abandoned(sink: &RecordingTelemetrySink) { + let batches = sink.batches.lock().expect("should lock telemetry batches"); + assert_eq!( + batches.len(), + 1, + "bodiless dispatched response should emit exactly one terminal batch" + ); + let reasons: Vec = batches[0] + .rows() + .iter() + .filter_map(|row| row.terminal_reason.clone()) + .collect(); + assert!( + reasons.iter().any(|reason| reason == "bodiless_response"), + "abandonment must carry the bodiless_response reason, got {reasons:?}" + ); + } + + // Streaming finalizer (Fastly). + let streaming_sink = Arc::new(RecordingTelemetrySink::default()); + let response = futures::executor::block_on(publisher_response_into_streaming_response( + make_stream_response(), + &Method::GET, + Arc::clone(&settings), + ®istry, + Arc::clone(&orchestrator), + noop_services_with_telemetry_sink(Arc::clone(&streaming_sink) as _), + )) + .expect("streaming finalize should succeed"); + assert_eq!(response.status(), StatusCode::NOT_MODIFIED); + assert_bodiless_abandoned(&streaming_sink); + + // Buffered finalizer (Axum/Cloudflare/Spin). + let buffered_sink = Arc::new(RecordingTelemetrySink::default()); + let buffered_services = noop_services_with_telemetry_sink(Arc::clone(&buffered_sink) as _); + let response = futures::executor::block_on(buffer_publisher_response_async( + make_stream_response(), + &Method::GET, + settings.as_ref(), + ®istry, + orchestrator.as_ref(), + &buffered_services, + )) + .expect("buffered finalize should succeed"); + assert_eq!(response.status(), StatusCode::NOT_MODIFIED); + assert_bodiless_abandoned(&buffered_sink); + } + #[test] fn publisher_response_streaming_finalize_processes_gzip_stream() { let compressed = From 5d6fc21f66aa9beb5f1915ec0e6bcc630a7fab7e Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Thu, 16 Jul 2026 17:07:57 +0530 Subject: [PATCH 09/15] Apply rustfmt after the edition-2024 merge The merge moved core to edition 2024, whose rustfmt sorts use-imports case-sensitively (lowercase items last) and rewraps long lines. Reformat the conflict-resolved imports and the new test assertions to match; no behavior change. --- .../trusted-server-adapter-fastly/src/app.rs | 4 ++-- crates/trusted-server-core/src/publisher.rs | 20 ++++++++++++------- .../src/streaming_processor.rs | 11 ++++++---- 3 files changed, 22 insertions(+), 13 deletions(-) diff --git a/crates/trusted-server-adapter-fastly/src/app.rs b/crates/trusted-server-adapter-fastly/src/app.rs index 7cdf640a3..ae4ca421f 100644 --- a/crates/trusted-server-adapter-fastly/src/app.rs +++ b/crates/trusted-server-adapter-fastly/src/app.rs @@ -116,8 +116,8 @@ use trusted_server_core::proxy::{ handle_first_party_proxy, handle_first_party_proxy_rebuild, handle_first_party_proxy_sign, }; use trusted_server_core::publisher::{ - handle_page_bids, handle_publisher_request, handle_tsjs_dynamic, page_bids_preflight_denied, - publisher_response_into_streaming_response, AuctionDispatch, + AuctionDispatch, handle_page_bids, handle_publisher_request, handle_tsjs_dynamic, + page_bids_preflight_denied, publisher_response_into_streaming_response, }; use trusted_server_core::request_signing::{ handle_deactivate_key, handle_rotate_key, handle_trusted_server_discovery, diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index 7046346f6..85b65a026 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -22,16 +22,16 @@ use std::io::Write; use std::sync::{Arc, Mutex}; use std::time::Duration; -use brotli::enc::writer::CompressorWriter; -use brotli::enc::BrotliEncoderParams; use brotli::Decompressor; +use brotli::enc::BrotliEncoderParams; +use brotli::enc::writer::CompressorWriter; use cookie::CookieJar; use edgezero_core::body::Body as EdgeBody; use error_stack::{Report, ResultExt}; use flate2::read::{GzDecoder, ZlibDecoder}; use flate2::write::{GzEncoder, ZlibEncoder}; use futures::StreamExt as _; -use http::{header, HeaderValue, Method, Request, Response, StatusCode, Uri}; +use http::{HeaderValue, Method, Request, Response, StatusCode, Uri, header}; use crate::auction::endpoints::{ merge_auction_eids, resolve_auction_eids, resolve_client_auction_eids, @@ -60,8 +60,8 @@ use crate::price_bucket::{PriceGranularity, price_bucket}; use crate::rsc_flight::RscFlightUrlRewriter; use crate::settings::Settings; use crate::streaming_processor::{ - BodyStreamDecoder, BodyStreamEncoder, Compression, PipelineConfig, StreamProcessor, - StreamingPipeline, STREAM_CHUNK_SIZE, + BodyStreamDecoder, BodyStreamEncoder, Compression, PipelineConfig, STREAM_CHUNK_SIZE, + StreamProcessor, StreamingPipeline, }; use crate::streaming_replacer::create_url_replacer; @@ -4134,7 +4134,10 @@ mod tests { "the prefix up to must be ready before collection" ); assert!( - ad_bids_state.lock().expect("should lock bid state").is_none(), + ad_bids_state + .lock() + .expect("should lock bid state") + .is_none(), "auction must not be collected while the ready prefix is emitted" ); @@ -4148,7 +4151,10 @@ mod tests { "the held close tail must be emitted after collection" ); assert!( - ad_bids_state.lock().expect("should lock bid state").is_some(), + ad_bids_state + .lock() + .expect("should lock bid state") + .is_some(), "collection must run when the held tail is emitted" ); } diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index 32532edbf..51b8d392c 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -787,14 +787,17 @@ mod tests { } Compression::Deflate => { let mut decoder = flate2::write::ZlibDecoder::new(Vec::new()); - decoder.write_all(encoded).expect("should write deflate bytes"); + decoder + .write_all(encoded) + .expect("should write deflate bytes"); decoder.flush().expect("should flush deflate decoder"); decoder.get_ref().clone() } Compression::Brotli => { - let mut decoder = - brotli::DecompressorWriter::new(Vec::new(), STREAM_CHUNK_SIZE); - decoder.write_all(encoded).expect("should write brotli bytes"); + let mut decoder = brotli::DecompressorWriter::new(Vec::new(), STREAM_CHUNK_SIZE); + decoder + .write_all(encoded) + .expect("should write brotli bytes"); decoder.flush().expect("should flush brotli decoder"); decoder.get_ref().clone() } From 33384cd87d4547ebfa827c75798128e532923f48 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Fri, 17 Jul 2026 13:07:45 +0530 Subject: [PATCH 10/15] Tolerate trailing garbage after the final gzip member MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace flate2's MultiGzDecoder with a shared GzipStreamDecoder used by both the streaming BodyStreamDecoder and the buffered pipeline (via a Read adapter). At each member boundary the next bytes are sniffed for the gzip magic number: a match starts the next member, anything else is dropped once at least one member has fully decoded — matching GNU gzip, main's single-member tolerance, and the deflate codec. Truncated or corrupt members still error. --- .../src/streaming_processor.rs | 420 +++++++++++++++++- 1 file changed, 400 insertions(+), 20 deletions(-) diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index 51b8d392c..ebf9312a6 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -27,7 +27,7 @@ use brotli::Decompressor; use brotli::enc::BrotliEncoderParams; use brotli::enc::writer::CompressorWriter; use error_stack::{Report, ResultExt as _}; -use flate2::read::{MultiGzDecoder, ZlibDecoder}; +use flate2::read::ZlibDecoder; use flate2::write::{GzEncoder, ZlibEncoder}; use crate::error::TrustedServerError; @@ -146,8 +146,10 @@ impl StreamingPipeline

{ (Compression::Gzip, Compression::Gzip) => { // Multi-member decoder: RFC 1952 permits concatenated gzip // members, so a single-member reader would stop after the first. - // Matches the streaming `BodyStreamDecoder` gzip codec. - let decoder = MultiGzDecoder::new(input); + // Shares `GzipStreamDecoder` with the streaming `BodyStreamDecoder` + // gzip codec, so both paths tolerate trailing garbage after the + // final member the same way. + let decoder = GzipDecodeReader::new(input); let mut encoder = GzEncoder::new(output, flate2::Compression::default()); self.process_chunks(decoder, &mut encoder)?; encoder.finish().change_context(TrustedServerError::Proxy { @@ -156,7 +158,7 @@ impl StreamingPipeline

{ Ok(()) } (Compression::Gzip, Compression::None) => { - self.process_chunks(MultiGzDecoder::new(input), output) + self.process_chunks(GzipDecodeReader::new(input), output) } (Compression::Deflate, Compression::Deflate) => { let decoder = ZlibDecoder::new(input); @@ -375,7 +377,8 @@ pub(crate) const STREAM_CHUNK_SIZE: usize = 8192; /// [`flate2::Decompress`] to its [`flate2::Status::StreamEnd`] marker (the /// `write`-based zlib decoder accepts truncated input silently, so the deflate /// arm drives [`flate2::Decompress`] directly). Concatenated gzip members -/// (RFC 1952) are decoded via [`flate2::write::MultiGzDecoder`]. +/// (RFC 1952) are decoded via [`GzipStreamDecoder`], which also tolerates +/// trailing garbage after the final member the way the deflate codec does. pub(crate) struct BodyStreamDecoder { codec: BodyStreamDecoderCodec, /// Cumulative decoded byte count, shared with the codec sinks so the cap is @@ -386,7 +389,7 @@ pub(crate) struct BodyStreamDecoder { enum BodyStreamDecoderCodec { None, - Gzip(flate2::write::MultiGzDecoder), + Gzip(GzipStreamDecoder), Deflate(DeflateStreamDecoder), Brotli(Box>), } @@ -440,6 +443,235 @@ impl Write for BoundedDecodeSink { } } +/// The RFC 1952 gzip member magic number. +const GZIP_MAGIC: [u8; 2] = [0x1f, 0x8b]; + +/// Push-style multi-member gzip decoder that tolerates trailing garbage. +/// +/// Decodes concatenated gzip members (RFC 1952) one [`flate2::write::GzDecoder`] +/// at a time. At each member boundary the next bytes are sniffed for the gzip +/// magic number: a match starts the next member, anything else is treated as +/// trailing garbage and silently dropped — matching GNU gzip, the deflate +/// codec's tolerance, and the single-member `read::GzDecoder` the buffered +/// pipeline used before multi-member support. `flate2`'s own `MultiGzDecoder` +/// instead errors on any post-member bytes that do not parse as a new header. +/// +/// Tolerance never weakens integrity checks *inside* a member: a truncated or +/// corrupt member (bad CRC, bad length, incomplete deflate stream) still fails +/// at [`Self::decode`] or [`Self::finish`]. Trailing bytes that happen to start +/// with the magic number are decoded as a member and fail if they are not one. +struct GzipStreamDecoder { + state: GzipStreamState, +} + +enum GzipStreamState { + /// Decoding a gzip member (header, deflate stream, or trailer). + Member(flate2::write::GzDecoder), + /// A member fully decoded and validated; sniffing whether the next bytes + /// start another member. `magic_prefix_seen` records that the first magic + /// byte (`0x1f`) arrived, possibly at the end of an earlier chunk. + Boundary { + sink: BoundedDecodeSink, + magic_prefix_seen: bool, + }, + /// Non-member bytes followed a completed member; all further input is + /// dropped. + TrailingGarbage(BoundedDecodeSink), + /// Transient placeholder while ownership moves between states. Never + /// observable by callers. + Poisoned, +} + +impl GzipStreamDecoder { + fn new(decoded_bytes: Rc>, max_decoded_bytes: usize) -> Self { + Self { + state: GzipStreamState::Member(flate2::write::GzDecoder::new(BoundedDecodeSink::new( + decoded_bytes, + max_decoded_bytes, + ))), + } + } + + /// Decode one compressed chunk, returning the decoded bytes it produced. + /// + /// # Errors + /// + /// Returns an error if a member is corrupt, its trailer fails validation, + /// or the decoded budget is exceeded. + fn decode(&mut self, chunk: &[u8]) -> io::Result> { + let mut input = chunk; + while !input.is_empty() { + match &mut self.state { + GzipStreamState::Member(decoder) => { + let consumed = decoder.write(input)?; + if consumed > 0 { + input = &input[consumed..]; + continue; + } + // A zero-byte write on non-empty input means the member + // (deflate stream plus trailer) is complete: validate the + // trailer and start sniffing for the next member. + let GzipStreamState::Member(decoder) = + std::mem::replace(&mut self.state, GzipStreamState::Poisoned) + else { + unreachable!("state was matched as Member above"); + }; + let sink = decoder.finish()?; + self.state = GzipStreamState::Boundary { + sink, + magic_prefix_seen: false, + }; + } + GzipStreamState::Boundary { + magic_prefix_seen, .. + } => { + let expected = if *magic_prefix_seen { + GZIP_MAGIC[1] + } else { + GZIP_MAGIC[0] + }; + if input[0] != expected { + self.enter_trailing_garbage(); + continue; + } + if !*magic_prefix_seen { + *magic_prefix_seen = true; + input = &input[1..]; + continue; + } + // Full magic number seen: start the next member, replaying + // the two magic bytes consumed during sniffing. + input = &input[1..]; + let GzipStreamState::Boundary { sink, .. } = + std::mem::replace(&mut self.state, GzipStreamState::Poisoned) + else { + unreachable!("state was matched as Boundary above"); + }; + let mut decoder = flate2::write::GzDecoder::new(sink); + decoder.write_all(&GZIP_MAGIC)?; + self.state = GzipStreamState::Member(decoder); + } + GzipStreamState::TrailingGarbage(_) => break, + GzipStreamState::Poisoned => { + unreachable!("poisoned state is never left in place across iterations") + } + } + } + Ok(self.take_decoded()) + } + + /// Finalize the stream at end of input, returning any remaining decoded bytes. + /// + /// # Errors + /// + /// Returns an error if the final member is truncated or its trailer fails + /// validation. Trailing garbage after a completed member is not an error. + fn finish(&mut self) -> io::Result> { + match &mut self.state { + GzipStreamState::Member(decoder) => { + decoder.try_finish()?; + Ok(std::mem::take(&mut decoder.get_mut().buffer)) + } + GzipStreamState::Boundary { + sink, + magic_prefix_seen, + } => { + if *magic_prefix_seen { + log::warn!( + "gzip body ends with a lone trailing byte after the final member; ignoring it" + ); + } + Ok(std::mem::take(&mut sink.buffer)) + } + GzipStreamState::TrailingGarbage(sink) => Ok(std::mem::take(&mut sink.buffer)), + GzipStreamState::Poisoned => { + unreachable!("poisoned state is never left in place across calls") + } + } + } + + /// Drop the rest of the input, keeping the sink so already-decoded bytes + /// can still be drained. + fn enter_trailing_garbage(&mut self) { + log::warn!("gzip body has trailing bytes after the final member; ignoring them"); + let sink = match std::mem::replace(&mut self.state, GzipStreamState::Poisoned) { + GzipStreamState::Boundary { sink, .. } => sink, + _ => unreachable!("trailing garbage is only entered from the boundary state"), + }; + self.state = GzipStreamState::TrailingGarbage(sink); + } + + /// Take the decoded bytes accumulated in the sink so far. + fn take_decoded(&mut self) -> Vec { + match &mut self.state { + GzipStreamState::Member(decoder) => std::mem::take(&mut decoder.get_mut().buffer), + GzipStreamState::Boundary { sink, .. } | GzipStreamState::TrailingGarbage(sink) => { + std::mem::take(&mut sink.buffer) + } + GzipStreamState::Poisoned => { + unreachable!("poisoned state is never left in place across calls") + } + } + } +} + +/// [`Read`] adapter that decodes multi-member gzip through [`GzipStreamDecoder`]. +/// +/// Used by the buffered [`StreamingPipeline`] so it shares the streaming +/// path's trailing-garbage tolerance instead of erroring like +/// `flate2::read::MultiGzDecoder`. The decode budget is unbounded here for +/// parity with the reader it replaces — the buffered path bounds output +/// downstream (e.g. via `BoundedWriter`). +struct GzipDecodeReader { + input: R, + decoder: GzipStreamDecoder, + raw: Vec, + decoded: Vec, + position: usize, + finished: bool, +} + +impl GzipDecodeReader { + fn new(input: R) -> Self { + Self { + input, + decoder: GzipStreamDecoder::new(Rc::new(Cell::new(0)), usize::MAX), + raw: vec![0_u8; STREAM_CHUNK_SIZE], + decoded: Vec::new(), + position: 0, + finished: false, + } + } +} + +impl Read for GzipDecodeReader { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + if buf.is_empty() { + return Ok(0); + } + loop { + if self.position < self.decoded.len() { + let available = self.decoded.len() - self.position; + let amount = available.min(buf.len()); + buf[..amount].copy_from_slice(&self.decoded[self.position..self.position + amount]); + self.position += amount; + return Ok(amount); + } + if self.finished { + return Ok(0); + } + let read = self.input.read(&mut self.raw)?; + if read == 0 { + self.decoded = self.decoder.finish()?; + self.finished = true; + } else { + self.decoded = self.decoder.decode(&self.raw[..read])?; + } + self.position = 0; + } + } +} + /// Charge `len` decoded bytes against `decoded_bytes`, erroring if the /// cumulative total would exceed `max_decoded_bytes`. fn charge_decoded( @@ -578,8 +810,9 @@ impl BodyStreamDecoder { let decoded_bytes = Rc::new(Cell::new(0usize)); let codec = match compression { Compression::None => BodyStreamDecoderCodec::None, - Compression::Gzip => BodyStreamDecoderCodec::Gzip(flate2::write::MultiGzDecoder::new( - BoundedDecodeSink::new(Rc::clone(&decoded_bytes), max_decoded_bytes), + Compression::Gzip => BodyStreamDecoderCodec::Gzip(GzipStreamDecoder::new( + Rc::clone(&decoded_bytes), + max_decoded_bytes, )), Compression::Deflate => BodyStreamDecoderCodec::Deflate(DeflateStreamDecoder::new( Rc::clone(&decoded_bytes), @@ -611,15 +844,13 @@ impl BodyStreamDecoder { Ok(chunk) } BodyStreamDecoderCodec::Gzip(decoder) => { - decoder - .write_all(&chunk) + // The sink charges the decoded bytes as the decoder writes them. + let decoded = decoder + .decode(&chunk) .change_context(TrustedServerError::Proxy { message: "Failed to decode gzip publisher body chunk".to_string(), })?; - // The sink charged the decoded bytes during `write_all`. - Ok(bytes::Bytes::from(std::mem::take( - &mut decoder.get_mut().buffer, - ))) + Ok(bytes::Bytes::from(decoded)) } BodyStreamDecoderCodec::Deflate(decoder) => { Ok(bytes::Bytes::from(decoder.decode(&chunk)?)) @@ -641,12 +872,9 @@ impl BodyStreamDecoder { match &mut self.codec { BodyStreamDecoderCodec::None => Ok(Vec::new()), BodyStreamDecoderCodec::Gzip(decoder) => { - decoder - .try_finish() - .change_context(TrustedServerError::Proxy { - message: "Failed to finalize gzip publisher body decoder".to_string(), - })?; - Ok(std::mem::take(&mut decoder.get_mut().buffer)) + decoder.finish().change_context(TrustedServerError::Proxy { + message: "Failed to finalize gzip publisher body decoder".to_string(), + }) } BodyStreamDecoderCodec::Deflate(decoder) => decoder.finish(), BodyStreamDecoderCodec::Brotli(decoder) => { @@ -1039,6 +1267,117 @@ mod tests { ); } + #[test] + fn body_stream_decoder_ignores_gzip_trailing_bytes() { + let mut with_trailing = gzip_member(b"gzip payload"); + with_trailing.extend_from_slice(b"junk"); + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, usize::MAX); + + let mut decoded = decoder + .decode_chunk(bytes::Bytes::from(with_trailing)) + .expect("complete gzip member should decode") + .to_vec(); + decoded.extend( + decoder + .finish() + .expect("trailing bytes after the final member should be ignored"), + ); + + assert_eq!( + decoded, b"gzip payload", + "should decode the payload and drop trailing junk" + ); + } + + #[test] + fn body_stream_decoder_ignores_gzip_trailing_bytes_split_across_chunks() { + let mut with_trailing = gzip_member(b"first member "); + with_trailing.extend(gzip_member(b"second member")); + with_trailing.extend_from_slice(b"trailing garbage"); + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, usize::MAX); + + let mut decoded = Vec::new(); + // Small pieces exercise every split point, including inside the + // member boundary sniff and inside the garbage itself. + for piece in with_trailing.chunks(3) { + decoded.extend( + decoder + .decode_chunk(bytes::Bytes::copy_from_slice(piece)) + .expect("gzip members followed by garbage should decode"), + ); + } + decoded.extend( + decoder + .finish() + .expect("trailing bytes after the final member should be ignored"), + ); + + assert_eq!( + decoded, b"first member second member", + "should decode every member and drop the trailing garbage" + ); + } + + #[test] + fn body_stream_decoder_ignores_gzip_lone_trailing_magic_prefix_byte() { + // A single 0x1f after the final member could be the start of another + // member; at end of input it must be treated as garbage, not an error. + let mut with_trailing = gzip_member(b"gzip payload"); + with_trailing.push(0x1f); + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, usize::MAX); + + let decoded = decoder + .decode_chunk(bytes::Bytes::from(with_trailing)) + .expect("complete gzip member should decode"); + decoder + .finish() + .expect("a lone trailing magic prefix byte should be ignored"); + + assert_eq!( + decoded.as_ref(), + b"gzip payload", + "should decode the payload and drop the lone trailing byte" + ); + } + + #[test] + fn body_stream_decoder_rejects_trailing_bytes_resembling_gzip_member() { + // Trailing bytes that start with the gzip magic number are decoded as + // a member, so a corrupt pseudo-member still fails: tolerance is + // best-effort and never accepts data that claims to be a member. + let mut with_trailing = gzip_member(b"gzip payload"); + with_trailing.extend_from_slice(&[0x1f, 0x8b, b'j', b'u', b'n', b'k']); + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, usize::MAX); + + let result = decoder + .decode_chunk(bytes::Bytes::from(with_trailing)) + .and_then(|_| decoder.finish().map(|_| ())); + + let err = result.expect_err("garbage that resembles a gzip member must fail"); + assert!( + format!("{err:?}").contains("gzip"), + "should report a gzip decode failure: {err:?}" + ); + } + + #[test] + fn body_stream_decoder_rejects_truncated_gzip_stream() { + let compressed = gzip_member(b"gzip payload that spans more than one deflate block"); + let truncated = &compressed[..compressed.len() / 2]; + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, usize::MAX); + decoder + .decode_chunk(bytes::Bytes::copy_from_slice(truncated)) + .expect("partial gzip input should decode incrementally"); + + let err = decoder + .finish() + .expect_err("a truncated gzip member must fail at finalization"); + assert!( + format!("{err:?}").contains("finalize gzip"), + "should report the gzip finalization failure: {err:?}" + ); + } + /// Verify that `lol_html` fragments text nodes when input chunks split /// mid-text-node. Script rewriters must be fragment-safe — they accumulate /// text fragments internally until `is_last_in_text_node` is true. @@ -1356,6 +1695,47 @@ mod tests { ); } + #[test] + fn test_gzip_pipeline_ignores_trailing_bytes_after_final_member() { + use flate2::read::GzDecoder; + use std::io::Read as _; + + // Arrange + let mut compressed_input = gzip_member(b"hello world"); + compressed_input.extend_from_slice(b"junk"); + + let replacer = StreamingReplacer::new(vec![Replacement { + find: "hello".to_owned(), + replace_with: "hi".to_owned(), + }]); + + let config = PipelineConfig { + input_compression: Compression::Gzip, + output_compression: Compression::Gzip, + chunk_size: 8192, + }; + + let mut pipeline = StreamingPipeline::new(config, replacer); + let mut output = Vec::new(); + + // Act + pipeline + .process(&*compressed_input, &mut output) + .expect("trailing bytes after the final gzip member should be ignored"); + + // Assert + let mut decompressed = Vec::new(); + GzDecoder::new(&*output) + .read_to_end(&mut decompressed) + .expect("should decompress output"); + + assert_eq!( + String::from_utf8(decompressed).expect("should be valid UTF-8"), + "hi world", + "should process the payload and drop trailing junk" + ); + } + #[test] fn test_gzip_to_none_produces_correct_output() { use flate2::write::GzEncoder; From 911268fb4815ba8a0629c88286a9f1144f8860a5 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Fri, 17 Jul 2026 13:08:02 +0530 Subject: [PATCH 11/15] Borrow post-release chunks in the auction hold step Once the close-body hold is released, decoded chunks were still copied via to_vec() before processing. Use Cow so the post-release path borrows the chunk and only the held path allocates. --- crates/trusted-server-core/src/publisher.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index 85b65a026..e03a96c4c 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -18,6 +18,7 @@ //! into any [`Write`] (a `Vec` for buffered routes, a streaming writer for //! the streaming route). It is not a content-rewriting concern. +use std::borrow::Cow; use std::io::Write; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -726,10 +727,11 @@ async fn hold_step_decoded_chunk( collect_refs: &AuctionHoldCollectRefs<'_>, ) -> Result> { let mut ready = Vec::new(); - let bytes = match state.hold.as_mut() { - // Once the hold has been released the chunk streams straight through. - None => chunk.to_vec(), - Some(hold_buffer) => hold_buffer.push(chunk), + let bytes: Cow<'_, [u8]> = match state.hold.as_mut() { + // Once the hold has been released the chunk streams straight through, + // borrowed rather than copied. + None => Cow::Borrowed(chunk), + Some(hold_buffer) => Cow::Owned(hold_buffer.push(chunk)), }; match process_and_encode_chunk(processor, encoder, &bytes, false, "Failed to process chunk") { Ok(Some(encoded)) => ready.push(encoded), From 666b3cb7816685aaef9f6daeb8dddc2231669778 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Fri, 17 Jul 2026 13:08:02 +0530 Subject: [PATCH 12/15] Document the post-processor buffering exception as out of scope HTML post-processor configs (the nextjs integration) accumulate the full rewritten document until origin EOF, so those pages do not stream body bytes early. Record the limitation and the intended follow-up (an up-front or streaming should_process gate) in the plan's Out of scope section. --- .../plans/2026-07-08-true-origin-streaming-fastly.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/docs/superpowers/plans/2026-07-08-true-origin-streaming-fastly.md b/docs/superpowers/plans/2026-07-08-true-origin-streaming-fastly.md index bc1983a97..677e68c85 100644 --- a/docs/superpowers/plans/2026-07-08-true-origin-streaming-fastly.md +++ b/docs/superpowers/plans/2026-07-08-true-origin-streaming-fastly.md @@ -22,6 +22,15 @@ In scope: Out of scope: +- HTML post-processor configs (the `nextjs` integration). When any + `IntegrationHtmlPostProcessor` is registered, `HtmlWithPostProcessing` + accumulates the full rewritten document and runs post-processors at origin + EOF, so the lazy stream emits no body bytes until the whole origin transfer + completes — even for pages where `should_process()` would return `false`. + Headers still commit early, but first byte/FCP tracks origin EOF for that + configuration; #849's objective is unmet there. The eventual fix is an + up-front or streaming `should_process` gate so non-RSC pages skip + accumulation entirely (follow-up issue). - Cloudflare origin streaming. Current adapter rejects `PlatformHttpRequest::stream_response`. - Spin streaming. Current adapter and upstream EdgeZero Spin conversion are buffered/blocking issues. - Axum client streaming. Axum is dev-only and has `LocalBoxStream`/`Send` constraints. From 69a2e70280d4f8d6c54c665673bb2ac5787d580d Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Wed, 22 Jul 2026 14:44:11 +0530 Subject: [PATCH 13/15] Bound streaming decode allocation and fix buffered gzip decoding Address the streaming-pipeline review findings. Deflate decode and finish now inflate through a fixed-size scratch buffer whose write window is bounded by the remaining decode budget, and charge produced bytes before appending them. `decompress_vec` previously let the output `Vec` double and be filled to capacity before the cap was checked, so a 16 MiB limit could peak at 32 MiB. `GzipDecodeReader` takes a decode budget instead of an unconditional `usize::MAX`; `StreamingPipeline::with_max_decoded_bytes` threads it and the buffered publisher path passes `max_buffered_body_bytes`, so a gzip bomb is rejected mid-decode rather than fully materialized before the bounded writer can act. The buffered auction hold path decodes concatenated gzip members through `GzipDecodeReader` instead of `flate2::read::GzDecoder`, which decoded only the first member and silently dropped the rest, potentially including `` and trailing markup on Axum, Cloudflare and Spin. `GzipStreamDecoder`'s poisoned-state arms return an error instead of `unreachable!()`, so a caller that retries after a mid-transition failure cannot abort the Wasm instance. Tests cover the deflate peak buffer capacity at the cap, a deflate bomb, a low-limit gzip bomb through the reader, a buffered auction body split across two gzip members, and early emission: the real lazy response body is polled once against an origin that yields one chunk and then stays pending, for both compressed HTML and the auction-hold prefix. Configuration docs now describe the raw and decoded streaming caps, mid-stream truncation after headers are committed, adapter differences, and the removal of the old 10 MiB Fastly origin-body ceiling. --- crates/trusted-server-core/src/publisher.rs | 260 +++++++++++++++++- .../src/streaming_processor.rs | 224 +++++++++++++-- docs/guide/configuration.md | 55 ++-- 3 files changed, 481 insertions(+), 58 deletions(-) diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index f0182c1f5..1d2f24759 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -29,7 +29,7 @@ use brotli::enc::writer::CompressorWriter; use cookie::CookieJar; use edgezero_core::body::Body as EdgeBody; use error_stack::{Report, ResultExt}; -use flate2::read::{GzDecoder, ZlibDecoder}; +use flate2::read::ZlibDecoder; use flate2::write::{GzEncoder, ZlibEncoder}; use futures::StreamExt as _; use http::{HeaderValue, Method, Request, Response, StatusCode, Uri, header}; @@ -61,8 +61,8 @@ use crate::price_bucket::{PriceGranularity, price_bucket}; use crate::rsc_flight::RscFlightUrlRewriter; use crate::settings::Settings; use crate::streaming_processor::{ - BodyStreamDecoder, BodyStreamEncoder, Compression, PipelineConfig, STREAM_CHUNK_SIZE, - StreamProcessor, StreamingPipeline, + BodyStreamDecoder, BodyStreamEncoder, Compression, GzipDecodeReader, PipelineConfig, + STREAM_CHUNK_SIZE, StreamProcessor, StreamingPipeline, }; use crate::streaming_replacer::create_url_replacer; @@ -434,6 +434,10 @@ fn process_response_streaming( output_compression: compression, chunk_size: 8192, }; + // Bound the gzip decode path to the same ceiling the buffered writer + // enforces, so a gzip bomb is rejected mid-decode instead of materializing + // its full expansion before the downstream writer can reject it. + let max_decoded_bytes = params.settings.publisher.max_buffered_body_bytes; if is_html { let processor = create_html_stream_processor( @@ -445,7 +449,9 @@ fn process_response_streaming( params.ad_slots_script.map(str::to_string), params.ad_bids_state.clone(), )?; - StreamingPipeline::new(config, processor).process(body_as_reader(body)?, output)?; + StreamingPipeline::new(config, processor) + .with_max_decoded_bytes(max_decoded_bytes) + .process(body_as_reader(body)?, output)?; } else if is_rsc_flight { // RSC Flight responses are length-prefixed (T rows). A naive string replacement will // corrupt the stream by changing byte lengths without updating the prefixes. @@ -455,7 +461,9 @@ fn process_response_streaming( params.request_host, params.request_scheme, ); - StreamingPipeline::new(config, processor).process(body_as_reader(body)?, output)?; + StreamingPipeline::new(config, processor) + .with_max_decoded_bytes(max_decoded_bytes) + .process(body_as_reader(body)?, output)?; } else { let replacer = create_url_replacer( params.origin_host, @@ -463,7 +471,9 @@ fn process_response_streaming( params.request_host, params.request_scheme, ); - StreamingPipeline::new(config, replacer).process(body_as_reader(body)?, output)?; + StreamingPipeline::new(config, replacer) + .with_max_decoded_bytes(max_decoded_bytes) + .process(body_as_reader(body)?, output)?; } Ok(()) @@ -1946,11 +1956,18 @@ async fn stream_html_with_auction_hold( .await; } + // Bound the gzip decode budget to the same ceiling the buffered writer + // enforces, matching the streaming arm above and the no-hold buffered path. + let max_body_bytes = ctx.settings.publisher.max_buffered_body_bytes; let body = body_as_reader(body)?; match compression { Compression::None => body_close_hold_loop(body, output, processor, ctx).await, Compression::Gzip => { - let decoder = GzDecoder::new(body); + // `GzipDecodeReader` decodes concatenated gzip members (RFC 1952) + // and bounds decoded output, unlike `flate2::read::GzDecoder`, which + // silently drops every member after the first — dropping trailing + // markup (potentially including ``) on buffered adapters. + let decoder = GzipDecodeReader::new(body, max_body_bytes); let mut encoder = GzEncoder::new(&mut *output, flate2::Compression::default()); body_close_hold_loop(decoder, &mut encoder, processor, ctx).await?; encoder.finish().change_context(TrustedServerError::Proxy { @@ -5795,6 +5812,72 @@ mod tests { }); } + #[test] + fn stream_publisher_body_async_auction_hold_decodes_multi_member_gzip_buffered() { + futures::executor::block_on(async { + let settings = create_test_settings(); + let registry = + IntegrationRegistry::new(&settings).expect("should create integration registry"); + let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); + let services = noop_services(); + let state = Arc::new(Mutex::new(None)); + let mut params = OwnedProcessResponseParams { + content_encoding: "gzip".to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/html; charset=utf-8".to_string(), + ad_slots_script: Some( + r#""# + .to_string(), + ), + ad_bids_state: state, + auction_observation: None, + auction_request: Some(test_auction_request()), + dispatched_auction: Some(DispatchedAuction::empty_for_test( + test_auction_request(), + 10, + )), + price_granularity: crate::price_bucket::PriceGranularity::default(), + }; + // The `` that triggers bid injection lives in the SECOND gzip + // member. `flate2::read::GzDecoder` decodes only the first member, so + // this buffered `Body::Once` body (the non-stream auction arm) proves + // the multi-member decoder now reads every member. + let mut compressed = gzip_encode(b"hello"); + compressed.extend(gzip_encode(b"")); + let body = EdgeBody::from(compressed); + let mut output = Vec::new(); + + stream_publisher_body_async( + body, + &mut output, + &mut params, + &settings, + ®istry, + &orchestrator, + &services, + ) + .await + .expect("buffered multi-member gzip auction body should process"); + + let html = String::from_utf8(gzip_decode(&output)).expect("should be valid UTF-8"); + assert!( + html.contains("hello"), + "should decode the first gzip member. Got: {html}" + ); + assert!( + html.contains(""), + "should decode the second gzip member that a single-member decoder drops. Got: {html}" + ); + assert!( + html.contains(".bids=JSON.parse"), + "should inject bids before the carried in the second member. Got: {html}" + ); + }); + } + #[test] fn stream_publisher_body_async_processes_non_html_stream_after_auction_collect() { futures::executor::block_on(async { @@ -5924,6 +6007,169 @@ mod tests { ); } + /// An origin body stream that yields `chunk` once and then stays `Pending` + /// forever without ever signalling EOF — modelling an origin that has sent + /// the document head but not yet finished the response. + fn origin_chunk_then_pending(chunk: bytes::Bytes) -> EdgeBody { + let mut sent = false; + EdgeBody::from_stream(futures::stream::poll_fn(move |_cx| { + if sent { + return std::task::Poll::Pending; + } + sent = true; + std::task::Poll::Ready(Some(Ok::<_, io::Error>(chunk.clone()))) + })) + } + + /// Poll a lazy `Body::Stream` exactly once and return the first emitted + /// chunk. Panics if the first poll is `Pending` (nothing emitted before the + /// origin would need its next chunk) or the stream ends. + fn first_lazy_body_chunk(body: EdgeBody) -> bytes::Bytes { + let mut stream = body + .into_stream() + .expect("streaming finalize should keep a lazy Body::Stream"); + let waker = futures::task::noop_waker(); + let mut cx = std::task::Context::from_waker(&waker); + let next = futures::StreamExt::next(&mut stream); + let mut next = std::pin::pin!(next); + match std::future::Future::poll(next.as_mut(), &mut cx) { + std::task::Poll::Ready(Some(Ok(bytes))) => bytes, + std::task::Poll::Ready(other) => { + panic!("first poll must emit a chunk, got Ready({other:?})") + } + std::task::Poll::Pending => { + panic!("first poll must emit rewritten content before origin EOF, got Pending") + } + } + } + + fn streaming_finalize_response(params: OwnedProcessResponseParams, body: EdgeBody) -> EdgeBody { + let settings = Arc::new(create_test_settings()); + let registry = Arc::new( + IntegrationRegistry::new(&settings).expect("should create integration registry"), + ); + let orchestrator = Arc::new(AuctionOrchestrator::new(settings.auction.clone())); + let services = noop_services(); + let response = Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, ¶ms.content_type) + .body(EdgeBody::empty()) + .expect("should build response"); + let publisher_response = PublisherResponse::Stream { + response, + body, + params: Box::new(params), + }; + futures::executor::block_on(publisher_response_into_streaming_response( + publisher_response, + &Method::GET, + Arc::clone(&settings), + registry.as_ref(), + orchestrator, + services, + )) + .expect("should build streaming response") + .into_body() + } + + fn html_stream_params( + content_encoding: &str, + dispatched_auction: Option, + ) -> OwnedProcessResponseParams { + OwnedProcessResponseParams { + content_encoding: content_encoding.to_string(), + origin_host: "origin.example.com".to_string(), + origin_url: "https://origin.example.com".to_string(), + request_host: "proxy.example.com".to_string(), + request_scheme: "https".to_string(), + content_type: "text/html; charset=utf-8".to_string(), + ad_slots_script: Some( + r#""# + .to_string(), + ), + ad_bids_state: Arc::new(Mutex::new(None)), + auction_observation: None, + auction_request: dispatched_auction.as_ref().map(|_| test_auction_request()), + dispatched_auction, + price_granularity: crate::price_bucket::PriceGranularity::default(), + } + } + + #[test] + fn streaming_finalize_emits_compressed_html_before_origin_eof() { + // The FCP regression from #849: the lazy body must emit its first + // rewritten, browser-decodable chunk as soon as the origin delivers one + // block — not buffer the whole (compressed) response until EOF. The + // origin here sends one gzip member and then stays Pending; the first + // poll must still yield gzip bytes that decode to the streamed content. + // The page exceeds the deflate decoder's internal output buffer so the + // first decoded block flushes downstream before the (never-arriving) + // EOF, exactly as a real page does. + let mut page = b"".to_vec(); + for i in 0..4000 { + page.extend(format!("

hello world paragraph {i}

").bytes()); + } + let params = html_stream_params("gzip", None); + let body = streaming_finalize_response( + params, + origin_chunk_then_pending(bytes::Bytes::from(gzip_encode(&page))), + ); + + let first = first_lazy_body_chunk(body); + assert!( + !first.is_empty(), + "first emitted gzip chunk must not be empty" + ); + + // Decode the flushed (not finished) gzip prefix the way a browser would + // decode bytes received so far while the response is still open. + let mut decoder = flate2::write::MultiGzDecoder::new(Vec::new()); + decoder + .write_all(&first) + .expect("first chunk must be valid gzip"); + decoder.flush().expect("should flush gzip decoder"); + let decoded = String::from_utf8(decoder.get_ref().clone()).expect("should be valid UTF-8"); + assert!( + decoded.contains("hello world"), + "first poll must emit decodable streamed content before EOF. Got: {decoded}" + ); + } + + #[test] + fn streaming_finalize_auction_hold_emits_prefix_before_origin_eof() { + // The auction-hold path must stream the document prefix (up to the held + // `` tail) before the origin finishes and before the auction is + // collected — otherwise the hold reintroduces the FCP regression. The + // origin sends the head/body prefix (no ``) then stays Pending. + let page = b"

hello

more streamed content here

"; + let params = html_stream_params( + "", + Some(DispatchedAuction::empty_for_test( + test_auction_request(), + 10, + )), + ); + let body = streaming_finalize_response( + params, + origin_chunk_then_pending(bytes::Bytes::from(&page[..])), + ); + + let first = first_lazy_body_chunk(body); + let html = String::from_utf8(first.to_vec()).expect("should be valid UTF-8"); + assert!( + html.contains("hello"), + "auction-hold path must stream the prefix before EOF. Got: {html}" + ); + assert!( + html.contains(".adSlots=JSON.parse"), + "prefix must carry the injected (rewritten) head before EOF. Got: {html}" + ); + assert!( + !html.contains(".bids=JSON.parse"), + "bids inject only at after collection, which the first poll must not wait for. Got: {html}" + ); + } + // (method, status, expected Content-Length after bodiless normalization). // 204 forbids Content-Length (removed); 205 must advertise a zero-length // body; HEAD and 304 legitimately advertise the GET representation length. diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index ebf9312a6..0d1cc7095 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -112,6 +112,11 @@ impl Default for PipelineConfig { pub struct StreamingPipeline { config: PipelineConfig, processor: P, + /// Cumulative decoded-byte ceiling for the gzip decode path. Defaults to + /// `usize::MAX` (unbounded); set via [`Self::with_max_decoded_bytes`] on the + /// buffered publisher path so a gzip bomb is rejected mid-decode instead of + /// materializing its full expansion. + max_decoded_bytes: usize, } impl StreamingPipeline

{ @@ -121,7 +126,25 @@ impl StreamingPipeline

{ /// /// No errors are returned by this constructor. pub fn new(config: PipelineConfig, processor: P) -> Self { - Self { config, processor } + Self { + config, + processor, + max_decoded_bytes: usize::MAX, + } + } + + /// Bound the cumulative decoded output of the gzip decode path. + /// + /// Only the gzip reader materializes decoded bytes ahead of the downstream + /// writer; the deflate and brotli read decoders already emit into the + /// caller's fixed-size buffer. Callers that buffer output under a configured + /// ceiling (e.g. `publisher.max_buffered_body_bytes`) should pass that + /// ceiling here so decode rejection cannot be preceded by a large + /// allocation. + #[must_use] + pub fn with_max_decoded_bytes(mut self, max_decoded_bytes: usize) -> Self { + self.max_decoded_bytes = max_decoded_bytes; + self } /// Process a stream from input to output @@ -149,7 +172,7 @@ impl StreamingPipeline

{ // Shares `GzipStreamDecoder` with the streaming `BodyStreamDecoder` // gzip codec, so both paths tolerate trailing garbage after the // final member the same way. - let decoder = GzipDecodeReader::new(input); + let decoder = GzipDecodeReader::new(input, self.max_decoded_bytes); let mut encoder = GzEncoder::new(output, flate2::Compression::default()); self.process_chunks(decoder, &mut encoder)?; encoder.finish().change_context(TrustedServerError::Proxy { @@ -158,7 +181,7 @@ impl StreamingPipeline

{ Ok(()) } (Compression::Gzip, Compression::None) => { - self.process_chunks(GzipDecodeReader::new(input), output) + self.process_chunks(GzipDecodeReader::new(input, self.max_decoded_bytes), output) } (Compression::Deflate, Compression::Deflate) => { let decoder = ZlibDecoder::new(input); @@ -553,7 +576,13 @@ impl GzipStreamDecoder { } GzipStreamState::TrailingGarbage(_) => break, GzipStreamState::Poisoned => { - unreachable!("poisoned state is never left in place across iterations") + // A prior call errored mid-transition, leaving the state + // poisoned. No current caller retries after an error, but + // returning an error (rather than `unreachable!`, which + // aborts the Wasm instance) keeps a re-entrant caller safe. + return Err(io::Error::other( + "gzip decoder previously failed; stream is unusable", + )); } } } @@ -585,7 +614,12 @@ impl GzipStreamDecoder { } GzipStreamState::TrailingGarbage(sink) => Ok(std::mem::take(&mut sink.buffer)), GzipStreamState::Poisoned => { - unreachable!("poisoned state is never left in place across calls") + // See the `Poisoned` arm in `decode`: a prior mid-transition + // error left the decoder unusable; error cleanly rather than + // aborting the Wasm instance via `unreachable!`. + Err(io::Error::other( + "gzip decoder previously failed; stream is unusable", + )) } } } @@ -617,12 +651,17 @@ impl GzipStreamDecoder { /// [`Read`] adapter that decodes multi-member gzip through [`GzipStreamDecoder`]. /// -/// Used by the buffered [`StreamingPipeline`] so it shares the streaming -/// path's trailing-garbage tolerance instead of erroring like -/// `flate2::read::MultiGzDecoder`. The decode budget is unbounded here for -/// parity with the reader it replaces — the buffered path bounds output -/// downstream (e.g. via `BoundedWriter`). -struct GzipDecodeReader { +/// Used by the buffered [`StreamingPipeline`] and the buffered auction hold path +/// so both share the streaming path's trailing-garbage tolerance and multi-member +/// support instead of erroring (or dropping members) like +/// `flate2::read::MultiGzDecoder`/`GzDecoder`. +/// +/// `max_decoded_bytes` bounds cumulative decoded output: each compressed block's +/// expansion is charged against the budget from inside [`GzipStreamDecoder`]'s +/// bounded sink, so a decompression bomb is rejected mid-decode rather than fully +/// materialized before a downstream writer (e.g. `BoundedWriter`) can act. Pass +/// `usize::MAX` where an upstream/downstream stage already bounds the size. +pub(crate) struct GzipDecodeReader { input: R, decoder: GzipStreamDecoder, raw: Vec, @@ -632,10 +671,10 @@ struct GzipDecodeReader { } impl GzipDecodeReader { - fn new(input: R) -> Self { + pub(crate) fn new(input: R, max_decoded_bytes: usize) -> Self { Self { input, - decoder: GzipStreamDecoder::new(Rc::new(Cell::new(0)), usize::MAX), + decoder: GzipStreamDecoder::new(Rc::new(Cell::new(0)), max_decoded_bytes), raw: vec![0_u8; STREAM_CHUNK_SIZE], decoded: Vec::new(), position: 0, @@ -695,6 +734,28 @@ fn charge_decoded( Ok(()) } +/// Append `data` to `output`, growing capacity with amortized doubling but +/// never past `max_decoded_bytes + 1`. +/// +/// A plain `Vec` doubles on growth, so accumulating a decode up to an N-byte +/// ceiling can leave the buffer with ~2N capacity. Capping the reservation at +/// the decoded limit keeps a decompression bomb from ballooning the accumulator +/// to roughly twice the configured ceiling before the cumulative-byte charge +/// rejects it. +fn extend_capped(output: &mut Vec, data: &[u8], max_decoded_bytes: usize) { + let needed = output.len() + data.len(); + if needed > output.capacity() { + let target = output + .capacity() + .saturating_mul(2) + .max(needed) + .min(max_decoded_bytes.saturating_add(1)) + .max(needed); + output.reserve_exact(target - output.len()); + } + output.extend_from_slice(data); +} + /// Streaming zlib decoder that tracks whether the stream reached its end /// marker, so truncated deflate bodies fail at finalization. struct DeflateStreamDecoder { @@ -719,6 +780,17 @@ impl DeflateStreamDecoder { charge_decoded(&self.decoded_bytes, self.max_decoded_bytes, len) } + /// The number of scratch bytes the inflater may write this step: the + /// remaining budget plus one — so a decompression bomb produces exactly one + /// byte past the ceiling, which [`Self::charge`] then rejects — capped at + /// `scratch_len`. + fn decode_window(&self, scratch_len: usize) -> usize { + let remaining = self + .max_decoded_bytes + .saturating_sub(self.decoded_bytes.get()); + remaining.saturating_add(1).min(scratch_len) + } + /// Decode as much of `chunk` as possible, draining any output the inflater /// can still produce once all input is consumed. /// @@ -729,34 +801,46 @@ impl DeflateStreamDecoder { /// makes no further progress, so those pending bytes are never stranded and /// a valid stream is not mistaken for a truncated one at `finish`. fn decode(&mut self, chunk: &[u8]) -> Result, Report> { - let mut output = Vec::with_capacity(STREAM_CHUNK_SIZE); + let mut output = Vec::new(); let mut offset = 0usize; + // The inflater writes into this fixed-size scratch buffer, never a + // growing `Vec`. `decompress_vec` fills a doubled `Vec` to its full + // capacity before the produced bytes can be charged, so a bomb could + // materialize ~2× the cap before rejection. A fixed scratch window + // bounded by the remaining budget produces at most one byte past the + // ceiling per step, which `charge` rejects before it is appended. + let mut scratch = [0_u8; STREAM_CHUNK_SIZE]; // Trailing bytes after the zlib end marker are ignored, matching the // read-based decoder used by the buffered pipeline. while !self.stream_ended { - if output.len() == output.capacity() { - output.reserve(STREAM_CHUNK_SIZE); - } + let window = self.decode_window(scratch.len()); let before_in = self.decompress.total_in(); let before_out = self.decompress.total_out(); let status = self .decompress - .decompress_vec(&chunk[offset..], &mut output, flate2::FlushDecompress::None) + .decompress( + &chunk[offset..], + &mut scratch[..window], + flate2::FlushDecompress::None, + ) .change_context(TrustedServerError::Proxy { message: "Failed to decode deflate publisher body chunk".to_string(), })?; let consumed = (self.decompress.total_in() - before_in) as usize; let produced = (self.decompress.total_out() - before_out) as usize; offset += consumed; + // Charge before appending so a bomb errors before its bytes are + // buffered. self.charge(produced)?; + extend_capped(&mut output, &scratch[..produced], self.max_decoded_bytes); match status { flate2::Status::StreamEnd => self.stream_ended = true, flate2::Status::Ok | flate2::Status::BufError => { - // Stop only when the inflater is starved for input: it made - // no progress and there is still spare output capacity, so - // the stall is missing input (arriving in a later chunk, or - // resolved at `finish`), not an exhausted output buffer. - if consumed == 0 && produced == 0 && output.len() < output.capacity() { + // The write window is always at least one byte, so no + // progress means the inflater is starved for input (arriving + // in a later chunk, or resolved at `finish`), not an + // exhausted output buffer. + if consumed == 0 && produced == 0 { break; } } @@ -773,19 +857,22 @@ impl DeflateStreamDecoder { /// truncated stream makes no further progress and errors. fn finish(&mut self) -> Result, Report> { let mut output = Vec::new(); + // Same fixed scratch buffer as `decode`: the terminal flush can still + // expand a withheld bomb, so it must be bounded by the remaining budget + // rather than draining into a doubling `Vec`. + let mut scratch = [0_u8; STREAM_CHUNK_SIZE]; while !self.stream_ended { - if output.len() == output.capacity() { - output.reserve(STREAM_CHUNK_SIZE); - } + let window = self.decode_window(scratch.len()); let before_out = self.decompress.total_out(); let status = self .decompress - .decompress_vec(&[], &mut output, flate2::FlushDecompress::Finish) + .decompress(&[], &mut scratch[..window], flate2::FlushDecompress::Finish) .change_context(TrustedServerError::Proxy { message: "Failed to finalize deflate publisher body decoder".to_string(), })?; let produced = (self.decompress.total_out() - before_out) as usize; self.charge(produced)?; + extend_capped(&mut output, &scratch[..produced], self.max_decoded_bytes); match status { flate2::Status::StreamEnd => self.stream_ended = true, flate2::Status::Ok | flate2::Status::BufError => { @@ -1211,6 +1298,60 @@ mod tests { ); } + fn zlib_compress(data: &[u8]) -> Vec { + let mut encoder = flate2::write::ZlibEncoder::new(Vec::new(), flate2::Compression::best()); + encoder + .write_all(data) + .expect("should write deflate test input"); + encoder.finish().expect("should finish deflate encoding") + } + + #[test] + fn deflate_decode_caps_buffer_capacity_at_decoded_limit() { + // Regression for the ~2× allocation: `decompress_vec` doubled the output + // `Vec` and let the inflater fill the doubled capacity before the cap + // check ran, so a 16 MiB limit peaked at 32 MiB. The fixed-scratch + // decoder must decode a body sitting exactly at the cap without letting + // its accumulator balloon toward twice the ceiling. + let cap = STREAM_CHUNK_SIZE * 16; + let payload = vec![b'a'; cap]; + let compressed = zlib_compress(&payload); + let mut decoder = DeflateStreamDecoder::new(Rc::new(Cell::new(0)), cap); + + let decoded = decoder + .decode(&compressed) + .expect("a body exactly at the cap must decode"); + + assert_eq!(decoded.len(), cap, "should decode the whole payload"); + assert!( + decoded.capacity() <= cap + STREAM_CHUNK_SIZE, + "decode buffer must not balloon toward ~2× the cap: cap={cap} capacity={}", + decoded.capacity() + ); + } + + #[test] + fn deflate_decode_rejects_bomb_before_expanding_past_limit() { + // A tiny compressed input that expands to far more than twice the cap + // must be rejected by the cumulative charge, not decoded in full first. + let cap = STREAM_CHUNK_SIZE * 16; + let bomb = zlib_compress(&vec![0_u8; cap * 64]); + assert!( + bomb.len() < cap, + "test precondition: the bomb's compressed form must stay well under the cap" + ); + let mut decoder = DeflateStreamDecoder::new(Rc::new(Cell::new(0)), cap); + + let err = decoder + .decode(&bomb) + .expect_err("a bomb expanding past the cap must error"); + + assert!( + format!("{err:?}").contains("decoded size exceeded"), + "should report the cumulative decoded cap: {err:?}" + ); + } + fn gzip_member(data: &[u8]) -> Vec { let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); encoder @@ -1219,6 +1360,35 @@ mod tests { encoder.finish().expect("should finish gzip encoding") } + #[test] + fn gzip_decode_reader_rejects_bomb_before_materializing_past_limit() { + // A tiny gzip member expanding far past the cap must be rejected by the + // reader's bounded sink mid-decode. Before this bound the reader used a + // `usize::MAX` budget and decoded the whole compressed block into its + // buffer (several MiB) before any downstream writer could reject it. + let cap = STREAM_CHUNK_SIZE; + let bomb = gzip_member(&vec![0_u8; cap * 512]); + assert!( + bomb.len() < cap, + "test precondition: the compressed bomb must stay under the cap" + ); + let mut reader = GzipDecodeReader::new(std::io::Cursor::new(bomb), cap); + + let mut sink = Vec::new(); + let err = std::io::copy(&mut reader, &mut sink) + .expect_err("a gzip bomb expanding past the cap must error"); + + assert!( + err.to_string().contains("decoded size exceeded"), + "should report the cumulative decoded cap: {err}" + ); + assert!( + sink.len() <= cap, + "no more than the cap may be emitted before rejection: {} bytes", + sink.len() + ); + } + #[test] fn body_stream_decoder_decodes_multi_member_gzip_single_chunk() { let mut compressed = gzip_member(b"first member "); diff --git a/docs/guide/configuration.md b/docs/guide/configuration.md index b975590c6..d686c0ee8 100644 --- a/docs/guide/configuration.md +++ b/docs/guide/configuration.md @@ -161,14 +161,14 @@ Core publisher settings for domain, origin, and proxy configuration. ### `[publisher]` -| Field | Type | Required | Description | -| ----------------------------- | ------- | -------- | --------------------------------------------------------------------------------------- | -| `domain` | String | Yes | Publisher's apex domain name | -| `cookie_domain` | String | Yes | Domain for non-EC cookies (typically with leading dot) | -| `origin_url` | String | Yes | Full URL of publisher origin server | -| `origin_host_header_override` | String | No | Outbound Host header to send while connecting to `origin_url` | -| `proxy_secret` | String | Yes | Secret key for encrypting/signing proxy URLs | -| `max_buffered_body_bytes` | Integer | No | Max bytes buffered when a publisher response is post-processed in full (default 16 MiB) | +| Field | Type | Required | Description | +| ----------------------------- | ------- | -------- | --------------------------------------------------------------------------- | +| `domain` | String | Yes | Publisher's apex domain name | +| `cookie_domain` | String | Yes | Domain for non-EC cookies (typically with leading dot) | +| `origin_url` | String | Yes | Full URL of publisher origin server | +| `origin_host_header_override` | String | No | Outbound Host header to send while connecting to `origin_url` | +| `proxy_secret` | String | Yes | Secret key for encrypting/signing proxy URLs | +| `max_buffered_body_bytes` | Integer | No | Buffered-body cap / Fastly stream raw+decoded byte ceiling (default 16 MiB) | > **Note:** EC cookies (`ts-ec`) derive their domain automatically as `.{domain}` and > do not use `cookie_domain`. The `cookie_domain` field is used by other cookie helpers. @@ -300,29 +300,36 @@ Changing `proxy_secret` invalidates all existing signed URLs. Plan rotations car #### `max_buffered_body_bytes` -**Purpose**: Upper bound on the in-memory buffer used when a publisher origin -response must be processed in full (HTML rewriting and integration injection) -instead of streamed. +**Purpose**: Upper bound on how much of a publisher origin body the rewrite +pipeline holds in memory — the post-rewrite output buffer on buffered adapters, +and the per-stream raw/decoded byte ceiling on the Fastly streaming path. **Usage**: -- Caps the _decoded, post-rewrite_ output buffer for any buffered publisher - response on both the legacy and EdgeZero paths. -- Exceeding the cap fails the response (mapped to a 5xx proxy error) rather than - allocating past the limit, preventing Wasm-heap exhaustion on highly - compressible documents. +- On **buffered adapters** (Axum, Cloudflare, Spin) it caps the _decoded, + post-rewrite_ output buffer for a publisher response processed in full. +- On the **Fastly streaming path** the origin body is preserved as a stream, so + the same value caps the stream twice over: the cumulative _raw_ (still + compressed) bytes pulled from origin, and the cumulative _decoded_ bytes + emitted by the decompressor. The decoded cap is enforced _during_ + decompression, so a decompression bomb is rejected before its expansion is + materialized rather than after. -**Default**: `16777216` (16 MiB). +**Behavior when exceeded**: -**Effective Fastly limit**: On Fastly the practical ceiling for a publisher page -is lower. The platform HTTP client rejects any origin response whose raw (still -compressed) body exceeds **10 MiB** before this buffer is filled, so raising the -value only helps highly compressible pages whose decoded size exceeds 16 MiB -while their compressed origin body stays under 10 MiB. Raising it above ~10 MiB -does not lift the platform cap for uncompressed pages. +- On **buffered adapters** the response fails before any bytes are committed. +- On the **streaming path** the response headers are already committed when + either cap trips, so the body is **truncated mid-stream** and the error is + logged — the client receives a short (incomplete) body rather than a `5xx`. + Size the cap above your largest expected decoded page so legitimate responses + are never truncated. + +**Default**: `16777216` (16 MiB). On the Fastly streaming path this is now the +sole ceiling: origin bodies are streamed rather than materialized in full, so +the previous ~10 MiB raw-body limit no longer applies. **Minimum**: Must be at least `1`. A value of `0` is rejected at startup because -a zero-byte cap fails every non-empty buffered response. +a zero-byte cap fails every non-empty publisher response. **Environment Override**: From 18fe060f1ca4e4c1070a45ec2c0c189ffa7ef043 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Mon, 27 Jul 2026 14:24:00 +0530 Subject: [PATCH 14/15] Cap the decode sink reservation and stop panicking on empty poisoned chunks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two consistency follow-ups to the streaming decode hardening. `BoundedDecodeSink::write` charged the cumulative decode budget before appending but then appended with a plain `extend_from_slice`, whose amortized doubling can commit roughly twice the configured ceiling of Wasm heap while the decoded total still sits under the cap. Route the append through `extend_capped`, the same helper the deflate decoder already uses, so the gzip and brotli decoders honour the documented "rejected before its expansion is materialized" guarantee too. `GzipStreamDecoder::decode`'s error-returning `Poisoned` arms live inside the `while !input.is_empty()` loop, so an empty chunk — which `BodyChunkSource` may legally yield — skipped them and fell through to `take_decoded`, whose `Poisoned` arm still aborted the Wasm instance via `unreachable!`. Drain no bytes instead; `finish` continues to error on a poisoned decoder, so a truncated body is never reported as success. Both regression tests were confirmed to fail against the previous behaviour. The capacity test drives the sink directly rather than through gzip, because the flate2 write decoder holds part of a member back until `finish` and makes the sink's length at drain time non-deterministic. --- .../src/streaming_processor.rs | 75 +++++++++++++++++-- 1 file changed, 69 insertions(+), 6 deletions(-) diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index 0d1cc7095..130d25751 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -457,7 +457,7 @@ impl Write for BoundedDecodeSink { ))); } self.decoded_bytes.set(next); - self.buffer.extend_from_slice(data); + extend_capped(&mut self.buffer, data, self.max_decoded_bytes); Ok(data.len()) } @@ -500,8 +500,9 @@ enum GzipStreamState { /// Non-member bytes followed a completed member; all further input is /// dropped. TrailingGarbage(BoundedDecodeSink), - /// Transient placeholder while ownership moves between states. Never - /// observable by callers. + /// Transient placeholder while ownership moves between states. An error + /// raised mid-transition can leave it in place, so every arm that reads it + /// fails (or drains empty) instead of panicking. Poisoned, } @@ -642,9 +643,11 @@ impl GzipStreamDecoder { GzipStreamState::Boundary { sink, .. } | GzipStreamState::TrailingGarbage(sink) => { std::mem::take(&mut sink.buffer) } - GzipStreamState::Poisoned => { - unreachable!("poisoned state is never left in place across calls") - } + // An empty chunk skips `decode`'s loop (and its error-returning + // `Poisoned` arm) and lands here, so a poisoned decoder must yield + // no bytes rather than abort the Wasm instance. `finish` still + // errors, so a truncated body can never be reported as success. + GzipStreamState::Poisoned => Vec::new(), } } } @@ -1389,6 +1392,66 @@ mod tests { ); } + #[test] + fn bounded_decode_sink_caps_buffer_capacity_at_decoded_limit() { + // Companion to `deflate_decode_caps_buffer_capacity_at_decoded_limit` + // for the sink the gzip and brotli decoders write through. It charged + // the budget correctly but appended with a plain `extend_from_slice`, so + // `Vec` doubling could commit ~2× the ceiling of Wasm heap while the + // decoded total still sat under the cap. The cap is deliberately off a + // doubling boundary: the last block lands past 128 KiB of capacity, so + // an uncapped reservation jumps to 256 KiB. + let cap = STREAM_CHUNK_SIZE * 20 + 1; + let mut sink = BoundedDecodeSink::new(Rc::new(Cell::new(0)), cap); + let block = vec![b'a'; STREAM_CHUNK_SIZE]; + + for _ in 0..20 { + let written = sink + .write(&block) + .expect("writes under the cap must succeed"); + assert_eq!( + written, STREAM_CHUNK_SIZE, + "a write under the cap must buffer the whole block" + ); + } + + assert_eq!( + sink.buffer.len(), + STREAM_CHUNK_SIZE * 20, + "every write under the cap should be buffered" + ); + assert!( + sink.buffer.capacity() <= cap + 1, + "decode sink must not balloon toward ~2× the cap: cap={cap} capacity={}", + sink.buffer.capacity() + ); + } + + #[test] + fn gzip_decode_of_empty_chunk_on_poisoned_decoder_yields_no_bytes() { + // `decode`'s error-returning `Poisoned` arm lives inside the + // `while !input.is_empty()` loop, so an empty chunk — which + // `BodyChunkSource` may legally yield — skipped it and reached + // `take_decoded`, whose `unreachable!` aborts the Wasm instance. + let mut decoder = GzipStreamDecoder::new(Rc::new(Cell::new(0)), usize::MAX); + decoder.state = GzipStreamState::Poisoned; + + let decoded = decoder + .decode(&[]) + .expect("an empty chunk must not fail on a poisoned decoder"); + + assert!( + decoded.is_empty(), + "a poisoned decoder must drain no bytes: {} byte(s)", + decoded.len() + ); + assert!( + decoder.finish().is_err(), + "finalizing a poisoned decoder must still fail, so a truncated body \ + is never reported as success" + ); + } + #[test] fn body_stream_decoder_decodes_multi_member_gzip_single_chunk() { let mut compressed = gzip_member(b"first member "); From 165aff841cf4dd3c9835287492f8c30dbb51eb43 Mon Sep 17 00:00:00 2001 From: prk-Jr Date: Mon, 27 Jul 2026 22:28:58 +0530 Subject: [PATCH 15/15] Flush streamed segments and bound gzip decode per step MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Resolve five review findings on the streaming pipeline. Flush Fastly's `StreamingBody` after each yielded chunk. It wraps a `BufWriter`, so a segment smaller than the write buffer stayed in the Wasm heap while the generator awaited the origin or the auction, leaving the client with committed headers and no body. Sync-flush the active gzip member per source chunk. `flate2`'s `write::GzDecoder` keeps decoded output in its own 32 KiB buffer and drains it only on the next write, so a member whose expansion fit there surfaced no bytes until `finish` and the caller polled the origin again instead of emitting content. `flush` is a sync flush, not a stream finish, so the trailer check still validates member integrity. Split the close-body hold finalization into a pre-collection ready stage and a tail stage, so bytes a decoder releases at its own finalization are emitted before the auction is collected. Only the held `` tail waits on collection, matching the mid-stream path. Bound the buffered gzip reader's pending decode buffer rather than its cumulative decoded input. As a cumulative cap it turned the buffered adapters' post-rewrite output ceiling into a gzip-only decoded-input limit, failing a valid body whose decode exceeded it even when the rewritten output fit — and only when gzip-encoded, since the deflate and brotli readers leave the total to the downstream cap. A bomb is still rejected mid-decode. Drop `Transfer-Encoding` and `Trailer` from normalized 204 and 205 responses. RFC 9112 §6.1 forbids `Transfer-Encoding` on a 204, and keeping it on a 205 conflicts with the inserted `Content-Length: 0`. `HEAD` and 304 keep both: their message length is fixed by the header section regardless of framing fields. --- crates/trusted-server-core/src/proxy.rs | 76 ++++++ crates/trusted-server-core/src/publisher.rs | 252 ++++++++++++++---- crates/trusted-server-core/src/settings.rs | 8 +- .../src/streaming_processor.rs | 165 ++++++++++-- docs/guide/configuration.md | 7 +- 5 files changed, 432 insertions(+), 76 deletions(-) diff --git a/crates/trusted-server-core/src/proxy.rs b/crates/trusted-server-core/src/proxy.rs index dad06c12d..a2022cf88 100644 --- a/crates/trusted-server-core/src/proxy.rs +++ b/crates/trusted-server-core/src/proxy.rs @@ -279,6 +279,18 @@ pub async fn stream_asset_body( .change_context(TrustedServerError::Proxy { message: "failed to write streaming platform response body".to_string(), })?; + // Flush per chunk: Fastly's `StreamingBody` is a + // `BufWriter`, so a segment smaller than + // the write buffer would sit in the Wasm heap until a later + // write filled it. The publisher stream yields compressed + // segments that are commonly smaller than that buffer, and the + // generator then awaits the origin (or the auction) before + // producing the next one — without this flush the client sees + // committed headers and no body, undoing the codec's per-chunk + // sync flush and delaying first paint. + output.flush().change_context(TrustedServerError::Proxy { + message: "failed to flush streaming platform response body".to_string(), + })?; } } } @@ -2027,8 +2039,10 @@ fn reconstruct_and_validate_signed_target( #[cfg(test)] mod tests { + use std::cell::RefCell; use std::collections::{HashMap, VecDeque}; use std::io; + use std::rc::Rc; use std::sync::{Arc, Mutex}; use super::{ @@ -4456,6 +4470,68 @@ mod tests { }); } + /// One observable step of [`stream_asset_body`] driving a buffered sink. + #[derive(Debug, PartialEq, Eq)] + enum StreamSinkEvent { + SourcePolled(usize), + Wrote(usize), + Flushed, + } + + /// Sink that records its calls so a test can assert how they interleave + /// with the source stream's polls. + struct RecordingSink(Rc>>); + + impl io::Write for RecordingSink { + fn write(&mut self, buf: &[u8]) -> io::Result { + self.0.borrow_mut().push(StreamSinkEvent::Wrote(buf.len())); + Ok(buf.len()) + } + + fn flush(&mut self) -> io::Result<()> { + self.0.borrow_mut().push(StreamSinkEvent::Flushed); + Ok(()) + } + } + + #[test] + fn stream_asset_body_flushes_each_chunk_before_polling_the_source_again() { + // Fastly's `StreamingBody` is a `BufWriter`, so a yielded segment only + // reaches the client once the sink is flushed. The publisher stream + // awaits the origin (and the auction) between segments, so a segment + // smaller than the write buffer must be flushed before the next poll or + // the client renders nothing behind committed headers. + let events = Rc::new(RefCell::new(Vec::new())); + let source_events = Rc::clone(&events); + let body = EdgeBody::from_stream(async_stream::stream! { + for (index, segment) in ["", ""].into_iter().enumerate() { + source_events + .borrow_mut() + .push(StreamSinkEvent::SourcePolled(index)); + yield Ok::(Bytes::from_static(segment.as_bytes())); + } + }); + + futures::executor::block_on(stream_asset_body( + body, + &mut RecordingSink(Rc::clone(&events)), + )) + .expect("should stream asset body"); + + assert_eq!( + *events.borrow(), + vec![ + StreamSinkEvent::SourcePolled(0), + StreamSinkEvent::Wrote(12), + StreamSinkEvent::Flushed, + StreamSinkEvent::SourcePolled(1), + StreamSinkEvent::Wrote(13), + StreamSinkEvent::Flushed, + ], + "each segment must be flushed to the client before the source is polled again" + ); + } + #[test] fn asset_proxy_response_into_response_rejects_stream_body() { futures::executor::block_on(async { diff --git a/crates/trusted-server-core/src/publisher.rs b/crates/trusted-server-core/src/publisher.rs index 8bbd487b7..600441f95 100644 --- a/crates/trusted-server-core/src/publisher.rs +++ b/crates/trusted-server-core/src/publisher.rs @@ -434,10 +434,13 @@ fn process_response_streaming( output_compression: compression, chunk_size: 8192, }; - // Bound the gzip decode path to the same ceiling the buffered writer - // enforces, so a gzip bomb is rejected mid-decode instead of materializing - // its full expansion before the downstream writer can reject it. - let max_decoded_bytes = params.settings.publisher.max_buffered_body_bytes; + // Bound how much decoded gzip output may sit in the heap at once, using the + // same ceiling the buffered writer enforces on the rewritten output: a gzip + // bomb is then rejected mid-decode instead of materializing its full + // expansion first. The bound is per-step, so a large honest body still + // streams through and only the buffered writer judges the total — otherwise + // gzip would reject bodies the identity, deflate and brotli paths accept. + let max_pending_decoded_bytes = params.settings.publisher.max_buffered_body_bytes; if is_html { let processor = create_html_stream_processor( @@ -450,7 +453,7 @@ fn process_response_streaming( params.ad_bids_state.clone(), )?; StreamingPipeline::new(config, processor) - .with_max_decoded_bytes(max_decoded_bytes) + .with_max_pending_decoded_bytes(max_pending_decoded_bytes) .process(body_as_reader(body)?, output)?; } else if is_rsc_flight { // RSC Flight responses are length-prefixed (T rows). A naive string replacement will @@ -462,7 +465,7 @@ fn process_response_streaming( params.request_scheme, ); StreamingPipeline::new(config, processor) - .with_max_decoded_bytes(max_decoded_bytes) + .with_max_pending_decoded_bytes(max_pending_decoded_bytes) .process(body_as_reader(body)?, output)?; } else { let replacer = create_url_replacer( @@ -472,7 +475,7 @@ fn process_response_streaming( params.request_scheme, ); StreamingPipeline::new(config, replacer) - .with_max_decoded_bytes(max_decoded_bytes) + .with_max_pending_decoded_bytes(max_pending_decoded_bytes) .process(body_as_reader(body)?, output)?; } @@ -760,8 +763,9 @@ async fn hold_step_decoded_chunk( /// Collect the dispatched auction and process the held `` tail. /// -/// Call only after [`hold_step_decoded_chunk`] (or [`hold_finish_segments`]) -/// reports `close_found` and the ready prefix has already been emitted: +/// Call only after [`hold_step_decoded_chunk`] (or +/// [`hold_finish_ready_segments`]) reports `close_found` and the ready prefix +/// has already been emitted: /// collecting here — after the prefix streams — is what keeps the auction /// riding alongside transfer instead of blocking it. Collection runs before the /// tail is processed so `lol_html` sees live bids at the injection point. @@ -802,7 +806,8 @@ async fn hold_collect_close_tail( /// through [`hold_step_decoded_chunk`]. /// /// Returns `Ok(None)` when the source is exhausted; the caller must then emit -/// [`hold_finish_segments`]. On read or decode failure the pending auction is +/// [`hold_finish_ready_segments`] followed by [`hold_finish_tail_segments`]. On +/// read or decode failure the pending auction is /// abandoned before the error is returned. Shared by the write-sink driver /// ([`body_close_hold_loop_stream`]) and the lazy publisher body stream so /// the two hold paths cannot drift apart. @@ -840,23 +845,24 @@ async fn hold_step_next_chunk( .map(Some) } -/// Finalize the close-body hold pipeline at end of the origin stream. +/// Drain the decoder tail at end of the origin stream, returning the prefix the +/// caller must emit before [`hold_finish_tail_segments`]. /// -/// Drains the decoder tail through the hold (or straight through when the -/// hold was already released mid-stream), collects the auction if the -/// close-body tag never streamed, processes the held tail plus the -/// processor's final chunk, and emits the encoder trailer. Returns the -/// encoded segments for the caller to emit. On decoder failure the pending -/// auction is abandoned before the error is returned. -async fn hold_finish_segments( +/// A codec can hold document bytes back until its own finalization — the gzip +/// decoder releases the remainder of the final member at `finish()` — and that +/// remainder may be the whole document for a small page. Returning it ahead of +/// collection keeps the invariant the mid-stream path already has: only the +/// closing `` tail waits for the auction, never renderable content. +/// +/// On decoder failure the pending auction is abandoned before the error is +/// returned. +async fn hold_finish_ready_segments( processor: &mut P, decoder: &mut BodyStreamDecoder, encoder: &mut BodyStreamEncoder, state: &mut AuctionHoldState, collect_refs: &AuctionCollectDeps<'_>, ) -> Result, Report> { - let mut segments = Vec::new(); - let decoded_tail = match decoder.finish() { Ok(decoded_tail) => decoded_tail, Err(err) => { @@ -864,15 +870,30 @@ async fn hold_finish_segments( return Err(err); } }; - if !decoded_tail.is_empty() { - let step = - hold_step_decoded_chunk(processor, encoder, &decoded_tail, state, collect_refs).await?; - segments.extend(step.ready); + if decoded_tail.is_empty() { + return Ok(Vec::new()); } + let step = + hold_step_decoded_chunk(processor, encoder, &decoded_tail, state, collect_refs).await?; + Ok(step.ready) +} + +/// Finalize the close-body hold pipeline after [`hold_finish_ready_segments`]. +/// +/// Collects the auction if the close-body tag never streamed, processes the held +/// tail plus the processor's final chunk, and emits the encoder trailer. Returns +/// the encoded segments for the caller to emit. +async fn hold_finish_tail_segments( + processor: &mut P, + encoder: &mut BodyStreamEncoder, + state: &mut AuctionHoldState, + collect_refs: &AuctionCollectDeps<'_>, +) -> Result, Report> { + let mut segments = Vec::new(); // If the hold is still armed the auction was never collected mid-stream: - // `` arrived only in this tail, or the document had none at all. - // Collect now and flush the held remainder before finalizing. + // `` arrived only in the decoder tail, or the document had none at + // all. Collect now and flush the held remainder before finalizing. if state.hold.is_some() { segments.extend(hold_collect_close_tail(processor, encoder, state, collect_refs).await?); } @@ -1321,7 +1342,11 @@ pub async fn publisher_response_into_streaming_response( } } - for encoded in hold_finish_segments( + // Whatever the decoder released at finalization is emitted + // before collection, for the same reason as the mid-stream + // prefix above: a small compressed page can surface its + // whole document here. + for encoded in hold_finish_ready_segments( &mut processor, &mut decoder, &mut encoder, @@ -1333,6 +1358,17 @@ pub async fn publisher_response_into_streaming_response( { yield encoded; } + for encoded in hold_finish_tail_segments( + &mut processor, + &mut encoder, + &mut state, + &collect_refs, + ) + .await + .map_err(publisher_stream_error)? + { + yield encoded; + } } else { while let Some(segments) = passthrough_step( &mut source, @@ -1385,21 +1421,38 @@ fn response_carries_body(method: &Method, status: StatusCode) -> bool { /// origin length is wrong (RFC 9110 §15.4.6); normalize it to `0`. /// - **HEAD and 304 Not Modified**: `Content-Length` legitimately advertises /// the `GET` representation length, so it is preserved untouched. +/// +/// The normalized 204/205 responses also drop chunked-framing metadata the +/// adapters would otherwise copy verbatim: RFC 9112 §6.1 forbids +/// `Transfer-Encoding` on a 204, and keeping it on a 205 alongside the +/// `Content-Length: 0` inserted above would advertise two conflicting framings. +/// `Trailer` describes fields that only a chunked body can carry, so it goes +/// with it. `HEAD` and 304 keep both for the same reason they keep +/// `Content-Length`: the fields describe the `GET` representation, not this +/// message. fn make_response_bodiless(response: &mut Response) { *response.body_mut() = EdgeBody::empty(); match response.status() { StatusCode::NO_CONTENT => { response.headers_mut().remove(header::CONTENT_LENGTH); + remove_chunked_framing_headers(response.headers_mut()); } StatusCode::RESET_CONTENT => { response .headers_mut() .insert(header::CONTENT_LENGTH, HeaderValue::from(0_u64)); + remove_chunked_framing_headers(response.headers_mut()); } _ => {} } } +/// Remove the framing fields that only apply to a message with a chunked body. +fn remove_chunked_framing_headers(headers: &mut header::HeaderMap) { + headers.remove(header::TRANSFER_ENCODING); + headers.remove(header::TRAILER); +} + /// A [`Write`] sink that buffers into a `Vec` but fails once the configured /// byte limit would be exceeded. /// @@ -2023,7 +2076,7 @@ async fn stream_html_with_auction_hold( /// Async-pull variant of [`body_close_hold_loop`] for live origin streams. /// -/// Shares [`hold_step_next_chunk`] and [`hold_finish_segments`] with the +/// Shares [`hold_step_next_chunk`] and the finish stages with the /// lazy streaming body built by [`publisher_response_into_streaming_response`], /// so the two async hold paths cannot drift apart. /// @@ -2074,7 +2127,9 @@ async fn body_close_hold_loop_stream( } } - for encoded in hold_finish_segments( + // Write the decoder-finalized prefix before collection, matching the lazy + // Fastly stream: only the held `` tail waits on the auction. + for encoded in hold_finish_ready_segments( processor, &mut decoder, &mut encoder, @@ -2085,6 +2140,11 @@ async fn body_close_hold_loop_stream( { write_encoded_segment(writer, &encoded)?; } + for encoded in + hold_finish_tail_segments(processor, &mut encoder, &mut state, &collect_refs).await? + { + write_encoded_segment(writer, &encoded)?; + } writer.flush().change_context(TrustedServerError::Proxy { message: "Failed to flush output".to_string(), })?; @@ -6247,16 +6307,102 @@ mod tests { ); } - // (method, status, expected Content-Length after bodiless normalization). - // 204 forbids Content-Length (removed); 205 must advertise a zero-length - // body; HEAD and 304 legitimately advertise the GET representation length. - const BODILESS_FRAMING_CASES: [(Method, StatusCode, Option<&str>); 4] = [ - (Method::HEAD, StatusCode::OK, Some("42")), - (Method::GET, StatusCode::NO_CONTENT, None), - (Method::GET, StatusCode::RESET_CONTENT, Some("0")), - (Method::GET, StatusCode::NOT_MODIFIED, Some("42")), + #[test] + fn streaming_finalize_auction_hold_emits_small_compressed_page_before_collection() { + // A small gzip page arrives as one origin chunk whose whole expansion + // fits inside `flate2`'s internal decode buffer, and its `` lands + // in that same chunk. Both stages that could withhold it must not: the + // decoder sync-flushes per source chunk instead of holding output until + // its own finalization, and the finish path emits the decoder tail + // before awaiting collection. Otherwise the client holds committed + // headers and an empty body until the auction resolves. + let page = b"

small compressed page

"; + let params = html_stream_params( + "gzip", + Some(DispatchedAuction::empty_for_test( + test_auction_request(), + 500, + )), + ); + let body = streaming_finalize_response( + params, + origin_chunk_then_pending(bytes::Bytes::from(gzip_encode(page))), + ); + + let first = first_lazy_body_chunk(body); + // Decode the flushed (not finished) prefix the way a browser decodes the + // bytes received while the response is still open. + let mut decoder = flate2::write::MultiGzDecoder::new(Vec::new()); + decoder + .write_all(&first) + .expect("first chunk must be valid gzip"); + decoder.flush().expect("should flush gzip decoder"); + let decoded = String::from_utf8(decoder.get_ref().clone()).expect("should be valid UTF-8"); + + assert!( + decoded.contains("small compressed page"), + "first poll must emit the decoded document prefix of a small gzip page. Got: {decoded}" + ); + assert!( + !decoded.contains(".bids=JSON.parse"), + "bids inject only at after collection, which the first poll must not wait for. Got: {decoded}" + ); + } + + // (method, status, expected Content-Length, expected Transfer-Encoding) + // after bodiless normalization. 204 forbids Content-Length (removed); 205 + // must advertise a zero-length body; HEAD and 304 legitimately advertise the + // GET representation length. Chunked framing is stripped from 204 (forbidden + // by RFC 9112 §6.1) and 205 (would conflict with the inserted + // `Content-Length: 0`), but preserved for HEAD and 304, whose message length + // is fixed by the header section regardless of framing fields + // (RFC 9112 §6.3, rule 1). + const BODILESS_FRAMING_CASES: [(Method, StatusCode, Option<&str>, Option<&str>); 4] = [ + (Method::HEAD, StatusCode::OK, Some("42"), Some("chunked")), + (Method::GET, StatusCode::NO_CONTENT, None, None), + (Method::GET, StatusCode::RESET_CONTENT, Some("0"), None), + ( + Method::GET, + StatusCode::NOT_MODIFIED, + Some("42"), + Some("chunked"), + ), ]; + /// Assert the framing fields left on a normalized bodiless response. + fn assert_bodiless_framing( + response: &Response, + method: &Method, + status: StatusCode, + expected_length: Option<&str>, + expected_transfer_encoding: Option<&str>, + ) { + let header_value = |name: header::HeaderName| { + response + .headers() + .get(name) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned) + }; + assert_eq!( + header_value(header::CONTENT_LENGTH).as_deref(), + expected_length, + "bodiless {method} {status} must carry the corrected Content-Length" + ); + assert_eq!( + header_value(header::TRANSFER_ENCODING).as_deref(), + expected_transfer_encoding, + "bodiless {method} {status} must carry the corrected Transfer-Encoding" + ); + assert_eq!( + header_value(header::TRAILER).as_deref(), + // `Trailer` only describes fields a chunked body can carry, so it + // survives exactly where chunked framing does. + expected_transfer_encoding.map(|_| "expires"), + "bodiless {method} {status} must drop Trailer with its chunked framing" + ); + } + #[test] fn publisher_response_streaming_finalize_normalizes_bodiless_framing() { // Fastly requests the origin body as a stream before classification, so @@ -6269,10 +6415,13 @@ mod tests { ); let orchestrator = Arc::new(AuctionOrchestrator::new(settings.auction.clone())); - for (method, status, expected_length) in BODILESS_FRAMING_CASES { + for (method, status, expected_length, expected_transfer_encoding) in BODILESS_FRAMING_CASES + { let response = Response::builder() .status(status) .header(header::CONTENT_LENGTH, "42") + .header(header::TRANSFER_ENCODING, "chunked") + .header(header::TRAILER, "expires") .body(EdgeBody::stream(futures::stream::iter(vec![ bytes::Bytes::from_static(b"origin body bytes that must not reach the client"), ]))) @@ -6293,13 +6442,12 @@ mod tests { !matches!(response.body(), EdgeBody::Stream(_)), "bodiless {method} {status} must not carry a streaming body" ); - assert_eq!( - response - .headers() - .get(header::CONTENT_LENGTH) - .and_then(|v| v.to_str().ok()), + assert_bodiless_framing( + &response, + &method, + status, expected_length, - "bodiless {method} {status} must carry the corrected Content-Length" + expected_transfer_encoding, ); let drained = futures::executor::block_on( @@ -6327,10 +6475,13 @@ mod tests { let orchestrator = AuctionOrchestrator::new(settings.auction.clone()); let services = noop_services(); - for (method, status, expected_length) in BODILESS_FRAMING_CASES { + for (method, status, expected_length, expected_transfer_encoding) in BODILESS_FRAMING_CASES + { let response = Response::builder() .status(status) .header(header::CONTENT_LENGTH, "42") + .header(header::TRANSFER_ENCODING, "chunked") + .header(header::TRAILER, "expires") .body(EdgeBody::from( b"origin body bytes that must not reach the client".to_vec(), )) @@ -6347,13 +6498,12 @@ mod tests { )) .expect("should finalize buffered response"); - assert_eq!( - response - .headers() - .get(header::CONTENT_LENGTH) - .and_then(|v| v.to_str().ok()), + assert_bodiless_framing( + &response, + &method, + status, expected_length, - "bodiless {method} {status} must carry the corrected Content-Length" + expected_transfer_encoding, ); let drained = futures::executor::block_on( diff --git a/crates/trusted-server-core/src/settings.rs b/crates/trusted-server-core/src/settings.rs index 6cf2009bb..61b8fb550 100644 --- a/crates/trusted-server-core/src/settings.rs +++ b/crates/trusted-server-core/src/settings.rs @@ -55,11 +55,17 @@ pub struct Publisher { /// raw (still compressed) bytes pulled from origin, and cumulative decoded /// bytes emitted by the decompressor — the latter so a decompression bomb /// cannot push an unbounded decoded volume through the rewrite pipeline. - /// Buffered adapters keep using it as the post-rewrite output buffer cap. /// On the streaming path headers are already committed when either cap /// trips, so the response is truncated mid-body (with the error logged) /// rather than replaced with a 5xx. /// + /// Buffered adapters keep using it as the post-rewrite output buffer cap. + /// There it additionally bounds how much decoded gzip output may sit in the + /// heap at once, so a bomb is rejected mid-decode instead of after its full + /// expansion; that bound is per-step, never cumulative, so a gzip-encoded + /// body is judged by the same post-rewrite total as an identity, deflate or + /// brotli one. + /// /// Must be at least 1: a zero-byte cap fails every non-empty buffered /// publisher response at request time, so it is rejected at config /// validation instead. diff --git a/crates/trusted-server-core/src/streaming_processor.rs b/crates/trusted-server-core/src/streaming_processor.rs index 130d25751..d95e6ebf8 100644 --- a/crates/trusted-server-core/src/streaming_processor.rs +++ b/crates/trusted-server-core/src/streaming_processor.rs @@ -112,11 +112,12 @@ impl Default for PipelineConfig { pub struct StreamingPipeline { config: PipelineConfig, processor: P, - /// Cumulative decoded-byte ceiling for the gzip decode path. Defaults to - /// `usize::MAX` (unbounded); set via [`Self::with_max_decoded_bytes`] on the - /// buffered publisher path so a gzip bomb is rejected mid-decode instead of - /// materializing its full expansion. - max_decoded_bytes: usize, + /// Ceiling on the decoded bytes the gzip decode path may hold at once. + /// Defaults to `usize::MAX` (unbounded); set via + /// [`Self::with_max_pending_decoded_bytes`] on the buffered publisher path so + /// a gzip bomb is rejected mid-decode instead of materializing its full + /// expansion. + max_pending_decoded_bytes: usize, } impl StreamingPipeline

{ @@ -129,21 +130,23 @@ impl StreamingPipeline

{ Self { config, processor, - max_decoded_bytes: usize::MAX, + max_pending_decoded_bytes: usize::MAX, } } - /// Bound the cumulative decoded output of the gzip decode path. + /// Bound how many decoded bytes the gzip decode path may hold at once. /// /// Only the gzip reader materializes decoded bytes ahead of the downstream /// writer; the deflate and brotli read decoders already emit into the /// caller's fixed-size buffer. Callers that buffer output under a configured /// ceiling (e.g. `publisher.max_buffered_body_bytes`) should pass that /// ceiling here so decode rejection cannot be preceded by a large - /// allocation. + /// allocation. The limit is per-step, not cumulative: the total decoded + /// volume stays the downstream output cap's responsibility, so gzip does not + /// reject a body the other encodings would accept. #[must_use] - pub fn with_max_decoded_bytes(mut self, max_decoded_bytes: usize) -> Self { - self.max_decoded_bytes = max_decoded_bytes; + pub fn with_max_pending_decoded_bytes(mut self, max_pending_decoded_bytes: usize) -> Self { + self.max_pending_decoded_bytes = max_pending_decoded_bytes; self } @@ -172,7 +175,7 @@ impl StreamingPipeline

{ // Shares `GzipStreamDecoder` with the streaming `BodyStreamDecoder` // gzip codec, so both paths tolerate trailing garbage after the // final member the same way. - let decoder = GzipDecodeReader::new(input, self.max_decoded_bytes); + let decoder = GzipDecodeReader::new(input, self.max_pending_decoded_bytes); let mut encoder = GzEncoder::new(output, flate2::Compression::default()); self.process_chunks(decoder, &mut encoder)?; encoder.finish().change_context(TrustedServerError::Proxy { @@ -180,9 +183,10 @@ impl StreamingPipeline

{ })?; Ok(()) } - (Compression::Gzip, Compression::None) => { - self.process_chunks(GzipDecodeReader::new(input, self.max_decoded_bytes), output) - } + (Compression::Gzip, Compression::None) => self.process_chunks( + GzipDecodeReader::new(input, self.max_pending_decoded_bytes), + output, + ), (Compression::Deflate, Compression::Deflate) => { let decoder = ZlibDecoder::new(input); let mut encoder = ZlibEncoder::new(output, flate2::Compression::default()); @@ -485,6 +489,11 @@ const GZIP_MAGIC: [u8; 2] = [0x1f, 0x8b]; /// with the magic number are decoded as a member and fail if they are not one. struct GzipStreamDecoder { state: GzipStreamState, + /// The charge counter shared with the active sink, so an owner that hands + /// decoded bytes onward can release them via [`Self::release_pending`] and + /// turn the sink's cap into a pending-bytes bound instead of a cumulative + /// one. + decoded_bytes: Rc>, } enum GzipStreamState { @@ -510,12 +519,24 @@ impl GzipStreamDecoder { fn new(decoded_bytes: Rc>, max_decoded_bytes: usize) -> Self { Self { state: GzipStreamState::Member(flate2::write::GzDecoder::new(BoundedDecodeSink::new( - decoded_bytes, + Rc::clone(&decoded_bytes), max_decoded_bytes, ))), + decoded_bytes, } } + /// Release `len` decoded bytes from the shared charge counter. + /// + /// Call this when decoded bytes leave the decoder's owner for good. The + /// counter then measures the bytes currently held rather than the total ever + /// produced, which turns the sink's ceiling into a bound on buffered decode + /// output instead of a cumulative decoded-input limit. + fn release_pending(&self, len: usize) { + self.decoded_bytes + .set(self.decoded_bytes.get().saturating_sub(len)); + } + /// Decode one compressed chunk, returning the decoded bytes it produced. /// /// # Errors @@ -587,6 +608,18 @@ impl GzipStreamDecoder { } } } + // `flate2::write::GzDecoder` keeps decoded output in its own 32 KiB + // buffer and only pushes it into the sink on the *next* write or on a + // flush. Without this flush a member whose expansion fits in that + // buffer — a small page, or the tail of a large one — surfaces no bytes + // until `finish`, so a streaming caller would poll the origin again + // instead of emitting renderable content. This is a sync flush, not a + // stream finish: member integrity is still validated by the trailer + // check in `finish`/`try_finish`. States other than `Member` already + // finished their decoder, so their output is in the sink. + if let GzipStreamState::Member(decoder) = &mut self.state { + decoder.flush()?; + } Ok(self.take_decoded()) } @@ -659,11 +692,19 @@ impl GzipStreamDecoder { /// support instead of erroring (or dropping members) like /// `flate2::read::MultiGzDecoder`/`GzDecoder`. /// -/// `max_decoded_bytes` bounds cumulative decoded output: each compressed block's -/// expansion is charged against the budget from inside [`GzipStreamDecoder`]'s -/// bounded sink, so a decompression bomb is rejected mid-decode rather than fully -/// materialized before a downstream writer (e.g. `BoundedWriter`) can act. Pass -/// `usize::MAX` where an upstream/downstream stage already bounds the size. +/// `max_pending_decoded_bytes` bounds the decoded bytes the reader may hold at +/// once — the expansion of the compressed block being decoded, plus whatever the +/// caller has not read out yet. Each block's expansion is charged against the +/// budget from inside [`GzipStreamDecoder`]'s bounded sink, so a decompression +/// bomb is rejected mid-decode rather than fully materialized before a +/// downstream writer (e.g. `BoundedWriter`) can act, while a large but honest +/// body still streams through in bounded steps. Deliberately *not* a cumulative +/// decoded-input limit: capping the total would reject a valid body whose +/// decoded input exceeds the ceiling even when the rewritten output stays under +/// it, and only the gzip encoding would be affected — the deflate and brotli +/// read decoders emit into the caller's fixed-size buffer and leave the total to +/// the downstream output cap. Pass `usize::MAX` where an upstream stage already +/// bounds the decoded size. pub(crate) struct GzipDecodeReader { input: R, decoder: GzipStreamDecoder, @@ -674,10 +715,10 @@ pub(crate) struct GzipDecodeReader { } impl GzipDecodeReader { - pub(crate) fn new(input: R, max_decoded_bytes: usize) -> Self { + pub(crate) fn new(input: R, max_pending_decoded_bytes: usize) -> Self { Self { input, - decoder: GzipStreamDecoder::new(Rc::new(Cell::new(0)), max_decoded_bytes), + decoder: GzipStreamDecoder::new(Rc::new(Cell::new(0)), max_pending_decoded_bytes), raw: vec![0_u8; STREAM_CHUNK_SIZE], decoded: Vec::new(), position: 0, @@ -697,6 +738,11 @@ impl Read for GzipDecodeReader { let amount = available.min(buf.len()); buf[..amount].copy_from_slice(&self.decoded[self.position..self.position + amount]); self.position += amount; + // These bytes have left the reader, so they stop counting + // against the pending-decode budget. Releasing here — rather + // than when the sink is drained — keeps the charge covering both + // the sink and this reader's undelivered remainder. + self.decoder.release_pending(amount); return Ok(amount); } if self.finished { @@ -1363,6 +1409,36 @@ mod tests { encoder.finish().expect("should finish gzip encoding") } + #[test] + fn body_stream_decoder_emits_small_gzip_member_before_finish() { + // `flate2::write::GzDecoder` holds decoded output in a 32 KiB internal + // buffer that drains only on the next write or on a flush, so a page + // whose whole expansion fits there surfaced zero bytes from + // `decode_chunk`. The streaming caller then polled the origin again + // instead of emitting content, and the document reached the client only + // at `finish` — after origin EOF, and on the publisher hold path after + // auction collection. + let document = b"small page"; + let mut decoder = BodyStreamDecoder::new(Compression::Gzip, usize::MAX); + + let decoded = decoder + .decode_chunk(bytes::Bytes::from(gzip_member(document))) + .expect("a complete small gzip member must decode"); + + assert_eq!( + decoded.as_ref(), + document, + "the first chunk must release the decoded document instead of withholding it until finish" + ); + assert!( + decoder + .finish() + .expect("a complete gzip member must finalize") + .is_empty(), + "finish must have nothing left to emit once the chunk released it" + ); + } + #[test] fn gzip_decode_reader_rejects_bomb_before_materializing_past_limit() { // A tiny gzip member expanding far past the cap must be rejected by the @@ -1383,7 +1459,7 @@ mod tests { assert!( err.to_string().contains("decoded size exceeded"), - "should report the cumulative decoded cap: {err}" + "should report the pending decoded cap: {err}" ); assert!( sink.len() <= cap, @@ -1392,6 +1468,49 @@ mod tests { ); } + /// Deterministic pseudo-random bytes (xorshift64). gzip cannot compress + /// these, so a member of them spans several source reads and each read + /// expands by roughly the read size. + fn incompressible_bytes(len: usize) -> Vec { + let mut state = 0x2545_f491_4f6c_dd1d_u64; + (0..len) + .map(|_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state >> 24) as u8 + }) + .collect() + } + + #[test] + fn gzip_decode_reader_streams_body_larger_than_the_pending_cap() { + // The cap bounds the decoded bytes held at once, not the cumulative + // decoded input. While it was cumulative it made the buffered adapters' + // configured output ceiling behave as a gzip-only decoded-input limit: a + // valid body whose decode exceeded it failed even when the rewritten + // output would have fit, and the identity, deflate and brotli paths + // accepted the same body. + let cap = STREAM_CHUNK_SIZE * 4; + let document = incompressible_bytes(STREAM_CHUNK_SIZE * 32); + let compressed = gzip_member(&document); + assert!( + compressed.len() > cap, + "test precondition: the member must span several source reads, got {} bytes", + compressed.len() + ); + let mut reader = GzipDecodeReader::new(std::io::Cursor::new(compressed), cap); + + let mut decoded = Vec::new(); + std::io::copy(&mut reader, &mut decoded) + .expect("a valid body decoding past the cap must still stream through"); + + assert_eq!( + decoded, document, + "should decode every byte of a body larger than the pending cap" + ); + } + #[test] fn bounded_decode_sink_caps_buffer_capacity_at_decoded_limit() { // Companion to `deflate_decode_caps_buffer_capacity_at_decoded_limit` diff --git a/docs/guide/configuration.md b/docs/guide/configuration.md index d686c0ee8..ad91a6026 100644 --- a/docs/guide/configuration.md +++ b/docs/guide/configuration.md @@ -307,7 +307,12 @@ and the per-stream raw/decoded byte ceiling on the Fastly streaming path. **Usage**: - On **buffered adapters** (Axum, Cloudflare, Spin) it caps the _decoded, - post-rewrite_ output buffer for a publisher response processed in full. + post-rewrite_ output buffer for a publisher response processed in full. It + also bounds how much decoded gzip output may sit in the heap at any one + moment, so a decompression bomb is rejected mid-decode rather than after its + full expansion. That second bound is per-step, not a total: a gzip-encoded + response passes or fails on the same post-rewrite output size as the identity, + deflate and brotli versions of the same body. - On the **Fastly streaming path** the origin body is preserved as a stream, so the same value caps the stream twice over: the cumulative _raw_ (still compressed) bytes pulled from origin, and the cumulative _decoded_ bytes