codetics-ocpp-ws 0.1.0

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Kevin Vavelin, Codetics (https://codetics.fr)

//! RFC 7692 permessage-deflate: negotiation and per-message codec.
//!
//! Memory policy, chosen for tens of thousands of mostly idle connections:
//! - Server → client: the server always declares `server_no_context_takeover`
//!   (allowed by RFC 7692 §7.1.1.1 regardless of the offer). Compression
//!   state is then not needed between messages, so compressors come from a
//!   process-wide pool instead of being held per connection.
//! - Client → server: a decompressor is allocated lazily, on the first
//!   compressed message. With `client_no_context_takeover` (requested by the
//!   client, or imposed by configuration per §7.1.1.2) it is reset after
//!   every message.
//! - `server_max_window_bits` < 15 cannot be honoured by the pure-Rust DEFLATE
//!   backend, so such offers are declined (the connection then simply runs
//!   uncompressed, which RFC 7692 and OCPP both allow).

use std::sync::Mutex;

use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress, Status};

/// Trailer removed after compression and re-appended before decompression (RFC 7692 §7.2.1).
const TRAILER: [u8; 4] = [0x00, 0x00, 0xFF, 0xFF];

#[derive(Debug, Clone, Copy)]
pub struct DeflateConfig {
    pub enabled: bool,
    /// Impose `client_no_context_takeover` even when the client did not offer it.
    /// Saves ~40 KiB per connection at the cost of client-side compression ratio.
    pub force_client_no_context_takeover: bool,
    /// Messages smaller than this are sent uncompressed.
    pub min_compress_size: usize,
    /// DEFLATE level 0–9.
    pub level: u32,
}

impl Default for DeflateConfig {
    fn default() -> Self {
        DeflateConfig {
            enabled: true,
            force_client_no_context_takeover: false,
            min_compress_size: 128,
            level: 6,
        }
    }
}

/// Parameters agreed during the handshake.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct DeflateParams {
    /// The sender of compressed client→server messages resets its context per message.
    pub client_no_context_takeover: bool,
    /// The server resets its context per message (always true for responses we produce).
    pub server_no_context_takeover: bool,
}

/// Server side: picks the first acceptable `permessage-deflate` offer and
/// returns the agreed parameters and the `Sec-WebSocket-Extensions` response value.
pub fn negotiate_server<'a, I>(offers: I, cfg: &DeflateConfig) -> Option<(DeflateParams, String)>
where
    I: IntoIterator<Item = &'a str>,
{
    if !cfg.enabled {
        return None;
    }
    for header in offers {
        for offer in header.split(',') {
            let mut parts = offer.split(';').map(str::trim);
            if !parts
                .next()
                .is_some_and(|n| n.eq_ignore_ascii_case("permessage-deflate"))
            {
                continue;
            }
            let mut seen = Vec::new();
            let mut client_nct = false;
            let mut acceptable = true;
            for p in parts.filter(|p| !p.is_empty()) {
                let (name, value) = match p.split_once('=') {
                    Some((n, v)) => (
                        n.trim().to_ascii_lowercase(),
                        Some(v.trim().trim_matches('"')),
                    ),
                    None => (p.to_ascii_lowercase(), None),
                };
                if seen.contains(&name) {
                    acceptable = false;
                    break;
                }
                match (name.as_str(), value) {
                    ("server_no_context_takeover", None) => {}
                    ("client_no_context_takeover", None) => client_nct = true,
                    ("server_max_window_bits", Some(v)) => acceptable &= v == "15",
                    ("client_max_window_bits", None) => {}
                    ("client_max_window_bits", Some(v)) => {
                        acceptable &= v.parse::<u8>().is_ok_and(|b| (8..=15).contains(&b))
                    }
                    _ => acceptable = false,
                }
                seen.push(name);
            }
            if !acceptable {
                continue;
            }
            let client_nct = client_nct || cfg.force_client_no_context_takeover;
            let mut response = String::from("permessage-deflate; server_no_context_takeover");
            if client_nct {
                response.push_str("; client_no_context_takeover");
            }
            return Some((
                DeflateParams {
                    client_no_context_takeover: client_nct,
                    server_no_context_takeover: true,
                },
                response,
            ));
        }
    }
    None
}

/// Client side: interprets the server's `Sec-WebSocket-Extensions` response.
pub fn accept_client_response(header: &str) -> Option<DeflateParams> {
    let mut parts = header.split(';').map(str::trim);
    if !parts.next()?.eq_ignore_ascii_case("permessage-deflate") {
        return None;
    }
    let mut params = DeflateParams {
        client_no_context_takeover: false,
        server_no_context_takeover: false,
    };
    for p in parts {
        match p
            .split('=')
            .next()
            .unwrap_or("")
            .trim()
            .to_ascii_lowercase()
            .as_str()
        {
            "client_no_context_takeover" => params.client_no_context_takeover = true,
            "server_no_context_takeover" => params.server_no_context_takeover = true,
            _ => {}
        }
    }
    Some(params)
}

static COMPRESSOR_POOL: Mutex<Vec<(u32, Compress)>> = Mutex::new(Vec::new());
const POOL_CAP: usize = 64;

/// Compresses one message without context takeover, using a pooled compressor.
pub fn compress_message(data: &[u8], level: u32) -> Vec<u8> {
    let mut c = {
        let mut pool = COMPRESSOR_POOL.lock().unwrap_or_else(|e| e.into_inner());
        match pool.iter().position(|(l, _)| *l == level) {
            Some(i) => pool.swap_remove(i).1,
            None => Compress::new(Compression::new(level), false),
        }
    };
    let start_in = c.total_in();
    let mut out = Vec::with_capacity(data.len() / 2 + 64);
    loop {
        if out.capacity() - out.len() < 64 {
            out.reserve(out.capacity().max(256));
        }
        let consumed = (c.total_in() - start_in) as usize;
        c.compress_vec(&data[consumed..], &mut out, FlushCompress::Sync)
            .expect("deflate never fails on valid state");
        let all_in = (c.total_in() - start_in) as usize == data.len();
        // A sync flush is complete once the output buffer was not filled up.
        if all_in && out.len() < out.capacity() {
            break;
        }
    }
    if out.ends_with(&TRAILER) {
        out.truncate(out.len() - TRAILER.len());
    }
    c.reset();
    let mut pool = COMPRESSOR_POOL.lock().unwrap_or_else(|e| e.into_inner());
    if pool.len() < POOL_CAP {
        pool.push((level, c));
    }
    out
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum InflateError {
    TooBig,
    Corrupt,
}

/// Per-connection decompressor for messages received from the peer.
pub struct Inflater {
    state: Option<Box<Decompress>>,
    no_context_takeover: bool,
}

impl Inflater {
    pub fn new(no_context_takeover: bool) -> Self {
        Inflater {
            state: None,
            no_context_takeover,
        }
    }

    pub fn decompress(&mut self, data: &[u8], max_size: usize) -> Result<Vec<u8>, InflateError> {
        let d = self
            .state
            .get_or_insert_with(|| Box::new(Decompress::new(false)));
        let result = inflate(d, data, max_size);
        match &result {
            // The sliding window is unusable after an error.
            Err(_) => self.state = None,
            Ok(_) if self.no_context_takeover => d.reset(false),
            Ok(_) => {}
        }
        result
    }
}

fn inflate(d: &mut Decompress, data: &[u8], max_size: usize) -> Result<Vec<u8>, InflateError> {
    let mut input = Vec::with_capacity(data.len() + TRAILER.len());
    input.extend_from_slice(data);
    input.extend_from_slice(&TRAILER);
    let start_in = d.total_in();
    let consumed = |d: &Decompress| (d.total_in() - start_in) as usize;
    let mut out = Vec::with_capacity((data.len() * 3).clamp(256, max_size.max(256)));
    loop {
        if out.capacity() - out.len() < 256 {
            if out.len() >= max_size {
                return Err(InflateError::TooBig);
            }
            out.reserve(out.capacity().min(max_size));
        }
        let (in_before, out_before) = (consumed(d), out.len());
        let status = d
            .decompress_vec(&input[in_before..], &mut out, FlushDecompress::Sync)
            .map_err(|_| InflateError::Corrupt)?;
        if out.len() > max_size {
            return Err(InflateError::TooBig);
        }
        let all_in = consumed(d) == input.len();
        let output_has_room = out.len() < out.capacity();
        if status == Status::StreamEnd || (all_in && output_has_room) {
            return Ok(out);
        }
        let progressed = consumed(d) != in_before || out.len() != out_before;
        if !progressed && output_has_room {
            // Input remains but the decoder cannot make progress: truncated or invalid stream.
            return Err(InflateError::Corrupt);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn negotiation() {
        let cfg = DeflateConfig::default();
        let (p, r) =
            negotiate_server(["permessage-deflate; client_max_window_bits"], &cfg).unwrap();
        assert_eq!(r, "permessage-deflate; server_no_context_takeover");
        assert!(!p.client_no_context_takeover);
        let (p, r) = negotiate_server(["permessage-deflate; server_max_window_bits=10, permessage-deflate; client_no_context_takeover"], &cfg).unwrap();
        assert!(p.client_no_context_takeover);
        assert_eq!(
            r,
            "permessage-deflate; server_no_context_takeover; client_no_context_takeover"
        );
        assert!(negotiate_server(["x-webkit-deflate-frame"], &cfg).is_none());
        assert!(negotiate_server(["permessage-deflate; foo"], &cfg).is_none());
        assert!(
            negotiate_server(
                ["permessage-deflate; server_no_context_takeover; server_no_context_takeover"],
                &cfg
            )
            .is_none()
        );
        assert!(
            negotiate_server(
                ["permessage-deflate"],
                &DeflateConfig {
                    enabled: false,
                    ..cfg
                }
            )
            .is_none()
        );
        let forced = DeflateConfig {
            force_client_no_context_takeover: true,
            ..cfg
        };
        assert!(
            negotiate_server(["permessage-deflate"], &forced)
                .unwrap()
                .0
                .client_no_context_takeover
        );
    }

    #[test]
    fn roundtrip_with_and_without_context() {
        for nct in [false, true] {
            let mut inf = Inflater::new(nct);
            for i in 0..20 {
                let msg = format!(
                    r#"[2,"{i}","MeterValues",{{"connectorId":1,"meterValue":[{{"timestamp":"2024-01-01T00:00:00Z","sampledValue":[{{"value":"{}"}}]}}]}}]"#,
                    i * 1000
                );
                let c = compress_message(msg.as_bytes(), 6);
                assert!(!c.ends_with(&TRAILER));
                assert_eq!(inf.decompress(&c, 1 << 20).unwrap(), msg.as_bytes());
            }
        }
    }

    #[test]
    fn large_and_bomb() {
        let big = "x".repeat(3_000_000);
        let c = compress_message(big.as_bytes(), 6);
        assert!(c.len() < 20_000);
        let mut inf = Inflater::new(true);
        assert_eq!(inf.decompress(&c, 1 << 20), Err(InflateError::TooBig));
        assert_eq!(inf.decompress(&c, 4 << 20).unwrap().len(), 3_000_000);
        assert_eq!(
            Inflater::new(true).decompress(&[0xFF, 0xFF, 0xFF], 1024),
            Err(InflateError::Corrupt)
        );
    }
}