rullst_core/
ws.rs

1use axum::extract::ws::{Message as AxumMessage, WebSocket as AxumWebSocket};
2
3/// Error type for Rullst WebSocket operations.
4#[derive(Debug)]
5pub enum WsError {
6    /// Returned when sending a message to the client fails (e.g. the connection was dropped).
7    SendError(String),
8    /// Returned when receiving a message from the client fails (e.g. the connection was reset).
9    RecvError(String),
10}
11
12impl std::fmt::Display for WsError {
13    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
14        match self {
15            WsError::SendError(err) => write!(f, "WebSocket send error: {}", err),
16            WsError::RecvError(err) => write!(f, "WebSocket receive error: {}", err),
17        }
18    }
19}
20
21impl std::error::Error for WsError {}
22
23/// A high-level, elegant wrapper for Axum WebSockets, built for premium DX and HTMX
24pub struct WebSocket {
25    inner: AxumWebSocket,
26}
27
28impl WebSocket {
29    /// Wraps a raw Axum WebSocket connection in the high-level `WebSocket` interface.
30    pub fn new(inner: AxumWebSocket) -> Self {
31        WebSocket { inner }
32    }
33
34    /// Send a text message to the WebSocket client
35    #[cfg_attr(mutants, mutants::skip)]
36    pub async fn send_text(&mut self, text: impl Into<String>) -> Result<(), WsError> {
37        self.inner
38            .send(AxumMessage::Text(text.into().into()))
39            .await
40            .map_err(|e| WsError::SendError(e.to_string()))
41    }
42
43    /// Send an HTML fragment to the client (perfect for out-of-band HTMX swapping)
44    #[cfg_attr(mutants, mutants::skip)]
45    pub async fn send_html(&mut self, html_content: String) -> Result<(), WsError> {
46        self.send_text(html_content).await
47    }
48
49    /// Receive the next text/binary frame from the client.
50    ///
51    /// Ping and Pong control frames are skipped (the underlying socket already
52    /// answers pings), so client keepalives never surface as errors. A Close
53    /// frame or the end of the stream returns `None`.
54    #[cfg_attr(mutants, mutants::skip)]
55    pub async fn recv(&mut self) -> Option<Result<String, WsError>> {
56        loop {
57            return match self.inner.recv().await {
58                Some(Ok(AxumMessage::Text(t))) => Some(Ok(t.to_string())),
59                Some(Ok(AxumMessage::Binary(b))) => {
60                    Some(Ok(String::from_utf8_lossy(&b).to_string()))
61                }
62                Some(Ok(AxumMessage::Ping(_) | AxumMessage::Pong(_))) => continue,
63                Some(Ok(AxumMessage::Close(_))) | None => None,
64                Some(Err(e)) => Some(Err(WsError::RecvError(e.to_string()))),
65            };
66        }
67    }
68}
69
70#[cfg(test)]
71#[allow(clippy::unwrap_used, clippy::expect_used)]
72mod tests {
73    use super::*;
74
75    #[test]
76    fn test_ws_error_display() {
77        let err1 = WsError::SendError("timeout".to_string());
78        let err2 = WsError::RecvError("closed".to_string());
79
80        assert_eq!(format!("{}", err1), "WebSocket send error: timeout");
81        assert_eq!(format!("{}", err2), "WebSocket receive error: closed");
82    }
83}
rullst_core/ws.rs

rullst_core/
ws.rs
