rustlavel_ws/
message.rs

1//! What an application sends and receives.
2//!
3//! A [`Message`] is a *whole* message: fragmentation is a transport detail that
4//! [`crate::connection::WebSocket`] reassembles before anything reaches here.
5//!
6//! `Ping`, `Pong` and `Close` carry their payloads because the protocol
7//! requires it — a pong must echo the ping's application data back byte for
8//! byte, and a close frame that drops its status code tells the peer nothing.
9
10use crate::error::{WsError, WsResult};
11use crate::frame::{CloseCode, CloseFrame, Frame, OpCode};
12use rustlavel_core::Json;
13
14#[derive(Debug, Clone, PartialEq, Eq)]
15pub enum Message {
16    Text(String),
17    Binary(Vec<u8>),
18    Ping(Vec<u8>),
19    Pong(Vec<u8>),
20    /// `None` when the peer closed without a status code, which is legal.
21    Close(Option<CloseFrame>),
22}
23
24impl Message {
25    pub fn text(body: impl Into<String>) -> Message {
26        Message::Text(body.into())
27    }
28
29    /// A JSON message. The broadcasting protocol is built entirely out of these.
30    pub fn json(value: &Json) -> Message {
31        Message::Text(value.to_string())
32    }
33
34    pub fn close(code: CloseCode, reason: impl Into<String>) -> Message {
35        Message::Close(Some(CloseFrame::new(code, reason)))
36    }
37
38    pub fn as_text(&self) -> Option<&str> {
39        match self {
40            Message::Text(body) => Some(body),
41            _ => None,
42        }
43    }
44
45    pub fn as_bytes(&self) -> &[u8] {
46        match self {
47            Message::Text(body) => body.as_bytes(),
48            Message::Binary(bytes) | Message::Ping(bytes) | Message::Pong(bytes) => bytes,
49            Message::Close(_) => &[],
50        }
51    }
52
53    /// Text or binary — the messages an application actually cares about.
54    pub fn is_data(&self) -> bool {
55        matches!(self, Message::Text(_) | Message::Binary(_))
56    }
57
58    pub fn len(&self) -> usize {
59        self.as_bytes().len()
60    }
61
62    pub fn is_empty(&self) -> bool {
63        self.len() == 0
64    }
65
66    /// A stable label for instrumentation. Never the payload itself: a socket
67    /// carrying chat or auth tokens must not leak them into Telescope.
68    pub fn kind(&self) -> &'static str {
69        match self {
70            Message::Text(_) => "text",
71            Message::Binary(_) => "binary",
72            Message::Ping(_) => "ping",
73            Message::Pong(_) => "pong",
74            Message::Close(_) => "close",
75        }
76    }
77
78    /// The frame a server writes for this message: final, and never masked.
79    pub fn into_frame(self) -> Frame {
80        match self {
81            Message::Text(body) => Frame::server(OpCode::Text, body.into_bytes()),
82            Message::Binary(bytes) => Frame::server(OpCode::Binary, bytes),
83            Message::Ping(bytes) => Frame::server(OpCode::Ping, bytes),
84            Message::Pong(bytes) => Frame::server(OpCode::Pong, bytes),
85            Message::Close(close) => Frame::server(
86                OpCode::Close,
87                close.map(|close| close.to_payload()).unwrap_or_default(),
88            ),
89        }
90    }
91
92    /// Turn a reassembled payload into a message.
93    ///
94    /// Text is validated here rather than at the frame layer, because a message
95    /// split across fragments is only valid UTF-8 once every fragment has
96    /// arrived — a multi-byte character is allowed to straddle the boundary.
97    pub fn from_data(opcode: OpCode, payload: Vec<u8>) -> WsResult<Message> {
98        match opcode {
99            OpCode::Text => {
100                String::from_utf8(payload).map(Message::Text).map_err(|_| WsError::InvalidUtf8)
101            }
102            OpCode::Binary => Ok(Message::Binary(payload)),
103            other => Err(WsError::protocol(format!("{other:?} is not a data frame"))),
104        }
105    }
106}
107
108impl From<String> for Message {
109    fn from(body: String) -> Message {
110        Message::Text(body)
111    }
112}
113
114impl From<&str> for Message {
115    fn from(body: &str) -> Message {
116        Message::Text(body.to_string())
117    }
118}
119
120impl From<Json> for Message {
121    fn from(value: Json) -> Message {
122        Message::Text(value.to_string())
123    }
124}
125
126impl From<Vec<u8>> for Message {
127    fn from(bytes: Vec<u8>) -> Message {
128        Message::Binary(bytes)
129    }
130}
131
132#[cfg(test)]
133mod tests {
134    use super::*;
135
136    #[test]
137    fn a_text_message_becomes_an_unmasked_text_frame() {
138        let frame = Message::text("hi").into_frame();
139
140        assert_eq!(frame.opcode, OpCode::Text);
141        assert!(frame.fin);
142        assert_eq!(frame.mask, None, "a server must never mask");
143        assert_eq!(frame.encode(), b"\x81\x02hi");
144    }
145
146    #[test]
147    fn a_close_message_encodes_its_code_and_reason() {
148        let frame = Message::close(CloseCode::POLICY, "no").into_frame();
149        assert_eq!(frame.encode(), b"\x88\x04\x03\xf0no");
150
151        // Closing with no status at all is legal and sends an empty payload.
152        assert_eq!(Message::Close(None).into_frame().encode(), b"\x88\x00");
153    }
154
155    #[test]
156    fn reassembled_text_must_be_valid_utf8() {
157        assert_eq!(
158            Message::from_data(OpCode::Text, "héllo".as_bytes().to_vec()).unwrap(),
159            Message::text("héllo")
160        );
161
162        // A lone continuation byte: the classic 1007 case.
163        let error = Message::from_data(OpCode::Text, vec![0xf0, 0x28, 0x8c, 0x28]).unwrap_err();
164        assert_eq!(error.close_code(), CloseCode::INVALID_PAYLOAD);
165    }
166
167    #[test]
168    fn binary_payloads_are_passed_through_untouched() {
169        let bytes = vec![0x00, 0xff, 0x80];
170        assert_eq!(
171            Message::from_data(OpCode::Binary, bytes.clone()).unwrap(),
172            Message::Binary(bytes)
173        );
174    }
175
176    #[test]
177    fn instrumentation_sees_a_kind_and_a_size_but_never_the_payload() {
178        let message = Message::text("a secret token");
179
180        assert_eq!(message.kind(), "text");
181        assert_eq!(message.len(), 14);
182        assert!(message.is_data());
183        assert!(!Message::Ping(Vec::new()).is_data());
184    }
185
186    #[test]
187    fn json_values_convert_straight_into_messages() {
188        let value = Json::object([("event", Json::from("ping"))]);
189        assert_eq!(Message::json(&value).as_text(), Some(r#"{"event":"ping"}"#));
190    }
191}
rustlavel_ws/message.rs

rustlavel_ws/
message.rs
