Files
hottub/src/util/requester.rs
2026-03-22 15:56:25 +00:00

576 lines
19 KiB
Rust

use serde::Serialize;
use std::env;
use std::fmt;
use std::sync::{Arc, OnceLock};
use wreq::Client;
use wreq::Proxy;
use wreq::Response;
use wreq::Uri;
use wreq::Version;
use wreq::cookie::{CookieStore, Cookies, Jar};
use wreq::header::{HeaderMap, HeaderValue, SET_COOKIE, USER_AGENT};
use wreq::multipart::Form;
use wreq::redirect::Policy;
use wreq_util::Emulation;
use crate::util::flaresolverr::FlareSolverrRequest;
use crate::util::flaresolverr::Flaresolverr;
use crate::util::proxy;
// A Send + Sync error type for all async paths
type AnyErr = Box<dyn std::error::Error + Send + Sync + 'static>;
#[derive(serde::Serialize, serde::Deserialize, Clone)]
pub struct Requester {
#[serde(skip)]
client: Client,
#[serde(skip)]
cookie_jar: Arc<Jar>,
#[serde(skip)]
debug_trace_id: Option<String>,
proxy: bool,
flaresolverr_session: Option<String>,
user_agent: Option<String>,
}
impl fmt::Debug for Requester {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Requester")
.field("proxy", &self.proxy)
.field("debug_trace_id", &self.debug_trace_id)
.field("flaresolverr_session", &self.flaresolverr_session)
.field("user_agent", &self.user_agent)
.finish()
}
}
impl Requester {
fn shared_cookie_jar() -> Arc<Jar> {
static SHARED_COOKIE_JAR: OnceLock<Arc<Jar>> = OnceLock::new();
SHARED_COOKIE_JAR
.get_or_init(|| Arc::new(Jar::default()))
.clone()
}
fn origin_url_for_cookie_scope(url: &str) -> Option<url::Url> {
let parsed = url::Url::parse(url).ok()?;
let host = parsed.host_str()?;
let scheme = parsed.scheme();
url::Url::parse(&format!("{scheme}://{host}/")).ok()
}
fn store_response_cookies(&self, url: &str, response: &Response) {
let Some(origin) = Self::origin_url_for_cookie_scope(url) else {
return;
};
for value in response.headers().get_all(SET_COOKIE).iter() {
if let Ok(cookie) = value.to_str() {
self.cookie_jar.add_cookie_str(cookie, &origin.to_string());
}
}
}
fn store_flaresolverr_cookies(
&mut self,
request_url: &str,
cookies: &[crate::util::flaresolverr::FlaresolverrCookie],
) {
let fallback_origin = Self::origin_url_for_cookie_scope(request_url);
for cookie in cookies {
let origin = if !cookie.domain.is_empty() {
let scheme = fallback_origin
.as_ref()
.map(|url| url.scheme())
.unwrap_or("https");
let host = cookie.domain.trim_start_matches('.');
url::Url::parse(&format!("{scheme}://{host}/"))
.ok()
.or_else(|| fallback_origin.clone())
} else {
fallback_origin.clone()
};
let Some(origin) = origin else {
continue;
};
let mut cookie_string =
format!("{}={}; Path={}", cookie.name, cookie.value, cookie.path);
if !cookie.domain.is_empty() {
cookie_string.push_str(&format!("; Domain={}", cookie.domain));
}
if cookie.secure {
cookie_string.push_str("; Secure");
}
if cookie.httpOnly {
cookie_string.push_str("; HttpOnly");
}
if let Some(same_site) = cookie.sameSite.as_deref() {
if !same_site.is_empty() {
cookie_string.push_str(&format!("; SameSite={same_site}"));
}
}
self.cookie_jar
.add_cookie_str(&cookie_string, &origin.to_string());
}
}
fn debug_cookie_preview_from_owned_headers(
&self,
url: &str,
headers: &[(String, String)],
) -> String {
if let Some((_, value)) = headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case("cookie"))
{
return crate::util::flow_debug::preview(value, 160);
}
self.cookie_header_for_url(url)
.map(|cookie| crate::util::flow_debug::preview(&cookie, 160))
.unwrap_or_else(|| "none".to_string())
}
#[cfg(any(not(hottub_single_provider), hottub_provider = "hypnotube"))]
fn debug_cookie_preview_from_borrowed_headers(
&self,
url: &str,
headers: &[(&str, &str)],
) -> String {
if let Some((_, value)) = headers
.iter()
.find(|(key, _)| key.eq_ignore_ascii_case("cookie"))
{
return crate::util::flow_debug::preview(value, 160);
}
self.cookie_header_for_url(url)
.map(|cookie| crate::util::flow_debug::preview(&cookie, 160))
.unwrap_or_else(|| "none".to_string())
}
fn build_client(cookie_jar: Arc<Jar>, user_agent: Option<&str>) -> Client {
let mut builder = Client::builder()
.cert_verification(false)
.emulation(Emulation::Firefox146)
.cookie_provider(cookie_jar)
.redirect(Policy::default());
if let Some(user_agent) = user_agent {
let mut headers = HeaderMap::new();
if let Ok(value) = HeaderValue::from_str(user_agent) {
headers.insert(USER_AGENT, value);
builder = builder.default_headers(headers);
}
}
builder.build().expect("Failed to create HTTP client")
}
pub fn new() -> Self {
let cookie_jar = Self::shared_cookie_jar();
let client = Self::build_client(cookie_jar.clone(), None);
let requester = Requester {
client,
cookie_jar,
debug_trace_id: None,
proxy: false,
flaresolverr_session: None,
user_agent: None,
};
proxy::init_all_proxies_background(requester.clone());
requester
}
pub fn set_proxy(&mut self, proxy: bool) {
if proxy {
println!("Proxy enabled");
}
self.proxy = proxy;
}
pub fn set_debug_trace_id(&mut self, debug_trace_id: Option<String>) {
self.debug_trace_id = debug_trace_id;
}
#[cfg(feature = "debug")]
pub fn debug_trace_id(&self) -> Option<&str> {
self.debug_trace_id.as_deref()
}
pub fn cookie_header_for_url(&self, url: &str) -> Option<String> {
let parsed = url.parse::<Uri>().ok()?;
match self.cookie_jar.cookies(&parsed) {
Cookies::Compressed(value) => value.to_str().ok().map(ToOwned::to_owned),
Cookies::Uncompressed(values) => {
let joined = values
.into_iter()
.filter_map(|value| value.to_str().ok().map(ToOwned::to_owned))
.collect::<Vec<_>>()
.join("; ");
(!joined.is_empty()).then_some(joined)
}
Cookies::Empty => None,
_ => None,
}
}
pub async fn get_raw(&mut self, url: &str) -> Result<Response, wreq::Error> {
let cookie_preview = self
.cookie_header_for_url(url)
.map(|cookie| crate::util::flow_debug::preview(&cookie, 160))
.unwrap_or_else(|| "none".to_string());
#[cfg(not(feature = "debug"))]
let _ = &cookie_preview;
crate::flow_debug!(
"trace={} requester get_raw url={} cookies={} proxy={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
cookie_preview,
self.proxy
);
let client = Self::build_client(self.cookie_jar.clone(), self.user_agent.as_deref());
let mut request = client.get(url).version(Version::HTTP_11);
if self.proxy {
if let Ok(proxy_url) = env::var("BURP_URL") {
let proxy = Proxy::all(&proxy_url).unwrap();
request = request.proxy(proxy);
}
}
let response = request.send().await?;
self.store_response_cookies(url, &response);
Ok(response)
}
pub async fn get_raw_with_headers(
&mut self,
url: &str,
headers: Vec<(String, String)>,
) -> Result<Response, wreq::Error> {
let cookie_preview = self.debug_cookie_preview_from_owned_headers(url, &headers);
#[cfg(not(feature = "debug"))]
let _ = &cookie_preview;
crate::flow_debug!(
"trace={} requester get_raw_with_headers url={} headers={} cookies={} proxy={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
headers.len(),
cookie_preview,
self.proxy
);
let client = Self::build_client(self.cookie_jar.clone(), self.user_agent.as_deref());
let mut request = client.get(url).version(Version::HTTP_11);
if self.proxy {
if let Ok(proxy_url) = env::var("BURP_URL") {
let proxy = Proxy::all(&proxy_url).unwrap();
request = request.proxy(proxy);
}
}
// Set custom headers
for (key, value) in headers.iter() {
request = request.header(key, value);
}
let response = request.send().await?;
self.store_response_cookies(url, &response);
Ok(response)
}
pub async fn post_json<S>(
&mut self,
url: &str,
data: &S,
headers: Vec<(String, String)>,
) -> Result<Response, wreq::Error>
where
S: Serialize + ?Sized,
{
let cookie_preview = self.debug_cookie_preview_from_owned_headers(url, &headers);
#[cfg(not(feature = "debug"))]
let _ = &cookie_preview;
crate::flow_debug!(
"trace={} requester post_json url={} headers={} cookies={} proxy={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
headers.len(),
cookie_preview,
self.proxy
);
let mut request = self.client.post(url).version(Version::HTTP_11).json(data);
// Set custom headers
for (key, value) in headers.iter() {
request = request.header(key, value);
}
if self.proxy {
if let Ok(proxy_url) = env::var("BURP_URL") {
let proxy = Proxy::all(&proxy_url).unwrap();
request = request.proxy(proxy);
}
}
let response = request.send().await?;
self.store_response_cookies(url, &response);
Ok(response)
}
#[cfg(any(not(hottub_single_provider), hottub_provider = "hypnotube"))]
pub async fn post(
&mut self,
url: &str,
data: &str,
headers: Vec<(&str, &str)>,
) -> Result<Response, wreq::Error> {
let cookie_preview = self.debug_cookie_preview_from_borrowed_headers(url, &headers);
#[cfg(not(feature = "debug"))]
let _ = &cookie_preview;
crate::flow_debug!(
"trace={} requester post url={} headers={} cookies={} body_len={} proxy={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
headers.len(),
cookie_preview,
data.len(),
self.proxy
);
let mut request = self
.client
.post(url)
.version(Version::HTTP_11)
.body(data.to_string());
// Set custom headers
for (key, value) in headers.iter() {
request = request.header(key.to_string(), value.to_string());
}
if self.proxy {
if let Ok(proxy_url) = env::var("BURP_URL") {
let proxy = Proxy::all(&proxy_url).unwrap();
request = request.proxy(proxy);
}
}
let response = request.send().await?;
self.store_response_cookies(url, &response);
Ok(response)
}
pub async fn post_multipart(
&mut self,
url: &str,
form: Form,
headers: Vec<(String, String)>,
_http_version: Option<Version>,
) -> Result<Response, wreq::Error> {
let cookie_preview = self.debug_cookie_preview_from_owned_headers(url, &headers);
#[cfg(not(feature = "debug"))]
let _ = &cookie_preview;
crate::flow_debug!(
"trace={} requester post_multipart url={} headers={} cookies={} proxy={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
headers.len(),
cookie_preview,
self.proxy
);
let http_version = match _http_version {
Some(v) => v,
None => Version::HTTP_11,
};
let mut request = self.client.post(url).multipart(form).version(http_version);
// Set custom headers
for (key, value) in headers.iter() {
request = request.header(key, value);
}
if self.proxy {
if let Ok(proxy_url) = env::var("BURP_URL") {
let proxy = Proxy::all(&proxy_url).unwrap();
request = request.proxy(proxy);
}
}
let response = request.send().await?;
self.store_response_cookies(url, &response);
Ok(response)
}
pub async fn get(
&mut self,
url: &str,
_http_version: Option<Version>,
) -> Result<String, AnyErr> {
self.get_with_headers(url, Vec::new(), _http_version).await
}
pub async fn get_with_headers(
&mut self,
url: &str,
headers: Vec<(String, String)>,
_http_version: Option<Version>,
) -> Result<String, AnyErr> {
let cookie_preview = self.debug_cookie_preview_from_owned_headers(url, &headers);
#[cfg(not(feature = "debug"))]
let _ = &cookie_preview;
crate::flow_debug!(
"trace={} requester get_with_headers start url={} headers={} cookies={} http_version={:?} proxy={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
headers.len(),
cookie_preview,
_http_version,
self.proxy
);
let http_version = match _http_version {
Some(v) => v,
None => Version::HTTP_11,
};
loop {
let mut request = self.client.get(url).version(http_version);
for (key, value) in headers.iter() {
request = request.header(key, value);
}
if self.proxy {
if let Ok(proxy_url) = env::var("BURP_URL") {
let proxy = Proxy::all(&proxy_url).unwrap();
request = request.proxy(proxy);
}
}
let response = request.send().await?;
self.store_response_cookies(url, &response);
crate::flow_debug!(
"trace={} requester direct response url={} status={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
response.status()
);
if response.status().is_success() || response.status().as_u16() == 404 {
return Ok(response.text().await?);
}
if response.status().as_u16() == 429 {
crate::flow_debug!(
"trace={} requester direct retry url={} status=429",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120)
);
tokio::time::sleep(std::time::Duration::from_secs(1)).await;
continue;
} else {
println!(
"Direct request to {} failed with status: {}",
url,
response.status()
);
break;
}
}
// If direct request failed, try FlareSolverr. Map its error to a Send+Sync error immediately,
// so no non-Send error value lives across later `.await`s.
let flare_url = match env::var("FLARE_URL") {
Ok(url) => url,
Err(e) => return Err(format!("FLARE_URL not set: {e}").into()),
};
let mut flare = Flaresolverr::new(flare_url);
if self.proxy && env::var("BURP_URL").is_ok() {
flare.set_proxy(true);
}
crate::flow_debug!(
"trace={} requester flaresolverr url={} proxy={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
self.proxy
);
let res = flare
.solve(FlareSolverrRequest {
cmd: "request.get".to_string(),
url: url.to_string(),
maxTimeout: 60000,
})
.await
.map_err(|e| -> AnyErr { format!("Failed to solve FlareSolverr: {e}").into() })?;
// Rebuild client and apply UA/cookies from FlareSolverr
let useragent = res.solution.userAgent;
self.user_agent = Some(useragent);
self.store_flaresolverr_cookies(url, &res.solution.cookies);
self.client = Self::build_client(self.cookie_jar.clone(), self.user_agent.as_deref());
crate::flow_debug!(
"trace={} requester flaresolverr solved url={} user_agent={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
crate::util::flow_debug::preview(self.user_agent.as_deref().unwrap_or("unknown"), 96)
);
// Retry the original URL with the updated client & (optional) proxy
let mut request = self.client.get(url).version(Version::HTTP_11);
for (key, value) in headers.iter() {
request = request.header(key, value);
}
if self.proxy {
if let Ok(proxy_url) = env::var("BURP_URL") {
let proxy = Proxy::all(&proxy_url).unwrap();
request = request.proxy(proxy);
}
}
let response = request.send().await?;
self.store_response_cookies(url, &response);
crate::flow_debug!(
"trace={} requester retry response url={} status={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120),
response.status()
);
if response.status().is_success() {
return Ok(response.text().await?);
}
// Fall back to FlareSolverr-provided body
crate::flow_debug!(
"trace={} requester fallback body url={}",
self.debug_trace_id().unwrap_or("none"),
crate::util::flow_debug::preview(url, 120)
);
Ok(res.solution.response)
}
}
#[cfg(test)]
mod tests {
use super::Requester;
#[test]
fn new_requesters_share_cookie_jar() {
let a = Requester::new();
let b = Requester::new();
let origin = "https://shared-cookie-requester-test.invalid/";
a.cookie_jar
.add_cookie_str("shared_cookie=1; Path=/; SameSite=Lax", origin);
let cookie_header = b
.cookie_header_for_url("https://shared-cookie-requester-test.invalid/path")
.unwrap_or_default();
assert!(cookie_header.contains("shared_cookie=1"));
}
}