1use anyhow::Context;
2use qmux::tokio_tungstenite;
3use std::collections::HashSet;
4use std::sync::{Arc, LazyLock, Mutex};
5use std::{net, time};
6use url::Url;
7
8static WEBSOCKET_WON: LazyLock<Mutex<HashSet<(String, u16)>>> = LazyLock::new(|| Mutex::new(HashSet::new()));
10
11#[derive(Clone, Debug, clap::Args, serde::Serialize, serde::Deserialize)]
13#[serde(default, deny_unknown_fields)]
14#[non_exhaustive]
15pub struct ClientWebSocket {
16 #[arg(
18 id = "websocket-enabled",
19 long = "websocket-enabled",
20 env = "MOQ_CLIENT_WEBSOCKET_ENABLED",
21 default_value = "true"
22 )]
23 pub enabled: bool,
24
25 #[arg(
28 id = "websocket-delay",
29 long = "websocket-delay",
30 env = "MOQ_CLIENT_WEBSOCKET_DELAY",
31 default_value = "200ms",
32 value_parser = humantime::parse_duration,
33 )]
34 #[serde(with = "humantime_serde")]
35 #[serde(skip_serializing_if = "Option::is_none")]
36 pub delay: Option<time::Duration>,
37}
38
39impl Default for ClientWebSocket {
40 fn default() -> Self {
41 Self {
42 enabled: true,
43 delay: Some(time::Duration::from_millis(200)),
44 }
45 }
46}
47
48pub(crate) async fn race_handle(
49 config: &ClientWebSocket,
50 tls: &rustls::ClientConfig,
51 url: Url,
52) -> Option<anyhow::Result<qmux::Session>> {
53 if !config.enabled {
54 return None;
55 }
56
57 match url.scheme() {
60 "http" | "https" | "ws" | "wss" => {}
61 _ => return None,
62 }
63
64 let res = connect(config, tls, url).await;
65 if let Err(err) = &res {
66 tracing::warn!(%err, "WebSocket connection failed");
67 }
68 Some(res)
69}
70
71pub(crate) async fn connect(
72 config: &ClientWebSocket,
73 tls: &rustls::ClientConfig,
74 mut url: Url,
75) -> anyhow::Result<qmux::Session> {
76 anyhow::ensure!(config.enabled, "WebSocket support is disabled");
77
78 let host = url.host_str().context("missing hostname")?.to_string();
79 let port = url.port().unwrap_or_else(|| match url.scheme() {
80 "https" | "wss" | "moql" | "moqt" => 443,
81 "http" | "ws" => 80,
82 _ => 443,
83 });
84 let key = (host, port);
85
86 match config.delay {
90 Some(delay) if !WEBSOCKET_WON.lock().unwrap().contains(&key) => {
91 tokio::time::sleep(delay).await;
92 tracing::debug!(%url, delay_ms = %delay.as_millis(), "QUIC not yet connected, attempting WebSocket fallback");
93 }
94 _ => {}
95 }
96
97 let needs_tls = match url.scheme() {
100 "http" => {
101 url.set_scheme("ws").expect("failed to set scheme");
102 false
103 }
104 "https" => {
105 url.set_scheme("wss").expect("failed to set scheme");
106 true
107 }
108 "ws" => false,
109 "wss" => true,
110 _ => anyhow::bail!("unsupported URL scheme for WebSocket: {}", url.scheme()),
111 };
112
113 tracing::debug!(%url, "connecting via WebSocket");
114
115 let connector = if needs_tls {
117 tokio_tungstenite::Connector::Rustls(Arc::new(tls.clone()))
118 } else {
119 tokio_tungstenite::Connector::Plain
120 };
121
122 let session = qmux::Client::new()
123 .with_config(qmux::tungstenite::protocol::WebSocketConfig {
124 max_message_size: Some(64 << 20), max_frame_size: Some(16 << 20), accept_unmasked_frames: false,
127 ..Default::default()
128 })
129 .with_connector(connector)
130 .connect(url.as_str())
131 .await
132 .context("failed to connect WebSocket")?;
133
134 tracing::warn!(%url, "using WebSocket fallback");
135 WEBSOCKET_WON.lock().unwrap().insert(key);
136
137 Ok(session)
138}
139
140pub struct WebSocketListener {
145 listener: tokio::net::TcpListener,
146 server: qmux::Server,
147}
148
149impl WebSocketListener {
150 pub async fn bind(addr: net::SocketAddr) -> anyhow::Result<Self> {
151 let listener = tokio::net::TcpListener::bind(addr).await?;
152 let server = qmux::Server::new();
153 Ok(Self { listener, server })
154 }
155
156 pub fn local_addr(&self) -> anyhow::Result<net::SocketAddr> {
157 Ok(self.listener.local_addr()?)
158 }
159
160 pub async fn accept(&self) -> Option<anyhow::Result<qmux::Session>> {
161 match self.listener.accept().await {
162 Ok((stream, addr)) => {
163 tracing::debug!(%addr, "accepted WebSocket TCP connection");
164 let server = self.server.clone();
165 Some(
166 server
167 .accept(stream)
168 .await
169 .map_err(|e| anyhow::anyhow!("WebSocket accept failed: {e}")),
170 )
171 }
172 Err(e) => Some(Err(e.into())),
173 }
174 }
175}