#![deny(missing_docs)]
use async_trait::async_trait;
use bytes::Bytes;
use churust_core::{AppBuilder, Body, Call, Middleware, Next, Phase, Plugin, Response};
use futures_util::StreamExt;
use http::header::{
HeaderValue, ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_RANGE, CONTENT_TYPE,
ETAG, VARY,
};
use http::{Method, StatusCode};
use std::sync::Arc;
use tokio::io::AsyncRead;
use tokio_util::io::{ReaderStream, StreamReader};
const CHUNK: usize = 16 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Encoding {
Brotli,
Gzip,
Deflate,
}
impl Encoding {
pub fn token(self) -> &'static str {
match self {
Encoding::Brotli => "br",
Encoding::Gzip => "gzip",
Encoding::Deflate => "deflate",
}
}
fn from_token(token: &str) -> Option<Self> {
match token {
"br" => Some(Encoding::Brotli),
"gzip" | "x-gzip" => Some(Encoding::Gzip),
"deflate" => Some(Encoding::Deflate),
_ => None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Level {
Fastest,
#[default]
Default,
Best,
}
impl From<Level> for async_compression::Level {
fn from(level: Level) -> Self {
match level {
Level::Fastest => async_compression::Level::Fastest,
Level::Default => async_compression::Level::Default,
Level::Best => async_compression::Level::Best,
}
}
}
pub fn compressible_by_default(content_type: &str) -> bool {
let essence = content_type
.split(';')
.next()
.unwrap_or("")
.trim()
.to_ascii_lowercase();
if essence.starts_with("text/") {
return true;
}
if essence.ends_with("+json") || essence.ends_with("+xml") || essence.ends_with("+text") {
return true;
}
matches!(
essence.as_str(),
"application/json"
| "application/javascript"
| "application/xml"
| "application/xhtml+xml"
| "application/rss+xml"
| "application/atom+xml"
| "application/wasm"
| "application/manifest+json"
| "application/x-ndjson"
| "image/svg+xml"
)
}
type CompressibleFn = Arc<dyn Fn(&str) -> bool + Send + Sync>;
#[derive(Clone)]
pub struct Compression {
min_size: usize,
level: Level,
preference: Vec<Encoding>,
compressible: CompressibleFn,
}
impl std::fmt::Debug for Compression {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Compression")
.field("min_size", &self.min_size)
.field("level", &self.level)
.field("preference", &self.preference)
.finish_non_exhaustive()
}
}
impl Default for Compression {
fn default() -> Self {
Self::new()
}
}
impl Compression {
pub fn new() -> Self {
Self {
min_size: 1024,
level: Level::Default,
preference: vec![Encoding::Brotli, Encoding::Gzip, Encoding::Deflate],
compressible: Arc::new(compressible_by_default),
}
}
pub fn min_size(mut self, bytes: usize) -> Self {
self.min_size = bytes;
self
}
pub fn level(mut self, level: Level) -> Self {
self.level = level;
self
}
pub fn encodings(mut self, encodings: impl IntoIterator<Item = Encoding>) -> Self {
let list: Vec<Encoding> = encodings.into_iter().collect();
assert!(
!list.is_empty(),
"Compression needs at least one encoding to offer"
);
self.preference = list;
self
}
pub fn compressible<F>(mut self, f: F) -> Self
where
F: Fn(&str) -> bool + Send + Sync + 'static,
{
self.compressible = Arc::new(f);
self
}
fn negotiate(&self, header: &str) -> Option<Encoding> {
let mut offers: Vec<(Encoding, f32, bool)> = Vec::new();
let mut wildcard: Option<f32> = None;
let mut explicit: Vec<(Encoding, f32)> = Vec::new();
for part in header.split(',') {
let mut bits = part.split(';');
let token = bits.next().unwrap_or("").trim().to_ascii_lowercase();
let mut q = 1.0f32;
for param in bits {
let param = param.trim().to_ascii_lowercase();
if let Some(value) = param.strip_prefix("q=") {
q = value.trim().parse().unwrap_or(0.0);
}
}
if token == "*" {
wildcard = Some(wildcard.map_or(q, |prev: f32| prev.max(q)));
} else if let Some(enc) = Encoding::from_token(&token) {
explicit.push((enc, q));
}
}
for (index, enc) in self.preference.iter().enumerate() {
let quality = match explicit.iter().find(|(e, _)| e == enc) {
Some((_, q)) => Some((*q, false)),
None => wildcard.map(|q| (q, true)),
};
if let Some((q, is_wildcard)) = quality {
if q > 0.0 {
let _ = index;
offers.push((*enc, q, is_wildcard));
}
}
}
offers
.into_iter()
.enumerate()
.max_by(|(ia, a), (ib, b)| {
a.1.partial_cmp(&b.1)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| (!a.2).cmp(&(!b.2)))
.then_with(|| ib.cmp(ia))
})
.map(|(_, (enc, _, _))| enc)
}
fn should_compress(&self, res: &Response, is_head: bool) -> bool {
if res.status.is_informational()
|| res.status == StatusCode::NO_CONTENT
|| res.status == StatusCode::NOT_MODIFIED
{
return false;
}
if res.status == StatusCode::PARTIAL_CONTENT || res.headers.contains_key(CONTENT_RANGE) {
return false;
}
if res.headers.contains_key(CONTENT_ENCODING) {
return false;
}
let content_type = res
.headers
.get(CONTENT_TYPE)
.and_then(|v| v.to_str().ok())
.unwrap_or("");
if !(self.compressible)(content_type) {
return false;
}
match res.body.as_bytes() {
Some(bytes) if is_head && bytes.is_empty() => {
head_identity_len(res).is_some_and(|len| len >= self.min_size)
}
Some(bytes) => bytes.len() >= self.min_size,
None => true,
}
}
}
fn head_identity_len(res: &Response) -> Option<usize> {
res.headers
.get(CONTENT_LENGTH)?
.to_str()
.ok()?
.trim()
.parse()
.ok()
}
fn encode(body: Body, encoding: Encoding, level: Level) -> Body {
let source = match body {
Body::Bytes(bytes) => {
futures_util::stream::unfold((bytes, false), |(mut rest, started)| async move {
if rest.is_empty() {
return None;
}
if started {
tokio::task::yield_now().await;
}
let take = CHUNK.min(rest.len());
let chunk = rest.split_to(take);
Some((Ok::<Bytes, std::io::Error>(chunk), (rest, true)))
})
.boxed()
}
Body::Stream(stream) => stream
.map(|chunk| chunk.map_err(|e| std::io::Error::other(e.to_string())))
.boxed(),
};
let reader = StreamReader::new(source);
let level = level.into();
let encoded: std::pin::Pin<Box<dyn AsyncRead + Send>> = match encoding {
Encoding::Brotli => {
Box::pin(async_compression::tokio::bufread::BrotliEncoder::with_quality(reader, level))
}
Encoding::Gzip => {
Box::pin(async_compression::tokio::bufread::GzipEncoder::with_quality(reader, level))
}
Encoding::Deflate => {
Box::pin(async_compression::tokio::bufread::ZlibEncoder::with_quality(reader, level))
}
};
Body::from_stream(futures_util::stream::unfold(
(ReaderStream::new(encoded), false),
|(mut out, started)| async move {
if started {
tokio::task::yield_now().await;
}
let chunk = out.next().await?;
Some((chunk, (out, true)))
},
))
}
fn vary_on_accept_encoding(res: &mut Response) {
let existing: Vec<String> = res
.headers
.get_all(VARY)
.iter()
.filter_map(|v| v.to_str().ok())
.flat_map(|v| v.split(','))
.map(|v| v.trim().to_ascii_lowercase())
.filter(|v| !v.is_empty())
.collect();
if existing.iter().any(|v| v == "*" || v == "accept-encoding") {
return;
}
let mut merged = existing;
merged.push("accept-encoding".to_string());
if let Ok(value) = HeaderValue::from_str(&merged.join(", ")) {
res.headers.insert(VARY, value);
}
}
fn weaken_etag(res: &mut Response) {
let Some(tag) = res.headers.get(ETAG).and_then(|v| v.to_str().ok()) else {
return;
};
if tag.starts_with("W/") {
return;
}
if let Ok(value) = HeaderValue::from_str(&format!("W/{tag}")) {
res.headers.insert(ETAG, value);
}
}
#[async_trait]
impl Middleware for Compression {
async fn handle(&self, call: Call, next: Next) -> Response {
let accept = call
.header(ACCEPT_ENCODING.as_str())
.unwrap_or_default()
.to_string();
let is_head = call.method() == Method::HEAD;
let mut res = next.run(call).await;
vary_on_accept_encoding(&mut res);
if accept.is_empty() {
return res;
}
let Some(encoding) = self.negotiate(&accept) else {
return res;
};
if !self.should_compress(&res, is_head) {
return res;
}
if !is_head {
let original = res.body.as_bytes().cloned();
let was_buffered = original.is_some();
let body = std::mem::take(&mut res.body);
let encoded = encode(body, encoding, self.level);
res.body = if was_buffered {
match encoded.into_bytes().await {
Ok(bytes) => Body::Bytes(bytes),
Err(_) => {
res.body = Body::Bytes(original.unwrap_or_default());
return res;
}
}
} else {
encoded
};
}
res.headers
.insert(CONTENT_ENCODING, HeaderValue::from_static(encoding.token()));
res.headers.remove(CONTENT_LENGTH);
res.headers.remove(http::header::ACCEPT_RANGES);
weaken_etag(&mut res);
res
}
}
impl Plugin for Compression {
fn install(self: Box<Self>, app: &mut AppBuilder) {
app.add_middleware_in(Phase::Plugins, Arc::new(*self));
}
}
#[cfg(test)]
mod tests {
use super::*;
fn plugin() -> Compression {
Compression::new()
}
#[test]
fn negotiation_prefers_the_client_quality() {
let c = plugin();
assert_eq!(c.negotiate("gzip;q=1.0, br;q=0.1"), Some(Encoding::Gzip));
assert_eq!(c.negotiate("gzip;q=0.1, br;q=1.0"), Some(Encoding::Brotli));
}
#[test]
fn server_order_breaks_a_quality_tie() {
let c = plugin();
assert_eq!(c.negotiate("gzip, br, deflate"), Some(Encoding::Brotli));
let gzip_first = plugin().encodings([Encoding::Gzip, Encoding::Brotli]);
assert_eq!(gzip_first.negotiate("gzip, br"), Some(Encoding::Gzip));
}
#[test]
fn a_zero_quality_refuses_that_coding() {
let c = plugin();
assert_eq!(c.negotiate("br;q=0, gzip"), Some(Encoding::Gzip));
assert_eq!(c.negotiate("br;q=0, gzip;q=0, deflate;q=0"), None);
}
#[test]
fn a_quality_parameter_is_matched_case_insensitively() {
let c = plugin();
assert_eq!(c.negotiate("gzip;Q=0, deflate"), Some(Encoding::Deflate));
assert_eq!(c.negotiate("br;Q=0, gzip;Q=0, deflate;Q=0"), None);
assert_eq!(c.negotiate("gzip;Q=1.0, br;Q=0.1"), Some(Encoding::Gzip));
}
#[test]
fn an_explicit_coding_beats_the_wildcard() {
let c = plugin();
assert_eq!(c.negotiate("gzip;q=0.5, *;q=0.4"), Some(Encoding::Gzip));
}
#[test]
fn the_wildcard_alone_selects_the_server_preference() {
assert_eq!(plugin().negotiate("*"), Some(Encoding::Brotli));
}
#[test]
fn unknown_codings_are_ignored() {
assert_eq!(plugin().negotiate("exi, sdch"), None);
}
#[test]
fn media_types_are_classified_conservatively() {
assert!(compressible_by_default("text/html"));
assert!(compressible_by_default("application/json"));
assert!(compressible_by_default("image/svg+xml"));
assert!(!compressible_by_default("image/jpeg"));
assert!(!compressible_by_default("application/zip"));
assert!(!compressible_by_default(""));
}
#[tokio::test]
async fn encoding_round_trips_through_gzip() {
use tokio::io::AsyncReadExt;
let input = Bytes::from("hello ".repeat(4000));
let body = encode(Body::Bytes(input.clone()), Encoding::Gzip, Level::Default);
let compressed = body.into_bytes().await.unwrap();
assert!(
compressed.len() < input.len(),
"repetitive text should shrink"
);
let mut decoded = Vec::new();
async_compression::tokio::bufread::GzipDecoder::new(std::io::Cursor::new(
compressed.to_vec(),
))
.read_to_end(&mut decoded)
.await
.unwrap();
assert_eq!(Bytes::from(decoded), input);
}
#[tokio::test]
async fn a_buffered_body_yields_the_worker_between_chunks() {
use std::sync::atomic::{AtomicBool, Ordering};
let ran = Arc::new(AtomicBool::new(false));
let flag = Arc::clone(&ran);
tokio::spawn(async move { flag.store(true, Ordering::SeqCst) });
let input = Bytes::from(vec![b'x'; CHUNK * 4]);
let encoded = encode(Body::Bytes(input), Encoding::Gzip, Level::Fastest)
.into_bytes()
.await
.unwrap();
assert!(
!encoded.is_empty(),
"the encode still has to produce output"
);
assert!(
ran.load(Ordering::SeqCst),
"a multi-chunk encode must let the runtime run something else"
);
}
#[test]
fn vary_is_appended_not_replaced() {
let mut res = Response::text("x");
res.headers.insert(VARY, HeaderValue::from_static("origin"));
vary_on_accept_encoding(&mut res);
assert_eq!(res.headers.get(VARY).unwrap(), "origin, accept-encoding");
}
#[test]
fn vary_is_not_duplicated() {
let mut res = Response::text("x");
vary_on_accept_encoding(&mut res);
vary_on_accept_encoding(&mut res);
assert_eq!(res.headers.get_all(VARY).iter().count(), 1);
assert_eq!(res.headers.get(VARY).unwrap(), "accept-encoding");
}
#[test]
fn a_strong_etag_becomes_weak() {
let mut res = Response::text("x");
res.headers
.insert(ETAG, HeaderValue::from_static("\"abc\""));
weaken_etag(&mut res);
assert_eq!(res.headers.get(ETAG).unwrap(), "W/\"abc\"");
}
#[test]
fn an_already_weak_etag_is_left_alone() {
let mut res = Response::text("x");
res.headers
.insert(ETAG, HeaderValue::from_static("W/\"abc\""));
weaken_etag(&mut res);
assert_eq!(res.headers.get(ETAG).unwrap(), "W/\"abc\"");
}
#[test]
fn partial_content_is_never_compressed() {
let mut res = Response::text("x".repeat(4096));
res.status = StatusCode::PARTIAL_CONTENT;
assert!(!plugin().should_compress(&res, false));
}
#[test]
fn an_already_encoded_response_is_left_alone() {
let mut res = Response::text("x".repeat(4096));
res.headers
.insert(CONTENT_ENCODING, HeaderValue::from_static("gzip"));
assert!(!plugin().should_compress(&res, false));
}
#[test]
fn small_bodies_are_left_alone() {
let res = Response::text("small");
assert!(!plugin().should_compress(&res, false));
assert!(
plugin().min_size(1).should_compress(&res, false),
"lowering the floor should admit it"
);
}
}