tachyon_web/ws/

deflate.rs

//! `permessage-deflate` (RFC 7692): negotiation of the `Sec-WebSocket-Extensions` offer/response,
//! and the actual per-message DEFLATE compressor/decompressor built on raw (headerless) deflate
//! streams via `flate2`.
//!
//! Compression is negotiated per RFC 7692 §7 and applied to the *payload of a full message* (all
//! fragments concatenated) rather than per-frame: RSV1 is only ever set on the first frame of a
//! message, and continuation frames carry no RSV1 of their own. The wire format for a compressed
//! message is produced with `Z_SYNC_FLUSH` and then has its trailing 4-byte empty-block marker
//! (`00 00 ff ff`) stripped per §7.2.1; the decompressor puts that marker back before inflating.

use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress};
use hyper::header::{HeaderMap, HeaderValue, SEC_WEBSOCKET_EXTENSIONS};

/// The 4 bytes a `Z_SYNC_FLUSH` block ends with, which RFC 7692 requires senders to strip and
/// receivers to restore before inflating.
const DEFLATE_TAIL: [u8; 4] = [0x00, 0x00, 0xff, 0xff];

/// Server-side tuning for the `permessage-deflate` extension.
///
/// Constructed via [`Default`] and passed to
/// [`WebSocketUpgrade::deflate_config`](super::WebSocketUpgrade::deflate_config).
#[derive(Debug, Clone, Copy)]
#[non_exhaustive]
pub struct DeflateConfig {
    /// Reset our own compression context after every message instead of reusing the sliding
    /// window across messages. Lowers compression ratio, lowers memory use. Default: `false`.
    pub server_no_context_takeover: bool,
    /// Require the client to reset its compression context after every message. This is only a
    /// request; correctness on our end does not depend on the client honoring it. Default: `false`.
    pub client_no_context_takeover: bool,
    /// The base-2 logarithm of the LZ77 window we use to compress outgoing messages, `9..=15`.
    /// Values are clamped into that range. Default: `15` (32 KiB window, maximum compression).
    pub server_max_window_bits: u8,
}

impl Default for DeflateConfig {
    fn default() -> Self {
        Self {
            server_no_context_takeover: false,
            client_no_context_takeover: false,
            server_max_window_bits: 15,
        }
    }
}

/// One `permessage-deflate` offer parsed out of a `Sec-WebSocket-Extensions` request header.
#[derive(Debug, Default, Clone, Copy)]
pub(super) struct Offer {
    server_no_context_takeover: bool,
    client_no_context_takeover: bool,
    /// Whether the client's offer named `server_max_window_bits` at all (with or without a
    /// value) — distinguishes "absent" from "present with the default window size". We don't
    /// honor `client_max_window_bits` (our decompressor always uses the maximum window, which
    /// is always compatible with whatever smaller window the client's compressor might use), so
    /// there's nothing to record for it beyond validating it during parsing.
    server_max_window_bits_offered: bool,
}

/// The agreed-upon parameters after negotiating an [`Offer`] against a [`DeflateConfig`].
#[derive(Debug, Clone, Copy)]
pub(super) struct Agreement {
    pub(super) server_no_context_takeover: bool,
    pub(super) client_no_context_takeover: bool,
    pub(super) server_max_window_bits: u8,
    /// Whether `server_max_window_bits` should be echoed in the response — only valid if the
    /// client's offer itself named the parameter.
    echo_server_max_window_bits: bool,
}

/// Extracts every `permessage-deflate` offer from the request's `Sec-WebSocket-Extensions`
/// header(s), in the order they appeared. Unrecognized extensions and unrecognized/malformed
/// parameters on an otherwise-recognized offer are skipped per-offer (per RFC 7692 §7, a server
/// declines an individual malformed offer rather than failing the whole negotiation).
pub(super) fn parse_offers(headers: &HeaderMap) -> Vec<Offer> {
    let mut offers = Vec::new();
    for header in headers.get_all(SEC_WEBSOCKET_EXTENSIONS) {
        let Ok(text) = header.to_str() else { continue };
        for extension in text.split(',') {
            let mut parts = extension.split(';').map(str::trim);
            let Some(name) = parts.next() else { continue };
            if !name.eq_ignore_ascii_case("permessage-deflate") {
                continue;
            }
            if let Some(offer) = parse_params(parts) {
                offers.push(offer);
            }
        }
    }
    offers
}

/// Parses one `window_bits` parameter's optional value, validating it against RFC 7692's
/// `9..=15` range when present. Returns `Err` if the offer should be declined outright.
fn parse_window_bits(value: Option<&str>) -> Result<(), ()> {
    match value.map(str::parse::<u8>) {
        None | Some(Ok(9..=15)) => Ok(()),
        Some(_) => Err(()),
    }
}

fn parse_params<'a>(params: impl Iterator<Item = &'a str>) -> Option<Offer> {
    let mut offer = Offer::default();
    for param in params {
        if param.is_empty() {
            continue;
        }
        let (key, value) = match param.split_once('=') {
            Some((k, v)) => (k.trim(), Some(v.trim().trim_matches('"'))),
            None => (param, None),
        };
        match key {
            "server_no_context_takeover" if value.is_none() => {
                offer.server_no_context_takeover = true;
            }
            "client_no_context_takeover" if value.is_none() => {
                offer.client_no_context_takeover = true;
            }
            "server_max_window_bits" => {
                parse_window_bits(value).ok()?;
                offer.server_max_window_bits_offered = true;
            }
            "client_max_window_bits" => {
                parse_window_bits(value).ok()?;
            }
            // Unrecognized parameter, or a value on a no-value-only flag: decline this offer.
            _ => return None,
        }
    }
    Some(offer)
}

/// Picks the first offer we can accept and applies `config`'s server-side preferences to it.
/// Returns `None` if there is nothing to negotiate (no offers at all).
pub(super) fn negotiate(offers: &[Offer], config: DeflateConfig) -> Option<Agreement> {
    let offer = offers.first()?;
    let server_max_window_bits = config.server_max_window_bits.clamp(9, 15);
    Some(Agreement {
        server_no_context_takeover: config.server_no_context_takeover
            || offer.server_no_context_takeover,
        client_no_context_takeover: config.client_no_context_takeover
            || offer.client_no_context_takeover,
        server_max_window_bits,
        echo_server_max_window_bits: offer.server_max_window_bits_offered
            && server_max_window_bits < 15,
    })
}

/// Builds the `Sec-WebSocket-Extensions` response header value for an accepted [`Agreement`].
pub(super) fn agreement_header_value(agreement: Agreement) -> HeaderValue {
    let mut value = String::from("permessage-deflate");
    if agreement.server_no_context_takeover {
        value.push_str("; server_no_context_takeover");
    }
    if agreement.client_no_context_takeover {
        value.push_str("; client_no_context_takeover");
    }
    if agreement.echo_server_max_window_bits {
        value.push_str("; server_max_window_bits=");
        value.push_str(&agreement.server_max_window_bits.to_string());
    }
    HeaderValue::from_str(&value).unwrap_or_else(|_| HeaderValue::from_static("permessage-deflate"))
}

/// Per-connection compressor/decompressor for an agreed `permessage-deflate` extension.
pub(super) struct PerMessageDeflate {
    compress: Compress,
    decompress: Decompress,
    server_no_context_takeover: bool,
    client_no_context_takeover: bool,
}

impl PerMessageDeflate {
    pub(super) fn new(agreement: Agreement) -> Self {
        Self {
            compress: Compress::new_with_window_bits(
                Compression::default(),
                false,
                agreement.server_max_window_bits,
            ),
            decompress: Decompress::new_with_window_bits(false, 15),
            server_no_context_takeover: agreement.server_no_context_takeover,
            client_no_context_takeover: agreement.client_no_context_takeover,
        }
    }

    /// Compresses `data`, but only when the result actually comes out smaller — RFC 7692 leaves
    /// per-message compression up to the sender (RSV1 is just left unset for the ones we skip),
    /// and deflate's per-block overhead means a small or already-dense payload can come out
    /// *larger* compressed. Returns `None` when the caller should send `data` verbatim instead.
    ///
    /// When we skip, the compressor's context is reset regardless of `server_no_context_takeover`:
    /// RFC 7692 requires an unsent-compressed message not to affect the compression context, but
    /// `flate2` gives no way to "undo" the trial `compress_vec` call already made while sizing up
    /// the candidate. Resetting is the only way back to a self-consistent state — deflate back-
    /// references are entirely encoder-side, so a reset simply means future messages don't
    /// reference data the client's decompressor never saw; the client's own (untouched, larger)
    /// window carrying unused extra history is harmless.
    pub(super) fn compress_if_smaller(
        &mut self,
        data: &[u8],
    ) -> Result<Option<Vec<u8>>, crate::http::error::Error> {
        let compressed = self.compress_raw(data)?;
        if compressed.len() < data.len() {
            if self.server_no_context_takeover {
                self.compress.reset();
            }
            Ok(Some(compressed))
        } else {
            self.compress.reset();
            Ok(None)
        }
    }

    /// Compresses one full message payload, stripping the trailing sync-flush marker per §7.2.1.
    fn compress_raw(&mut self, data: &[u8]) -> Result<Vec<u8>, crate::http::error::Error> {
        let total_in_before = self.compress.total_in();
        let mut out = Vec::with_capacity(data.len() + 32);
        loop {
            grow(&mut out, 1024.max(data.len()));
            self.compress
                .compress_vec(data, &mut out, FlushCompress::Sync)
                .map_err(|e| crate::http::error::Error::Internal(e.to_string()))?;
            let consumed =
                usize::try_from(self.compress.total_in() - total_in_before).unwrap_or(usize::MAX);
            if consumed >= data.len() {
                break;
            }
        }
        out.truncate(out.len().saturating_sub(DEFLATE_TAIL.len()));
        Ok(out)
    }

    /// Decompresses one full message payload, restoring the trailing sync-flush marker first.
    pub(super) fn decompress(
        &mut self,
        data: &[u8],
        max_size: Option<usize>,
    ) -> Result<Vec<u8>, crate::http::error::Error> {
        let max_size = max_size.unwrap_or(usize::MAX);
        let mut input = Vec::with_capacity(data.len() + DEFLATE_TAIL.len());
        input.extend_from_slice(data);
        input.extend_from_slice(&DEFLATE_TAIL);

        let total_in_before = self.decompress.total_in();
        let mut out = Vec::with_capacity((data.len() * 3 + 32).min(max_size));
        loop {
            grow(&mut out, 1024);
            let consumed_before =
                usize::try_from(self.decompress.total_in() - total_in_before).unwrap_or(usize::MAX);
            self.decompress
                .decompress_vec(&input[consumed_before..], &mut out, FlushDecompress::Sync)
                .map_err(|e| crate::http::error::Error::Internal(e.to_string()))?;
            if out.len() > max_size {
                return Err(crate::http::error::Error::Internal(
                    "decompressed message exceeds the configured maximum size".to_string(),
                ));
            }
            let consumed =
                usize::try_from(self.decompress.total_in() - total_in_before).unwrap_or(usize::MAX);
            if consumed >= input.len() {
                break;
            }
        }
        if self.client_no_context_takeover {
            self.decompress.reset(false);
        }
        Ok(out)
    }
}

/// Reserves more spare capacity in `out`, doubling what's already there (or `min_initial` on the
/// first call) — geometric growth means a large payload needing several `compress_vec`/
/// `decompress_vec` rounds costs O(log n) reallocations instead of O(n).
fn grow(out: &mut Vec<u8>, min_initial: usize) {
    let additional = out.capacity().max(min_initial);
    out.reserve(additional);
}