RustProxy/src/transform/apply.rs

pub async fn apply(
    response: Response,
    path: &str,
    settings: &Settings,
    origin_no_transform: bool,
) -> Result<Response> {
    let skip = |mut response: Response, kind: Option<Kind>, 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(::http::header::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(::http::header::CONTENT_LENGTH)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.parse::<usize>().ok())
    else {
        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 {
        return Ok(skip(response, Some(kind), "memory_pressure"));
    };

    let (mut parts, body) = response.into_parts();
    let raw = Limited::new(body, settings.max_body_bytes)
        .collect()
        .await
        .map_err(|error| {
            anyhow::anyhow!("collect bounded static response for transformation: {error}")
        })?
        .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() {
        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());
    drop(decoded);
    drop(raw);
    trim_allocator_after_large_transform(released_bytes);

    let validators = parts
        .extensions
        .get::<OriginValidators>()
        .cloned()
        .unwrap_or_else(|| OriginValidators {
            etag: parts.headers.get(::http::header::ETAG).cloned(),
            last_modified: parts.headers.get(::http::header::LAST_MODIFIED).cloned(),
        });
    parts.headers.remove(::http::header::CONTENT_ENCODING);
    parts.headers.remove(::http::header::CONTENT_LENGTH);
    parts.headers.remove("content-md5");
    parts.headers.remove("digest");
    parts.headers.remove("content-digest");
    parts.headers.remove(::http::header::ACCEPT_RANGES);
    parts.headers.insert(::http::header::ETAG, transformed_etag(&transformed));
    parts.headers.insert(::http::header::CONTENT_LENGTH,
        HeaderValue::from(transformed.len()),
    );
    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)))
}