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 std::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                        eprintln!("[WSCONN-DIAG] received {} bytes from wire", s.len());
99                        if s.len() < 5000 {
100                            eprintln!("[WSCONN-DIAG] content: {}", s);
101                        }
102                        debug!("WsConn received: {}", s);
103                        match Message::try_from(s, ctx2.addr.clone(), allow_public_space) {
104                            Ok(msgs) => {
105                                eprintln!(
106                                    "[WSCONN-DIAG] parsed {} messages, forwarding to router",
107                                    msgs.len()
108                                );
109                                for msg in msgs.into_iter() {
110                                    if ctx2.router.send(msg).is_err() {
111                                        error!("failed to forward incoming message to router");
112                                    }
113                                }
114                            }
115                            Err(e) => {
116                                eprintln!(
117                                    "[WSCONN-DIAG] try_from FAILED — err: '{}' (len={})",
118                                    e,
119                                    s.len()
120                                );
121                            }
122                        }
123                    }
124                    future::ok(())
125                })
126                .await;
127            ctx2.stop();
128        });
129    }
130
131    async fn stopping(&mut self, _context: &ActorContext) {
132        info!("WsConn stopping — sending WebSocket Close frame");
133        // Send a WebSocket Close frame to initiate a clean close handshake.
134        // The remote peer should respond with its own Close frame, after
135        // which the connection is fully closed. We use a timeout so a
136        // non-responsive peer doesn't block shutdown.
137        let close_result = tokio::time::timeout(Duration::from_secs(2), self.sender.close()).await;
138
139        match close_result {
140            Ok(Ok(())) => debug!("WsConn Close frame acknowledged"),
141            Ok(Err(e)) => debug!("WsConn Close error (non-fatal): {}", e),
142            Err(_) => debug!("WsConn Close timed out — connection dropped"),
143        }
144    }
145}
beam/adapters/ws_conn.rs

beam/adapters/
ws_conn.rs
