RustProxy/src/transform/compile.rs

pub trait Transformer {
    fn transform(&self, body: &[u8]) -> Result<Vec<u8>>;
}

struct JsTransformer;
impl Transformer for JsTransformer {
    fn transform(&self, body: &[u8]) -> Result<Vec<u8>> {
        js::minify(body)
    }
}
struct CssTransformer;
impl Transformer for CssTransformer {
    fn transform(&self, body: &[u8]) -> Result<Vec<u8>> {
        css::minify(body)
    }
}
struct HtmlTransformer;
impl Transformer for HtmlTransformer {
    fn transform(&self, body: &[u8]) -> Result<Vec<u8>> {
        html::minify(body)
    }
}
struct SvgTransformer;
impl Transformer for SvgTransformer {
    fn transform(&self, body: &[u8]) -> Result<Vec<u8>> {
        svg::optimize(body)
    }
}
struct JsonTransformer;
impl Transformer for JsonTransformer {
    fn transform(&self, body: &[u8]) -> Result<Vec<u8>> {
        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<Kind> {
    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<String> {
    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<Kind> {
    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<Option<&'static str>> {
    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<Cow<'a, [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 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::<String>::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<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)
}