1use std::sync::Arc;
2
3use futures_util::{SinkExt, StreamExt};
4use tokio_tungstenite::WebSocketStream;
5use tokio_tungstenite::tungstenite::Message;
6use tokio_tungstenite::tungstenite::protocol::{CloseFrame as WireClose, Role};
7
8use crate::core::headers::reorder;
9use crate::h2::client::H2ConnectStream;
10use crate::h2::config::H2Config;
11use crate::h2::connection::PseudoHeaders;
12use crate::header_str::HeaderStr;
13use crate::pool::Pool;
14use crate::profile::preset::HeaderPair;
15use crate::tls::{FingerprintConnector, TlsIo};
16
17use crate::core::WebSocketConfig;
18use crate::core::error::{Error, Kind, Result};
19
20mod debug;
21mod handshake;
22mod split;
23
24use handshake::{
25 H2_NO_CONNECT_PROTOCOL, check_subprotocol, check_upgrade_response, handshake_target,
26 is_reserved_ws_header, overlay_headers, random_sec_ws_key, response_header, tungstenite_config,
27 wire_error, ws_header_pair,
28};
29use http::header::SEC_WEBSOCKET_PROTOCOL;
30pub use split::{WsSink, WsStream};
31use split::{WsSinkInner, WsStreamInner};
32
33#[derive(Debug, Clone, PartialEq, Eq)]
34#[non_exhaustive]
35pub enum WsMessage {
36 Text(String),
37 Binary(Vec<u8>),
38 Ping,
39 Pong,
40 Close(Option<CloseFrame>),
41}
42
43#[derive(Debug, Clone, PartialEq, Eq)]
44#[non_exhaustive]
45pub struct CloseFrame {
46 pub code: u16,
47 pub reason: String,
48}
49
50impl CloseFrame {
51 pub fn new(code: u16, reason: impl Into<String>) -> Self {
52 Self {
53 code,
54 reason: reason.into(),
55 }
56 }
57}
58
59impl WsMessage {
60 fn wire(msg: Message) -> Self {
61 match msg {
62 Message::Text(s) => Self::Text(s.as_str().to_owned()),
63 Message::Binary(b) => Self::Binary(b.to_vec()),
64 Message::Ping(_) => Self::Ping,
65 Message::Pong(_) => Self::Pong,
66 Message::Close(frame) => Self::Close(frame.map(|f| CloseFrame {
67 code: f.code.into(),
68 reason: f.reason.as_str().to_owned(),
69 })),
70 Message::Frame(f) => Self::Binary(f.into_payload().to_vec()),
71 }
72 }
73
74 fn into_wire(self) -> Message {
75 match self {
76 Self::Text(s) => Message::Text(s.into()),
77 Self::Binary(b) => Message::Binary(b.into()),
78 Self::Ping => Message::Ping(Vec::new().into()),
79 Self::Pong => Message::Pong(Vec::new().into()),
80 Self::Close(frame) => Message::Close(frame.map(|f| WireClose {
81 code: f.code.into(),
82 reason: f.reason.into(),
83 })),
84 }
85 }
86}
87
88enum WsInner {
89 H1(WebSocketStream<TlsIo>),
90 H2(WebSocketStream<H2ConnectStream>),
91}
92
93pub struct WsConnection {
94 inner: WsInner,
95 protocol: Option<String>,
96 headers: Vec<(String, String)>,
97}
98
99impl WsConnection {
100 #[expect(
101 clippy::too_many_arguments,
102 reason = "flat per-request wire fields across one internal call path"
103 )]
104 pub(crate) async fn connect_h1(
105 connector: &FingerprintConnector,
106 url: &str,
107 proxy: Option<&str>,
108 user_agent: &str,
109 origin: &str,
110 extra_headers: &[(String, String)],
111 header_order: Option<&[String]>,
112 ws_config: &WebSocketConfig,
113 ) -> Result<Self> {
114 let (parsed, host, port) = handshake_target(url)?;
115
116 let mut stream = connector
117 .connect_h1(&host, port, proxy)
118 .await
119 .map_err(Error::from)?
120 .stream;
121
122 let sec_key = random_sec_ws_key();
123 let mut request: Vec<HeaderPair> = vec![
124 ("Connection".into(), "Upgrade".into()),
125 ("Upgrade".into(), "websocket".into()),
126 ("User-Agent".into(), user_agent.to_owned().into()),
127 ("Origin".into(), origin.to_owned().into()),
128 ("Sec-WebSocket-Version".into(), "13".into()),
129 ("Sec-WebSocket-Key".into(), sec_key.clone().into()),
130 ];
131 overlay_headers(&mut request, extra_headers);
132 if let Some(order) = header_order {
133 reorder(&mut request, order);
134 }
135 let request = request
136 .into_iter()
137 .map(|(n, v)| (n.into_owned(), v.into_owned()))
138 .collect();
139
140 let ((status, headers, _), leftover) =
141 crate::pool::upgrade_on_stream(&mut stream, &parsed, request)
142 .await
143 .map_err(crate::core::transport::h1_error_to_core)?;
144 check_upgrade_response(status, &headers, &sec_key, extra_headers)?;
145
146 let protocol =
147 response_header(&headers, SEC_WEBSOCKET_PROTOCOL.as_str()).map(str::to_owned);
148 let ws_stream = WebSocketStream::from_partially_read(
149 stream,
150 leftover,
151 Role::Client,
152 Some(tungstenite_config(ws_config)),
153 )
154 .await;
155
156 Ok(Self {
157 inner: WsInner::H1(ws_stream),
158 protocol,
159 headers,
160 })
161 }
162
163 #[expect(
164 clippy::too_many_arguments,
165 reason = "flat per-request wire fields across one internal call path"
166 )]
167 pub(crate) async fn connect_h2(
168 pool: &Arc<Pool>,
169 connector: &FingerprintConnector,
170 h2_config: &H2Config,
171 url: &str,
172 proxy: Option<&str>,
173 user_agent: &str,
174 origin: &str,
175 extra_headers: &[(String, String)],
176 ws_config: &WebSocketConfig,
177 ) -> Result<Self> {
178 let (parsed, host, port) = handshake_target(url)?;
179 let path = crate::util::request_target(&parsed);
180 let authority = crate::util::authority(&parsed, &host, port);
181
182 let (h2_client, _tls) =
183 crate::pool::checkout_handle(pool, connector, h2_config, &host, port, proxy).await?;
184
185 if !h2_client.peer_enables_connect_protocol() {
186 return Err(Error::new(Kind::Request).with_message(H2_NO_CONNECT_PROTOCOL));
187 }
188
189 for (name, value) in extra_headers
190 .iter()
191 .filter(|(name, _)| !is_reserved_ws_header(name))
192 {
193 drop(ws_header_pair(name, value)?);
194 }
195 let sec_key = random_sec_ws_key();
196 let mut headers: Vec<HeaderPair> = vec![
197 ("sec-websocket-version".into(), "13".into()),
198 ("sec-websocket-key".into(), sec_key.into()),
199 ("user-agent".into(), user_agent.to_owned().into()),
200 ("origin".into(), origin.to_owned().into()),
201 ];
202 overlay_headers(&mut headers, extra_headers);
203
204 let pseudo = PseudoHeaders {
205 method: HeaderStr::from_static("CONNECT"),
206 scheme: HeaderStr::from_static("https"),
207 authority: HeaderStr::from(authority),
208 path: HeaderStr::from(path),
209 protocol: Some(HeaderStr::from_static("websocket")),
210 };
211
212 let stream = h2_client
213 .open_extended_connect(
214 pseudo,
215 headers
216 .into_iter()
217 .map(|(name, value)| {
218 (std::borrow::Cow::Owned(name.to_ascii_lowercase()), value)
219 })
220 .collect(),
221 )
222 .await
223 .map_err(|e| Error::new(Kind::Request).with_message(format!("h2 ws: {e}")))?;
224
225 if stream.status() != 200 {
226 return Err(Error::new(Kind::Request).with_message(format!(
227 "ws handshake: h2 extended CONNECT returned :status {}",
228 stream.status()
229 )));
230 }
231
232 check_subprotocol(stream.response_headers(), extra_headers)?;
233
234 let protocol = response_header(stream.response_headers(), SEC_WEBSOCKET_PROTOCOL.as_str())
235 .map(str::to_owned);
236 let headers = stream.response_headers().to_vec();
237
238 let ws_stream = WebSocketStream::from_raw_socket(
239 stream,
240 Role::Client,
241 Some(tungstenite_config(ws_config)),
242 )
243 .await;
244 Ok(Self {
245 inner: WsInner::H2(ws_stream),
246 protocol,
247 headers,
248 })
249 }
250
251 pub(crate) fn is_h2_fallback_trigger(err: &Error) -> bool {
252 crate::core::transport::is_h2_alpn_mismatch(err)
253 || err.kind() == Kind::Request
254 && err
255 .message()
256 .is_some_and(|s| s.contains(H2_NO_CONNECT_PROTOCOL))
257 }
258
259 pub async fn send(&mut self, msg: WsMessage) -> Result<()> {
260 let msg = msg.into_wire();
261 match &mut self.inner {
262 WsInner::H1(s) => s.send(msg).await,
263 WsInner::H2(s) => s.send(msg).await,
264 }
265 .map_err(|e| wire_error("ws send", e))
266 }
267
268 pub async fn recv(&mut self) -> Result<Option<WsMessage>> {
269 let next = match &mut self.inner {
270 WsInner::H1(s) => s.next().await,
271 WsInner::H2(s) => s.next().await,
272 };
273 match next {
274 Some(Ok(msg)) => Ok(Some(WsMessage::wire(msg))),
275 Some(Err(e)) => Err(wire_error("ws recv", e)),
276 None => Ok(None),
277 }
278 }
279
280 pub async fn close(&mut self) -> Result<()> {
281 match &mut self.inner {
282 WsInner::H1(s) => s.close(None).await.map_err(|e| wire_error("ws close", e)),
283 WsInner::H2(s) => s.close(None).await.map_err(|e| wire_error("ws close", e)),
284 }
285 }
286
287 pub fn protocol(&self) -> Option<&str> {
288 self.protocol.as_deref()
289 }
290
291 pub fn header(&self, name: &str) -> Option<&str> {
292 self.headers
293 .iter()
294 .find(|(n, _)| n.eq_ignore_ascii_case(name))
295 .map(|(_, v)| v.as_str())
296 }
297
298 pub fn split(self) -> (WsSink, WsStream) {
299 match self.inner {
300 WsInner::H1(s) => {
301 let (tx, rx) = s.split();
302 (
303 WsSink {
304 inner: WsSinkInner::H1(tx),
305 },
306 WsStream {
307 inner: WsStreamInner::H1(rx),
308 },
309 )
310 }
311 WsInner::H2(s) => {
312 let (tx, rx) = s.split();
313 (
314 WsSink {
315 inner: WsSinkInner::H2(tx),
316 },
317 WsStream {
318 inner: WsStreamInner::H2(rx),
319 },
320 )
321 }
322 }
323 }
324}
325
326#[cfg(test)]
327mod tests;