rpfm_server/server_websocket.rs
1//---------------------------------------------------------------------------//
2// Copyright (c) 2017-2026 Ismael Gutiérrez González. All rights reserved.
3//
4// This file is part of the Rusted PackFile Manager (RPFM) project,
5// which can be found here: https://github.com/Frodo45127/rpfm.
6//
7// This file is licensed under the MIT license, which can be found here:
8// https://github.com/Frodo45127/rpfm/blob/master/LICENSE.
9//---------------------------------------------------------------------------//
10
11//! WebSocket upgrade handler and message multiplexer for the `/ws` endpoint.
12//!
13//! On upgrade, the handler either reuses an existing [`Session`] (when the
14//! client supplies `?session_id=N`) or creates a new one. From then on the
15//! socket carries a stream of JSON-encoded [`IpcMessage<Command>`] frames
16//! from the client and [`IpcMessage<Response>`] frames back. Each command
17//! is dispatched into the session's dedicated background thread, whose
18//! responses are forwarded back over the same socket with the originating
19//! request `id` preserved so the client can correlate them.
20//!
21//! Graceful disconnect (`Command::ClientDisconnecting`) tears the session
22//! down immediately and flushes telemetry. Hard disconnects (socket close
23//! without that command) leave the session in a 5-minute grace period so
24//! the client can reconnect with the same `session_id` and pick up where it
25//! left off.
26//!
27//! [`Session`]: crate::session::Session
28//! [`IpcMessage<Command>`]: rpfm_ipc::messages::Message
29//! [`IpcMessage<Response>`]: rpfm_ipc::messages::Message
30
31use axum::{
32 extract::ws::{Message, WebSocket, WebSocketUpgrade},
33 extract::{Query, State},
34 response::IntoResponse
35};
36use futures::stream::StreamExt;
37use futures::sink::SinkExt;
38use serde::Deserialize;
39use tokio::sync::{broadcast, mpsc};
40
41use std::sync::Arc;
42
43use rpfm_ipc::messages::{Command, Message as IpcMessage, Response};
44use rpfm_telemetry::{info, warn};
45
46use crate::session::{DEFAULT_SESSION_TIMEOUT_SECS, SessionId, SessionManager, recv_response};
47use crate::settings::SETTINGS_CHANGED;
48
49//-------------------------------------------------------------------------------//
50// Enums & Structs
51//-------------------------------------------------------------------------------//
52
53
54/// Query parameters for WebSocket connection.
55#[derive(Debug, Deserialize)]
56pub struct WsQueryParams {
57
58 /// Optional session ID to connect to an existing session.
59 pub session_id: Option<SessionId>,
60}
61
62//-------------------------------------------------------------------------------//
63// Implementations
64//-------------------------------------------------------------------------------//
65
66/// WebSocket handler to upgrade the connection and handle messages.
67///
68/// Accepts an optional `session_id` query parameter to reconnect to an existing session.
69/// Example: `ws://localhost:45127/ws?session_id=123`
70pub(crate) async fn ws_handler(
71 State(session_manager): State<Arc<SessionManager>>,
72 Query(params): Query<WsQueryParams>,
73 ws: WebSocketUpgrade,
74) -> impl IntoResponse {
75 ws.max_message_size(usize::MAX)
76 .max_frame_size(usize::MAX)
77 .on_upgrade(move |socket| handle_socket(socket, session_manager, params.session_id))
78}
79
80/// Function to handle a WebSocket connection.
81///
82/// Each WebSocket connection gets its own session with an isolated background thread.
83/// If a session_id is provided and that session exists, the client reconnects to it.
84async fn handle_socket(socket: WebSocket, session_manager: Arc<SessionManager>, requested_session_id: Option<SessionId>) {
85
86 // Get or create a session for this client connection.
87 let (session, is_new) = session_manager.get_or_create_session(requested_session_id);
88 let session_id = session.id();
89
90 if is_new {
91 info!("New WebSocket client connected, created session ID: {}", session_id);
92 } else {
93 info!("WebSocket client reconnected to existing session ID: {}", session_id);
94 }
95
96 let (mut sink, mut receiver) = socket.split();
97 let (tx, mut rx) = mpsc::unbounded_channel::<IpcMessage<Response>>();
98
99 // Send the session ID to the client immediately after connection.
100 let session_connected_msg = IpcMessage {
101 id: 0, // Special ID for connection message
102 data: Response::SessionConnected(session_id),
103 };
104 if let Ok(json) = serde_json::to_string(&session_connected_msg) {
105 let _ = sink.send(Message::Text(json.into())).await;
106 }
107
108 // Task to send responses back to the client.
109 let sender_task = tokio::spawn(async move {
110 while let Some(response_msg) = rx.recv().await {
111 match serde_json::to_string(&response_msg) {
112 Ok(json) => {
113 if sink.send(Message::Text(json.into())).await.is_err() {
114 break;
115 }
116 }
117 Err(error) => {
118 let error_msg = IpcMessage {
119 id: response_msg.id,
120 data: Response::Error(format!("Serialization error: {}", error)),
121 };
122
123 if let Ok(json) = serde_json::to_string(&error_msg) {
124 let _ = sink.send(Message::Text(json.into())).await;
125 }
126 }
127 }
128 }
129 });
130
131 // Task to forward settings changes made by any session to this client, so
132 // other open UI instances refresh without needing to reconnect.
133 let mut settings_changed_rx = SETTINGS_CHANGED.subscribe();
134 let settings_tx = tx.clone();
135 let settings_forward_task = tokio::spawn(async move {
136 loop {
137 match settings_changed_rx.recv().await {
138 Ok(snapshot) => {
139 let msg = IpcMessage { id: 0, data: Response::SettingsChanged(snapshot) };
140 let _ = settings_tx.send(msg);
141 }
142
143 // We missed some updates because we were too slow, but there's always a newer
144 // one coming right after, so just keep going instead of tearing down the task.
145 Err(broadcast::error::RecvError::Lagged(_)) => continue,
146 Err(broadcast::error::RecvError::Closed) => break,
147 }
148 }
149 });
150
151 // Track whether the client requested a graceful disconnect.
152 let mut graceful_disconnect = false;
153
154 // Loop to receive commands from the client.
155 while let Some(msg) = receiver.next().await {
156 if let Ok(msg) = msg {
157 match msg {
158 Message::Text(t) => {
159 // Try to parse the message to check for ClientDisconnecting.
160 match serde_json::from_str::<IpcMessage<Command>>(&t) {
161 Ok(msg) => {
162 info!("Session {}: Received command [ID {}]: {:?}", session.id(), msg.id, msg.data);
163
164 // Handle ClientDisconnecting specially - it needs access to session_manager.
165 if matches!(msg.data, Command::ClientDisconnecting) {
166 // Send success response before cleanup.
167 let response_msg = IpcMessage {
168 id: msg.id,
169 data: Response::Success,
170 };
171 let _ = tx.send(response_msg);
172 graceful_disconnect = true;
173 break;
174 }
175
176 // Route other commands through the session's background thread.
177 let tx = tx.clone();
178 let session = session.clone();
179 tokio::spawn(async move {
180 let mut receiver = session.send(msg.data);
181 let response = recv_response(&mut receiver).await;
182 let response_msg = IpcMessage {
183 id: msg.id,
184 data: response,
185 };
186 let _ = tx.send(response_msg);
187 });
188 }
189 Err(error) => {
190 warn!("Session {}: Deserialization error: {}", session.id(), error);
191
192 // Try to extract the message ID from the malformed message so we can
193 // send an error response back to the client.
194 if let Some(id) = serde_json::from_str::<serde_json::Value>(&t)
195 .ok()
196 .and_then(|v| v.get("id")?.as_u64()) {
197 let error_msg = IpcMessage {
198 id,
199 data: Response::Error(format!("Server failed to deserialize command: {}", error)),
200 };
201 let _ = tx.send(error_msg);
202 }
203
204 // TODO: Handle the error case when the message ID cannot be extracted.
205 }
206 }
207 }
208 Message::Close(_) => {
209 info!("Session {}: Client disconnected", session_id);
210 break;
211 }
212 _ => {}
213 }
214 } else {
215 info!("Session {}: Client disconnected (error)", session_id);
216 break;
217 }
218 }
219
220 sender_task.abort();
221 settings_forward_task.abort();
222
223 // Client requested graceful disconnect - remove session immediately.
224 if graceful_disconnect {
225 info!("Session {}: Client requested graceful disconnect, removing session immediately", session_id);
226 session_manager.remove_session(session_id);
227
228 // Check if this was the last session and shutdown the server if so.
229 if session_manager.session_count() == 0 {
230 info!("No more active sessions, shutting down server...");
231 rpfm_telemetry::flush("Server Action Telemetry");
232 std::process::exit(0);
233 }
234 }
235
236 // Unexpected disconnect - mark session for timeout cleanup.
237 else {
238 SessionManager::client_disconnected(session_manager.clone(),session_id);
239 info!("Session {} client disconnected, session will timeout in {} minutes if not reconnected", session_id, DEFAULT_SESSION_TIMEOUT_SECS / 60);
240 }
241}