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 session = &self.session;
147 let handshake = async {
148 let _pass = session
149 .inner
150 .host_limits
151 .admit(&http_equivalent(&url))
152 .await;
153 session
154 .websocket_with_options(url.as_str(), self.config, self.proxy.as_ref(), &headers)
155 .await
156 };
157 session
158 .unless_shut_down(session.deadline(None).total(handshake))
159 .await
160 }
161}
162
163impl std::future::IntoFuture for WebSocketBuilder {
164 type Output = Result<crate::core::websocket::WsConnection>;
165 type IntoFuture = std::pin::Pin<Box<dyn std::future::Future<Output = Self::Output> + Send>>;
166
167 fn into_future(self) -> Self::IntoFuture {
168 Box::pin(async move { self.connect().await })
169 }
170}
171
172impl Session {
173 fn websocket_headers(&self, url: &url::Url, caller: &HeaderList) -> Vec<(String, String)> {
174 let caller_has = |name: &str| caller.get(name).is_some();
175 let cookie_url = http_equivalent(url);
176 let lookup = &cookie_url;
177 let mut headers = Vec::new();
178 self.merge_session_headers(&mut headers, &caller_has, false, lookup);
179 apply_extra_headers(&mut headers, caller, false, &sensitive_header);
180 let site = super::execute::SiteContext {
181 fetch_site: crate::FetchSite::SameOrigin,
182 initiator: Some(lookup),
183 chain: &[],
184 };
185 self.add_jar_cookie(&mut headers, lookup, &site, "GET");
186 headers
187 .into_iter()
188 .map(|(k, v)| (k.into_owned(), v.into_owned()))
189 .collect()
190 }
191}
192
193fn http_equivalent(url: &url::Url) -> url::Url {
194 let mut http = url.clone();
195 let scheme = match url.scheme() {
196 "ws" => "http",
197 _ => "https",
198 };
199 if http.set_scheme(scheme).is_err() {
200 return url.clone();
201 }
202 http
203}
204
205fn ws_origin(url: &str) -> Result<String> {
206 let parsed = url::Url::parse(url).map_err(crate::core::Error::from_url_parse)?;
207 let scheme = match parsed.scheme() {
208 "wss" => "https",
209 "ws" => {
210 return Err(Error::new(Kind::Request)
211 .with_message("plaintext ws:// is not supported; use wss://"));
212 }
213 other => {
214 return Err(Error::new(Kind::Request)
215 .with_message(format!("websocket URL must use wss://, not {other}")));
216 }
217 };
218 let host = parsed.host_str().unwrap_or("");
219 Ok(match parsed.port() {
220 Some(port) => format!("{scheme}://{host}:{port}"),
221 None => format!("{scheme}://{host}"),
222 })
223}
224
225#[cfg(test)]
226#[path = "websocket_tests.rs"]
227mod tests;