//! Static-response transformation stage. Transformations run on the identity //! representation before shared-cache storage and before downstream compression. use anyhow::{ensure, Context as _, Result}; use axum::{body::Body, response::Response}; use http::{HeaderMap, HeaderValue}; use http_body_util::{BodyExt, Limited}; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use std::borrow::Cow; use std::collections::BTreeMap; use std::io::Read; use std::sync::{ atomic::{AtomicUsize, Ordering}, Mutex, }; use std::time::Duration; mod css; mod html; mod js; mod json; mod svg; const DEFAULT_MAX_BODY_BYTES: usize = 2 * 1024 * 1024; const DEFAULT_MEMORY_BUDGET_BYTES: usize = 128 * 1024 * 1024; const DEFAULT_MAX_INFLIGHT: usize = 32; const TRANSFORM_MEMORY_UNIT_BYTES: usize = 64 * 1024; static TRANSFORM_MEMORY_IN_USE_BYTES: AtomicUsize = AtomicUsize::new(0); static TRANSFORM_INFLIGHT: AtomicUsize = AtomicUsize::new(0); #[cfg(all(target_os = "linux", target_env = "gnu"))] const ALLOCATOR_TRIM_MIN_BODY_BYTES: usize = 256 * 1024; #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] #[serde(rename_all = "snake_case")] pub enum Preset { Conservative, Recommended, Performance, Uber, } /// Large transformable responses can temporarily own the collected origin body, /// a decoded buffer and the transformed output at the same time. On glibc, freed /// pages may remain in malloc arenas for a long time. The systemd unit also caps /// arena proliferation with `MALLOC_ARENA_MAX=2`; this low-frequency trim handles /// the large transient buffers without adding a trim syscall to every request. fn trim_allocator_after_large_transform(released_bytes: usize) { #[cfg(all(target_os = "linux", target_env = "gnu"))] { use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; if released_bytes < ALLOCATOR_TRIM_MIN_BODY_BYTES { return; } static LAST_TRIM_SECONDS: AtomicU64 = AtomicU64::new(0); let now = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs(); let previous = LAST_TRIM_SECONDS.load(Ordering::Relaxed); if now.saturating_sub(previous) < 30 { return; } if LAST_TRIM_SECONDS .compare_exchange(previous, now, Ordering::Relaxed, Ordering::Relaxed) .is_ok() { unsafe { libc::malloc_trim(0); } } } #[cfg(not(all(target_os = "linux", target_env = "gnu")))] { let _ = released_bytes; } } #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(default, deny_unknown_fields)] pub struct Settings { /// Global kill switch. Per-vhost opt-in is still required. pub enabled: bool, /// Last resource profile explicitly selected in the administrator UI. The /// numeric limits below are authoritative; this field only preserves UI /// intent across reloads. pub preset: Option, pub javascript: bool, pub css: bool, pub html: bool, pub svg: bool, pub json: bool, /// Optional extension allow-list. Content-Type recognition remains available, /// but an extension outside this list is never transformed by extension alone. pub extensions: Vec, /// Hard aggregation/decode ceiling for transformable static responses. pub max_body_bytes: usize, /// Process-wide ceiling for estimated temporary response buffers owned by /// active transformations. This is a limit, not a pre-allocation. pub memory_budget_bytes: usize, /// Process-wide number of transformation attempts allowed to own an active /// transform slot at the same time. Excess work fails open unchanged. pub max_inflight: usize, } impl Default for Settings { fn default() -> Self { Self { enabled: false, preset: None, javascript: true, css: true, html: false, svg: false, json: false, extensions: vec![ "js", "mjs", "cjs", "css", "html", "htm", "svg", "json", ] .into_iter() .map(str::to_owned) .collect(), max_body_bytes: DEFAULT_MAX_BODY_BYTES, memory_budget_bytes: DEFAULT_MEMORY_BUDGET_BYTES, max_inflight: DEFAULT_MAX_INFLIGHT, } } } impl Settings { pub fn validate(&self) -> Result<()> { ensure!( (1024..=32 * 1024 * 1024).contains(&self.max_body_bytes), "transform.max_body_bytes must be 1 KiB..32 MiB" ); ensure!( (1024 * 1024..=2 * 1024 * 1024 * 1024).contains(&self.memory_budget_bytes), "transform.memory_budget_bytes must be 1 MiB..2 GiB" ); ensure!( (1..=1024).contains(&self.max_inflight), "transform.max_inflight must be 1..1024" ); ensure!(self.extensions.len() <= 64, "at most 64 transform extensions"); let mut seen = std::collections::HashSet::new(); for extension in &self.extensions { ensure!( !extension.is_empty() && extension.len() <= 16 && extension.bytes().all(|b| b.is_ascii_alphanumeric()), "transform extensions must be 1..16 ASCII letters or numbers without a dot" ); ensure!( seen.insert(extension.to_ascii_lowercase()), "duplicate transform extension" ); } Ok(()) } } /// Per-vhost transformation policy. Unlike cache overrides, transformation is /// deliberately opt-in per vhost even when the global engine is enabled. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(default, deny_unknown_fields)] pub struct Overrides { pub enabled: bool, pub javascript: Option, pub css: Option, pub html: Option, pub svg: Option, pub json: Option, pub extensions: Option>, } impl Default for Overrides { fn default() -> Self { Self { enabled: false, javascript: None, css: None, html: None, svg: None, json: None, extensions: None, } } } impl Overrides { pub fn resolve(&self, global: &Settings) -> Result { let mut resolved = global.clone(); resolved.enabled = global.enabled && self.enabled; if let Some(value) = self.javascript { resolved.javascript = value; } if let Some(value) = self.css { resolved.css = value; } if let Some(value) = self.html { resolved.html = value; } if let Some(value) = self.svg { resolved.svg = value; } if let Some(value) = self.json { resolved.json = value; } if let Some(value) = &self.extensions { resolved.extensions = value.clone(); } resolved.validate()?; Ok(resolved) } } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] #[serde(rename_all = "snake_case")] pub enum Kind { Javascript, Css, Html, Svg, Json, } impl Kind { pub fn as_str(self) -> &'static str { match self { Self::Javascript => "javascript", Self::Css => "css", Self::Html => "html", Self::Svg => "svg", Self::Json => "json", } } } /// Origin validators are retained separately because transformed bytes receive /// a RustProxy validator that must never be sent back to the origin. #[derive(Debug, Clone)] pub struct OriginValidators { pub etag: Option, pub last_modified: Option, } #[derive(Debug, Clone, Copy)] pub struct Applied { pub kind: Kind, pub original_bytes: usize, pub transformed_bytes: usize, } /// Internal observability marker for a response which was examined by the /// transform stage but intentionally left unchanged. Reasons are fixed tokens /// and never contain request or response data. #[derive(Debug, Clone, Copy)] pub struct Skipped { pub kind: Option, pub reason: &'static str, } /// Marker restored on responses served from a cached transformed /// representation. It lets diagnostics distinguish a transformed HIT from an /// ordinary cached object without re-running the transformer. #[derive(Debug, Clone, Copy)] pub struct CachedTransformed; #[derive(Debug, Clone, Default, Serialize)] pub struct Counters { pub examined: u64, pub applied: u64, pub skipped: u64, pub errors: u64, pub cache_hits: u64, pub original_bytes: u64, pub transformed_bytes: u64, pub saved_bytes: u64, pub duration_us: u64, pub by_kind: BTreeMap, pub reasons: BTreeMap, } impl Counters { fn record_applied(&mut self, applied: Applied, elapsed: Duration) { self.examined = self.examined.saturating_add(1); self.applied = self.applied.saturating_add(1); self.duration_us = self .duration_us .saturating_add(elapsed.as_micros().min(u64::MAX as u128) as u64); self.original_bytes = self .original_bytes .saturating_add(applied.original_bytes as u64); self.transformed_bytes = self .transformed_bytes .saturating_add(applied.transformed_bytes as u64); self.saved_bytes = self.saved_bytes.saturating_add( applied .original_bytes .saturating_sub(applied.transformed_bytes) as u64, ); *self.by_kind.entry(applied.kind.as_str().into()).or_default() += 1; } fn record_skipped(&mut self, skipped: Skipped, elapsed: Duration) { self.examined = self.examined.saturating_add(1); self.skipped = self.skipped.saturating_add(1); self.duration_us = self .duration_us .saturating_add(elapsed.as_micros().min(u64::MAX as u128) as u64); *self.reasons.entry(skipped.reason.into()).or_default() += 1; } fn record_error(&mut self, elapsed: Duration) { self.examined = self.examined.saturating_add(1); self.errors = self.errors.saturating_add(1); self.duration_us = self .duration_us .saturating_add(elapsed.as_micros().min(u64::MAX as u128) as u64); *self.reasons.entry("body_error".into()).or_default() += 1; } fn record_cache_hit(&mut self) { self.cache_hits = self.cache_hits.saturating_add(1); } } #[derive(Debug, Clone, Serialize)] pub struct ScopeCounters { pub vhost: String, #[serde(flatten)] pub counters: Counters, } #[derive(Debug, Clone, Default, Serialize)] pub struct MetricsView { pub total: Counters, pub scopes: Vec, } #[derive(Debug, Clone, Copy, Serialize)] pub struct ResourceUsage { /// Estimated temporary transform bytes currently reserved by active work. pub memory_reserved_bytes: usize, pub memory_budget_bytes: usize, pub inflight: usize, pub max_inflight: usize, } pub fn resource_usage(settings: &Settings) -> ResourceUsage { ResourceUsage { memory_reserved_bytes: TRANSFORM_MEMORY_IN_USE_BYTES.load(Ordering::Acquire), memory_budget_bytes: settings.memory_budget_bytes, inflight: TRANSFORM_INFLIGHT.load(Ordering::Acquire), max_inflight: settings.max_inflight, } } #[derive(Default)] struct MetricsStore { total: Counters, scopes: BTreeMap, } #[derive(Default)] pub struct Metrics { inner: Mutex, } impl Metrics { fn with_scope(&self, scope: &str, mut update: impl FnMut(&mut Counters)) { let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); update(&mut inner.total); if inner.scopes.len() < 4096 || inner.scopes.contains_key(scope) { update(inner.scopes.entry(scope.to_owned()).or_default()); } } pub fn record_response(&self, scope: &str, response: &Response, elapsed: Duration) { if let Some(applied) = response.extensions().get::().copied() { self.with_scope(scope, |counters| counters.record_applied(applied, elapsed)); return; } let skipped = response .extensions() .get::() .copied() .unwrap_or(Skipped { kind: None, reason: "not_applied", }); self.with_scope(scope, |counters| counters.record_skipped(skipped, elapsed)); } pub fn record_error(&self, scope: &str, elapsed: Duration) { self.with_scope(scope, |counters| counters.record_error(elapsed)); } pub fn record_cache_hit(&self, scope: &str) { self.with_scope(scope, Counters::record_cache_hit); } pub fn total(&self) -> Counters { self.inner .lock() .unwrap_or_else(|e| e.into_inner()) .total .clone() } pub fn view(&self) -> MetricsView { let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); MetricsView { total: inner.total.clone(), scopes: inner .scopes .iter() .map(|(vhost, counters)| ScopeCounters { vhost: vhost.clone(), counters: counters.clone(), }) .collect(), } } } /// Marks a response whose upstream content-negotiation dimension has been /// normalized to the identity representation. The shared cache can therefore /// ignore origin `Vary: Accept-Encoding`; downstream compression recreates the /// client-specific encoding after the cache lookup. #[derive(Debug, Clone, Copy)] pub struct CanonicalIdentity; pub trait Transformer { fn transform(&self, body: &[u8]) -> Result>; } struct JsTransformer; impl Transformer for JsTransformer { fn transform(&self, body: &[u8]) -> Result> { js::minify(body) } } struct CssTransformer; impl Transformer for CssTransformer { fn transform(&self, body: &[u8]) -> Result> { css::minify(body) } } struct HtmlTransformer; impl Transformer for HtmlTransformer { fn transform(&self, body: &[u8]) -> Result> { html::minify(body) } } struct SvgTransformer; impl Transformer for SvgTransformer { fn transform(&self, body: &[u8]) -> Result> { svg::optimize(body) } } struct JsonTransformer; impl Transformer for JsonTransformer { fn transform(&self, body: &[u8]) -> Result> { json::compact(body) } } fn transformer(kind: Kind) -> &'static dyn Transformer { static JS: JsTransformer = JsTransformer; static CSS: CssTransformer = CssTransformer; static HTML: HtmlTransformer = HtmlTransformer; static SVG: SvgTransformer = SvgTransformer; static JSON: JsonTransformer = JsonTransformer; match kind { Kind::Javascript => &JS, Kind::Css => &CSS, Kind::Html => &HTML, Kind::Svg => &SVG, Kind::Json => &JSON, } } pub fn no_transform(headers: &HeaderMap) -> bool { headers .get_all("cache-control") .iter() .filter_map(|value| value.to_str().ok()) .any(|value| { value.split(',').any(|directive| { directive .trim() .split('=') .next() .is_some_and(|name| name.trim().eq_ignore_ascii_case("no-transform")) }) }) } fn content_type_kind(headers: &HeaderMap) -> Option { let value = headers.get("content-type")?.to_str().ok()?; let mime = value.split(';').next()?.trim().to_ascii_lowercase(); match mime.as_str() { "application/javascript" | "text/javascript" | "application/ecmascript" | "text/ecmascript" => Some(Kind::Javascript), "text/css" => Some(Kind::Css), "text/html" | "application/xhtml+xml" => Some(Kind::Html), "image/svg+xml" => Some(Kind::Svg), "application/json" | "text/json" => Some(Kind::Json), _ if mime.ends_with("+json") => Some(Kind::Json), _ => None, } } fn extension(path: &str) -> Option { let path = path.split('?').next().unwrap_or(path); let segment = path.rsplit('/').next().unwrap_or(path); let (_, extension) = segment.rsplit_once('.')?; (!extension.is_empty()).then(|| extension.to_ascii_lowercase()) } fn extension_kind(path: &str, settings: &Settings) -> Option { let extension = extension(path)?; if !settings .extensions .iter() .any(|allowed| allowed.eq_ignore_ascii_case(&extension)) { return None; } match extension.as_str() { "js" | "mjs" | "cjs" => Some(Kind::Javascript), "css" => Some(Kind::Css), "html" | "htm" => Some(Kind::Html), "svg" => Some(Kind::Svg), "json" => Some(Kind::Json), _ => None, } } fn kind_enabled(kind: Kind, settings: &Settings) -> bool { match kind { Kind::Javascript => settings.javascript, Kind::Css => settings.css, Kind::Html => settings.html, Kind::Svg => settings.svg, Kind::Json => settings.json, } } fn supported_encoding(headers: &HeaderMap) -> Option> { let Some(value) = headers.get("content-encoding") else { return Some(None); }; let value = value.to_str().ok()?.trim(); if value.eq_ignore_ascii_case("identity") { return Some(None); } for known in ["gzip", "deflate", "br", "zstd"] { if value.eq_ignore_ascii_case(known) { return Some(Some(known)); } } None } fn decode<'a>(body: &'a [u8], encoding: Option<&str>, limit: usize) -> Result> { fn read_bounded(mut reader: impl Read, limit: usize, label: &'static str) -> Result> { let mut output = Vec::with_capacity(limit.min(64 * 1024)); let mut chunk = [0u8; 16 * 1024]; loop { let read = reader .read(&mut chunk) .with_context(|| format!("decode {label}"))?; if read == 0 { break; } ensure!( output.len().saturating_add(read) <= limit, "decoded transform body exceeds configured limit" ); if output.capacity().saturating_sub(output.len()) < read { output .try_reserve_exact(read) .map_err(|_| anyhow::anyhow!("decoded transform allocation failed"))?; } output.extend_from_slice(&chunk[..read]); } Ok(output) } match encoding { None => { ensure!( body.len() <= limit, "transform body exceeds configured limit" ); // Identity bodies are already owned by `Bytes` after collection. // Borrow them instead of creating a second full-size allocation. Ok(Cow::Borrowed(body)) } Some("gzip") => read_bounded(flate2::read::GzDecoder::new(body), limit, "gzip") .map(Cow::Owned), Some("deflate") => { let zlib = flate2::read::ZlibDecoder::new(body); match read_bounded(zlib, limit, "deflate") { Ok(value) => Ok(Cow::Owned(value)), Err(_) => read_bounded( flate2::read::DeflateDecoder::new(body), limit, "deflate", ) .map(Cow::Owned), } } Some("br") => read_bounded( brotli::Decompressor::new(body, 16 * 1024), limit, "brotli", ) .map(Cow::Owned), Some("zstd") => { let decoder = zstd::stream::read::Decoder::new(body).context("open zstd decoder")?; read_bounded(decoder, limit, "zstd").map(Cow::Owned) } Some(_) => unreachable!(), } } fn drop_vary_accept_encoding(headers: &mut HeaderMap) { let values: Vec<_> = headers.get_all("vary").iter().cloned().collect(); if values.is_empty() { return; } let mut retained = Vec::::new(); for value in values { let Ok(value) = value.to_str() else { return; }; for token in value.split(',').map(str::trim).filter(|token| !token.is_empty()) { if token == "*" { // `Vary: *` has special semantics and is rejected by the shared // cache. Do not rewrite it into a different policy. return; } if token.eq_ignore_ascii_case("accept-encoding") { continue; } if !retained.iter().any(|existing| existing.eq_ignore_ascii_case(token)) { retained.push(token.to_owned()); } } } headers.remove("vary"); if !retained.is_empty() { if let Ok(value) = HeaderValue::from_str(&retained.join(", ")) { headers.insert("vary", value); } } } fn transformed_etag(body: &[u8]) -> HeaderValue { let digest = Sha256::digest(body); let mut value = String::with_capacity(76); value.push_str("\"rustproxy-"); for byte in digest { use std::fmt::Write as _; let _ = write!(&mut value, "{byte:02x}"); } value.push('"'); HeaderValue::from_str(&value) .expect("sha256 validator is always a valid header") } fn transform_memory_required_bytes( content_length: usize, encoding: Option<&str>, max_body_bytes: usize, ) -> Option { let estimated = if encoding.is_none() { content_length.checked_mul(2)? } else { content_length.checked_add(max_body_bytes.checked_mul(2)?)? }; let units = estimated.checked_add(TRANSFORM_MEMORY_UNIT_BYTES - 1)? / TRANSFORM_MEMORY_UNIT_BYTES; units.max(1).checked_mul(TRANSFORM_MEMORY_UNIT_BYTES) } struct TransformMemoryReservation<'a> { counter: &'a AtomicUsize, bytes: usize, } impl Drop for TransformMemoryReservation<'_> { fn drop(&mut self) { self.counter.fetch_sub(self.bytes, Ordering::AcqRel); } } fn try_reserve_memory_from( counter: &AtomicUsize, required_bytes: usize, budget_bytes: usize, ) -> Option> { if required_bytes > budget_bytes { return None; } let mut current = counter.load(Ordering::Acquire); loop { let next = current.checked_add(required_bytes)?; if next > budget_bytes { return None; } match counter.compare_exchange_weak( current, next, Ordering::AcqRel, Ordering::Acquire, ) { Ok(_) => { return Some(TransformMemoryReservation { counter, bytes: required_bytes, }) } Err(actual) => current = actual, } } } fn try_reserve_transform_memory( required_bytes: usize, budget_bytes: usize, ) -> Option> { try_reserve_memory_from( &TRANSFORM_MEMORY_IN_USE_BYTES, required_bytes, budget_bytes, ) } struct TransformInflightReservation<'a> { counter: &'a AtomicUsize, } impl Drop for TransformInflightReservation<'_> { fn drop(&mut self) { self.counter.fetch_sub(1, Ordering::AcqRel); } } fn try_reserve_slot_from( counter: &AtomicUsize, max_inflight: usize, ) -> Option> { let mut current = counter.load(Ordering::Acquire); loop { if current >= max_inflight { return None; } let next = current.checked_add(1)?; match counter.compare_exchange_weak( current, next, Ordering::AcqRel, Ordering::Acquire, ) { Ok(_) => return Some(TransformInflightReservation { counter }), Err(actual) => current = actual, } } } fn try_reserve_transform_slot( max_inflight: usize, ) -> Option> { try_reserve_slot_from(&TRANSFORM_INFLIGHT, max_inflight) } /// Transform an eligible response. The call is fail-open for unsupported media, /// encodings and malformed minifiable content; transport/body collection errors /// still propagate because the original stream can no longer be replayed safely. pub async fn apply( response: Response, path: &str, settings: &Settings, origin_no_transform: bool, ) -> Result { let skip = |mut response: Response, kind: Option, reason: &'static str| { response.extensions_mut().insert(Skipped { kind, reason }); response }; if !settings.enabled { return Ok(skip(response, None, "disabled")); } if origin_no_transform { return Ok(skip(response, None, "origin_no_transform")); } if no_transform(response.headers()) { return Ok(skip(response, None, "response_no_transform")); } if response.status().is_informational() || response.status().as_u16() == 204 || response.status().as_u16() == 304 { return Ok(skip(response, None, "status_ineligible")); } if response.headers().contains_key("content-range") { return Ok(skip(response, None, "content_range")); } let Some(kind) = content_type_kind(response.headers()).or_else(|| extension_kind(path, settings)) else { return Ok(skip(response, None, "unsupported_type")); }; if !kind_enabled(kind, settings) { return Ok(skip(response, Some(kind), "type_disabled")); } let Some(encoding) = supported_encoding(response.headers()) else { return Ok(skip(response, Some(kind), "unsupported_encoding")); }; let Some(content_length) = response .headers() .get("content-length") .and_then(|value| value.to_str().ok()) .and_then(|value| value.parse::().ok()) else { // Keep streaming semantics for unknown-size bodies instead of introducing // unbounded buffering into the dataplane. return Ok(skip(response, Some(kind), "unknown_length")); }; if content_length > settings.max_body_bytes { return Ok(skip(response, Some(kind), "too_large")); } let Some(required_memory) = transform_memory_required_bytes(content_length, encoding, settings.max_body_bytes) else { return Ok(skip(response, Some(kind), "memory_budget")); }; let Some(_inflight_slot) = try_reserve_transform_slot(settings.max_inflight) else { return Ok(skip(response, Some(kind), "concurrency_pressure")); }; let Some(_memory_reservation) = try_reserve_transform_memory(required_memory, settings.memory_budget_bytes) else { // Transformations are an optimization. Under concurrent allocation // pressure, preserve streaming/origin bytes instead of queueing more // full-response buffers behind the memory budget. return Ok(skip(response, Some(kind), "memory_pressure")); }; let (mut parts, body) = response.into_parts(); // Do not trust Content-Length as a memory bound. A broken or malicious // upstream can send more bytes than declared; Limited stops aggregation at // the configured ceiling instead of letting BodyExt::collect grow without // bound. Once aggregation starts an over-limit stream cannot be replayed // safely, so it is treated like the other body/transport collection errors. let raw = Limited::new(body, settings.max_body_bytes) .collect() .await .context("collect bounded static response for transformation")? .to_bytes(); if raw.len() > settings.max_body_bytes { let mut response = Response::from_parts(parts, Body::from(raw)); response.extensions_mut().insert(Skipped { kind: Some(kind), reason: "too_large", }); return Ok(response); } let decoded = match decode(&raw, encoding, settings.max_body_bytes) { Ok(value) => value, Err(_) => { let mut response = Response::from_parts(parts, Body::from(raw)); response.extensions_mut().insert(Skipped { kind: Some(kind), reason: "decode_fail_open", }); return Ok(response); } }; let transformed = match transformer(kind).transform(decoded.as_ref()) { Ok(value) if !value.is_empty() => value, _ => { let mut response = Response::from_parts(parts, Body::from(raw)); response.extensions_mut().insert(Skipped { kind: Some(kind), reason: "transform_fail_open", }); return Ok(response); } }; if transformed.as_slice() == decoded.as_ref() && encoding.is_none() { // The origin was explicitly requested without Accept-Encoding and the // eligible response is already identity. The bytes/validator stay intact, // but the cache no longer needs an Accept-Encoding variant dimension. drop_vary_accept_encoding(&mut parts.headers); parts.extensions.insert(CanonicalIdentity); parts.extensions.insert(Skipped { kind: Some(kind), reason: "unchanged", }); return Ok(Response::from_parts(parts, Body::from(raw))); } let original_bytes = decoded.len(); let released_bytes = raw.len().saturating_add(decoded.len()); // The successful transformed response no longer needs the collected origin // body or decoded temporary buffer. Drop both before building the response so // large short-lived allocations do not overlap longer than necessary. drop(decoded); drop(raw); trim_allocator_after_large_transform(released_bytes); let validators = parts .extensions .get::() .cloned() .unwrap_or_else(|| OriginValidators { etag: parts.headers.get("etag").cloned(), last_modified: parts.headers.get("last-modified").cloned(), }); parts.headers.remove("content-encoding"); parts.headers.remove("content-length"); parts.headers.remove("content-md5"); parts.headers.remove("digest"); parts.headers.remove("content-digest"); parts.headers.remove("accept-ranges"); parts.headers.insert("etag", transformed_etag(&transformed)); parts.headers.insert( "content-length", HeaderValue::from_str(&transformed.len().to_string()).expect("integer header"), ); drop_vary_accept_encoding(&mut parts.headers); parts.extensions.insert(validators); parts.extensions.insert(CanonicalIdentity); parts.extensions.insert(Applied { kind, original_bytes, transformed_bytes: transformed.len(), }); Ok(Response::from_parts(parts, Body::from(transformed))) } #[cfg(test)] mod tests { use super::*; #[test] fn identity_decode_borrows_the_collected_body() { let body = b"const value = 1;"; let decoded = decode(body, None, 1024).unwrap(); assert!(matches!(decoded, Cow::Borrowed(_))); } #[test] fn transform_memory_budget_accounts_for_identity_and_decoded_buffers() { assert_eq!( transform_memory_required_bytes(2 * 1024 * 1024, None, 2 * 1024 * 1024), Some(4 * 1024 * 1024) ); assert_eq!( transform_memory_required_bytes( 2 * 1024 * 1024, Some("gzip"), 2 * 1024 * 1024 ), Some(6 * 1024 * 1024) ); assert_eq!( transform_memory_required_bytes( DEFAULT_MEMORY_BUDGET_BYTES, Some("gzip"), DEFAULT_MEMORY_BUDGET_BYTES ), Some(3 * DEFAULT_MEMORY_BUDGET_BYTES) ); } #[test] fn transform_memory_reservation_obeys_runtime_budget() { let counter = AtomicUsize::new(0); let required = 4 * 1024 * 1024; let reservation = try_reserve_memory_from(&counter, required, required).unwrap(); assert!(try_reserve_memory_from(&counter, required, required).is_none()); drop(reservation); assert!(try_reserve_memory_from(&counter, required, required).is_some()); } #[test] fn transform_inflight_reservation_obeys_runtime_limit() { let counter = AtomicUsize::new(0); let slot = try_reserve_slot_from(&counter, 1).unwrap(); assert!(try_reserve_slot_from(&counter, 1).is_none()); drop(slot); assert!(try_reserve_slot_from(&counter, 1).is_some()); } #[test] fn resource_usage_exposes_configured_limits() { let settings = Settings { memory_budget_bytes: 512 * 1024 * 1024, max_inflight: 64, ..Default::default() }; let usage = resource_usage(&settings); assert_eq!(usage.memory_budget_bytes, 512 * 1024 * 1024); assert_eq!(usage.max_inflight, 64); } #[test] fn resource_preset_metadata_round_trips() { let mut settings = Settings::default(); settings.preset = Some(Preset::Uber); settings.memory_budget_bytes = 1024 * 1024 * 1024; settings.max_body_bytes = 32 * 1024 * 1024; settings.max_inflight = 128; let encoded = serde_json::to_value(&settings).unwrap(); assert_eq!(encoded["preset"], "uber"); let decoded: Settings = serde_json::from_value(encoded).unwrap(); assert_eq!(decoded.preset, Some(Preset::Uber)); assert_eq!(decoded.memory_budget_bytes, 1024 * 1024 * 1024); assert_eq!(decoded.max_inflight, 128); } #[test] fn legacy_transform_settings_receive_resource_defaults() { let decoded: Settings = serde_json::from_str(r#"{"enabled":true}"#).unwrap(); assert!(decoded.enabled); assert_eq!(decoded.preset, None); assert_eq!(decoded.memory_budget_bytes, DEFAULT_MEMORY_BUDGET_BYTES); assert_eq!(decoded.max_body_bytes, DEFAULT_MAX_BODY_BYTES); assert_eq!(decoded.max_inflight, DEFAULT_MAX_INFLIGHT); } #[test] fn bounded_decoder_rejects_expansion_past_limit() { use std::io::Write as _; let identity = vec![b'x'; 128 * 1024]; let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); encoder.write_all(&identity).unwrap(); let encoded = encoder.finish().unwrap(); assert!(decode(&encoded, Some("gzip"), 64 * 1024).is_err()); } #[tokio::test] async fn aggregation_is_bounded_even_when_content_length_is_false() { let settings = Settings { enabled: true, json: true, max_body_bytes: 1024, ..Default::default() }; let mut response = Response::new(Body::from(vec![b' '; 2048])); response .headers_mut() .insert("content-type", HeaderValue::from_static("application/json")); response .headers_mut() .insert("content-length", HeaderValue::from_static("128")); assert!(apply(response, "/data.json", &settings, false).await.is_err()); } #[test] fn metrics_track_applied_skipped_and_cached_transformed_responses() { let metrics = Metrics::default(); let mut applied = Response::new(Body::empty()); applied.extensions_mut().insert(Applied { kind: Kind::Css, original_bytes: 100, transformed_bytes: 70, }); metrics.record_response("app", &applied, Duration::from_micros(500)); let mut skipped = Response::new(Body::empty()); skipped.extensions_mut().insert(Skipped { kind: Some(Kind::Css), reason: "response_no_transform", }); metrics.record_response("app", &skipped, Duration::from_micros(100)); metrics.record_cache_hit("app"); let total = metrics.total(); assert_eq!(total.applied, 1); assert_eq!(total.cache_hits, 1); let view = metrics.view(); assert_eq!(view.total.examined, 2); assert_eq!(view.total.applied, 1); assert_eq!(view.total.skipped, 1); assert_eq!(view.total.cache_hits, 1); assert_eq!(view.total.saved_bytes, 30); assert_eq!(view.total.by_kind.get("css"), Some(&1)); assert_eq!( view.total.reasons.get("response_no_transform"), Some(&1) ); assert_eq!(view.scopes.len(), 1); assert_eq!(view.scopes[0].vhost, "app"); } #[test] fn per_vhost_is_opt_in_even_when_global_is_enabled() { let global = Settings { enabled: true, ..Default::default() }; assert!(!Overrides::default().resolve(&global).unwrap().enabled); let enabled = Overrides { enabled: true, ..Default::default() }; assert!(enabled.resolve(&global).unwrap().enabled); } #[tokio::test] async fn json_is_compacted_and_gets_a_new_validator() { let settings = Settings { enabled: true, json: true, ..Default::default() }; let mut response = Response::new(Body::from("{ \"a\": 1, \"b\": [ true ] }")); response .headers_mut() .insert("content-type", HeaderValue::from_static("application/json")); response .headers_mut() .insert("content-length", HeaderValue::from_static("30")); response .headers_mut() .insert("etag", HeaderValue::from_static("\"origin\"")); response .headers_mut() .insert("vary", HeaderValue::from_static("Accept-Encoding")); let response = apply(response, "/data.json", &settings, false).await.unwrap(); assert!(response.headers()["etag"].to_str().unwrap().starts_with("\"rustproxy-")); assert!(!response.headers().contains_key("vary")); assert!(response.extensions().get::().is_some()); let validators = response.extensions().get::().unwrap(); assert_eq!(validators.etag.as_ref().unwrap(), "\"origin\""); let body = response.into_body().collect().await.unwrap().to_bytes(); assert_eq!(&body[..], br#"{"a":1,"b":[true]}"#); } #[tokio::test] async fn unchanged_body_keeps_origin_validator() { let settings = Settings { enabled: true, json: true, ..Default::default() }; let body = br#"{"a":1}"#; let mut response = Response::new(Body::from(body.as_slice())); response .headers_mut() .insert("content-type", HeaderValue::from_static("application/json")); response .headers_mut() .insert("content-length", HeaderValue::from_static("7")); response .headers_mut() .insert("etag", HeaderValue::from_static("\"origin\"")); response.headers_mut().insert( "vary", HeaderValue::from_static("Accept-Encoding, X-Theme"), ); let response = apply(response, "/data.json", &settings, false).await.unwrap(); assert_eq!(response.headers()["etag"], "\"origin\""); assert_eq!(response.headers()["vary"], "X-Theme"); assert!(response.extensions().get::().is_some()); assert!(response.extensions().get::().is_none()); } #[tokio::test] async fn unchanged_compressed_body_is_canonicalized_to_identity() { use std::io::Write as _; let settings = Settings { enabled: true, json: true, ..Default::default() }; let identity = br#"{"a":1}"#; let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); encoder.write_all(identity).unwrap(); let encoded = encoder.finish().unwrap(); let mut response = Response::new(Body::from(encoded.clone())); response .headers_mut() .insert("content-type", HeaderValue::from_static("application/json")); response.headers_mut().insert( "content-length", HeaderValue::from_str(&encoded.len().to_string()).unwrap(), ); response .headers_mut() .insert("content-encoding", HeaderValue::from_static("gzip")); response .headers_mut() .insert("etag", HeaderValue::from_static("\"origin\"")); let response = apply(response, "/data.json", &settings, false).await.unwrap(); assert!(!response.headers().contains_key("content-encoding")); assert_eq!( response.headers()["content-length"].to_str().unwrap(), identity.len().to_string() ); assert!(response.headers()["etag"].to_str().unwrap().starts_with("\"rustproxy-")); let validators = response.extensions().get::().unwrap(); assert_eq!(validators.etag.as_ref().unwrap(), "\"origin\""); let applied = response.extensions().get::().unwrap(); assert_eq!(applied.original_bytes, identity.len()); assert_eq!(applied.transformed_bytes, identity.len()); let body = response.into_body().collect().await.unwrap().to_bytes(); assert_eq!(&body[..], identity); } #[tokio::test] async fn origin_no_transform_blocks_body_changes_and_keeps_validator() { let settings = Settings { enabled: true, json: true, ..Default::default() }; let body = br#"{ "a": 1 }"#; let mut response = Response::new(Body::from(body.as_slice())); response .headers_mut() .insert("content-type", HeaderValue::from_static("application/json")); response.headers_mut().insert( "content-length", HeaderValue::from_str(&body.len().to_string()).unwrap(), ); response .headers_mut() .insert("etag", HeaderValue::from_static("\"origin\"")); let response = apply(response, "/data.json", &settings, true).await.unwrap(); assert_eq!(response.headers()["etag"], "\"origin\""); assert!(response.extensions().get::().is_none()); let received = response.into_body().collect().await.unwrap().to_bytes(); assert_eq!(&received[..], body); } #[tokio::test] async fn final_no_transform_blocks_body_changes() { let settings = Settings { enabled: true, json: true, ..Default::default() }; let body = br#"{ "a": 1 }"#; let mut response = Response::new(Body::from(body.as_slice())); response .headers_mut() .insert("content-type", HeaderValue::from_static("application/json")); response.headers_mut().insert( "content-length", HeaderValue::from_str(&body.len().to_string()).unwrap(), ); response.headers_mut().insert( "cache-control", HeaderValue::from_static("public, max-age=60, no-transform"), ); let response = apply(response, "/data.json", &settings, false).await.unwrap(); assert!(response.extensions().get::().is_none()); let received = response.into_body().collect().await.unwrap().to_bytes(); assert_eq!(&received[..], body); } #[tokio::test] async fn empty_head_body_with_origin_length_keeps_origin_representation_headers() { let settings = Settings { enabled: true, json: true, ..Default::default() }; let mut response = Response::new(Body::empty()); response .headers_mut() .insert("content-type", HeaderValue::from_static("application/json")); response .headers_mut() .insert("content-length", HeaderValue::from_static("76")); response .headers_mut() .insert("etag", HeaderValue::from_static("\"origin-head\"")); let response = apply(response, "/data.json", &settings, false).await.unwrap(); assert_eq!(response.headers()["content-length"], "76"); assert_eq!(response.headers()["etag"], "\"origin-head\""); assert!(response.extensions().get::().is_none()); assert!(response.extensions().get::().is_none()); let body = response.into_body().collect().await.unwrap().to_bytes(); assert!(body.is_empty()); } #[tokio::test] async fn compressed_no_transform_preserves_origin_encoding_and_bytes() { use std::io::Write as _; let settings = Settings { enabled: true, json: true, ..Default::default() }; let identity = br#"{ "a": 1 }"#; let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default()); encoder.write_all(identity).unwrap(); let encoded = encoder.finish().unwrap(); let mut response = Response::new(Body::from(encoded.clone())); response .headers_mut() .insert("content-type", HeaderValue::from_static("application/json")); response.headers_mut().insert( "content-length", HeaderValue::from_str(&encoded.len().to_string()).unwrap(), ); response .headers_mut() .insert("content-encoding", HeaderValue::from_static("gzip")); response.headers_mut().insert( "cache-control", HeaderValue::from_static("public, max-age=60, no-transform"), ); response .headers_mut() .insert("vary", HeaderValue::from_static("Accept-Encoding")); response .headers_mut() .insert("etag", HeaderValue::from_static("\"origin\"")); let response = apply(response, "/data.json", &settings, false).await.unwrap(); assert_eq!(response.headers()["content-encoding"], "gzip"); assert_eq!(response.headers()["etag"], "\"origin\""); assert_eq!(response.headers()["vary"], "Accept-Encoding"); assert!(response.extensions().get::().is_none()); let received = response.into_body().collect().await.unwrap().to_bytes(); assert_eq!(&received[..], encoded.as_slice()); } }