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

caelix_core/
websocket.rs
