livevue-rs/src/brotli_layer.rs

//! 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<u8>`.
///
/// 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<u8>);

impl std::io::Write for DrainWriter {
    fn write(&mut self, data: &[u8]) -> std::io::Result<usize> {
        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::<Vec<_>>())
    }
}

// ---------------------------------------------------------------------------
// 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<brotli::CompressorWriter<DrainWriter>>,
}

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, &params);
        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<Option<Result<http_body::Frame<Self::Data>, 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<SharedConfig>,
    request: axum::http::Request<Body>,
    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))
}