v0.24.7

src/transform/mod.rs

@@ -3,13 +3,16 @@
 use anyhow::{ensure, Context as _, Result};
 use axum::{body::Body, response::Response};
 use http::{HeaderMap, HeaderValue};
-use http_body_util::BodyExt;
+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::Mutex;
+use std::sync::{
+    atomic::{AtomicUsize, Ordering},
+    Mutex,
+};
 use std::time::Duration;
 
 mod css;
@@ -19,9 +22,23 @@ 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
@@ -65,6 +82,10 @@ fn trim_allocator_after_large_transform(released_bytes: usize) {
 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<Preset>,
     pub javascript: bool,
     pub css: bool,
     pub html: bool,
@@ -75,11 +96,18 @@ pub struct Settings {
     pub extensions: Vec<String>,
     /// 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,
@@ -92,6 +120,8 @@ impl Default for Settings {
             .map(str::to_owned)
             .collect(),
             max_body_bytes: DEFAULT_MAX_BODY_BYTES,
+            memory_budget_bytes: DEFAULT_MEMORY_BUDGET_BYTES,
+            max_inflight: DEFAULT_MAX_INFLIGHT,
         }
     }
 }
@@ -101,6 +131,14 @@ impl Settings {
             (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 {
@@ -296,6 +334,24 @@ pub struct MetricsView {
     pub scopes: Vec<ScopeCounters>,
 }
 
+#[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,
@@ -505,14 +561,27 @@ fn supported_encoding(headers: &HeaderMap) -> Option<Option<&'static str>> {
 }
 
 fn decode<'a>(body: &'a [u8], encoding: Option<&str>, limit: usize) -> Result<Cow<'a, [u8]>> {
-    fn read_bounded(reader: impl Read, limit: usize, label: &'static str) -> Result<Vec<u8>> {
+    fn read_bounded(mut reader: impl Read, limit: usize, label: &'static str) -> Result<Vec<u8>> {
         let mut output = Vec::with_capacity(limit.min(64 * 1024));
-        let mut bounded = reader.take(limit.saturating_add(1) as u64);
-        bounded.read_to_end(&mut output).with_context(|| format!("decode {label}"))?;
-        ensure!(
-            output.len() <= limit,
-            "decoded transform body exceeds configured limit"
-        );
+        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)
     }
 
@@ -588,15 +657,123 @@ fn drop_vary_accept_encoding(headers: &mut HeaderMap) {
 
 fn transformed_etag(body: &[u8]) -> HeaderValue {
     let digest = Sha256::digest(body);
-    let mut hex = String::with_capacity(64);
+    let mut value = String::with_capacity(76);
+    value.push_str("\"rustproxy-");
     for byte in digest {
         use std::fmt::Write as _;
-        let _ = write!(&mut hex, "{byte:02x}");
+        let _ = write!(&mut value, "{byte:02x}");
     }
-    HeaderValue::from_str(&format!("\"rustproxy-{hex}\""))
+    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<usize> {
+    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<TransformMemoryReservation<'_>> {
+    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<TransformMemoryReservation<'static>> {
+    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<TransformInflightReservation<'_>> {
+    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<TransformInflightReservation<'static>> {
+    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.
@@ -651,12 +828,33 @@ pub async fn apply(
     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();
-    let raw = body
+    // 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 static response for transformation")?
+        .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));
@@ -751,6 +949,114 @@ mod tests {
         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();