1use std::fmt;
2
3use http::{HeaderName, HeaderValue};
4
5use super::Session;
6use super::header_merge::apply_extra_headers;
7use crate::core::error::{Error, Kind, Result};
8use crate::core::headers::HeaderList;
9use crate::core::{IntoParamPair, IntoUrl, ProxyConfig, WebSocketConfig};
10use crate::util::{redact, sensitive_header};
11
12impl Session {
13 pub fn websocket(&self, url: impl IntoUrl) -> WebSocketBuilder {
14 WebSocketBuilder {
15 session: self.clone(),
16 url: url.into_url(),
17 config: self.inner.websocket_config,
18 proxy: None,
19 headers: Ok(HeaderList::new()),
20 }
21 }
22
23 async fn websocket_with_options(
24 &self,
25 url: &str,
26 config: WebSocketConfig,
27 request_proxy: Option<&ProxyConfig>,
28 extra_headers: &HeaderList,
29 ) -> Result<crate::core::websocket::WsConnection> {
30 let origin = ws_origin(url)?;
31 let parsed = url::Url::parse(url).map_err(crate::core::Error::from_url_parse)?;
32 let proxy = self.proxy_for(&parsed, request_proxy)?;
33 let extra_headers = &self.websocket_headers(&parsed, extra_headers);
34
35 let h1_only = parsed.host_str().is_some_and(|host| {
36 self.inner
37 .pool
38 .is_h1_only(host, parsed.port_or_known_default().unwrap_or(443), proxy)
39 });
40
41 if config.prefer_http2 && !h1_only {
42 match crate::core::websocket::WsConnection::connect_h2(
43 &self.inner.pool,
44 &self.inner.connector,
45 &self.inner.h2_config,
46 url,
47 proxy,
48 &self.inner.user_agent,
49 &origin,
50 extra_headers,
51 &config,
52 )
53 .await
54 {
55 Ok(conn) => return Ok(conn),
56 Err(e) if crate::core::websocket::WsConnection::is_h2_fallback_trigger(&e) => {
57 if crate::core::transport::is_h2_alpn_mismatch(&e)
58 && let Some(host) = parsed.host_str()
59 {
60 let port = parsed.port_or_known_default().unwrap_or(443);
61 self.inner.pool.note_h1_only(host, port, proxy);
62 }
63 tracing::debug!(
64 error = %e,
65 "H2 extended CONNECT not available, falling back to H1 Upgrade"
66 );
67 }
68 Err(e) => return Err(e),
69 }
70 }
71
72 crate::core::websocket::WsConnection::connect_h1(
73 &self.inner.connector,
74 url,
75 proxy,
76 &self.inner.user_agent,
77 &origin,
78 extra_headers,
79 self.session_header_order().as_deref(),
80 &config,
81 )
82 .await
83 }
84}
85
86#[must_use = "builders are lazy: nothing happens until `.connect()` / await"]
87pub struct WebSocketBuilder {
88 session: Session,
89 url: Result<url::Url>,
90 config: WebSocketConfig,
91 proxy: Option<ProxyConfig>,
92 headers: Result<HeaderList>,
93}
94
95impl fmt::Debug for WebSocketBuilder {
96 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
97 let url = self.url.as_ref().map(|url| redact(url.as_str()));
98 f.debug_struct("WebSocketBuilder")
99 .field("url", &url)
100 .field("config", &self.config)
101 .field("proxy", &self.proxy)
102 .field("headers", &self.headers)
103 .finish_non_exhaustive()
104 }
105}
106
107impl WebSocketBuilder {
108 pub fn config(mut self, config: WebSocketConfig) -> Self {
109 self.config = config;
110 self
111 }
112
113 pub fn proxy(mut self, config: impl Into<ProxyConfig>) -> Self {
114 self.proxy = Some(config.into());
115 self
116 }
117
118 pub fn headers<I, P>(mut self, headers: I) -> Self
119 where
120 I: IntoIterator<Item = P>,
121 P: IntoParamPair,
122 {
123 for pair in headers {
124 let (name, value) = pair.into_param_pair();
125 self = self.header(name, value);
126 }
127 self
128 }
129
130 pub fn header(
131 mut self,
132 name: impl TryInto<HeaderName>,
133 value: impl TryInto<HeaderValue>,
134 ) -> Self {
135 if let Ok(headers) = &mut self.headers
136 && let Err(err) = headers.append(name, value)
137 {
138 self.headers = Err(err);
139 }
140 self
141 }
142
143 pub async fn connect(self) -> Result<crate::core::websocket::WsConnection> {
144 let url = self.url?;
145 let headers = self.headers?;
146 let handshake = self.session.websocket_with_options(
147 url.as_str(),
148 self.config,
149 self.proxy.as_ref(),
150 &headers,
151 );
152 self.session.deadline(None).total(handshake).await
153 }
154}
155
156impl std::future::IntoFuture for WebSocketBuilder {
157 type Output = Result<crate::core::websocket::WsConnection>;
158 type IntoFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output> + Send>>;
159
160 fn into_future(self) -> Self::IntoFuture {
161 Box::pin(async move { self.connect().await })
162 }
163}
164
165impl Session {
166 fn websocket_headers(&self, url: &url::Url, caller: &HeaderList) -> Vec<(String, String)> {
167 let caller_has = |name: &str| caller.get(name).is_some();
168 let mut cookie_url = url.clone();
169 let http_scheme = match url.scheme() {
170 "ws" => "http",
171 _ => "https",
172 };
173 let lookup = if cookie_url.set_scheme(http_scheme).is_ok() {
174 &cookie_url
175 } else {
176 url
177 };
178 let mut headers = Vec::new();
179 self.merge_session_headers(&mut headers, &caller_has, false, lookup);
180 apply_extra_headers(&mut headers, caller, false, &sensitive_header);
181 let site = super::execute::SiteContext {
182 fetch_site: crate::FetchSite::SameOrigin,
183 initiator: Some(lookup),
184 chain: &[],
185 };
186 self.add_jar_cookie(&mut headers, lookup, &site, "GET");
187 headers
188 .into_iter()
189 .map(|(k, v)| (k.into_owned(), v.into_owned()))
190 .collect()
191 }
192}
193
194fn ws_origin(url: &str) -> Result<String> {
195 let parsed = url::Url::parse(url).map_err(crate::core::Error::from_url_parse)?;
196 let scheme = match parsed.scheme() {
197 "wss" => "https",
198 "ws" => {
199 return Err(Error::new(Kind::Request)
200 .with_message("plaintext ws:// is not supported; use wss://"));
201 }
202 other => {
203 return Err(Error::new(Kind::Request)
204 .with_message(format!("websocket URL must use wss://, not {other}")));
205 }
206 };
207 let host = parsed.host_str().unwrap_or("");
208 Ok(match parsed.port() {
209 Some(port) => format!("{scheme}://{host}:{port}"),
210 None => format!("{scheme}://{host}"),
211 })
212}
213
214#[cfg(test)]
215#[path = "websocket_tests.rs"]
216mod tests;