steamkit/networking/steam3/
web_socket_context.rs

1//! One WebSocket connection attempt: URI construction, message framing, the
2//! read loop and the send path.
3//!
4//! The socket itself is behind the [`WebSocketSocket`] and [`WebSocketFactory`]
5//! traits. [`TungsteniteWebSocketSocket`] is the real implementation; tests can
6//! plug in their own.
7
8use std::net::{IpAddr, SocketAddr};
9use std::sync::atomic::{AtomicBool, Ordering};
10use std::sync::Arc;
11use std::time::Duration;
12
13use async_trait::async_trait;
14use bytes::Bytes;
15use futures::{Sink, SinkExt, Stream, StreamExt};
16use tokio_tungstenite::tungstenite::client::IntoClientRequest;
17use tokio_tungstenite::tungstenite::error::Error as WebSocketError;
18use tokio_tungstenite::tungstenite::http::header::{HeaderName, HeaderValue};
19use tokio_tungstenite::tungstenite::http::Request;
20use tokio_tungstenite::tungstenite::Message;
21use tokio_util::sync::CancellationToken;
22
23use crate::error::SteamError;
24
25use super::disconnected_event_args::DisconnectedEventArgs;
26use super::i_connection::ConnectionEventSenders;
27use super::net_msg_event_args::NetMsgEventArgs;
28
29/// The path of the Steam websocket endpoint.
30pub const WEBSOCKET_PATH: &str = "/cmsocket/";
31
32/// The scheme of the Steam websocket endpoint.
33pub const WEBSOCKET_SCHEME: &str = "wss";
34
35/// The name of the header that carries the `User-Agent` of the upgrade request.
36const USER_AGENT: HeaderName = HeaderName::from_static("user-agent");
37
38/// The type of a websocket message.
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40pub enum WebSocketMessageType {
41    /// A UTF-8 text message.
42    Text,
43    /// A binary message.
44    Binary,
45    /// The close frame.
46    Close,
47}
48
49/// The outcome of a single websocket receive.
50#[derive(Clone, Debug, PartialEq, Eq)]
51pub struct WebSocketReceive {
52    /// The number of bytes carried by the message.
53    pub data: Bytes,
54    /// Whether this was the last chunk of the message.
55    pub end_of_message: bool,
56    /// The type of the message.
57    pub message_type: WebSocketMessageType,
58}
59
60/// A WebSocket socket.
61#[async_trait]
62pub trait WebSocketSocket: Send + Sync {
63    /// Sends a binary message.
64    async fn send_binary(&self, data: Bytes) -> Result<(), SteamError>;
65
66    /// Receives the next chunk of a message.
67    async fn receive(&self) -> Result<WebSocketReceive, SteamError>;
68
69    /// Closes the socket.
70    async fn close(&self) -> Result<(), SteamError>;
71
72    /// Whether the socket is still open.
73    fn is_open(&self) -> bool;
74}
75
76/// Creates websocket sockets.
77#[async_trait]
78pub trait WebSocketFactory: Send + Sync {
79    /// Connects to `uri`, giving up after `timeout`.
80    async fn connect(
81        &self,
82        uri: &str,
83        timeout: Duration,
84    ) -> Result<Arc<dyn WebSocketSocket>, SteamError>;
85}
86
87/// A single WebSocket connection attempt, with its own read loop.
88#[derive(Debug)]
89pub struct WebSocketContext {
90    end_point: SocketAddr,
91    uri: String,
92    events: ConnectionEventSenders,
93    cancel: CancellationToken,
94    disconnect_notified: AtomicBool,
95}
96
97impl WebSocketContext {
98    /// Creates a context for the given endpoint.
99    pub fn new(end_point: SocketAddr, events: ConnectionEventSenders) -> Self {
100        Self::new_with_uri(end_point, construct_uri(end_point), events)
101    }
102
103    /// Creates a context for the given endpoint and URI. The URI is passed in
104    /// because it needs the host name, which `end_point` doesn't have.
105    pub fn new_with_uri(
106        end_point: SocketAddr,
107        uri: String,
108        events: ConnectionEventSenders,
109    ) -> Self {
110        Self {
111            end_point,
112            uri,
113            events,
114            cancel: CancellationToken::new(),
115            disconnect_notified: AtomicBool::new(false),
116        }
117    }
118
119    /// The endpoint this context is connected to.
120    pub fn end_point(&self) -> SocketAddr {
121        self.end_point
122    }
123
124    /// The websocket URI this context connects to.
125    pub fn uri(&self) -> &str {
126        &self.uri
127    }
128
129    /// Whether the context has been cancelled.
130    pub fn is_cancelled(&self) -> bool {
131        self.cancel.is_cancelled()
132    }
133
134    /// Cancels the run loop.
135    pub fn cancel(&self) {
136        self.cancel.cancel();
137    }
138
139    /// Raises the `Disconnected` event exactly once.
140    pub fn notify_disconnected(&self, user_initiated: bool) {
141        if !self.disconnect_notified.swap(true, Ordering::SeqCst) {
142            self.events
143                .emit_disconnected(DisconnectedEventArgs::new(user_initiated));
144        }
145    }
146
147    /// Starts the read loop.
148    pub fn start(self: Arc<Self>, socket: Arc<dyn WebSocketSocket>) {
149        tokio::spawn(async move {
150            self.run_loop(socket).await;
151        });
152    }
153
154    async fn run_loop(self: Arc<Self>, socket: Arc<dyn WebSocketSocket>) {
155        self.events.emit_connected();
156
157        loop {
158            if self.cancel.is_cancelled() {
159                break;
160            }
161
162            let received = tokio::select! {
163                biased;
164                _ = self.cancel.cancelled() => break,
165                result = socket.receive() => result,
166            };
167
168            let Ok(received) = received else {
169                break;
170            };
171
172            match received.message_type {
173                WebSocketMessageType::Binary => {
174                    let mut message = received.data;
175
176                    while !received.end_of_message {
177                        let Ok(next) = socket.receive().await else {
178                            self.notify_disconnected(false);
179                            return;
180                        };
181
182                        if next.message_type != WebSocketMessageType::Binary {
183                            break;
184                        }
185
186                        message = extend(message, next.data);
187
188                        if next.end_of_message {
189                            break;
190                        }
191                    }
192
193                    if !message.is_empty() {
194                        self.events.emit_net_msg_received(NetMsgEventArgs {
195                            data: message.to_vec(),
196                            end_point: self.end_point,
197                        });
198                    }
199                }
200                WebSocketMessageType::Text => {
201                    tracing::debug!(
202                        "WebSocketContext: Received websocket text message: \"{}\"",
203                        String::from_utf8_lossy(&received.data)
204                    );
205                }
206                WebSocketMessageType::Close => {
207                    self.notify_disconnected(false);
208                    return;
209                }
210            }
211        }
212
213        let _ = socket.close().await;
214        self.notify_disconnected(false);
215    }
216}
217
218/// Appends the payload of a message chunk to the message that is being read.
219fn extend(message: Bytes, chunk: Bytes) -> Bytes {
220    let mut buffer = Vec::with_capacity(message.len() + chunk.len());
221    buffer.extend_from_slice(&message);
222    buffer.extend_from_slice(&chunk);
223    buffer.into()
224}
225
226/// Builds the WebSocket URI for an endpoint.
227pub fn construct_uri(end_point: SocketAddr) -> String {
228    construct_uri_for_host(&format_host(end_point.ip()), end_point.port())
229}
230
231/// Builds the WebSocket URI for a host name and port.
232pub fn construct_uri_for_host(host: &str, port: u16) -> String {
233    format!("{WEBSOCKET_SCHEME}://{host}:{port}{WEBSOCKET_PATH}")
234}
235
236fn format_host(ip: IpAddr) -> String {
237    match ip {
238        IpAddr::V4(v4) => v4.to_string(),
239        IpAddr::V6(v6) => format!("[{v6}]"),
240    }
241}
242
243type TungsteniteSink =
244    Box<dyn Sink<Message, Error = WebSocketError> + Unpin + Send>;
245
246type TungsteniteStream =
247    Box<dyn Stream<Item = Result<Message, WebSocketError>> + Unpin + Send>;
248
249/// A WebSocket socket backed by `tokio-tungstenite`.
250pub struct TungsteniteWebSocketSocket {
251    sink: tokio::sync::Mutex<Option<TungsteniteSink>>,
252    stream: tokio::sync::Mutex<Option<TungsteniteStream>>,
253    open: AtomicBool,
254}
255
256impl std::fmt::Debug for TungsteniteWebSocketSocket {
257    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
258        f.debug_struct("TungsteniteWebSocketSocket")
259            .field("open", &self.open.load(Ordering::SeqCst))
260            .finish()
261    }
262}
263
264impl TungsteniteWebSocketSocket {
265    fn new(socket: tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>) -> Self {
266        let (sink, stream) = socket.split();
267
268        Self {
269            sink: tokio::sync::Mutex::new(Some(Box::new(sink))),
270            stream: tokio::sync::Mutex::new(Some(Box::new(stream))),
271            open: AtomicBool::new(true),
272        }
273    }
274}
275
276#[async_trait]
277impl WebSocketSocket for TungsteniteWebSocketSocket {
278    async fn send_binary(&self, data: Bytes) -> Result<(), SteamError> {
279        let mut guard = self.sink.lock().await;
280
281        let Some(sink) = guard.as_mut() else {
282            return Err(SteamError::Connection("websocket is closed".to_string()));
283        };
284
285        sink.send(Message::Binary(data.to_vec().into()))
286            .await
287            .map_err(|err| SteamError::Connection(format!("websocket send failed: {err}")))
288    }
289
290    async fn receive(&self) -> Result<WebSocketReceive, SteamError> {
291        let mut guard = self.stream.lock().await;
292
293        let Some(stream) = guard.as_mut() else {
294            return Err(SteamError::Connection("websocket is closed".to_string()));
295        };
296
297        match stream.next().await {
298            Some(Ok(Message::Binary(payload))) => Ok(WebSocketReceive {
299                data: payload.into(),
300                end_of_message: true,
301                message_type: WebSocketMessageType::Binary,
302            }),
303            Some(Ok(Message::Text(text))) => Ok(WebSocketReceive {
304                data: text.into_bytes().into(),
305                end_of_message: true,
306                message_type: WebSocketMessageType::Text,
307            }),
308            Some(Ok(Message::Close(_))) => {
309                self.open.store(false, Ordering::SeqCst);
310
311                Ok(WebSocketReceive {
312                    data: Bytes::new(),
313                    end_of_message: true,
314                    message_type: WebSocketMessageType::Close,
315                })
316            }
317            Some(Ok(_)) => Ok(WebSocketReceive {
318                data: Bytes::new(),
319                end_of_message: true,
320                message_type: WebSocketMessageType::Binary,
321            }),
322            Some(Err(err)) => Err(SteamError::Connection(format!(
323                "websocket read failed: {err}"
324            ))),
325            None => {
326                self.open.store(false, Ordering::SeqCst);
327
328                Err(SteamError::Connection(
329                    "websocket stream ended".to_string(),
330                ))
331            }
332        }
333    }
334
335    async fn close(&self) -> Result<(), SteamError> {
336        self.open.store(false, Ordering::SeqCst);
337
338        let sink = self.sink.lock().await.take();
339        if let Some(mut sink) = sink {
340            let _ = sink.close().await;
341        }
342
343        Ok(())
344    }
345
346    fn is_open(&self) -> bool {
347        self.open.load(Ordering::SeqCst)
348    }
349}
350
351/// Creates `tokio-tungstenite` backed websocket sockets.
352#[derive(Clone, Debug, Default)]
353pub struct TungsteniteWebSocketFactory {
354    user_agent: Option<String>,
355}
356
357impl TungsteniteWebSocketFactory {
358    /// Creates a factory that sends no `User-Agent` header.
359    pub fn new() -> Self {
360        Self::default()
361    }
362
363    /// Creates a factory that sends the given `User-Agent` with the upgrade
364    /// request.
365    pub fn with_user_agent(user_agent: impl Into<String>) -> Self {
366        Self {
367            user_agent: Some(user_agent.into()),
368        }
369    }
370}
371
372#[async_trait]
373impl WebSocketFactory for TungsteniteWebSocketFactory {
374    async fn connect(
375        &self,
376        uri: &str,
377        timeout: Duration,
378    ) -> Result<Arc<dyn WebSocketSocket>, SteamError> {
379        let request = match &self.user_agent {
380            Some(user_agent) => {
381                let mut request: Request<()> = uri
382                    .into_client_request()
383                    .map_err(|error| SteamError::Connection(format!("invalid websocket URI: {error}")))?;
384
385                request.headers_mut().insert(
386                    USER_AGENT,
387                    HeaderValue::from_str(user_agent).map_err(|error| {
388                        SteamError::Connection(format!("invalid user agent: {error}"))
389                    })?,
390                );
391
392                request
393            }
394            None => uri
395                .into_client_request()
396                .map_err(|error| SteamError::Connection(format!("invalid websocket URI: {error}")))?,
397        };
398
399        let connect = tokio_tungstenite::connect_async(request);
400
401        let (socket, _response) = match tokio::time::timeout(timeout, connect).await {
402            Ok(Ok(pair)) => pair,
403            Ok(Err(err)) => {
404                return Err(SteamError::Connection(format!(
405                    "websocket handshake failed: {err}"
406                )))
407            }
408            Err(_) => return Err(SteamError::Timeout(timeout)),
409        };
410
411        Ok(Arc::new(TungsteniteWebSocketSocket::new(socket)))
412    }
413}
steamkit/networking/steam3/web_socket_context.rs

steamkit/networking/steam3/
web_socket_context.rs
