beam/adapters/
ws_conn.rs

1//! Per-connection WebSocket actor — bridges Gun protocol messages over WebSocket.
2//!
3//! [`WsConn`] wraps a single WebSocket connection (either inbound via
4//! [`crate::adapters::WsServer`] or outbound via
5//! [`crate::adapters::OutgoingWebsocketManager`]). It:
6//!
7//! - Sends outgoing messages as `WsMessage::Text` on the WebSocket
8//! - Receives incoming messages, parses them via [`Message::try_from`], and
9//!   forwards to the [`crate::router::Router`]
10//! - Sends a `Message::Hi` on startup to register with the router
11//!
12//! # Lifecycle
13//!
14//! 1. `pre_start`: Send `Hi` message, spawn receive loop
15//! 2. `handle`: Forward outgoing messages to WebSocket
16//! 3. `stopping`: Log shutdown
17//!
18//! The receive loop runs as a child task and calls `ctx.stop()` when the
19//! WebSocket closes, triggering actor shutdown.
20
21use crate::actor::{Actor, ActorContext};
22use crate::message::Message;
23use futures_util::SinkExt;
24use futures_util::stream::{SplitSink, SplitStream};
25
26use async_trait::async_trait;
27
28use futures_util::{TryStreamExt, future};
29use log::{debug, error, info};
30use web_time::Duration;
31
32use tokio_tungstenite::tungstenite::Message as WsMessage;
33use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
34
35/// Type alias for the WebSocket stream over a TLS or plain TCP connection.
36type WsStream = WebSocketStream<MaybeTlsStream<tokio::net::TcpStream>>;
37/// Type alias for the sending half of a split WebSocket stream.
38type WsSender = SplitSink<WsStream, WsMessage>;
39/// Type alias for the receiving half of a split WebSocket stream.
40type WsReceiver = SplitStream<WsStream>;
41
42/// A per-connection WebSocket actor that bridges Gun protocol messages.
43///
44/// Created by [`crate::adapters::WsServer`] (inbound) or
45/// [`crate::adapters::OutgoingWebsocketManager`] (outbound). Each `WsConn`
46/// manages a single WebSocket connection and translates between the
47/// Gun wire format (text) and [`Message`] enum values.
48pub struct WsConn {
49    sender: WsSender,
50    receiver: Option<WsReceiver>,
51    allow_public_space: bool,
52}
53
54impl WsConn {
55    /// Creates a new `WsConn` from the split halves of a WebSocket stream.
56    ///
57    /// # Arguments
58    ///
59    /// * `sender` - The writing half of the WebSocket stream
60    /// * `receiver` - The reading half of the WebSocket stream
61    /// * `allow_public_space` - Whether to accept public space writes (forwarded
62    ///   to `Message::try_from` for inbound message parsing)
63    pub fn new(sender: WsSender, receiver: WsReceiver, allow_public_space: bool) -> Self {
64        Self {
65            sender,
66            receiver: Some(receiver),
67            allow_public_space,
68        }
69    }
70}
71
72#[async_trait]
73impl Actor for WsConn {
74    async fn handle(&mut self, msg: Message, _ctx: &ActorContext) {
75        let _ = self
76            .sender
77            .send(WsMessage::Text(msg.to_string().into()))
78            .await;
79    }
80
81    async fn pre_start(&mut self, ctx: &ActorContext) {
82        info!("WsConn starting");
83        let hi = Message::Hi {
84            from: ctx.addr.clone(),
85            peer_id: ctx.peer_id.read().clone(),
86        };
87        let _ = self
88            .sender
89            .send(WsMessage::Text(hi.to_string().into()))
90            .await;
91        let receiver = self.receiver.take().unwrap();
92        let mut ctx2 = ctx.clone();
93        let allow_public_space = self.allow_public_space;
94        ctx.child_task(async move {
95            let _ = receiver
96                .try_for_each(|msg| {
97                    if let Ok(s) = msg.to_text() {
98                        debug!("WsConn received: {}", s);
99                        match Message::try_from(s, ctx2.addr.clone(), allow_public_space) {
100                            Ok(msgs) => {
101                                for msg in msgs.into_iter() {
102                                    if ctx2.router.send(msg).is_err() {
103                                        error!("failed to forward incoming message to router");
104                                    }
105                                }
106                            }
107                            Err(e) => {
108                                debug!("WsConn parse error: {} (len={})", e, s.len());
109                            }
110                        }
111                    }
112                    future::ok(())
113                })
114                .await;
115            ctx2.stop();
116        });
117    }
118
119    async fn stopping(&mut self, _context: &ActorContext) {
120        info!("WsConn stopping — sending WebSocket Close frame");
121        // Send a WebSocket Close frame to initiate a clean close handshake.
122        // The remote peer should respond with its own Close frame, after
123        // which the connection is fully closed. We use a timeout so a
124        // non-responsive peer doesn't block shutdown.
125        let close_result = crate::tokio_time::timeout(Duration::from_secs(2), self.sender.close()).await;
126
127        match close_result {
128            Ok(Ok(())) => debug!("WsConn Close frame acknowledged"),
129            Ok(Err(e)) => debug!("WsConn Close error (non-fatal): {}", e),
130            Err(_) => debug!("WsConn Close timed out — connection dropped"),
131        }
132    }
133}
beam/adapters/ws_conn.rs

beam/adapters/
ws_conn.rs
