workflow_websocket/server/
mod.rs

1//!
2//! async WebSocket server functionality (requires tokio executor)
3//!
4use async_trait::async_trait;
5use cfg_if::cfg_if;
6use downcast_rs::*;
7use futures::{future::FutureExt, select};
8use futures_util::{
9    SinkExt, StreamExt,
10    stream::{SplitSink, SplitStream},
11};
12use std::net::SocketAddr;
13use std::pin::Pin;
14use std::sync::Arc;
15use std::sync::atomic::{AtomicUsize, Ordering};
16use std::time::Duration;
17pub use tokio::net::TcpListener;
18use tokio::net::TcpStream;
19use tokio::sync::mpsc::{
20    UnboundedReceiver as TokioUnboundedReceiver, UnboundedSender as TokioUnboundedSender,
21};
22use tokio_tungstenite::{WebSocketStream, accept_async_with_config};
23use tungstenite::Error as WebSocketError;
24use workflow_core::channel::DuplexChannel;
25use workflow_log::*;
26pub mod error;
27pub mod result;
28
29pub use error::Error;
30pub use result::Result;
31pub use tungstenite::Message;
32pub use tungstenite::protocol::WebSocketConfig;
33/// WebSocket stream sender for dispatching [`tungstenite::Message`].
34/// This stream object must have a mutable reference and can not be cloned.
35pub type WebSocketSender = SplitSink<WebSocketStream<TcpStream>, Message>;
36/// WebSocket stream receiver for receiving [`tungstenite::Message`].
37/// This stream object must have a mutable reference and can not be cloned.
38pub type WebSocketReceiver = SplitStream<WebSocketStream<TcpStream>>;
39/// WebSocketSink [`tokio::sync::mpsc::UnboundedSender`] for dispatching
40/// messages from within the [`WebSocketHandler::message`]. This is an
41/// `MPSC` channel that can be cloned and retained externally for the
42/// lifetime of the WebSocket connection.
43pub type WebSocketSink = TokioUnboundedSender<Message>;
44
45/// Atomic counters that allow tracking connection counts
46/// and cumulative message sizes in bytes (bandwidth consumption
47/// without accounting for the websocket framing overhead).
48/// These counters can be created and supplied externally or
49/// supplied as `None`.
50pub struct WebSocketCounters {
51    /// Cumulative number of connections accepted since startup.
52    pub total_connections: Arc<AtomicUsize>,
53    /// Number of currently active (open) connections.
54    pub active_connections: Arc<AtomicUsize>,
55    /// Cumulative number of connections that failed during handshake.
56    pub handshake_failures: Arc<AtomicUsize>,
57    /// Total bytes received across all connections (excluding framing overhead).
58    pub rx_bytes: Arc<AtomicUsize>,
59    /// Total bytes sent across all connections (excluding framing overhead).
60    pub tx_bytes: Arc<AtomicUsize>,
61}
62
63impl Default for WebSocketCounters {
64    fn default() -> Self {
65        WebSocketCounters {
66            total_connections: Arc::new(AtomicUsize::new(0)),
67            active_connections: Arc::new(AtomicUsize::new(0)),
68            handshake_failures: Arc::new(AtomicUsize::new(0)),
69            rx_bytes: Arc::new(AtomicUsize::new(0)),
70            tx_bytes: Arc::new(AtomicUsize::new(0)),
71        }
72    }
73}
74
75/// WebSocketHandler trait that represents the WebSocket processor
76/// functionality.  This trait is supplied to the WebSocket
77/// which subsequently invokes it's functions during websocket
78/// connection and messages.  The trait can override `with_handshake()` method
79/// to enable invocation of the `handshake()` method upon receipt of the
80/// first valid websocket message from the incoming connection.
81#[async_trait]
82pub trait WebSocketHandler
83where
84    Arc<Self>: Sync,
85{
86    /// Context type used by impl trait to represent websocket connection
87    type Context: Send + Sync;
88
89    /// Called to determine if the connection should be accepted.
90    fn accept(&self, _peer: &SocketAddr) -> bool {
91        true
92    }
93
94    /// Called immediately when connection is established.
95    /// This function should return an error to terminate the connection.
96    /// If the server manages a client ban list, it should process it
97    /// in this function and return an [`Error`] to prevent further processing.
98    async fn connect(self: &Arc<Self>, _peer: &SocketAddr) -> Result<()> {
99        Ok(())
100    }
101
102    /// Called upon websocket disconnection
103    async fn disconnect(self: &Arc<Self>, _ctx: Self::Context, _result: Result<()>) {}
104
105    /// Called after [`Self::connect()`], after creating the [`tokio::sync::mpsc`] sender `sink`
106    /// channel, allowing the server to execute additional handshake communication phase,
107    /// or retain the sink for external message dispatch (such as server-side notifications).
108    async fn handshake(
109        self: &Arc<Self>,
110        peer: &SocketAddr,
111        sender: &mut WebSocketSender,
112        receiver: &mut WebSocketReceiver,
113        sink: &WebSocketSink,
114    ) -> Result<Self::Context>;
115
116    /// Called for every websocket message
117    /// This function can return an error to terminate the connection
118    async fn message(
119        self: &Arc<Self>,
120        ctx: &Self::Context,
121        msg: Message,
122        sink: &WebSocketSink,
123    ) -> Result<()>;
124
125    /// Called to handle control messages (such as `Ping`/`Pong`) when the
126    /// `ping-pong` feature is enabled. The default implementation replies to
127    /// a `Ping` with a matching `Pong`.
128    async fn ctl(self: &Arc<Self>, msg: Message, sender: &mut WebSocketSender) -> Result<()> {
129        if let Message::Ping(data) = msg {
130            sender.send(Message::Pong(data)).await?;
131        }
132        Ok(())
133    }
134}
135
136/// WebSocketServer that provides the main websocket connection
137/// and message processing loop that delivers messages to the
138/// installed WebSocketHandler trait.
139pub struct WebSocketServer<T>
140where
141    T: WebSocketHandler + Send + Sync + 'static + Sized,
142{
143    // pub connections: AtomicU64,
144    /// Connection and bandwidth counters tracked by this server.
145    pub counters: Arc<WebSocketCounters>,
146    /// The user-supplied handler driving connection and message processing.
147    pub handler: Arc<T>,
148    /// Duplex channel used to signal and await server shutdown.
149    pub stop: DuplexChannel,
150}
151
152impl<T> WebSocketServer<T>
153where
154    T: WebSocketHandler + Send + Sync + 'static,
155{
156    /// Creates a new server with the given handler, optionally sharing externally
157    /// supplied [`WebSocketCounters`] (new default counters are created if `None`).
158    pub fn new(handler: Arc<T>, counters: Option<Arc<WebSocketCounters>>) -> Arc<Self> {
159        Arc::new(WebSocketServer {
160            counters: counters.unwrap_or_default(),
161            handler,
162            stop: DuplexChannel::oneshot(),
163        })
164    }
165
166    async fn handle_connection(
167        self: &Arc<Self>,
168        peer: SocketAddr,
169        stream: TcpStream,
170        config: Option<WebSocketConfig>,
171    ) -> Result<()> {
172        let ws_stream = accept_async_with_config(stream, config).await?;
173        self.handler.connect(&peer).await?;
174        // log_trace!("WebSocket connected: {}", peer);
175
176        let (mut ws_sender, mut ws_receiver) = ws_stream.split();
177        let (sink_sender, sink_receiver) = tokio::sync::mpsc::unbounded_channel::<Message>();
178
179        let ctx = match self
180            .handler
181            .handshake(&peer, &mut ws_sender, &mut ws_receiver, &sink_sender)
182            .await
183        {
184            Ok(ctx) => ctx,
185            Err(err) => {
186                self.counters
187                    .handshake_failures
188                    .fetch_add(1, Ordering::Relaxed);
189                return Err(err);
190            }
191        };
192
193        let result = self
194            .connection_task(&ctx, ws_sender, ws_receiver, sink_sender, sink_receiver)
195            .await;
196        self.handler.disconnect(ctx, result).await;
197        // log_trace!("WebSocket disconnected: {}", peer);
198
199        Ok(())
200    }
201
202    async fn connection_task(
203        self: &Arc<Self>,
204        ctx: &T::Context,
205        mut ws_sender: WebSocketSender,
206        mut ws_receiver: WebSocketReceiver,
207        sink_sender: TokioUnboundedSender<Message>,
208        mut sink_receiver: TokioUnboundedReceiver<Message>,
209    ) -> Result<()> {
210        loop {
211            tokio::select! {
212                msg = sink_receiver.recv() => {
213                    let msg = msg.unwrap();
214                    match msg {
215                        Message::Binary(data)  => {
216                            self.counters.tx_bytes.fetch_add(data.len(), Ordering::Relaxed);
217                            ws_sender.send(Message::Binary(data)).await?;
218                        },
219                        Message::Text(text)  => {
220                            self.counters.tx_bytes.fetch_add(text.len(), Ordering::Relaxed);
221                            ws_sender.send(Message::Text(text)).await?;
222                        },
223                        Message::Close(_) => {
224                            ws_sender.send(msg).await?;
225                            break;
226                        },
227                        Message::Ping(data) => {
228                            self.counters.tx_bytes.fetch_add(data.len(), Ordering::Relaxed);
229                            ws_sender.send(Message::Ping(data)).await?;
230                        },
231                        Message::Pong(data) => {
232                            self.counters.tx_bytes.fetch_add(data.len(), Ordering::Relaxed);
233                            ws_sender.send(Message::Pong(data)).await?;
234                        },
235                        msg => {
236                            ws_sender.send(msg).await?;
237                        }
238                    }
239                },
240                msg = ws_receiver.next() => {
241                    match msg {
242                        Some(msg) => {
243                            let msg = msg?;
244                            match msg {
245                                Message::Binary(data)  => {
246                                    self.counters.rx_bytes.fetch_add(data.len(), Ordering::Relaxed);
247                                    self.handler.message(ctx, Message::Binary(data), &sink_sender).await?;
248                                },
249                                Message::Text(text)  => {
250                                    self.counters.rx_bytes.fetch_add(text.len(), Ordering::Relaxed);
251                                    self.handler.message(ctx, Message::Text(text), &sink_sender).await?;
252                                },
253                                Message::Close(_) => {
254                                    self.handler.message(ctx, msg, &sink_sender).await?;
255                                    break;
256                                },
257                                Message::Ping(data) => {
258                                    self.counters.rx_bytes.fetch_add(data.len(), Ordering::Relaxed);
259                                    cfg_if! {
260                                        if #[cfg(feature = "ping-pong")] {
261                                            self.handler.ctl(Message::Ping(data), &mut ws_sender).await?;
262                                        } else {
263                                            ws_sender.send(Message::Pong(data)).await?;
264                                        }
265                                    }
266                                },
267                                Message::Pong(data) => {
268                                    self.counters.rx_bytes.fetch_add(data.len(), Ordering::Relaxed);
269                                    cfg_if! {
270                                        if #[cfg(feature = "ping-pong")] {
271                                            self.handler.ctl(Message::Pong(data), &mut ws_sender).await?;
272                                        } else {
273                                            // ignore pong
274                                        }
275                                    }
276                                },
277                                _ => {
278                                }
279                            }
280                        }
281                        None => {
282                            return Err(Error::AbnormalClose);
283                        }
284                    }
285                }
286            }
287        }
288
289        Ok(())
290    }
291
292    /// Binds a TCP listener to the given address, returning an [`Error::Listen`]
293    /// if binding fails.
294    pub async fn bind(self: &Arc<Self>, addr: &str) -> Result<TcpListener> {
295        let listener = TcpListener::bind(&addr).await.map_err(|err| {
296            Error::Listen(format!(
297                "WebSocket server unable to listen on `{addr}`: {err}",
298            ))
299        })?;
300        // log_trace!("WebSocket server listening on: {}", addr);
301        Ok(listener)
302    }
303
304    async fn accept(self: &Arc<Self>, stream: TcpStream, config: Option<WebSocketConfig>) {
305        let peer = match stream.peer_addr() {
306            Ok(peer_address) => peer_address,
307            Err(_) => {
308                self.counters
309                    .handshake_failures
310                    .fetch_add(1, Ordering::Relaxed);
311                return;
312            }
313        };
314
315        self.counters
316            .total_connections
317            .fetch_add(1, Ordering::Relaxed);
318        self.counters
319            .active_connections
320            .fetch_add(1, Ordering::Relaxed);
321
322        let self_ = self.clone();
323        tokio::spawn(async move {
324            if let Err(e) = self_.handle_connection(peer, stream, config).await {
325                match &e {
326                    Error::WebSocketError(wse) => match **wse {
327                        WebSocketError::ConnectionClosed
328                        | WebSocketError::Protocol(_)
329                        | WebSocketError::Utf8(_) => {}
330                        _ => {
331                            log_error!("Error processing connection: {}", e);
332                        }
333                    },
334                    err => log_error!("Error processing connection: {}", err),
335                }
336            }
337            self_
338                .counters
339                .active_connections
340                .fetch_sub(1, Ordering::Relaxed)
341        });
342    }
343
344    /// Runs the accept loop on the given listener, dispatching each accepted
345    /// connection to the handler until [`stop()`](Self::stop) is signaled.
346    pub async fn listen(
347        self: &Arc<Self>,
348        listener: TcpListener,
349        config: Option<WebSocketConfig>,
350    ) -> Result<()> {
351        loop {
352            select! {
353                stream = listener.accept().fuse() => {
354                    if let Ok((stream,socket_addr)) = stream
355                        && self.handler.accept(&socket_addr) {
356                            self.accept(stream, config).await;
357                        }
358                },
359                _ = self.stop.request.receiver.recv().fuse() => break,
360            }
361        }
362
363        self.stop
364            .response
365            .sender
366            .send(())
367            .await
368            .map_err(|err| Error::Done(err.to_string()))
369    }
370
371    /// Signals the running [`listen()`](Self::listen) loop to shut down.
372    pub fn stop(&self) -> Result<()> {
373        self.stop
374            .request
375            .sender
376            .try_send(())
377            .map_err(|err| Error::Stop(err.to_string()))
378    }
379
380    /// Awaits completion of the [`listen()`](Self::listen) loop after a stop signal.
381    pub async fn join(&self) -> Result<()> {
382        self.stop
383            .response
384            .receiver
385            .recv()
386            .await
387            .map_err(|err| Error::Join(err.to_string()))
388    }
389
390    /// Signals shutdown and waits for the listen loop to finish.
391    pub async fn stop_and_join(&self) -> Result<()> {
392        self.stop()?;
393        self.join().await
394    }
395}
396
397/// Base WebSocketServer trait allows the [`WebSocketServer<T>`] struct
398/// to be retained by the trait reference by casting it to the trait
399/// as follows:
400///
401/// ```rust
402/// use std::sync::Arc;
403/// use async_trait::async_trait;
404/// use workflow_websocket::server::{Result,WebSocketServerTrait,WebSocketConfig,TcpListener};
405///
406/// struct Server{}
407///
408/// #[async_trait]
409/// impl WebSocketServerTrait for Server {
410///     async fn bind(self: Arc<Self>, addr: &str) -> Result<TcpListener>{
411///         unimplemented!()
412///     }
413///     async fn listen(self: Arc<Self>, listener : TcpListener, config: Option<WebSocketConfig>) -> Result<()>{
414///         unimplemented!()
415///     }
416///     fn stop(&self) -> Result<()>{
417///         unimplemented!()
418///     }
419///     async fn join(&self) -> Result<()>{
420///         unimplemented!()
421///     }
422///     async fn stop_and_join(&self) -> Result<()>{
423///         unimplemented!()
424///     }
425/// }
426/// let server_trait: Arc<dyn WebSocketServerTrait> = Arc::new(Server{});
427/// let server = server_trait.downcast_arc::<Server>();
428/// ```
429/// This can help simplify web socket handling in case the supplied
430/// `T` generic contains complex generic types that typically
431/// results in generics propagating up into the ownership type chain.
432///
433/// This trait is used in the [`workflow-rpc`](https://docs.rs/workflow-rpc)
434/// crate to isolate `RpcHandler` generics from the RpcServer owning the WebSocket.
435///
436#[async_trait]
437pub trait WebSocketServerTrait: DowncastSync {
438    /// Binds a TCP listener to the given address.
439    async fn bind(self: Arc<Self>, addr: &str) -> Result<TcpListener>;
440    /// Runs the accept loop on the given listener until stopped.
441    async fn listen(
442        self: Arc<Self>,
443        listener: TcpListener,
444        config: Option<WebSocketConfig>,
445    ) -> Result<()>;
446    /// Signals the running listen loop to shut down.
447    fn stop(&self) -> Result<()>;
448    /// Awaits completion of the listen loop after a stop signal.
449    async fn join(&self) -> Result<()>;
450    /// Signals shutdown and waits for the listen loop to finish.
451    async fn stop_and_join(&self) -> Result<()>;
452}
453impl_downcast!(sync WebSocketServerTrait);
454
455#[async_trait]
456impl<T> WebSocketServerTrait for WebSocketServer<T>
457where
458    T: WebSocketHandler + Send + Sync + 'static + Sized,
459{
460    async fn bind(self: Arc<Self>, addr: &str) -> Result<TcpListener> {
461        WebSocketServer::<T>::bind(&self, addr).await
462    }
463
464    async fn listen(
465        self: Arc<Self>,
466        listener: TcpListener,
467        config: Option<WebSocketConfig>,
468    ) -> Result<()> {
469        WebSocketServer::<T>::listen(&self, listener, config).await
470    }
471
472    fn stop(&self) -> Result<()> {
473        WebSocketServer::<T>::stop(self)
474    }
475
476    async fn join(&self) -> Result<()> {
477        WebSocketServer::<T>::join(self).await
478    }
479
480    async fn stop_and_join(&self) -> Result<()> {
481        WebSocketServer::<T>::stop_and_join(self).await
482    }
483}
484
485pub mod handshake {
486    //!
487    //! Module containing simple convenience handshake functions
488    //! such as `greeting()`
489    //!     
490
491    use super::*;
492
493    /// Handshake closure function type for [`greeting()`] handshake
494    pub type HandshakeFn = Pin<Box<dyn Send + Sync + Fn(&str) -> Result<()>>>;
495
496    /// Simple greeting handshake where supplied closure receives
497    /// the first message from the client and should return
498    /// `Ok(())` to proceed or [`Error`] to abort the connection.
499    pub async fn greeting<'ws>(
500        timeout_duration: Duration,
501        _sender: &'ws mut WebSocketSender,
502        receiver: &'ws mut WebSocketReceiver,
503        handler: HandshakeFn,
504    ) -> Result<()> {
505        let delay = tokio::time::sleep(timeout_duration);
506        tokio::select! {
507            msg = receiver.next() => {
508                if let Some(Ok(msg)) = msg
509                    && (msg.is_text() || msg.is_binary()) {
510                        return handler(msg.to_text()?);
511                    }
512                Err(Error::MalformedHandshake)
513            }
514            _ = delay => {
515                Err(Error::ConnectionTimeout)
516            }
517        }
518    }
519}
workflow_websocket/server/mod.rs

workflow_websocket/server/
mod.rs
