diff --git a/impit-node/index.wrapper.js b/impit-node/index.wrapper.js index 91b74fd4..a50a0ca4 100644 --- a/impit-node/index.wrapper.js +++ b/impit-node/index.wrapper.js @@ -23,8 +23,8 @@ Run your script with IMPIT_VERBOSE=1 environment variable to get more informatio class ResponsePatches { static async text() { - const buffer = await this.bytes(); - return this.decodeBuffer(buffer); + const buffer = await this.arrayBuffer(); + return this.decodeBuffer(Buffer.from(buffer)); } } @@ -39,14 +39,18 @@ function canonicalizeHeaders(headers) { return []; } +const REDIRECT_STATUSES = new Set([301, 302, 303, 307, 308]); + function isRedirectStatus(status) { - return [301, 302, 303, 307, 308].includes(status); + return REDIRECT_STATUSES.has(status); } +const textEncoder = new TextEncoder(); + // Stream chunks may be strings (e.g. a ReadableStream or Node stream that wasn't // fed bytes); encode them so they aren't coerced to zero bytes by Uint8Array.set. function toUint8Array(chunk) { - return typeof chunk === 'string' ? new TextEncoder().encode(chunk) : new Uint8Array(chunk); + return typeof chunk === 'string' ? textEncoder.encode(chunk) : new Uint8Array(chunk); } // Streamed bodies are passed to the native layer as a ReadableStream of bytes, so that the chunks @@ -69,7 +73,7 @@ function toByteStream(source) { function shouldRewriteRedirectToGet(httpStatus, method) { // See https://github.com/mozilla-firefox/firefox/blob/911b3eec6c5e58a9a49e23aa105e49aa76e00f9c/netwerk/protocol/http/HttpBaseChannel.cpp#L4801 - if ([301, 302].includes(httpStatus)) { + if (httpStatus === 301 || httpStatus === 302) { return method === 'POST'; } @@ -138,7 +142,6 @@ class Impit extends native.Impit { const blobParts = []; const rn = new Uint8Array([13, 10]); - const textEncoder = new TextEncoder(); for (const [name, value] of formData) { if (typeof value === 'string') { @@ -189,9 +192,9 @@ class Impit extends native.Impit { // Based on https://github.com/nodejs/undici/blob/14e62db0d0cff4bea27357aa5bd14881459b27c7/lib/web/fetch/body.js#L90 async #serializeBody(body) { if (typeof body === 'string') { - return { body: new TextEncoder().encode(body), type: 'text/plain;charset=UTF-8' }; + return { body: textEncoder.encode(body), type: 'text/plain;charset=UTF-8' }; } else if (body instanceof URLSearchParams) { - return { body: new TextEncoder().encode(body.toString()), type: 'application/x-www-form-urlencoded;charset=UTF-8' }; + return { body: textEncoder.encode(body.toString()), type: 'application/x-www-form-urlencoded;charset=UTF-8' }; } else if (body instanceof ArrayBuffer) { return { body: new Uint8Array(body.slice()), type: '' }; } else if (ArrayBuffer.isView(body)) { @@ -261,9 +264,9 @@ class Impit extends native.Impit { signal?.throwIfAborted(); let abortHandler; - const waitForAbort = new Promise((_, reject) => { + const waitForAbort = signal && new Promise((_, reject) => { abortHandler = () => reject(signal.reason); - signal?.addEventListener?.("abort", abortHandler, { once: true }); + signal.addEventListener?.("abort", abortHandler, { once: true }); }); try { @@ -296,13 +299,12 @@ class Impit extends native.Impit { while (true) { signal?.throwIfAborted(); - const headers = [...(options.headers || [])]; - const hasUserCookie = headers.some(([k]) => k.toLowerCase() === 'cookie'); + let headers = options.headers || []; - if (this.#cookieJar && !hasUserCookie) { + if (this.#cookieJar && !headers.some(([k]) => k.toLowerCase() === 'cookie')) { const cookieHeader = await this.#getCookies(url); if (cookieHeader) { - headers.push(['Cookie', cookieHeader]); + headers = [...headers, ['Cookie', cookieHeader]]; } } @@ -314,10 +316,9 @@ class Impit extends native.Impit { bodyStream: method === 'GET' ? undefined : options.bodyStream, }); - const originalResponse = await Promise.race([ - response, - waitForAbort - ]); + const originalResponse = waitForAbort + ? await Promise.race([response, waitForAbort]) + : await response; const responseHeaders = new Headers(originalResponse.headers); @@ -334,7 +335,7 @@ class Impit extends native.Impit { const location = responseHeaders.get('location'); if (!location) { - return this.#wrapResponse(originalResponse, signal); + return this.#wrapResponse(originalResponse, signal, responseHeaders); } redirectCount++; @@ -353,7 +354,7 @@ class Impit extends native.Impit { } } - return this.#wrapResponse(originalResponse, signal); + return this.#wrapResponse(originalResponse, signal, responseHeaders); } } @@ -361,9 +362,10 @@ class Impit extends native.Impit { * Wrap a native response with JS enhancements * @param {object} originalResponse * @param {AbortSignal} signal + * @param {Headers} headers Response headers, already parsed by the caller * @returns {object} */ - #wrapResponse(originalResponse, signal) { + #wrapResponse(originalResponse, signal, headers) { signal?.throwIfAborted(); let abortHandler; @@ -411,7 +413,7 @@ class Impit extends native.Impit { }); Object.defineProperty(originalResponse, 'headers', { - value: new Headers(originalResponse.headers) + value: headers }); Object.defineProperty(originalResponse, 'clone', { diff --git a/impit-node/src/lib.rs b/impit-node/src/lib.rs index 5fa5be65..57fa4820 100644 --- a/impit-node/src/lib.rs +++ b/impit-node/src/lib.rs @@ -120,30 +120,27 @@ impl ImpitWrapper { url: String, request_init: Option, ) -> Result { + let RequestInit { + method, + headers, + body, + body_stream, + timeout, + force_http3, + .. + } = request_init.unwrap_or_default(); + let request_options = Some(RequestOptions { - headers: request_init - .as_ref() - .and_then(|init| init.headers.as_ref()) - .cloned() - .unwrap_or_default(), - timeout: request_init - .as_ref() - .and_then(|init| init.timeout) - .map(|timeout| Some(Duration::from_millis(timeout.into()))), - http3_prior_knowledge: request_init - .as_ref() - .and_then(|init| init.force_http3) - .unwrap_or_default(), + headers: headers.unwrap_or_default(), + timeout: timeout.map(|timeout| Some(Duration::from_millis(timeout.into()))), + http3_prior_knowledge: force_http3.unwrap_or_default(), }); - let method = request_init - .as_ref() - .and_then(|init| init.method.to_owned()) - .unwrap_or_default(); - let body = request_init.and_then(|init| match (init.body, init.body_stream) { + let method = method.unwrap_or_default(); + let body = match (body, body_stream) { (_, Some(stream)) => Some(ImpitBody::from_stream(stream.into_byte_stream())), (bytes, None) => bytes.map(|array| array.to_vec().into()), - }); + }; let response = if matches!(method, HttpMethod::Get | HttpMethod::Head) && body.is_some() { Err(ImpitError::BindingPassthroughError( diff --git a/impit-node/src/response.rs b/impit-node/src/response.rs index 9959afe6..3d30eb8c 100644 --- a/impit-node/src/response.rs +++ b/impit-node/src/response.rs @@ -92,7 +92,7 @@ impl<'env> ImpitResponse { // JS Fetch semantics: header values are decoded as ISO-8859-1 (each byte 0x00..=0xFF maps to // the code point U+0000..=U+00FF). Since that mapping is a bijection, the string stays // byte-recoverable via `Buffer.from(value, 'latin1')`. - let mut headers_vec: Vec<(String, String)> = Vec::new(); + let mut headers_vec: Vec<(String, String)> = Vec::with_capacity(response.headers().len()); for (k, v) in response.headers().iter() { headers_vec.push(( k.as_str().to_string(), diff --git a/impit-python/src/async_client.rs b/impit-python/src/async_client.rs index 646c0903..9a5a5004 100644 --- a/impit-python/src/async_client.rs +++ b/impit-python/src/async_client.rs @@ -376,13 +376,11 @@ impl AsyncClient { url: String, content: Option, mut data: Option, - headers: Option>, + mut headers: Option>, timeout: Option>, force_http3: Option, stream: Option, ) -> Result, PyErr> { - let mut headers = headers.clone(); - if let Some(content) = content { data = Some(content); } @@ -392,11 +390,7 @@ impl AsyncClient { let timeout = parse_timeout(timeout)?; let options = RequestOptions { - headers: headers - .unwrap_or_default() - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), + headers: headers.unwrap_or_default().into_iter().collect(), timeout, http3_prior_knowledge: force_http3.unwrap_or(false), }; diff --git a/impit-python/src/client.rs b/impit-python/src/client.rs index e691f93b..708589b9 100644 --- a/impit-python/src/client.rs +++ b/impit-python/src/client.rs @@ -373,13 +373,11 @@ impl Client { url: String, content: Option, mut data: Option, - headers: Option>, + mut headers: Option>, timeout: Option>, force_http3: Option, stream: Option, ) -> Result { - let mut headers = headers.clone(); - if let Some(content) = content { data = Some(content); } @@ -391,11 +389,7 @@ impl Client { .map_err(|e| ImpitPyError(ImpitError::BindingPassthroughError(e.to_string())))?; let options = RequestOptions { - headers: headers - .unwrap_or_default() - .iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect(), + headers: headers.unwrap_or_default().into_iter().collect(), timeout, http3_prior_knowledge: force_http3.unwrap_or(false), }; diff --git a/impit-python/src/response.rs b/impit-python/src/response.rs index 803d5930..c23c6163 100644 --- a/impit-python/src/response.rs +++ b/impit-python/src/response.rs @@ -7,7 +7,7 @@ use encoding::label::encoding_from_whatwg_label; use futures::{Stream, StreamExt}; use impit::{errors::ImpitError, utils::ContentType}; use pyo3::prelude::*; -use pyo3::types::PyBytes; +use pyo3::types::{PyBytes, PyString}; use reqwest::{Response, StatusCode, Version}; use std::pin::Pin; @@ -18,7 +18,7 @@ type SharedStream = #[pyclass] pub struct PyResponseBytesIterator { - ready_content: Option>, + ready_content: Option, stream: Option> + Send + Sync>>>, runtime: tokio::runtime::Handle, content_returned: bool, @@ -31,10 +31,11 @@ impl PyResponseBytesIterator { slf } - fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult>> { + fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult>> { if let Some(content) = slf.ready_content.take() { slf.content_returned = true; - return Ok(Some(content)); + let py = slf.py(); + return Ok(Some(PyBytes::new(py, &content).unbind())); } if slf.content_returned { @@ -58,7 +59,7 @@ impl PyResponseBytesIterator { let result = py.detach(|| runtime.block_on(stream.next())); match result { - Some(Ok(chunk)) => Ok(Some(chunk.to_vec())), + Some(Ok(chunk)) => Ok(Some(PyBytes::new(py, &chunk).unbind())), Some(Err(e)) => { slf.content_returned = true; if let Some(parent) = &slf.parent_response { @@ -91,7 +92,7 @@ impl PyResponseBytesIterator { #[pyclass] pub struct PyResponseAsyncBytesIterator { - ready_content: Option>, + ready_content: Option, stream: Option, content_returned: bool, parent_response: Option>, @@ -106,8 +107,9 @@ impl PyResponseAsyncBytesIterator { fn __anext__<'py>(&mut self, py: Python<'py>) -> PyResult> { if let Some(content) = self.ready_content.take() { self.content_returned = true; + let content = PyBytes::new(py, &content).unbind(); let future = - pyo3_async_runtimes::tokio::future_into_py::<_, Vec>(py, async move { + pyo3_async_runtimes::tokio::future_into_py::<_, Py>(py, async move { Ok(content) })?; return Ok(future); @@ -139,43 +141,44 @@ impl PyResponseAsyncBytesIterator { let parent_response = self.parent_response.as_ref().map(|p| p.clone_ref(py)); - let future = pyo3_async_runtimes::tokio::future_into_py::<_, Vec>(py, async move { - let chunk_result = { - let mut stream_guard = stream_arc.lock().await; - if let Some(stream) = stream_guard.as_mut() { - stream.next().await - } else { - None - } - }; - match chunk_result { - Some(Ok(chunk)) => Ok(chunk.to_vec()), - Some(Err(e)) => { - if let Some(parent) = parent_response { - Python::attach(|py| { - if let Ok(mut parent_ref) = parent.try_borrow_mut(py) { - parent_ref.inner_state = InnerResponseState::StreamingClosed; - parent_ref.is_stream_consumed = true; - parent_ref.is_closed = true; - } - }); + let future = + pyo3_async_runtimes::tokio::future_into_py::<_, Py>(py, async move { + let chunk_result = { + let mut stream_guard = stream_arc.lock().await; + if let Some(stream) = stream_guard.as_mut() { + stream.next().await + } else { + None } - Err(ImpitPyError(ImpitError::from(e, None)).into()) - } - None => { - if let Some(parent) = parent_response { - Python::attach(|py| { - if let Ok(mut parent_ref) = parent.try_borrow_mut(py) { - parent_ref.inner_state = InnerResponseState::StreamingClosed; - parent_ref.is_stream_consumed = true; - parent_ref.is_closed = true; - } - }); + }; + match chunk_result { + Some(Ok(chunk)) => Ok(Python::attach(|py| PyBytes::new(py, &chunk).unbind())), + Some(Err(e)) => { + if let Some(parent) = parent_response { + Python::attach(|py| { + if let Ok(mut parent_ref) = parent.try_borrow_mut(py) { + parent_ref.inner_state = InnerResponseState::StreamingClosed; + parent_ref.is_stream_consumed = true; + parent_ref.is_closed = true; + } + }); + } + Err(ImpitPyError(ImpitError::from(e, None)).into()) + } + None => { + if let Some(parent) = parent_response { + Python::attach(|py| { + if let Ok(mut parent_ref) = parent.try_borrow_mut(py) { + parent_ref.inner_state = InnerResponseState::StreamingClosed; + parent_ref.is_stream_consumed = true; + parent_ref.is_closed = true; + } + }); + } + Err(pyo3::exceptions::PyStopAsyncIteration::new_err("")) } - Err(pyo3::exceptions::PyStopAsyncIteration::new_err("")) } - } - })?; + })?; Ok(future) } @@ -219,7 +222,7 @@ pub struct ImpitPyResponse { // #[pyo3(get)] // elapsed: Duration, text: Option, - content: Option>, + content: Option, inner: Option, inner_state: InnerResponseState, // Raw, undecoded header name/value byte pairs. The `headers` getter builds the Python-side @@ -271,7 +274,7 @@ impl ImpitPyResponse { is_closed: true, is_stream_consumed: true, text: None, - content: Some(content.unwrap_or_default()), + content: Some(content.map(Bytes::from).unwrap_or_default()), inner: None, inner_state: InnerResponseState::Read, raw_headers, @@ -400,23 +403,26 @@ impl ImpitPyResponse { }); match state { - InnerResponseState::Read => Ok(content_option.unwrap_or_default()), + InnerResponseState::Read => Ok(Python::attach(|py| { + PyBytes::new(py, &content_option.unwrap_or_default()).unbind() + })), InnerResponseState::Unread => { if let Some(response) = response_option { - let content = - response.bytes().await.map(|b| b.to_vec()).map_err(|_| { - ImpitPyError(impit::errors::ImpitError::NetworkError) - })?; + let content = response + .bytes() + .await + .map_err(|_| ImpitPyError(impit::errors::ImpitError::NetworkError))?; - Python::attach(|py| { + Ok(Python::attach(|py| { let mut slf_ref = slf.borrow_mut(py); slf_ref.content = Some(content.clone()); slf_ref.inner_state = InnerResponseState::Read; slf_ref.is_stream_consumed = true; slf_ref.is_closed = true; - }); + drop(slf_ref); - Ok(content) + PyBytes::new(py, &content).unbind() + })) } else { Err(ImpitPyError(impit::errors::ImpitError::StreamClosed).into()) } @@ -466,22 +472,20 @@ impl ImpitPyResponse { } #[getter] - fn content(&mut self, py: Python<'_>) -> PyResult> { + fn content<'py>(&mut self, py: Python<'py>) -> PyResult> { self.read(py) } #[getter] - fn text(&mut self, py: Python<'_>) -> PyResult { - if let Some(cached_text) = &self.text { - return Ok(cached_text.clone()); - } + fn text<'py>(&mut self, py: Python<'py>) -> PyResult> { + if self.text.is_none() { + let decoder = encoding_from_whatwg_label(&self.encoding); + let decoded_text = impit::utils::decode(self.read_bytes(py)?, decoder); - let content_bytes = self.read(py)?; - let decoder = encoding_from_whatwg_label(&self.encoding); - let decoded_text = impit::utils::decode(&content_bytes, decoder); + self.text = Some(decoded_text); + } - self.text = Some(decoded_text.clone()); - Ok(decoded_text) + Ok(PyString::new(py, self.text.as_deref().unwrap_or_default())) } fn json(&mut self, py: Python<'_>) -> PyResult> { @@ -491,12 +495,19 @@ impl ImpitPyResponse { Ok(parsed.into()) } - fn read(&mut self, py: Python<'_>) -> PyResult> { + fn read<'py>(&mut self, py: Python<'py>) -> PyResult> { + Ok(PyBytes::new(py, self.read_bytes(py)?)) + } +} + +impl ImpitPyResponse { + /// Buffers the response body if it hasn't been read yet and borrows it, so that the callers + /// that only need to look at the bytes don't have to copy the whole body first. + fn read_bytes(&mut self, py: Python<'_>) -> PyResult<&[u8]> { match self.inner_state { InnerResponseState::Read => self .content - .as_ref() - .cloned() + .as_deref() .ok_or_else(|| ImpitPyError(impit::errors::ImpitError::ResponseNotRead).into()), InnerResponseState::Streaming | InnerResponseState::StreamingClosed => { if self.is_stream_consumed { @@ -517,23 +528,19 @@ impl ImpitPyResponse { response .bytes() .await - .map(|b| b.to_vec()) .map_err(|_| ImpitPyError(impit::errors::ImpitError::NetworkError)) }) })?; - self.content = Some(content.clone()); self.inner_state = InnerResponseState::Read; self.is_stream_consumed = true; self.is_closed = true; - Ok(content) + Ok(self.content.insert(content)) } } } -} -impl ImpitPyResponse { fn close_async_impl(slf: Py) -> PyResult<()> { Python::attach(|py| { let mut slf_ref = slf.borrow_mut(py); @@ -580,15 +587,11 @@ impl ImpitPyResponse { .and_then(|ct| ct.into()); let (content, inner_state, encoding, inner, is_closed, is_stream_consumed) = if !stream { - let content = val - .bytes() - .await - .map(|b| b.to_vec()) - .map_err(|e| ImpitError::from(e, None))?; + let content = val.bytes().await.map_err(|e| ImpitError::from(e, None))?; let encoding = preferred_encoding .and_then(|e| encoding_from_whatwg_label(&e)) .or(content_type_charset) - .or(impit::utils::determine_encoding(content.as_slice())) + .or(impit::utils::determine_encoding(&content)) .unwrap_or(impit::utils::encodings::UTF_8); ( diff --git a/impit/src/http3.rs b/impit/src/http3.rs index 7391460c..b3f6acf6 100644 --- a/impit/src/http3.rs +++ b/impit/src/http3.rs @@ -37,7 +37,7 @@ impl H3Engine { } } - pub async fn host_supports_h3(&self, host: &String) -> bool { + pub async fn host_supports_h3(&self, host: &str) -> bool { { let cache = self.h3_alt_svc.read().await; if let Some(&supports_h3) = cache.get(host) { @@ -76,7 +76,7 @@ impl H3Engine { false } - pub async fn set_h3_support(&self, host: &String, supports_h3: bool) { + pub async fn set_h3_support(&self, host: &str, supports_h3: bool) { let mut cache = self.h3_alt_svc.write().await; cache.insert(host.to_owned(), supports_h3); } diff --git a/impit/src/http_headers/mod.rs b/impit/src/http_headers/mod.rs index 4702c157..b1651ae5 100644 --- a/impit/src/http_headers/mod.rs +++ b/impit/src/http_headers/mod.rs @@ -1,135 +1,68 @@ use crate::{errors::ImpitError, fingerprint::BrowserFingerprint}; use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; -use std::{collections::HashSet, str::FromStr}; - -pub struct HttpHeaders { - context: HttpHeadersBuilder, +use std::str::FromStr; + +/// Merges the impersonated fingerprint headers with the caller-supplied ones into a [`HeaderMap`]. +/// +/// Header names are matched case-insensitively and the first source that provides a name wins. +/// Custom header sources are consulted in reverse registration order (the last +/// [`with_custom_headers`](HttpHeadersBuilder::with_custom_headers) call has the highest priority), +/// followed by the fingerprint headers. An empty value drops the header instead of sending it. +#[derive(Default)] +pub struct HttpHeadersBuilder<'a> { + fingerprint: Option<&'a BrowserFingerprint>, + custom_headers: Vec<&'a [(String, String)]>, } -impl HttpHeaders { - pub fn new(options: &HttpHeadersBuilder) -> HttpHeaders { - HttpHeaders { - context: options.clone(), - } - } - - pub fn get_builder() -> HttpHeadersBuilder { - HttpHeadersBuilder::default() +impl<'a> HttpHeadersBuilder<'a> { + pub fn with_fingerprint(mut self, fingerprint: Option<&'a BrowserFingerprint>) -> Self { + self.fingerprint = fingerprint; + self } -} - -impl HttpHeaders { - pub fn iter(&self) -> impl Iterator + '_ { - // Use fingerprint headers if available, otherwise fall back to browser enum - let impersonated_headers: Vec<(String, String)> = - if let Some(ref fp) = self.context.fingerprint { - fp.headers.to_vec() - } else { - vec![] - }; - - let custom_headers = self - .context - .custom_headers - .iter() - .map(|(k, v)| (k.clone(), v.clone())); - let mut used_header_names: HashSet = HashSet::new(); - - custom_headers - .chain(impersonated_headers) - .filter_map(move |(name, value)| { - if used_header_names.contains(&name.to_lowercase()) { - None - } else { - used_header_names.insert(name.to_lowercase()); - if value.is_empty() { - None - } else { - Some((name, value)) - } - } - }) + pub fn with_custom_headers(mut self, custom_headers: &'a [(String, String)]) -> Self { + self.custom_headers.push(custom_headers); + self } -} -impl From> for HttpHeaders { - fn from(val: Vec<(String, String)>) -> Self { - let mut builder = HttpHeaders::get_builder(); - builder.with_custom_headers(Some(val)); - builder.build() - } -} + pub fn build(self) -> Result { + let fingerprint_headers = self + .fingerprint + .map(|fingerprint| fingerprint.headers.as_slice()) + .unwrap_or_default(); -impl From for Result { - fn from(val: HttpHeaders) -> Self { let mut headers = HeaderMap::new(); + let mut dropped: Vec = Vec::new(); - for (name, value) in val.iter() { - let header_name = HeaderName::from_str(&name); - let header_value = HeaderValue::from_str(&value); - - match (header_name, header_value) { - (Err(_), _) => { - return Err(ImpitError::InvalidHeaderName(name)); - } - (_, Err(_)) => { - return Err(ImpitError::InvalidHeaderValue(value)); - } - (Ok(header_name), Ok(header_value)) => { - headers.append(header_name, header_value); + for (name, value) in self + .custom_headers + .iter() + .rev() + .copied() + .chain(std::iter::once(fingerprint_headers)) + .flatten() + { + if value.is_empty() { + // An empty value removes the header, but still shadows the lower-priority sources. + if let Ok(name) = HeaderName::from_str(name) { + dropped.push(name); } + continue; } - } - Ok(headers) - } -} - -#[derive(Default, Clone)] -pub struct HttpHeadersBuilder { - host: String, - fingerprint: Option, - https: bool, - custom_headers: Vec<(String, String)>, -} -impl HttpHeadersBuilder { - // TODO: Enforce `with_host` to be called before `build` - pub fn with_host(&mut self, host: &String) -> &mut Self { - self.host = host.to_owned(); - self - } + let name = HeaderName::from_str(name) + .map_err(|_| ImpitError::InvalidHeaderName(name.clone()))?; - pub fn with_fingerprint(&mut self, fingerprint: &Option) -> &mut Self { - self.fingerprint = fingerprint.clone(); - self - } + if headers.contains_key(&name) || dropped.contains(&name) { + continue; + } - pub fn with_https(&mut self, https: bool) -> &mut Self { - self.https = https; - self - } + let value = HeaderValue::from_str(value) + .map_err(|_| ImpitError::InvalidHeaderValue(value.clone()))?; - pub fn with_custom_headers( - &mut self, - custom_headers: Option>, - ) -> &mut Self { - match custom_headers { - Some(headers) => { - // Later call to with_custom_headers should override existing headers. - // We need to prepend the new headers (higher prio) to the existing ones (lower prio). - self.custom_headers = headers - .iter() - .chain(self.custom_headers.iter()) - .map(|(k, v)| (k.to_owned(), v.to_owned())) - .collect(); - self - } - None => self, + headers.insert(name, value); } - } - pub fn build(&self) -> HttpHeaders { - HttpHeaders::new(self) + Ok(headers) } } diff --git a/impit/src/impit.rs b/impit/src/impit.rs index 1dbf3bc9..4c45610c 100644 --- a/impit/src/impit.rs +++ b/impit/src/impit.rs @@ -2,14 +2,14 @@ use tokio::sync::RwLock; use log::debug; use reqwest::{cookie::CookieStore, header::HeaderMap, Method, Response, Version}; -use std::{fmt::Debug, net::IpAddr, str::FromStr, sync::Arc, time::Duration}; +use std::{fmt::Debug, net::IpAddr, sync::Arc, time::Duration}; use url::Url; use crate::{ errors::{ErrorContext, ImpitError}, fingerprint::BrowserFingerprint, http3::H3Engine, - http_headers::HttpHeaders, + http_headers::HttpHeadersBuilder, request::{ImpitBody, ImpitRequest, RequestOptions}, tls, }; @@ -371,7 +371,7 @@ impl Impit { } } - async fn should_use_h3(&self, host: &String) -> bool { + async fn should_use_h3(&self, host: &str) -> bool { if self.config.max_http_version < Version::HTTP_3 { debug!("HTTP/3 is disabled, falling back to TCP-based requests."); return false; @@ -402,24 +402,20 @@ impl Impit { method: Method, url: Url, body: Option, - headers: Vec<(String, String)>, - ) -> ImpitRequest { - let host = url.host_str().unwrap_or_default().to_string(); - - let headers = HttpHeaders::get_builder() - .with_fingerprint(&self.config.fingerprint) - .with_host(&host) - .with_https(url.scheme() == "https") - .with_custom_headers(self.config.headers.to_owned()) - .with_custom_headers(Some(headers)) - .build(); - - ImpitRequest { + headers: &[(String, String)], + ) -> Result { + let headers = HttpHeadersBuilder::default() + .with_fingerprint(self.config.fingerprint.as_ref()) + .with_custom_headers(self.config.headers.as_deref().unwrap_or_default()) + .with_custom_headers(headers) + .build()?; + + Ok(ImpitRequest { url, body: body.unwrap_or_default(), - headers: headers.iter().collect(), - method: method.to_string(), - } + headers, + method, + }) } async fn execute_request( @@ -459,37 +455,27 @@ impl Impit { return Err(ImpitError::Http3Disabled); } - let url = request.url.to_string(); - let host = request.url.host_str().unwrap_or_default().to_string(); let h3 = http3_prior_knowledge || self - .should_use_h3(&request.url.host_str().unwrap_or_default().to_string()) + .should_use_h3(request.url.host_str().unwrap_or_default()) .await; let client = if h3 { - debug!("Using QUIC for request to {url}"); + debug!("Using QUIC for request to {}", request.url); self.h3_client.as_ref().unwrap_or(&self.base_client) } else { - debug!("{url} doesn't seem to have HTTP3 support"); + debug!("{} doesn't seem to have HTTP3 support", request.url); &self.base_client }; - let header_map_result: Result = - HttpHeaders::from(request.headers).into(); - let header_map = header_map_result?; - - let method = Method::from_str(&request.method).map_err(|_| { - ImpitError::InvalidMethod(format!("Invalid HTTP method: {}", request.method)) - })?; - let max_redirects = match self.config.redirect { RedirectBehavior::FollowRedirect(max) => max, RedirectBehavior::ManualRedirect => 0, }; let mut prepared = PreparedRequest { - method: method.clone(), - url: request.url.clone(), - headers: header_map, + method: request.method, + url: request.url, + headers: request.headers, body: request.body, }; @@ -505,9 +491,9 @@ impl Impit { Some(ErrorContext { timeout: Some(timeout.unwrap_or(self.config.request_timeout)), max_redirects: Some(max_redirects), - method: Some(method.to_string()), - protocol: Some(request.url.scheme().to_string()), - url: Some(url.clone()), + method: Some(prepared.method.to_string()), + protocol: Some(prepared.url.scheme().to_string()), + url: Some(prepared.url.to_string()), }), ); @@ -520,7 +506,8 @@ impl Impit { }; debug!( - "Primary request to {url} failed with {primary_error}, retrying with vanilla client" + "Primary request to {} failed with {primary_error}, retrying with vanilla client", + prepared.url ); match self .execute_request(vanilla_client, &mut prepared, timeout, false) @@ -535,7 +522,8 @@ impl Impit { if !h3 { let engine_guard = self.h3_engine.read().await; if let Some(h3_engine) = engine_guard.as_ref() { - h3_engine.set_h3_support(&host, false).await; + let host = prepared.url.host_str().unwrap_or_default(); + h3_engine.set_h3_support(host, false).await; if let Some(alt_svc) = response.headers().get("Alt-Svc") { if let Ok(alt_svc_str) = alt_svc.to_str() { @@ -543,7 +531,7 @@ impl Impit { debug!( "{host} supports HTTP/3 (alt-svc header), adding to Alt-Svc cache" ); - h3_engine.set_h3_support(&host, true).await; + h3_engine.set_h3_support(host, true).await; } } } @@ -563,8 +551,7 @@ impl Impit { let url = self.parse_url(url)?; let request_options = options.unwrap_or_default(); - let headers = request_options.headers; - let request = self.build_request(method, url, body, headers); + let request = self.build_request(method, url, body, &request_options.headers)?; let timeout = match request_options.timeout { None => None, diff --git a/impit/src/request.rs b/impit/src/request.rs index 1f440534..12e878aa 100644 --- a/impit/src/request.rs +++ b/impit/src/request.rs @@ -2,6 +2,7 @@ use std::time::Duration; use bytes::Bytes; use futures_core::TryStream; +use reqwest::{header::HeaderMap, Method}; use url::Url; /// A struct that holds the request options. @@ -90,8 +91,8 @@ impl> From for ImpitBody { pub struct ImpitRequest { pub url: Url, pub body: ImpitBody, - pub headers: Vec<(String, String)>, - pub method: String, + pub headers: HeaderMap, + pub method: Method, } #[cfg(test)] diff --git a/impit/src/response_parsing/mod.rs b/impit/src/response_parsing/mod.rs index 9874e1de..db67ff62 100644 --- a/impit/src/response_parsing/mod.rs +++ b/impit/src/response_parsing/mod.rs @@ -1,5 +1,7 @@ use encoding::Encoding; +use lol_html::{ElementContentHandlers, HandlerResult, Selector}; use mime::{Mime, TEXT_PLAIN}; +use std::{borrow::Cow, sync::LazyLock}; /// Implements the BOM sniffing algorithm to detect the encoding of the response. /// If the BOM sniffing algorithm fails, the function returns `None`. @@ -10,19 +12,29 @@ fn bom_sniffing(bytes: &[u8]) -> Option { return None; } - if [0xEF, 0xBB, 0xBF].to_vec() == bytes[0..3].to_vec() { - return Some(encoding::all::UTF_8); - } - - if [0xFE, 0xFF].to_vec() == bytes[0..2].to_vec() { - return Some(encoding::all::UTF_16BE); - } - - if [0xFF, 0xFE].to_vec() == bytes[0..2].to_vec() { - return Some(encoding::all::UTF_16LE); + match bytes { + [0xEF, 0xBB, 0xBF, ..] => Some(encoding::all::UTF_8), + [0xFE, 0xFF, ..] => Some(encoding::all::UTF_16BE), + [0xFF, 0xFE, ..] => Some(encoding::all::UTF_16LE), + _ => None, } +} - None +const META_TAG_START: &[u8] = b" = LazyLock::new(|| "meta[charset]".parse().unwrap()); +static HTTP_EQUIV_SELECTOR: LazyLock = + LazyLock::new(|| "meta[http-equiv]".parse().unwrap()); + +fn on_element<'h>( + selector: &'static Selector, + handler: impl FnMut(&mut lol_html::html_content::Element<'_, '_>) -> HandlerResult + 'h, +) -> (Cow<'static, Selector>, ElementContentHandlers<'h>) { + ( + Cow::Borrowed(selector), + ElementContentHandlers::default().element(handler), + ) } /// A lazy implementation of the prescan algorithm, using `lol_html` to parse the HTML and extract the encoding. @@ -35,6 +47,15 @@ fn prescan_bytestream(bytes: &[u8]) -> Option { let limit = std::cmp::min(1024, bytes.len()); + // Both handlers below can only fire on a ` Option { let mut rewriter = lol_html::HtmlRewriter::new( lol_html::Settings { element_content_handlers: vec![ - lol_html::element!("meta[charset]", move |el| { + on_element(&CHARSET_SELECTOR, move |el| { if found_charset.borrow().is_none() { if let Some(charset) = el.get_attribute("charset") { *found_charset.borrow_mut() = @@ -55,7 +76,7 @@ fn prescan_bytestream(bytes: &[u8]) -> Option { } Ok(()) }), - lol_html::element!("meta[http-equiv]", move |el| { + on_element(&HTTP_EQUIV_SELECTOR, move |el| { if found_http_equiv.borrow().is_none() { let is_content_type = el .get_attribute("http-equiv") @@ -173,3 +194,51 @@ impl From for Option { encoding::label::encoding_from_whatwg_label(val.charset.as_str()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn prescan_finds_the_meta_charset() { + let html = br#""#; + + assert_eq!( + prescan_bytestream(html).map(|encoding| encoding.name()), + Some("windows-1250") + ); + } + + #[test] + fn prescan_matches_the_meta_tag_case_insensitively() { + let html = br#""#; + + assert_eq!( + prescan_bytestream(html).map(|encoding| encoding.name()), + Some("windows-1250") + ); + } + + #[test] + fn prescan_reads_the_charset_from_the_content_type_meta() { + let html = br#""#; + + assert_eq!( + prescan_bytestream(html).map(|encoding| encoding.name()), + Some("windows-1250") + ); + } + + #[test] + fn prescan_ignores_meta_tags_past_the_first_kilobyte() { + let mut html = vec![b' '; 1024]; + html.extend_from_slice(br#""#); + + assert!(prescan_bytestream(&html).is_none()); + } + + #[test] + fn prescan_skips_documents_without_a_meta_tag() { + assert!(prescan_bytestream(br#"{"hello":"world"}"#).is_none()); + } +}