//! Streaming brotli compression middleware for axum. //! //! # Why a custom encoder instead of `tower-http`'s `CompressionLayer`? //! //! `tower-http` does support brotli (via `async-compression`), but it does not //! expose the `lgwin` (window size) parameter. For Datastar SSE streams the //! window size is the key tuning knob: all SSE events on a connection travel //! through a **single persistent brotli encoder**, so later events benefit from //! the full history of earlier ones. A 4 MB window (`lgwin = 22`) means the //! encoder can reference up to 4 MB of previously transmitted HTML when //! compressing a new fragment — dramatically better ratios than per-chunk gzip. //! //! This module implements a [`BrotliBody`] wrapper and an axum //! [`brotli_compression`] middleware that honour the [`BrotliConfig`] (quality //! + window size) loaded from `livevue.toml`. //! //! # Behaviour //! //! - Only activates when the client sends `Accept-Encoding: br`. //! - Reads config from `SharedConfig` on **each new request**, so a hot-reload //! applies to the next connection without a restart. //! - Calls `encoder.flush()` after every body frame so SSE clients receive //! data promptly (brotli metablock boundary = decodable checkpoint). //! - The encoder's sliding-window context is preserved across flushes, giving //! progressively better compression as the SSE stream grows. use crate::config::SharedConfig; use axum::body::Body; use axum::extract::State; use axum::http::header::{ACCEPT_ENCODING, CONTENT_ENCODING, CONTENT_LENGTH}; use axum::http::HeaderValue; use axum::middleware::Next; use axum::response::Response; use bytes::Bytes; use http_body::Body as HttpBody; use pin_project::pin_project; use std::io::Write as _; use std::pin::Pin; use std::task::{Context, Poll}; // --------------------------------------------------------------------------- // DrainWriter — a `Write` impl whose output can be drained incrementally // --------------------------------------------------------------------------- /// A `std::io::Write` target that accumulates bytes in a `Vec`. /// /// After flushing the brotli encoder we call [`DrainWriter::drain`] to collect /// only the bytes written since the last drain. This lets us send one /// compressed chunk per SSE event while keeping the encoder (and its sliding /// window) alive across the entire connection. struct DrainWriter(Vec); impl std::io::Write for DrainWriter { fn write(&mut self, data: &[u8]) -> std::io::Result { self.0.extend_from_slice(data); Ok(data.len()) } /// No-op: the actual flushing is driven by `CompressorWriter::flush()`. fn flush(&mut self) -> std::io::Result<()> { Ok(()) } } impl DrainWriter { fn drain(&mut self) -> Bytes { Bytes::from(self.0.drain(..).collect::>()) } } // --------------------------------------------------------------------------- // BrotliBody — streaming Body wrapper // --------------------------------------------------------------------------- /// An `http_body::Body` that brotli-compresses a wrapped [`axum::body::Body`]. /// /// The encoder is constructed once at stream creation and persists for the /// lifetime of the connection, maintaining its sliding-window context across /// every SSE event. [`flush()`] is called after each data frame to create a /// metablock boundary so the browser can decode partial data immediately. /// /// [`flush()`]: std::io::Write::flush #[pin_project] pub struct BrotliBody { #[pin] inner: Body, /// `None` once we have consumed the encoder to finalize the stream. encoder: Option>, } impl BrotliBody { pub fn new(inner: Body, quality: u32, lgwin: u32) -> Self { let params = brotli::enc::BrotliEncoderParams { quality: quality as i32, lgwin: lgwin as i32, ..Default::default() }; // Buffer size (4 KiB) is the encoder's internal I/O buffer — smaller // means lower latency per flush, which matters for SSE. let encoder = brotli::CompressorWriter::with_params(DrainWriter(Vec::new()), 4096, ¶ms); Self { inner, encoder: Some(encoder), } } } impl HttpBody for BrotliBody { type Data = Bytes; type Error = axum::Error; fn poll_frame( self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll, Self::Error>>> { let this = self.project(); match this.inner.poll_frame(cx) { // ---------------------------------------------------------------- // Got a data frame — compress it and yield immediately. // ---------------------------------------------------------------- Poll::Ready(Some(Ok(frame))) => { if let Ok(data) = frame.into_data() { match this.encoder { Some(enc) => { if let Err(e) = enc.write_all(&data) { return Poll::Ready(Some(Err(axum::Error::new(e)))); } if let Err(e) = enc.flush() { return Poll::Ready(Some(Err(axum::Error::new(e)))); } let compressed = enc.get_mut().drain(); Poll::Ready(Some(Ok(http_body::Frame::data(compressed)))) } None => Poll::Ready(None), } } else { // Trailers / unknown frame type — pass through as-is. // (Re-construct the original frame; trailers don't need compression.) Poll::Ready(None) } } // ---------------------------------------------------------------- // Inner body is exhausted — finalise the encoder. // ---------------------------------------------------------------- Poll::Ready(None) => { if let Some(enc) = this.encoder.take() { // into_inner() flushes any pending bytes and writes the // final brotli stream terminator, returning the inner // DrainWriter directly (not wrapped in Result). let mut drain = enc.into_inner(); let final_bytes = drain.drain(); if !final_bytes.is_empty() { return Poll::Ready(Some(Ok(http_body::Frame::data(final_bytes)))); } } Poll::Ready(None) } // ---------------------------------------------------------------- // Pass-through. // ---------------------------------------------------------------- Poll::Ready(Some(Err(e))) => Poll::Ready(Some(Err(e))), Poll::Pending => Poll::Pending, } } } // --------------------------------------------------------------------------- // axum middleware // --------------------------------------------------------------------------- /// Axum middleware that wraps SSE (and any other) responses with streaming /// brotli compression, obeying `Accept-Encoding: br`. /// /// Reads [`BrotliConfig`] from [`SharedConfig`] on each new request, so /// quality and window-size changes take effect for new connections immediately /// after a config reload. /// /// # Usage /// /// ```rust,ignore /// use axum::{middleware, Router}; /// use livevue_rs::brotli_layer::brotli_compression; /// /// let app = Router::new() /// /* ... routes ... */ /// .layer(middleware::from_fn_with_state(config.clone(), brotli_compression)); /// ``` pub async fn brotli_compression( State(config): State, request: axum::http::Request, next: Next, ) -> Response { // Check whether the client signals brotli support. let accepts_brotli = request .headers() .get(ACCEPT_ENCODING) .and_then(|v| v.to_str().ok()) .map(|v| v.contains("br")) .unwrap_or(false); let response = next.run(request).await; if !accepts_brotli { return response; } // Snapshot the brotli config; new values apply to future connections. let (enabled, quality, window_size) = { let cfg = config.read().await; (cfg.brotli.enabled, cfg.brotli.quality, cfg.brotli.window_size) }; if !enabled { return response; } let (mut parts, body) = response.into_parts(); // Remove Content-Length — the compressed length will differ. parts.headers.remove(CONTENT_LENGTH); // Advertise brotli encoding. parts .headers .insert(CONTENT_ENCODING, HeaderValue::from_static("br")); let brotli_body = BrotliBody::new(body, quality, window_size); Response::from_parts(parts, Body::new(brotli_body)) }