arc_web/websocket/
connection.rs

1use actix::prelude::*;
2use actix_session::Session;
3use actix_web::{web, Error, HttpRequest, HttpResponse};
4use actix_web_actors::ws;
5use std::time::{Duration, Instant};
6use tracing::{info, warn};
7use uuid::Uuid;
8
9use crate::websocket::server::{Connect, Disconnect, Subscribe, Unsubscribe, WsMessage, WsServer};
10
11/// How often heartbeat pings are sent
12const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
13
14/// How long before lack of client response causes a timeout
15const CLIENT_TIMEOUT: Duration = Duration::from_secs(10);
16
17/// WebSocket connection actor
18pub struct WsConnection {
19    /// Unique connection ID
20    id: Uuid,
21    /// User ID if authenticated (optional)
22    user_id: Option<String>,
23    /// Client heartbeat tracking
24    heartbeat: Instant,
25    /// Address of the broadcast server
26    server_addr: Addr<WsServer>,
27}
28
29impl WsConnection {
30    pub fn new(user_id: Option<String>, server_addr: Addr<WsServer>) -> Self {
31        WsConnection {
32            id: Uuid::new_v4(),
33            user_id,
34            heartbeat: Instant::now(),
35            server_addr,
36        }
37    }
38
39    /// Start heartbeat process
40    fn start_heartbeat(&self, ctx: &mut ws::WebsocketContext<Self>) {
41        ctx.run_interval(HEARTBEAT_INTERVAL, |act, ctx| {
42            // Check client heartbeat
43            if Instant::now().duration_since(act.heartbeat) > CLIENT_TIMEOUT {
44                warn!(connection_id = %act.id, "WebSocket client heartbeat timeout, disconnecting");
45                ctx.stop();
46                return;
47            }
48            ctx.ping(b"");
49        });
50    }
51
52    /// Handle text messages from client
53    fn handle_text_message(&mut self, text: &str, ctx: &mut ws::WebsocketContext<Self>) {
54        // Parse JSON commands from client
55        if let Ok(cmd) = serde_json::from_str::<ClientCommand>(text) {
56            match cmd {
57                ClientCommand::Subscribe { room } => {
58                    self.server_addr.do_send(Subscribe { id: self.id, room });
59                }
60                ClientCommand::Unsubscribe { room } => {
61                    self.server_addr.do_send(Unsubscribe { id: self.id, room });
62                }
63                ClientCommand::Ping => {
64                    ctx.text(r#"{"type":"pong"}"#);
65                }
66            }
67        }
68    }
69}
70
71/// Client commands that can be sent via WebSocket
72#[derive(serde::Deserialize)]
73#[serde(tag = "type", rename_all = "snake_case")]
74enum ClientCommand {
75    Subscribe { room: String },
76    Unsubscribe { room: String },
77    Ping,
78}
79
80impl Actor for WsConnection {
81    type Context = ws::WebsocketContext<Self>;
82
83    fn started(&mut self, ctx: &mut Self::Context) {
84        // Start heartbeat
85        self.start_heartbeat(ctx);
86
87        // Register with server
88        let addr = ctx.address();
89        self.server_addr.do_send(Connect {
90            id: self.id,
91            user_id: self.user_id.clone(),
92            addr: addr.recipient(),
93        });
94
95        info!(connection_id = %self.id, user_id = ?self.user_id, "WebSocket connection actor started");
96    }
97
98    fn stopping(&mut self, _: &mut Self::Context) -> Running {
99        // Notify server of disconnect
100        self.server_addr.do_send(Disconnect { id: self.id });
101        Running::Stop
102    }
103}
104
105/// Handle messages from WebSocket server
106impl Handler<WsMessage> for WsConnection {
107    type Result = ();
108
109    fn handle(&mut self, msg: WsMessage, ctx: &mut Self::Context) {
110        ctx.text(msg.0);
111    }
112}
113
114/// Handle WebSocket messages
115impl StreamHandler<Result<ws::Message, ws::ProtocolError>> for WsConnection {
116    fn handle(&mut self, msg: Result<ws::Message, ws::ProtocolError>, ctx: &mut Self::Context) {
117        match msg {
118            Ok(ws::Message::Ping(msg)) => {
119                self.heartbeat = Instant::now();
120                ctx.pong(&msg);
121            }
122            Ok(ws::Message::Pong(_)) => {
123                self.heartbeat = Instant::now();
124            }
125            Ok(ws::Message::Text(text)) => {
126                self.heartbeat = Instant::now();
127                self.handle_text_message(&text, ctx);
128            }
129            Ok(ws::Message::Binary(bin)) => {
130                // Binary messages not used for Turbo Streams, but handle gracefully
131                ctx.binary(bin);
132            }
133            Ok(ws::Message::Close(reason)) => {
134                ctx.close(reason);
135                ctx.stop();
136            }
137            _ => ctx.stop(),
138        }
139    }
140}
141
142/// WebSocket handler endpoint
143pub async fn ws_handler(
144    req: HttpRequest,
145    stream: web::Payload,
146    session: Session,
147    server: web::Data<Addr<WsServer>>,
148) -> Result<HttpResponse, Error> {
149    // Pull the aggregate UUID out of the cached `SessionUser`, if any.
150    let user_id = session
151        .get::<crate::helpers::session::SessionUser>("user")
152        .ok()
153        .flatten()
154        .map(|u| u.id);
155
156    let ws_connection = WsConnection::new(user_id, server.get_ref().clone());
157
158    ws::start(ws_connection, &req, stream)
159}
arc_web/websocket/connection.rs

arc_web/websocket/
connection.rs
