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, and greets the client
15//! with a `session.connected` notification. From then on the socket carries
16//! JSON-RPC 2.0 requests from the client, and their responses and the job
17//! notifications back. Each request is queued into the session's background
18//! thread in the order it arrived, and its response is sent back with the
19//! request's `id` so the client can correlate them.
20//!
21//! Graceful disconnect (`session.disconnect`) tears the session down
22//! immediately and flushes telemetry. Hard disconnects (socket close without
23//! that request) leave the session in a 5-minute grace period so the client
24//! can reconnect with the same `session_id` and pick up where it left off.
25//!
26//! [`Session`]: crate::session::Session
27
28use axum::{
29    extract::ws::{Message, WebSocket, WebSocketUpgrade},
30    extract::{Query, State},
31    http::{header::ORIGIN, HeaderMap, StatusCode},
32    response::{IntoResponse, Response}
33};
34use futures::stream::StreamExt;
35use futures::sink::SinkExt;
36use serde::Deserialize;
37use tokio::sync::{broadcast, mpsc};
38
39use std::sync::Arc;
40
41use rpfm_ipc::api::{ApiError, Done, JOB_UPDATED_NOTIFICATION, JSONRPC_VERSION, Request, RpcNotification, RpcRequest, RpcResponse};
42use rpfm_ipc::api::session::{Disconnect, SESSION_CONNECTED_NOTIFICATION, SessionConnected};
43use rpfm_telemetry::{info, warn};
44
45use crate::allowed_origins;
46use crate::session::{DEFAULT_SESSION_TIMEOUT_SECS, Session, SessionId, SessionManager};
47
48//-------------------------------------------------------------------------------//
49//                              Enums & Structs
50//-------------------------------------------------------------------------------//
51
52
53/// Query parameters for WebSocket connection.
54#[derive(Debug, Deserialize)]
55pub struct WsQueryParams {
56
57    /// Optional session ID to connect to an existing session.
58    pub session_id: Option<SessionId>,
59}
60
61/// A message for the client.
62enum Outgoing {
63
64    /// A response to a request.
65    Response(RpcResponse),
66
67    /// A notification.
68    Notification(RpcNotification),
69
70    /// Closes the connection, once the messages queued before it are sent.
71    Close,
72}
73
74//-------------------------------------------------------------------------------//
75//                             Implementations
76//-------------------------------------------------------------------------------//
77
78/// WebSocket handler to upgrade the connection and handle messages.
79///
80/// Accepts an optional `session_id` query parameter to reconnect to an existing session.
81/// Example: `ws://localhost:45127/ws?session_id=123`
82///
83/// Rejects connections from browser pages of other origins: browsers don't apply CORS to WebSockets,
84/// so any page could otherwise drive the server.
85pub(crate) async fn ws_handler(
86    State(session_manager): State<Arc<SessionManager>>,
87    Query(params): Query<WsQueryParams>,
88    headers: HeaderMap,
89    ws: WebSocketUpgrade,
90) -> Response {
91    if let Some(origin) = headers.get(ORIGIN) {
92        if !allowed_origins().iter().any(|allowed| origin.as_bytes() == allowed.as_bytes()) {
93            warn!("Rejected a WebSocket connection from the origin {:?}.", origin);
94            return StatusCode::FORBIDDEN.into_response();
95        }
96    }
97
98    ws.max_message_size(usize::MAX)
99        .max_frame_size(usize::MAX)
100        .on_upgrade(move |socket| handle_socket(socket, session_manager, params.session_id))
101        .into_response()
102}
103
104/// Function to handle a WebSocket connection.
105///
106/// Each WebSocket connection gets its own session with an isolated background thread.
107/// If a session_id is provided and that session exists, the client reconnects to it.
108async fn handle_socket(socket: WebSocket, session_manager: Arc<SessionManager>, requested_session_id: Option<SessionId>) {
109
110    // Get or create a session for this client connection.
111    let (session, is_new) = session_manager.get_or_create_session(requested_session_id);
112    let session_id = session.id();
113
114    if is_new {
115        info!("New WebSocket client connected, created session ID: {}", session_id);
116    } else {
117        info!("WebSocket client reconnected to existing session ID: {}", session_id);
118    }
119
120    let (mut sink, mut receiver) = socket.split();
121    let (tx, mut rx) = mpsc::unbounded_channel::<Outgoing>();
122
123    // Send the session ID to the client immediately after connection.
124    if let Ok(params) = serde_json::to_value(SessionConnected { session_id }) {
125        if let Ok(json) = serde_json::to_string(&RpcNotification::new(SESSION_CONNECTED_NOTIFICATION, params)) {
126            let _ = sink.send(Message::Text(json.into())).await;
127        }
128    }
129
130    // Task to send responses back to the client.
131    let sender_task = tokio::spawn(async move {
132        while let Some(outgoing) = rx.recv().await {
133            let json = match outgoing {
134                Outgoing::Response(response) => serde_json::to_string(&response)
135                    .unwrap_or_else(|error| serde_json::to_string(&RpcResponse::new(response.id, Err(ApiError::Internal(format!("Serialization error: {error}"))))).unwrap_or_default()),
136                Outgoing::Notification(notification) => match serde_json::to_string(&notification) {
137                    Ok(json) => json,
138                    Err(_) => continue,
139                },
140                Outgoing::Close => {
141                    let _ = sink.close().await;
142                    break;
143                }
144            };
145
146            if sink.send(Message::Text(json.into())).await.is_err() {
147                break;
148            }
149        }
150    });
151
152    // Task to notify the client of every change of the session's jobs.
153    let mut job_updates = session.jobs().subscribe();
154    let jobs_tx = tx.clone();
155    let jobs_forward_task = tokio::spawn(async move {
156        loop {
157            match job_updates.recv().await {
158                Ok(status) => {
159                    if let Ok(params) = serde_json::to_value(&status) {
160                        let _ = jobs_tx.send(Outgoing::Notification(RpcNotification::new(JOB_UPDATED_NOTIFICATION, params)));
161                    }
162                }
163
164                // Clients can always ask for the current state of a job, so missed updates are skipped.
165                Err(broadcast::error::RecvError::Lagged(_)) => continue,
166                Err(broadcast::error::RecvError::Closed) => break,
167            }
168        }
169    });
170
171    // Track whether the client requested a graceful disconnect.
172    let mut graceful_disconnect = false;
173
174    // Loop to receive requests from the client.
175    while let Some(msg) = receiver.next().await {
176        if let Ok(msg) = msg {
177            match msg {
178                Message::Text(t) => {
179
180                    let request = match serde_json::from_str::<RpcRequest>(&t) {
181                        Ok(request) => request,
182                        Err(error) => {
183                            warn!("Session {}: Deserialization error: {}", session.id(), error);
184
185                            // Answer with the ID of the malformed request, if it can be found.
186                            let id = serde_json::from_str::<serde_json::Value>(&t).ok()
187                                .and_then(|value| value.get("id")?.as_u64())
188                                .unwrap_or_default();
189
190                            let _ = tx.send(Outgoing::Response(RpcResponse::invalid_request(id, format!("Invalid request: {error}"))));
191                            continue;
192                        }
193                    };
194
195                    // Disconnecting needs the session manager, so it's answered here, before cleaning up.
196                    if request.method == Disconnect::METHOD {
197                        let done = serde_json::to_value(Done {}).map_err(|error| ApiError::Internal(error.to_string()));
198                        let _ = tx.send(Outgoing::Response(RpcResponse::new(request.id, done)));
199                        let _ = tx.send(Outgoing::Close);
200                        graceful_disconnect = true;
201                        break;
202                    }
203
204                    handle_api_request(request, &session, &tx);
205                }
206                Message::Close(_) => {
207                    info!("Session {}: Client disconnected", session_id);
208                    break;
209                }
210                _ => {}
211            }
212        } else {
213            info!("Session {}: Client disconnected (error)", session_id);
214            break;
215        }
216    }
217
218    jobs_forward_task.abort();
219
220    // On a graceful disconnect, the sender ends after sending the disconnect's response.
221    if graceful_disconnect {
222        let _ = sender_task.await;
223    } else {
224        sender_task.abort();
225    }
226
227    // Client requested graceful disconnect - remove session immediately.
228    if graceful_disconnect {
229        info!("Session {}: Client requested graceful disconnect, removing session immediately", session_id);
230        session_manager.remove_session(session_id);
231        session_manager.exit_if_no_sessions();
232    }
233
234    // Unexpected disconnect - mark session for timeout cleanup.
235    else {
236        session_manager.client_disconnected(session_id);
237        info!("Session {} client disconnected, session will timeout in {} minutes if not reconnected", session_id, DEFAULT_SESSION_TIMEOUT_SECS / 60);
238    }
239}
240
241/// Sends a request to the session's background thread, and its response to the client when it's done.
242fn handle_api_request(request: RpcRequest, session: &Arc<Session>, tx: &mpsc::UnboundedSender<Outgoing>) {
243    let id = request.id;
244    if request.jsonrpc != JSONRPC_VERSION {
245        let _ = tx.send(Outgoing::Response(RpcResponse::invalid_request(id, format!("Unsupported JSON-RPC version: {}", request.jsonrpc))));
246        return;
247    }
248
249    // Enqueue the request before spawning, so requests run in the order they arrived.
250    let mut receiver = session.call(request);
251    let tx = tx.clone();
252    tokio::spawn(async move {
253        let response = receiver.recv().await
254            .unwrap_or_else(|| RpcResponse::new(id, Err(ApiError::Internal("Session response channel closed unexpectedly".to_owned()))));
255
256        let _ = tx.send(Outgoing::Response(response));
257    });
258}
rpfm_server/server_websocket.rs

rpfm_server/
server_websocket.rs
