1use std::error::Error;
2use std::sync::{LazyLock, OnceLock};
3use std::{any::type_name, borrow::Cow, mem, pin::Pin, task::Poll, time::Duration};
4
5use anyhow::anyhow;
6use bytes::{BufMut, Bytes, BytesMut};
7use futures::{AsyncRead, TryStreamExt as _};
8use http_client::{RedirectPolicy, Url, http};
9use regex::Regex;
10use reqwest::{
11 header::{HeaderMap, HeaderValue},
12 redirect,
13};
14use smol::future::FutureExt;
15
16const DEFAULT_CAPACITY: usize = 4096;
17static RUNTIME: OnceLock<tokio::runtime::Runtime> = OnceLock::new();
18static REDACT_REGEX: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"key=[^&]+").unwrap());
19
20pub struct ReqwestClient {
21 client: reqwest::Client,
22 proxy: Option<Url>,
23 handle: tokio::runtime::Handle,
24}
25
26impl ReqwestClient {
27 fn builder() -> reqwest::ClientBuilder {
28 reqwest::Client::builder()
29 .use_rustls_tls()
30 .connect_timeout(Duration::from_secs(10))
31 }
32
33 pub fn new() -> Self {
34 Self::builder()
35 .build()
36 .expect("Failed to initialize HTTP client")
37 .into()
38 }
39
40 pub fn user_agent(agent: &str) -> anyhow::Result<Self> {
41 let mut map = HeaderMap::new();
42 map.insert(http::header::USER_AGENT, HeaderValue::from_str(agent)?);
43 let client = Self::builder().default_headers(map).build()?;
44 Ok(client.into())
45 }
46
47 pub fn proxy_and_user_agent(proxy: Option<Url>, agent: &str) -> anyhow::Result<Self> {
48 let mut map = HeaderMap::new();
49 map.insert(http::header::USER_AGENT, HeaderValue::from_str(agent)?);
50 let mut client = Self::builder().default_headers(map);
51 let client_has_proxy;
52
53 if let Some(proxy) = proxy.as_ref().and_then(|proxy_url| {
54 reqwest::Proxy::all(proxy_url.clone())
55 .inspect_err(|e| {
56 log::error!(
57 "Failed to parse proxy URL '{}': {}",
58 proxy_url,
59 e.source().unwrap_or(&e as &_)
60 )
61 })
62 .ok()
63 }) {
64 // Respect NO_PROXY env var
65 client = client.proxy(proxy.no_proxy(reqwest::NoProxy::from_env()));
66 client_has_proxy = true;
67 } else {
68 client_has_proxy = false;
69 };
70
71 let client = client
72 .use_preconfigured_tls(http_client_tls::tls_config())
73 .build()?;
74 let mut client: ReqwestClient = client.into();
75 client.proxy = client_has_proxy.then_some(proxy).flatten();
76 Ok(client)
77 }
78}
79
80impl From<reqwest::Client> for ReqwestClient {
81 fn from(client: reqwest::Client) -> Self {
82 let handle = tokio::runtime::Handle::try_current().unwrap_or_else(|_| {
83 log::debug!("no tokio runtime found, creating one for Reqwest...");
84 let runtime = RUNTIME.get_or_init(|| {
85 tokio::runtime::Builder::new_multi_thread()
86 // Since we now have two executors, let's try to keep our footprint small
87 .worker_threads(1)
88 .enable_all()
89 .build()
90 .expect("Failed to initialize HTTP client")
91 });
92
93 runtime.handle().clone()
94 });
95 Self {
96 client,
97 handle,
98 proxy: None,
99 }
100 }
101}
102
103// This struct is essentially a re-implementation of
104// https://docs.rs/tokio-util/0.7.12/tokio_util/io/struct.ReaderStream.html
105// except outside of Tokio's aegis
106struct StreamReader {
107 reader: Option<Pin<Box<dyn futures::AsyncRead + Send + Sync>>>,
108 buf: BytesMut,
109 capacity: usize,
110}
111
112impl StreamReader {
113 fn new(reader: Pin<Box<dyn futures::AsyncRead + Send + Sync>>) -> Self {
114 Self {
115 reader: Some(reader),
116 buf: BytesMut::new(),
117 capacity: DEFAULT_CAPACITY,
118 }
119 }
120}
121
122impl futures::Stream for StreamReader {
123 type Item = std::io::Result<Bytes>;
124
125 fn poll_next(
126 mut self: Pin<&mut Self>,
127 cx: &mut std::task::Context<'_>,
128 ) -> Poll<Option<Self::Item>> {
129 let mut this = self.as_mut();
130
131 let mut reader = match this.reader.take() {
132 Some(r) => r,
133 None => return Poll::Ready(None),
134 };
135
136 if this.buf.capacity() == 0 {
137 let capacity = this.capacity;
138 this.buf.reserve(capacity);
139 }
140
141 match poll_read_buf(&mut reader, cx, &mut this.buf) {
142 Poll::Pending => Poll::Pending,
143 Poll::Ready(Err(err)) => {
144 self.reader = None;
145
146 Poll::Ready(Some(Err(err)))
147 }
148 Poll::Ready(Ok(0)) => {
149 self.reader = None;
150 Poll::Ready(None)
151 }
152 Poll::Ready(Ok(_)) => {
153 let chunk = this.buf.split();
154 self.reader = Some(reader);
155 Poll::Ready(Some(Ok(chunk.freeze())))
156 }
157 }
158 }
159}
160
161/// Implementation from <https://docs.rs/tokio-util/0.7.12/src/tokio_util/util/poll_buf.rs.html>
162/// Specialized for this use case
163pub fn poll_read_buf(
164 io: &mut Pin<Box<dyn futures::AsyncRead + Send + Sync>>,
165 cx: &mut std::task::Context<'_>,
166 buf: &mut BytesMut,
167) -> Poll<std::io::Result<usize>> {
168 if !buf.has_remaining_mut() {
169 return Poll::Ready(Ok(0));
170 }
171
172 let n = {
173 let dst = buf.chunk_mut();
174
175 // Safety: `chunk_mut()` returns a `&mut UninitSlice`, and `UninitSlice` is a
176 // transparent wrapper around `[MaybeUninit<u8>]`.
177 let dst = unsafe { &mut *(dst as *mut _ as *mut [std::mem::MaybeUninit<u8>]) };
178 let mut buf = tokio::io::ReadBuf::uninit(dst);
179 let ptr = buf.filled().as_ptr();
180 let unfilled_portion = buf.initialize_unfilled();
181 // SAFETY: Pin projection
182 let io_pin = unsafe { Pin::new_unchecked(io) };
183 std::task::ready!(io_pin.poll_read(cx, unfilled_portion)?);
184
185 // Ensure the pointer does not change from under us
186 assert_eq!(ptr, buf.filled().as_ptr());
187 buf.filled().len()
188 };
189
190 // Safety: This is guaranteed to be the number of initialized (and read)
191 // bytes due to the invariants provided by `ReadBuf::filled`.
192 unsafe {
193 buf.advance_mut(n);
194 }
195
196 Poll::Ready(Ok(n))
197}
198
199fn redact_error(mut error: reqwest::Error) -> reqwest::Error {
200 if let Some(url) = error.url_mut() {
201 if let Some(query) = url.query() {
202 if let Cow::Owned(redacted) = REDACT_REGEX.replace_all(query, "key=REDACTED") {
203 url.set_query(Some(redacted.as_str()));
204 }
205 }
206 }
207 error
208}
209
210impl http_client::HttpClient for ReqwestClient {
211 fn proxy(&self) -> Option<&Url> {
212 self.proxy.as_ref()
213 }
214
215 fn type_name(&self) -> &'static str {
216 type_name::<Self>()
217 }
218
219 fn send(
220 &self,
221 req: http::Request<http_client::AsyncBody>,
222 ) -> futures::future::BoxFuture<
223 'static,
224 anyhow::Result<http_client::Response<http_client::AsyncBody>>,
225 > {
226 let (parts, body) = req.into_parts();
227
228 let mut request = self.client.request(parts.method, parts.uri.to_string());
229 request = request.headers(parts.headers);
230 if let Some(redirect_policy) = parts.extensions.get::<RedirectPolicy>() {
231 request = request.redirect_policy(match redirect_policy {
232 RedirectPolicy::NoFollow => redirect::Policy::none(),
233 RedirectPolicy::FollowLimit(limit) => redirect::Policy::limited(*limit as usize),
234 RedirectPolicy::FollowAll => redirect::Policy::limited(100),
235 });
236 }
237 let request = request.body(match body.0 {
238 http_client::Inner::Empty => reqwest::Body::default(),
239 http_client::Inner::Bytes(cursor) => cursor.into_inner().into(),
240 http_client::Inner::AsyncReader(stream) => {
241 reqwest::Body::wrap_stream(StreamReader::new(stream))
242 }
243 });
244
245 let handle = self.handle.clone();
246 async move {
247 let mut response = handle
248 .spawn(async { request.send().await })
249 .await?
250 .map_err(redact_error)?;
251
252 let headers = mem::take(response.headers_mut());
253 let mut builder = http::Response::builder()
254 .status(response.status().as_u16())
255 .version(response.version());
256 *builder.headers_mut().unwrap() = headers;
257
258 let bytes = response
259 .bytes_stream()
260 .map_err(|e| futures::io::Error::new(futures::io::ErrorKind::Other, e))
261 .into_async_read();
262 let body = http_client::AsyncBody::from_reader(bytes);
263
264 builder.body(body).map_err(|e| anyhow!(e))
265 }
266 .boxed()
267 }
268}
269
270#[cfg(test)]
271mod tests {
272 use http_client::{HttpClient, Url};
273
274 use crate::ReqwestClient;
275
276 #[test]
277 fn test_proxy_uri() {
278 let client = ReqwestClient::new();
279 assert_eq!(client.proxy(), None);
280
281 let proxy = Url::parse("http://localhost:10809").unwrap();
282 let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
283 assert_eq!(client.proxy(), Some(&proxy));
284
285 let proxy = Url::parse("https://localhost:10809").unwrap();
286 let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
287 assert_eq!(client.proxy(), Some(&proxy));
288
289 let proxy = Url::parse("socks4://localhost:10808").unwrap();
290 let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
291 assert_eq!(client.proxy(), Some(&proxy));
292
293 let proxy = Url::parse("socks4a://localhost:10808").unwrap();
294 let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
295 assert_eq!(client.proxy(), Some(&proxy));
296
297 let proxy = Url::parse("socks5://localhost:10808").unwrap();
298 let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
299 assert_eq!(client.proxy(), Some(&proxy));
300
301 let proxy = Url::parse("socks5h://localhost:10808").unwrap();
302 let client = ReqwestClient::proxy_and_user_agent(Some(proxy.clone()), "test").unwrap();
303 assert_eq!(client.proxy(), Some(&proxy));
304 }
305
306 #[test]
307 fn test_invalid_proxy_uri() {
308 let proxy = Url::parse("socks://127.0.0.1:20170").unwrap();
309 let client = ReqwestClient::proxy_and_user_agent(Some(proxy), "test").unwrap();
310 assert!(
311 client.proxy.is_none(),
312 "An invalid proxy URL should add no proxy to the client!"
313 )
314 }
315}