codetics-ocpp-ws 0.1.0

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

//! Message-level reader and writer halves of a WebSocket connection.

use std::fmt;
use std::io;

use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

use crate::deflate::{self, DeflateParams, InflateError, Inflater};
use crate::frame::{self, OpCode};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Role {
    Server,
    Client,
}

/// Standard close codes (RFC 6455 §7.4.1).
pub mod close_code {
    pub const NORMAL: u16 = 1000;
    pub const GOING_AWAY: u16 = 1001;
    pub const PROTOCOL_ERROR: u16 = 1002;
    pub const UNSUPPORTED_DATA: u16 = 1003;
    pub const INVALID_PAYLOAD: u16 = 1007;
    pub const POLICY_VIOLATION: u16 = 1008;
    pub const MESSAGE_TOO_BIG: u16 = 1009;
    pub const INTERNAL_ERROR: u16 = 1011;
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Message {
    Text(String),
    Binary(Vec<u8>),
    Ping(Vec<u8>),
    Pong(Vec<u8>),
    /// Close frame with optional (code, reason).
    Close(Option<(u16, String)>),
}

#[derive(Debug)]
pub enum WsError {
    Io(io::Error),
    Protocol(&'static str),
    InvalidUtf8,
    MessageTooBig,
}

impl WsError {
    /// Close code to send to the peer for this error.
    pub fn close_code(&self) -> u16 {
        match self {
            WsError::Io(_) => close_code::INTERNAL_ERROR,
            WsError::Protocol(_) => close_code::PROTOCOL_ERROR,
            WsError::InvalidUtf8 => close_code::INVALID_PAYLOAD,
            WsError::MessageTooBig => close_code::MESSAGE_TOO_BIG,
        }
    }
}

impl fmt::Display for WsError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            WsError::Io(e) => write!(f, "I/O error: {e}"),
            WsError::Protocol(m) => write!(f, "WebSocket protocol error: {m}"),
            WsError::InvalidUtf8 => f.write_str("text message is not valid UTF-8"),
            WsError::MessageTooBig => f.write_str("message exceeds the maximum size"),
        }
    }
}

impl std::error::Error for WsError {}

impl From<io::Error> for WsError {
    fn from(e: io::Error) -> Self {
        WsError::Io(e)
    }
}

#[derive(Debug, Clone, Copy)]
pub struct WsConfig {
    /// Maximum size of a (decompressed) message.
    pub max_message_size: usize,
    /// Buffer capacity kept for idle connections; larger buffers shrink back to this.
    pub idle_buffer: usize,
}

impl Default for WsConfig {
    fn default() -> Self {
        WsConfig {
            max_message_size: 4 << 20,
            idle_buffer: 4 << 10,
        }
    }
}

struct Partial {
    opcode: OpCode,
    compressed: bool,
    data: Vec<u8>,
}

/// Reading half. [`WsReader::recv`] is cancel-safe: all progress is kept in
/// `self`, so it can be used as a `tokio::select!` branch.
pub struct WsReader<R> {
    io: R,
    role: Role,
    cfg: WsConfig,
    buf: Vec<u8>,
    pos: usize,
    partial: Option<Partial>,
    inflater: Option<Inflater>,
    eof: bool,
}

impl<R: AsyncRead + Unpin> WsReader<R> {
    pub fn new(
        io: R,
        role: Role,
        cfg: WsConfig,
        deflate: Option<DeflateParams>,
        initial: Vec<u8>,
    ) -> Self {
        let inflater = deflate.map(|p| {
            Inflater::new(match role {
                Role::Server => p.client_no_context_takeover,
                Role::Client => p.server_no_context_takeover,
            })
        });
        WsReader {
            io,
            role,
            cfg,
            buf: initial,
            pos: 0,
            partial: None,
            inflater,
            eof: false,
        }
    }

    /// Receives the next message. `Ok(None)` means the TCP stream ended.
    pub async fn recv(&mut self) -> Result<Option<Message>, WsError> {
        loop {
            if let Some(msg) = self.try_parse()? {
                return Ok(Some(msg));
            }
            if self.eof {
                return Ok(None);
            }
            self.compact();
            if self.buf.capacity() - self.buf.len() < 1024 {
                self.buf.reserve(self.cfg.idle_buffer.max(1024));
            }
            // Cancel-safe: if cancelled, no data has been read.
            let n = self.io.read_buf(&mut self.buf).await?;
            if n == 0 {
                self.eof = true;
                if self.buf.len() > self.pos {
                    return Err(WsError::Protocol(
                        "connection closed in the middle of a frame",
                    ));
                }
            }
        }
    }

    fn compact(&mut self) {
        if self.pos > 0 {
            self.buf.drain(..self.pos);
            self.pos = 0;
        }
        if self.buf.is_empty() && self.buf.capacity() > self.cfg.idle_buffer * 4 {
            self.buf = Vec::with_capacity(self.cfg.idle_buffer);
        }
    }

    fn try_parse(&mut self) -> Result<Option<Message>, WsError> {
        loop {
            let avail = &self.buf[self.pos..];
            let Some(h) = frame::parse_header(avail).map_err(WsError::Protocol)? else {
                return Ok(None);
            };
            let max_frame = self.cfg.max_message_size as u64 + 1024;
            if h.payload_len > max_frame {
                return Err(WsError::MessageTooBig);
            }
            let total = h.header_len + h.payload_len as usize;
            if avail.len() < total {
                return Ok(None);
            }
            match (self.role, h.mask) {
                (Role::Server, None) => {
                    return Err(WsError::Protocol("client frames must be masked"));
                }
                (Role::Client, Some(_)) => {
                    return Err(WsError::Protocol("server frames must not be masked"));
                }
                _ => {}
            }
            if h.rsv2 || h.rsv3 {
                return Err(WsError::Protocol("reserved bits set"));
            }
            let opcode = OpCode::from_u8(h.opcode).ok_or(WsError::Protocol("unknown opcode"))?;
            let start = self.pos + h.header_len;
            let end = self.pos + total;
            if let Some(m) = h.mask {
                frame::apply_mask(&mut self.buf[start..end], m);
            }
            self.pos = end;
            let payload = &self.buf[start..end];

            if opcode.is_control() {
                if !h.fin || payload.len() > 125 {
                    return Err(WsError::Protocol("invalid control frame"));
                }
                if h.rsv1 {
                    return Err(WsError::Protocol("RSV1 set on control frame"));
                }
                return Ok(Some(match opcode {
                    OpCode::Ping => Message::Ping(payload.to_vec()),
                    OpCode::Pong => Message::Pong(payload.to_vec()),
                    _ => Message::Close(parse_close(payload)?),
                }));
            }

            match (opcode, &mut self.partial) {
                (OpCode::Continuation, None) => {
                    return Err(WsError::Protocol("unexpected continuation frame"));
                }
                (OpCode::Continuation, Some(p)) => {
                    if h.rsv1 {
                        return Err(WsError::Protocol("RSV1 set on continuation frame"));
                    }
                    if p.data.len() + payload.len() > self.cfg.max_message_size {
                        return Err(WsError::MessageTooBig);
                    }
                    p.data.extend_from_slice(payload);
                }
                (_, Some(_)) => {
                    return Err(WsError::Protocol(
                        "new data frame before previous message finished",
                    ));
                }
                (_, None) => {
                    if h.rsv1 && self.inflater.is_none() {
                        return Err(WsError::Protocol("RSV1 set without negotiated compression"));
                    }
                    self.partial = Some(Partial {
                        opcode,
                        compressed: h.rsv1,
                        data: payload.to_vec(),
                    });
                }
            }
            if h.fin {
                let p = self.partial.take().expect("partial message present");
                let data = if p.compressed {
                    let inflater = self.inflater.as_mut().expect("checked above");
                    inflater
                        .decompress(&p.data, self.cfg.max_message_size)
                        .map_err(|e| match e {
                            InflateError::TooBig => WsError::MessageTooBig,
                            InflateError::Corrupt => WsError::Protocol("invalid compressed data"),
                        })?
                } else {
                    p.data
                };
                return Ok(Some(match p.opcode {
                    OpCode::Text => {
                        Message::Text(String::from_utf8(data).map_err(|_| WsError::InvalidUtf8)?)
                    }
                    _ => Message::Binary(data),
                }));
            }
        }
    }
}

fn parse_close(payload: &[u8]) -> Result<Option<(u16, String)>, WsError> {
    match payload.len() {
        0 => Ok(None),
        1 => Err(WsError::Protocol("invalid close payload")),
        _ => {
            let code = u16::from_be_bytes([payload[0], payload[1]]);
            let reason = std::str::from_utf8(&payload[2..]).map_err(|_| WsError::InvalidUtf8)?;
            Ok(Some((code, reason.to_owned())))
        }
    }
}

/// Writing half. Not cancel-safe: do not race `send_*` futures in `select!`.
pub struct WsWriter<W> {
    io: W,
    role: Role,
    compress: Option<(usize, u32)>,
    rng: u64,
    out: Vec<u8>,
    closed: bool,
}

impl<W: AsyncWrite + Unpin> WsWriter<W> {
    pub fn new(
        io: W,
        role: Role,
        deflate: Option<(DeflateParams, deflate::DeflateConfig)>,
    ) -> Self {
        // Each RandomState has fresh keys, so this differs per writer.
        let seed = {
            use std::hash::{BuildHasher, Hasher};
            std::collections::hash_map::RandomState::new()
                .build_hasher()
                .finish()
                | 1
        };
        WsWriter {
            io,
            role,
            compress: deflate.map(|(_, c)| (c.min_compress_size, c.level)),
            rng: seed,
            out: Vec::new(),
            closed: false,
        }
    }

    fn mask(&mut self) -> Option<[u8; 4]> {
        match self.role {
            Role::Server => None,
            Role::Client => {
                // xorshift64*; masking keys only need to be unpredictable to intermediaries.
                self.rng ^= self.rng >> 12;
                self.rng ^= self.rng << 25;
                self.rng ^= self.rng >> 27;
                Some((self.rng.wrapping_mul(0x2545_F491_4F6C_DD1D) as u32).to_le_bytes())
            }
        }
    }

    async fn write_frame(
        &mut self,
        opcode: OpCode,
        rsv1: bool,
        payload: &[u8],
    ) -> Result<(), WsError> {
        if self.closed {
            return Err(WsError::Io(io::Error::new(
                io::ErrorKind::NotConnected,
                "close frame already sent",
            )));
        }
        let mask = self.mask();
        self.out.clear();
        frame::encode(&mut self.out, opcode, rsv1, payload, mask);
        self.io.write_all(&self.out).await?;
        self.io.flush().await?;
        if self.out.capacity() > 64 << 10 {
            self.out = Vec::new();
        }
        Ok(())
    }

    pub async fn send_text(&mut self, text: &str) -> Result<(), WsError> {
        match self.compress {
            Some((min, level)) if text.len() >= min => {
                let c = deflate::compress_message(text.as_bytes(), level);
                self.write_frame(OpCode::Text, true, &c).await
            }
            _ => self.write_frame(OpCode::Text, false, text.as_bytes()).await,
        }
    }

    pub async fn send_ping(&mut self, data: &[u8]) -> Result<(), WsError> {
        self.write_frame(OpCode::Ping, false, &data[..data.len().min(125)])
            .await
    }

    pub async fn send_pong(&mut self, data: &[u8]) -> Result<(), WsError> {
        self.write_frame(OpCode::Pong, false, &data[..data.len().min(125)])
            .await
    }

    /// Sends a close frame. Further sends fail.
    pub async fn send_close(&mut self, code: u16, reason: &str) -> Result<(), WsError> {
        let mut p = code.to_be_bytes().to_vec();
        let mut cut = reason.len().min(123);
        while !reason.is_char_boundary(cut) {
            cut -= 1;
        }
        p.extend_from_slice(&reason.as_bytes()[..cut]);
        let r = self.write_frame(OpCode::Close, false, &p).await;
        self.closed = true;
        r
    }

    pub fn is_closed(&self) -> bool {
        self.closed
    }

    pub async fn shutdown(&mut self) -> Result<(), WsError> {
        self.io.shutdown().await.map_err(WsError::Io)
    }
}