caelix_core/
websocket.rs

1use crate::{BoxFuture, Gateway, HttpException, Injectable, Result};
2use bytes::Bytes;
3use std::{
4    collections::BTreeMap,
5    fmt,
6    net::SocketAddr,
7    sync::{
8        Arc,
9        atomic::{AtomicBool, Ordering},
10    },
11};
12
13#[derive(Clone, Copy, Debug, Eq, PartialEq)]
14/// Public Caelix enumeration `WebSocketCloseCode`.
15pub enum WebSocketCloseCode {
16    /// Public Caelix API.
17    Normal,
18    /// Public Caelix API.
19    GoingAway,
20    /// Public Caelix API.
21    Protocol,
22    /// Public Caelix API.
23    Unsupported,
24    /// Public Caelix API.
25    InvalidData,
26    /// Public Caelix API.
27    Policy,
28    /// Public Caelix API.
29    MessageTooBig,
30    /// Public Caelix API.
31    Internal,
32    /// Public Caelix API.
33    Restart,
34    /// Public Caelix API.
35    TryAgainLater,
36    /// Public Caelix API.
37    MandatoryExtension,
38    /// Public Caelix API.
39    BadGateway,
40    /// Public Caelix API.
41    Other(u16),
42}
43
44impl WebSocketCloseCode {
45    /// Runs the `as_u16` public API operation.
46    pub fn as_u16(self) -> u16 {
47        match self {
48            Self::Normal => 1000,
49            Self::GoingAway => 1001,
50            Self::Protocol => 1002,
51            Self::Unsupported => 1003,
52            Self::InvalidData => 1007,
53            Self::Policy => 1008,
54            Self::MessageTooBig => 1009,
55            Self::Internal => 1011,
56            Self::Restart => 1012,
57            Self::TryAgainLater => 1013,
58            Self::MandatoryExtension => 1010,
59            Self::BadGateway => 1014,
60            Self::Other(code) => code,
61        }
62    }
63    /// Runs the `from_u16` public API operation.
64    pub fn from_u16(code: u16) -> Option<Self> {
65        Some(match code {
66            1000 => Self::Normal,
67            1001 => Self::GoingAway,
68            1002 => Self::Protocol,
69            1003 => Self::Unsupported,
70            1007 => Self::InvalidData,
71            1008 => Self::Policy,
72            1009 => Self::MessageTooBig,
73            1010 => Self::MandatoryExtension,
74            1011 => Self::Internal,
75            1012 => Self::Restart,
76            1013 => Self::TryAgainLater,
77            1014 => Self::BadGateway,
78            3000..=4999 => Self::Other(code),
79            _ => return None,
80        })
81    }
82    /// Runs the `is_valid` public API operation.
83    pub fn is_valid(self) -> bool {
84        Self::from_u16(self.as_u16()).is_some()
85    }
86    /// Runs the `is_server_sendable` public API operation.
87    pub fn is_server_sendable(self) -> bool {
88        self.is_valid() && self != Self::MandatoryExtension
89    }
90}
91
92#[derive(Clone, Debug, Eq, PartialEq)]
93/// Public Caelix type `WebSocketCloseFrame`.
94pub struct WebSocketCloseFrame {
95    /// The `code` value.
96    pub code: WebSocketCloseCode,
97    /// The `reason` value.
98    pub reason: String,
99}
100impl WebSocketCloseFrame {
101    /// Runs the `new` public API operation.
102    pub fn new(code: WebSocketCloseCode, reason: impl Into<String>) -> Self {
103        Self {
104            code,
105            reason: reason.into(),
106        }
107    }
108}
109
110#[derive(Clone, Debug)]
111/// Public Caelix type `WebSocketRequest`.
112pub struct WebSocketRequest {
113    path: String,
114    query: String,
115    peer_addr: Option<SocketAddr>,
116    headers: BTreeMap<String, String>,
117}
118impl WebSocketRequest {
119    /// Runs the `new` public API operation.
120    pub fn new(
121        path: impl Into<String>,
122        query: impl Into<String>,
123        peer_addr: Option<SocketAddr>,
124        headers: BTreeMap<String, String>,
125    ) -> Self {
126        Self {
127            path: path.into(),
128            query: query.into(),
129            peer_addr,
130            headers,
131        }
132    }
133    /// Runs the `path` public API operation.
134    pub fn path(&self) -> &str {
135        &self.path
136    }
137    /// Runs the `query_string` public API operation.
138    pub fn query_string(&self) -> &str {
139        &self.query
140    }
141    /// Runs the `peer_addr` public API operation.
142    pub fn peer_addr(&self) -> Option<SocketAddr> {
143        self.peer_addr
144    }
145    /// Runs the `header` public API operation.
146    pub fn header(&self, name: &str) -> Option<&str> {
147        self.headers
148            .get(&name.to_ascii_lowercase())
149            .map(String::as_str)
150    }
151    /// Runs the `headers` public API operation.
152    pub fn headers(&self) -> &BTreeMap<String, String> {
153        &self.headers
154    }
155}
156
157#[derive(Clone, Debug)]
158/// Public Caelix type `WebSocketError`.
159pub struct WebSocketError {
160    message: String,
161}
162impl WebSocketError {
163    /// Runs the `new` public API operation.
164    pub fn new(message: impl Into<String>) -> Self {
165        Self {
166            message: message.into(),
167        }
168    }
169}
170impl fmt::Display for WebSocketError {
171    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
172        f.write_str(&self.message)
173    }
174}
175impl std::error::Error for WebSocketError {}
176
177#[doc(hidden)]
178pub trait WebSocketTransport: Send + Sync {
179    /// Public Caelix API.
180    fn send_text(&self, text: String) -> BoxFuture<'_, Result<()>>;
181    /// Public Caelix API.
182    fn send_binary(&self, data: Bytes) -> BoxFuture<'_, Result<()>>;
183    /// Public Caelix API.
184    fn ping(&self, data: Bytes) -> BoxFuture<'_, Result<()>>;
185    /// Public Caelix API.
186    fn pong(&self, data: Bytes) -> BoxFuture<'_, Result<()>>;
187    /// Public Caelix API.
188    fn close(&self, frame: Option<WebSocketCloseFrame>) -> BoxFuture<'_, Result<()>>;
189}
190
191#[derive(Clone)]
192/// Public Caelix type `WebSocketSession`.
193pub struct WebSocketSession {
194    id: String,
195    open: Arc<AtomicBool>,
196    transport: Arc<dyn WebSocketTransport>,
197    close_frame: Arc<std::sync::Mutex<Option<WebSocketCloseFrame>>>,
198}
199impl WebSocketSession {
200    #[doc(hidden)]
201    pub fn new(
202        id: impl Into<String>,
203        open: Arc<AtomicBool>,
204        transport: Arc<dyn WebSocketTransport>,
205    ) -> Self {
206        Self {
207            id: id.into(),
208            open,
209            transport,
210            close_frame: Arc::new(std::sync::Mutex::new(None)),
211        }
212    }
213    /// Runs the `id` public API operation.
214    pub fn id(&self) -> &str {
215        &self.id
216    }
217    /// Runs the `is_open` public API operation.
218    pub fn is_open(&self) -> bool {
219        self.open.load(Ordering::Acquire)
220    }
221    /// Runs the `send_text` public API operation.
222    pub async fn send_text(&self, text: impl Into<String>) -> Result<()> {
223        self.ensure_open()?;
224        self.transport.send_text(text.into()).await
225    }
226    /// Runs the `send_binary` public API operation.
227    pub async fn send_binary(&self, data: impl Into<Bytes>) -> Result<()> {
228        self.ensure_open()?;
229        self.transport.send_binary(data.into()).await
230    }
231    /// Runs the `ping` public API operation.
232    pub async fn ping(&self, data: impl Into<Bytes>) -> Result<()> {
233        self.ensure_open()?;
234        self.transport.ping(data.into()).await
235    }
236    /// Runs the `pong` public API operation.
237    pub async fn pong(&self, data: impl Into<Bytes>) -> Result<()> {
238        self.ensure_open()?;
239        self.transport.pong(data.into()).await
240    }
241    /// Runs the `close` public API operation.
242    pub async fn close(&self, frame: Option<WebSocketCloseFrame>) -> Result<()> {
243        *self
244            .close_frame
245            .lock()
246            .expect("websocket close state lock poisoned") = frame.clone();
247        self.transport.close(frame).await
248    }
249    #[doc(hidden)]
250    pub fn take_local_close_frame(&self) -> Option<WebSocketCloseFrame> {
251        self.close_frame
252            .lock()
253            .expect("websocket close state lock poisoned")
254            .take()
255    }
256    fn ensure_open(&self) -> Result<()> {
257        if self.is_open() {
258            Ok(())
259        } else {
260            Err(HttpException::new(
261                http::StatusCode::INTERNAL_SERVER_ERROR,
262                "Internal Server Error",
263                "websocket session is closed",
264            ))
265        }
266    }
267}
268
269/// Public Caelix extension trait `WebSocketGateway`.
270pub trait WebSocketGateway: Injectable {
271    /// Public Caelix API.
272    fn on_connect(
273        &self,
274        _session: Arc<WebSocketSession>,
275        _request: WebSocketRequest,
276    ) -> BoxFuture<'_, Result<()>> {
277        Box::pin(async { Ok(()) })
278    }
279    /// Public Caelix API.
280    fn on_text(&self, _session: Arc<WebSocketSession>, _text: String) -> BoxFuture<'_, Result<()>> {
281        Box::pin(async { Ok(()) })
282    }
283    /// Public Caelix API.
284    fn on_binary(
285        &self,
286        _session: Arc<WebSocketSession>,
287        _data: Bytes,
288    ) -> BoxFuture<'_, Result<()>> {
289        Box::pin(async { Ok(()) })
290    }
291    /// Public Caelix API.
292    fn on_error(
293        &self,
294        _session: Arc<WebSocketSession>,
295        _error: WebSocketError,
296    ) -> BoxFuture<'_, ()> {
297        Box::pin(async {})
298    }
299    /// Public Caelix API.
300    fn on_close(
301        &self,
302        _session: Arc<WebSocketSession>,
303        _frame: Option<WebSocketCloseFrame>,
304    ) -> BoxFuture<'_, ()> {
305        Box::pin(async {})
306    }
307}
308
309/// Marker implemented by [`#[gateway]`](https://docs.rs/caelix) for RFC 6455
310/// gateways. It separates path metadata from the transport callbacks so the
311/// same `ModuleMetadata::gateway` registration API can also support optional
312/// transports such as Socket.IO.
313pub trait RegisteredWebSocketGateway: WebSocketGateway + Gateway {}
314
315impl<T: WebSocketGateway + Gateway> RegisteredWebSocketGateway for T {}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320
321    #[test]
322    /// Public Caelix API.
323    fn close_codes_round_trip_without_corruption() {
324        for code in [1000, 1012, 1013, 3000, 3999, 4000, 4999] {
325            assert_eq!(WebSocketCloseCode::from_u16(code).unwrap().as_u16(), code);
326        }
327    }
328
329    #[test]
330    /// Public Caelix API.
331    fn prohibited_wire_close_codes_are_rejected() {
332        for code in [0, 999, 1004, 1005, 1006, 1015, 2000, 2999, 5000] {
333            assert_eq!(WebSocketCloseCode::from_u16(code), None);
334        }
335        assert!(!WebSocketCloseCode::Other(1005).is_valid());
336        assert!(!WebSocketCloseCode::Other(5000).is_valid());
337        assert!(WebSocketCloseCode::MandatoryExtension.is_valid());
338        assert!(!WebSocketCloseCode::MandatoryExtension.is_server_sendable());
339    }
340}
caelix_core/websocket.rs

caelix_core/
websocket.rs
