Skip to main content

rig_http/
ws_client.rs

1//! Transport-independent websocket handshakes, connections, and frames.
2//! Backends preserve rejected upgrades as HTTP errors with status, headers, and body.
3//!
4//! ```
5//! use rig_http::ws_client::websocket_url;
6//!
7//! assert_eq!(websocket_url("https://example.com/v1", "responses")?,
8//!            "wss://example.com/v1/responses");
9//! # Ok::<(), rig_http::http_client::Error>(())
10//! ```
11
12use crate::http_client::{Error, NoBody, Request, Result};
13use crate::wasm_compat::{WasmBoxedFuture, WasmCompatSend, WasmCompatSync};
14use bytes::Bytes;
15use std::time::Duration;
16
17/// One websocket data or control frame, in either direction.
18#[derive(Clone, Debug, PartialEq, Eq)]
19pub enum Frame {
20    /// A UTF-8 text frame.
21    Text(String),
22    /// A binary frame.
23    Binary(Bytes),
24    /// A ping, with its application payload.
25    Ping(Bytes),
26    /// A pong, with its application payload.
27    Pong(Bytes),
28    /// A close frame, with the peer's status and reason when it sent one.
29    Close(Option<CloseFrame>),
30}
31
32/// The status code and reason carried by a websocket close frame.
33#[derive(Clone, Debug, PartialEq, Eq)]
34pub struct CloseFrame {
35    /// The RFC 6455 close code.
36    pub code: u16,
37    /// The peer's reason, empty when it sent none.
38    pub reason: String,
39}
40
41/// Options enforced by the backend during the handshake.
42#[derive(Clone, Debug, Default)]
43pub struct ConnectOptions {
44    /// Backend-enforced handshake timeout, separate from session event timeouts.
45    /// `None` imposes no handshake deadline.
46    pub timeout: Option<Duration>,
47}
48
49impl ConnectOptions {
50    /// Options with no connect timeout.
51    #[must_use]
52    pub fn new() -> Self {
53        Self::default()
54    }
55
56    /// Sets the handshake timeout.
57    #[must_use]
58    pub fn with_timeout(mut self, timeout: Option<Duration>) -> Self {
59        self.timeout = timeout;
60        self
61    }
62}
63
64/// Opens websocket connections from requests with WS(S) URIs and authentication
65/// headers. Backends supply websocket handshake headers such as
66/// `Sec-WebSocket-Key`; callers must not supply those headers.
67pub trait WebSocketClientExt: Clone + WasmCompatSend + WasmCompatSync + 'static {
68    /// Opens a connection, preserving rejected upgrades with their HTTP status,
69    /// headers, and response body.
70    fn connect(
71        &self,
72        request: Request<NoBody>,
73        options: ConnectOptions,
74    ) -> impl Future<Output = Result<BoxedWebSocketConnection>> + WasmCompatSend;
75}
76
77/// One open websocket connection, usable as a trait object.
78/// Calls are sequential: sessions must not poll send and receive concurrently.
79/// WASM-compatible bounds preserve the containing session's thread-safety contract.
80pub trait WebSocketConnection: WasmCompatSend + WasmCompatSync {
81    /// Write one frame.
82    fn send(&mut self, frame: Frame) -> WasmBoxedFuture<'_, Result<()>>;
83
84    /// Read the next frame; `Ok(None)` means the peer ended the stream.
85    fn recv(&mut self) -> WasmBoxedFuture<'_, Result<Option<Frame>>>;
86
87    /// Completes a close handshake. Callers must avoid repeated closes;
88    /// backends may return an error for an already closed socket.
89    fn close(&mut self, frame: Option<CloseFrame>) -> WasmBoxedFuture<'_, Result<()>>;
90}
91
92/// A type-erased [`WebSocketConnection`].
93pub type BoxedWebSocketConnection = Box<dyn WebSocketConnection>;
94
95impl WebSocketConnection for BoxedWebSocketConnection {
96    fn send(&mut self, frame: Frame) -> WasmBoxedFuture<'_, Result<()>> {
97        (**self).send(frame)
98    }
99
100    fn recv(&mut self) -> WasmBoxedFuture<'_, Result<Option<Frame>>> {
101        (**self).recv()
102    }
103
104    fn close(&mut self, frame: Option<CloseFrame>) -> WasmBoxedFuture<'_, Result<()>> {
105        (**self).close(frame)
106    }
107}
108
109/// A base URL that cannot be turned into a websocket URL.
110#[derive(Debug, thiserror::Error)]
111#[error("invalid websocket base URL: {0}")]
112pub struct InvalidWebSocketUrl(String);
113
114/// Appends `path` to a base URL, converting HTTP(S) to WS(S) and retaining
115/// existing WS(S) schemes, query, and fragment. Trims boundary slashes from the
116/// appended path. Returns an error for invalid URLs or unsupported schemes.
117pub fn websocket_url(base_url: &str, path: &str) -> Result<String> {
118    fn invalid(message: impl Into<String>) -> Error {
119        Error::instance(InvalidWebSocketUrl(message.into()))
120    }
121
122    let mut url =
123        url::Url::parse(base_url).map_err(|error| invalid(format!("{base_url}: {error}")))?;
124
125    let scheme = match url.scheme() {
126        "https" | "wss" => "wss",
127        "http" | "ws" => "ws",
128        other => {
129            return Err(invalid(format!(
130                "unsupported base URL scheme for websocket mode: {other}"
131            )));
132        }
133    };
134    url.set_scheme(scheme)
135        .map_err(|()| invalid(format!("failed to convert {base_url} to a websocket URL")))?;
136
137    let path = format!(
138        "{}/{}",
139        url.path().trim_end_matches('/'),
140        path.trim_matches('/')
141    );
142    url.set_path(&path);
143    Ok(url.to_string())
144}
145
146#[cfg(test)]
147mod tests;