codetics_ocpp_ws/
conn.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2026 Kevin Vavelin, Codetics (https://codetics.fr)
3
4//! Message-level reader and writer halves of a WebSocket connection.
5
6use std::fmt;
7use std::io;
8
9use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
10
11use crate::deflate::{self, DeflateParams, InflateError, Inflater};
12use crate::frame::{self, OpCode};
13
14#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum Role {
16    Server,
17    Client,
18}
19
20/// Standard close codes (RFC 6455 §7.4.1).
21pub mod close_code {
22    pub const NORMAL: u16 = 1000;
23    pub const GOING_AWAY: u16 = 1001;
24    pub const PROTOCOL_ERROR: u16 = 1002;
25    pub const UNSUPPORTED_DATA: u16 = 1003;
26    pub const INVALID_PAYLOAD: u16 = 1007;
27    pub const POLICY_VIOLATION: u16 = 1008;
28    pub const MESSAGE_TOO_BIG: u16 = 1009;
29    pub const INTERNAL_ERROR: u16 = 1011;
30}
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub enum Message {
34    Text(String),
35    Binary(Vec<u8>),
36    Ping(Vec<u8>),
37    Pong(Vec<u8>),
38    /// Close frame with optional (code, reason).
39    Close(Option<(u16, String)>),
40}
41
42#[derive(Debug)]
43pub enum WsError {
44    Io(io::Error),
45    Protocol(&'static str),
46    InvalidUtf8,
47    MessageTooBig,
48}
49
50impl WsError {
51    /// Close code to send to the peer for this error.
52    pub fn close_code(&self) -> u16 {
53        match self {
54            WsError::Io(_) => close_code::INTERNAL_ERROR,
55            WsError::Protocol(_) => close_code::PROTOCOL_ERROR,
56            WsError::InvalidUtf8 => close_code::INVALID_PAYLOAD,
57            WsError::MessageTooBig => close_code::MESSAGE_TOO_BIG,
58        }
59    }
60}
61
62impl fmt::Display for WsError {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        match self {
65            WsError::Io(e) => write!(f, "I/O error: {e}"),
66            WsError::Protocol(m) => write!(f, "WebSocket protocol error: {m}"),
67            WsError::InvalidUtf8 => f.write_str("text message is not valid UTF-8"),
68            WsError::MessageTooBig => f.write_str("message exceeds the maximum size"),
69        }
70    }
71}
72
73impl std::error::Error for WsError {}
74
75impl From<io::Error> for WsError {
76    fn from(e: io::Error) -> Self {
77        WsError::Io(e)
78    }
79}
80
81#[derive(Debug, Clone, Copy)]
82pub struct WsConfig {
83    /// Maximum size of a (decompressed) message.
84    pub max_message_size: usize,
85    /// Buffer capacity kept for idle connections; larger buffers shrink back to this.
86    pub idle_buffer: usize,
87}
88
89impl Default for WsConfig {
90    fn default() -> Self {
91        WsConfig {
92            max_message_size: 4 << 20,
93            idle_buffer: 4 << 10,
94        }
95    }
96}
97
98struct Partial {
99    opcode: OpCode,
100    compressed: bool,
101    data: Vec<u8>,
102}
103
104/// Reading half. [`WsReader::recv`] is cancel-safe: all progress is kept in
105/// `self`, so it can be used as a `tokio::select!` branch.
106pub struct WsReader<R> {
107    io: R,
108    role: Role,
109    cfg: WsConfig,
110    buf: Vec<u8>,
111    pos: usize,
112    partial: Option<Partial>,
113    inflater: Option<Inflater>,
114    eof: bool,
115}
116
117impl<R: AsyncRead + Unpin> WsReader<R> {
118    pub fn new(
119        io: R,
120        role: Role,
121        cfg: WsConfig,
122        deflate: Option<DeflateParams>,
123        initial: Vec<u8>,
124    ) -> Self {
125        let inflater = deflate.map(|p| {
126            Inflater::new(match role {
127                Role::Server => p.client_no_context_takeover,
128                Role::Client => p.server_no_context_takeover,
129            })
130        });
131        WsReader {
132            io,
133            role,
134            cfg,
135            buf: initial,
136            pos: 0,
137            partial: None,
138            inflater,
139            eof: false,
140        }
141    }
142
143    /// Receives the next message. `Ok(None)` means the TCP stream ended.
144    pub async fn recv(&mut self) -> Result<Option<Message>, WsError> {
145        loop {
146            if let Some(msg) = self.try_parse()? {
147                return Ok(Some(msg));
148            }
149            if self.eof {
150                return Ok(None);
151            }
152            self.compact();
153            if self.buf.capacity() - self.buf.len() < 1024 {
154                self.buf.reserve(self.cfg.idle_buffer.max(1024));
155            }
156            // Cancel-safe: if cancelled, no data has been read.
157            let n = self.io.read_buf(&mut self.buf).await?;
158            if n == 0 {
159                self.eof = true;
160                if self.buf.len() > self.pos {
161                    return Err(WsError::Protocol(
162                        "connection closed in the middle of a frame",
163                    ));
164                }
165            }
166        }
167    }
168
169    fn compact(&mut self) {
170        if self.pos > 0 {
171            self.buf.drain(..self.pos);
172            self.pos = 0;
173        }
174        if self.buf.is_empty() && self.buf.capacity() > self.cfg.idle_buffer * 4 {
175            self.buf = Vec::with_capacity(self.cfg.idle_buffer);
176        }
177    }
178
179    fn try_parse(&mut self) -> Result<Option<Message>, WsError> {
180        loop {
181            let avail = &self.buf[self.pos..];
182            let Some(h) = frame::parse_header(avail).map_err(WsError::Protocol)? else {
183                return Ok(None);
184            };
185            let max_frame = self.cfg.max_message_size as u64 + 1024;
186            if h.payload_len > max_frame {
187                return Err(WsError::MessageTooBig);
188            }
189            let total = h.header_len + h.payload_len as usize;
190            if avail.len() < total {
191                return Ok(None);
192            }
193            match (self.role, h.mask) {
194                (Role::Server, None) => {
195                    return Err(WsError::Protocol("client frames must be masked"));
196                }
197                (Role::Client, Some(_)) => {
198                    return Err(WsError::Protocol("server frames must not be masked"));
199                }
200                _ => {}
201            }
202            if h.rsv2 || h.rsv3 {
203                return Err(WsError::Protocol("reserved bits set"));
204            }
205            let opcode = OpCode::from_u8(h.opcode).ok_or(WsError::Protocol("unknown opcode"))?;
206            let start = self.pos + h.header_len;
207            let end = self.pos + total;
208            if let Some(m) = h.mask {
209                frame::apply_mask(&mut self.buf[start..end], m);
210            }
211            self.pos = end;
212            let payload = &self.buf[start..end];
213
214            if opcode.is_control() {
215                if !h.fin || payload.len() > 125 {
216                    return Err(WsError::Protocol("invalid control frame"));
217                }
218                if h.rsv1 {
219                    return Err(WsError::Protocol("RSV1 set on control frame"));
220                }
221                return Ok(Some(match opcode {
222                    OpCode::Ping => Message::Ping(payload.to_vec()),
223                    OpCode::Pong => Message::Pong(payload.to_vec()),
224                    _ => Message::Close(parse_close(payload)?),
225                }));
226            }
227
228            match (opcode, &mut self.partial) {
229                (OpCode::Continuation, None) => {
230                    return Err(WsError::Protocol("unexpected continuation frame"));
231                }
232                (OpCode::Continuation, Some(p)) => {
233                    if h.rsv1 {
234                        return Err(WsError::Protocol("RSV1 set on continuation frame"));
235                    }
236                    if p.data.len() + payload.len() > self.cfg.max_message_size {
237                        return Err(WsError::MessageTooBig);
238                    }
239                    p.data.extend_from_slice(payload);
240                }
241                (_, Some(_)) => {
242                    return Err(WsError::Protocol(
243                        "new data frame before previous message finished",
244                    ));
245                }
246                (_, None) => {
247                    if h.rsv1 && self.inflater.is_none() {
248                        return Err(WsError::Protocol("RSV1 set without negotiated compression"));
249                    }
250                    self.partial = Some(Partial {
251                        opcode,
252                        compressed: h.rsv1,
253                        data: payload.to_vec(),
254                    });
255                }
256            }
257            if h.fin {
258                let p = self.partial.take().expect("partial message present");
259                let data = if p.compressed {
260                    let inflater = self.inflater.as_mut().expect("checked above");
261                    inflater
262                        .decompress(&p.data, self.cfg.max_message_size)
263                        .map_err(|e| match e {
264                            InflateError::TooBig => WsError::MessageTooBig,
265                            InflateError::Corrupt => WsError::Protocol("invalid compressed data"),
266                        })?
267                } else {
268                    p.data
269                };
270                return Ok(Some(match p.opcode {
271                    OpCode::Text => {
272                        Message::Text(String::from_utf8(data).map_err(|_| WsError::InvalidUtf8)?)
273                    }
274                    _ => Message::Binary(data),
275                }));
276            }
277        }
278    }
279}
280
281fn parse_close(payload: &[u8]) -> Result<Option<(u16, String)>, WsError> {
282    match payload.len() {
283        0 => Ok(None),
284        1 => Err(WsError::Protocol("invalid close payload")),
285        _ => {
286            let code = u16::from_be_bytes([payload[0], payload[1]]);
287            let reason = std::str::from_utf8(&payload[2..]).map_err(|_| WsError::InvalidUtf8)?;
288            Ok(Some((code, reason.to_owned())))
289        }
290    }
291}
292
293/// Writing half. Not cancel-safe: do not race `send_*` futures in `select!`.
294pub struct WsWriter<W> {
295    io: W,
296    role: Role,
297    compress: Option<(usize, u32)>,
298    rng: u64,
299    out: Vec<u8>,
300    closed: bool,
301}
302
303impl<W: AsyncWrite + Unpin> WsWriter<W> {
304    pub fn new(
305        io: W,
306        role: Role,
307        deflate: Option<(DeflateParams, deflate::DeflateConfig)>,
308    ) -> Self {
309        // Each RandomState has fresh keys, so this differs per writer.
310        let seed = {
311            use std::hash::{BuildHasher, Hasher};
312            std::collections::hash_map::RandomState::new()
313                .build_hasher()
314                .finish()
315                | 1
316        };
317        WsWriter {
318            io,
319            role,
320            compress: deflate.map(|(_, c)| (c.min_compress_size, c.level)),
321            rng: seed,
322            out: Vec::new(),
323            closed: false,
324        }
325    }
326
327    fn mask(&mut self) -> Option<[u8; 4]> {
328        match self.role {
329            Role::Server => None,
330            Role::Client => {
331                // xorshift64*; masking keys only need to be unpredictable to intermediaries.
332                self.rng ^= self.rng >> 12;
333                self.rng ^= self.rng << 25;
334                self.rng ^= self.rng >> 27;
335                Some((self.rng.wrapping_mul(0x2545_F491_4F6C_DD1D) as u32).to_le_bytes())
336            }
337        }
338    }
339
340    async fn write_frame(
341        &mut self,
342        opcode: OpCode,
343        rsv1: bool,
344        payload: &[u8],
345    ) -> Result<(), WsError> {
346        if self.closed {
347            return Err(WsError::Io(io::Error::new(
348                io::ErrorKind::NotConnected,
349                "close frame already sent",
350            )));
351        }
352        let mask = self.mask();
353        self.out.clear();
354        frame::encode(&mut self.out, opcode, rsv1, payload, mask);
355        self.io.write_all(&self.out).await?;
356        self.io.flush().await?;
357        if self.out.capacity() > 64 << 10 {
358            self.out = Vec::new();
359        }
360        Ok(())
361    }
362
363    pub async fn send_text(&mut self, text: &str) -> Result<(), WsError> {
364        match self.compress {
365            Some((min, level)) if text.len() >= min => {
366                let c = deflate::compress_message(text.as_bytes(), level);
367                self.write_frame(OpCode::Text, true, &c).await
368            }
369            _ => self.write_frame(OpCode::Text, false, text.as_bytes()).await,
370        }
371    }
372
373    pub async fn send_ping(&mut self, data: &[u8]) -> Result<(), WsError> {
374        self.write_frame(OpCode::Ping, false, &data[..data.len().min(125)])
375            .await
376    }
377
378    pub async fn send_pong(&mut self, data: &[u8]) -> Result<(), WsError> {
379        self.write_frame(OpCode::Pong, false, &data[..data.len().min(125)])
380            .await
381    }
382
383    /// Sends a close frame. Further sends fail.
384    pub async fn send_close(&mut self, code: u16, reason: &str) -> Result<(), WsError> {
385        let mut p = code.to_be_bytes().to_vec();
386        let mut cut = reason.len().min(123);
387        while !reason.is_char_boundary(cut) {
388            cut -= 1;
389        }
390        p.extend_from_slice(&reason.as_bytes()[..cut]);
391        let r = self.write_frame(OpCode::Close, false, &p).await;
392        self.closed = true;
393        r
394    }
395
396    pub fn is_closed(&self) -> bool {
397        self.closed
398    }
399
400    pub async fn shutdown(&mut self) -> Result<(), WsError> {
401        self.io.shutdown().await.map_err(WsError::Io)
402    }
403}
codetics_ocpp_ws/conn.rs

codetics_ocpp_ws/
conn.rs
