aerosocket_server/
handler.rs

1//! WebSocket connection handlers
2//!
3//! This module provides handler abstractions for processing WebSocket connections.
4
5use aerosocket_core::{Message, Result};
6use std::future::Future;
7use std::pin::Pin;
8
9/// Trait for handling WebSocket connections
10pub trait Handler: Send + Sync + 'static {
11    /// Handle a new connection
12    fn handle<'a>(
13        &'a self,
14        connection: crate::connection::ConnectionHandle,
15    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>>;
16
17    /// Clone the handler
18    fn clone_box(&self) -> Box<dyn Handler>;
19}
20
21impl Clone for Box<dyn Handler> {
22    fn clone(&self) -> Box<dyn Handler> {
23        self.clone_box()
24    }
25}
26
27/// Boxed handler type
28pub type BoxedHandler = Box<dyn Handler>;
29
30/// Default handler implementation
31#[derive(Debug, Clone)]
32pub struct DefaultHandler;
33
34impl DefaultHandler {
35    /// Create a new default handler
36    pub fn new() -> Self {
37        Self
38    }
39}
40
41impl Default for DefaultHandler {
42    fn default() -> Self {
43        Self::new()
44    }
45}
46
47impl Handler for DefaultHandler {
48    fn handle<'a>(
49        &'a self,
50        connection: crate::connection::ConnectionHandle,
51    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
52        Box::pin(async move {
53            // Get the connection from the handle
54            let mut conn = connection.try_lock().await.map_err(|_| {
55                aerosocket_core::Error::Other("Failed to lock connection".to_string())
56            })?;
57
58            while let Some(msg) = conn.next().await? {
59                match msg {
60                    Message::Text(text) => {
61                        conn.send_text(text.as_str()).await?;
62                    }
63                    Message::Binary(data) => {
64                        conn.send_binary(data.as_bytes().to_vec()).await?;
65                    }
66                    Message::Ping(_) => {
67                        conn.pong(None).await?;
68                    }
69                    Message::Close(close_msg) => {
70                        let code = close_msg.code();
71                        let reason = close_msg.reason();
72                        conn.close(code, Some(reason)).await?;
73                        break;
74                    }
75                    Message::Pong(_) => {
76                        // Ignore pong messages
77                    }
78                }
79            }
80
81            Ok(())
82        })
83    }
84
85    fn clone_box(&self) -> Box<dyn Handler> {
86        Box::new(self.clone())
87    }
88}
89
90/// Echo handler implementation
91#[derive(Debug, Clone)]
92pub struct EchoHandler;
93
94impl EchoHandler {
95    /// Create a new echo handler
96    pub fn new() -> Self {
97        Self
98    }
99}
100
101impl Default for EchoHandler {
102    fn default() -> Self {
103        Self::new()
104    }
105}
106
107impl Handler for EchoHandler {
108    fn handle<'a>(
109        &'a self,
110        connection: crate::connection::ConnectionHandle,
111    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
112        Box::pin(async move {
113            // Get the connection from the handle
114            let mut conn = connection.try_lock().await.map_err(|_| {
115                aerosocket_core::Error::Other("Failed to lock connection".to_string())
116            })?;
117
118            while let Some(msg) = conn.next().await? {
119                match msg {
120                    Message::Text(text) => {
121                        let echo_text = format!("Echo: {}", text.as_str());
122                        conn.send_text(&echo_text).await?;
123                    }
124                    Message::Binary(data) => {
125                        conn.send_binary(data.as_bytes().to_vec()).await?;
126                    }
127                    Message::Ping(_) => {
128                        conn.pong(None).await?;
129                    }
130                    Message::Close(close_msg) => {
131                        let code = close_msg.code();
132                        let reason = close_msg.reason();
133                        conn.close(code, Some(reason)).await?;
134                        break;
135                    }
136                    Message::Pong(_) => {
137                        // Ignore pong messages
138                    }
139                }
140            }
141
142            Ok(())
143        })
144    }
145
146    fn clone_box(&self) -> Box<dyn Handler> {
147        Box::new(self.clone())
148    }
149}
150
151/// Function-based handler
152#[derive(Clone)]
153pub struct FnHandler<F> {
154    f: F,
155}
156
157impl<F> FnHandler<F> {
158    /// Create a new function-based handler
159    pub fn new(f: F) -> Self {
160        Self { f }
161    }
162}
163
164impl<F> Handler for FnHandler<F>
165where
166    F: Fn(crate::connection::ConnectionHandle) -> Pin<Box<dyn Future<Output = Result<()>> + Send>>
167        + Send
168        + Sync
169        + Clone
170        + 'static,
171{
172    fn handle<'a>(
173        &'a self,
174        connection: crate::connection::ConnectionHandle,
175    ) -> Pin<Box<dyn Future<Output = Result<()>> + Send + 'a>> {
176        Box::pin((self.f)(connection))
177    }
178
179    fn clone_box(&self) -> Box<dyn Handler> {
180        Box::new(self.clone())
181    }
182}
183
184impl<F> std::fmt::Debug for FnHandler<F> {
185    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
186        f.debug_struct("FnHandler")
187            .field("f", &"<function>")
188            .finish()
189    }
190}
191
192/// Create a handler from a function
193pub fn from_fn<F, Fut>(f: F) -> FnHandler<F>
194where
195    F: Fn(crate::connection::ConnectionHandle) -> Fut + Send + Sync + Clone + 'static,
196    Fut: Future<Output = Result<()>> + Send + 'static,
197{
198    FnHandler::new(f)
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204    use bytes::Bytes;
205
206    #[tokio::test]
207    async fn test_default_handler() {
208        let handler = DefaultHandler::new();
209        let remote = "127.0.0.1:12345".parse().unwrap();
210        let local = "127.0.0.1:8080".parse().unwrap();
211        let connection = crate::connection::Connection::new(remote, local);
212
213        // Note: This test will fail until Connection::next and send are implemented
214        // For now, we just test that the handler can be created
215    }
216
217    #[tokio::test]
218    async fn test_echo_handler() {
219        let handler = EchoHandler::new();
220        let remote = "127.0.0.1:12345".parse().unwrap();
221        let local = "127.0.0.1:8080".parse().unwrap();
222        let connection = crate::connection::Connection::new(remote, local);
223
224        // Note: This test will fail until Connection::next and send are implemented
225        // For now, we just test that the handler can be created
226    }
227
228    #[tokio::test]
229    async fn test_fn_handler() {
230        let handler = from_fn(|_conn| async { Ok(()) });
231
232        let remote = "127.0.0.1:12345".parse().unwrap();
233        let local = "127.0.0.1:8080".parse().unwrap();
234        let connection = crate::connection::Connection::new(remote, local);
235
236        // Note: This test will fail until Connection::next and send are implemented
237        // For now, we just test that the handler can be created
238    }
239}