use std::time::Duration;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::browser::{Cookie, CookieParam, StaticElement, Tab};
use crate::launcher::Proxy;
use crate::{Error, Result};
mod http;
use http::{HttpBackend, RawResponse};
pub use http::BrowserProfile;
const DEFAULT_UA: &str =
"Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:135.0) Gecko/20100101 Firefox/135.0";
#[derive(Debug, Clone)]
pub struct SessionOptions {
pub user_agent: String,
pub headers: Vec<(String, String)>,
pub proxy: Option<Proxy>,
pub timeout: Duration,
pub ignore_https_errors: bool,
pub max_redirects: usize,
pub profile: BrowserProfile,
}
impl Default for SessionOptions {
fn default() -> Self {
Self {
user_agent: DEFAULT_UA.to_string(),
headers: Vec::new(),
proxy: None,
timeout: Duration::from_secs(30),
ignore_https_errors: false,
max_redirects: 10,
profile: BrowserProfile::None,
}
}
}
impl SessionOptions {
pub fn new() -> Self {
Self::default()
}
pub fn user_agent(mut self, ua: impl Into<String>) -> Self {
self.user_agent = ua.into();
self
}
pub fn header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
self.headers.push((name.into(), value.into()));
self
}
pub fn proxy(mut self, proxy: Proxy) -> Self {
self.proxy = Some(proxy);
self
}
pub fn timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
pub fn ignore_https_errors(mut self, yes: bool) -> Self {
self.ignore_https_errors = yes;
self
}
pub fn max_redirects(mut self, n: usize) -> Self {
self.max_redirects = n;
self
}
pub fn profile(mut self, profile: BrowserProfile) -> Self {
self.profile = profile;
self
}
}
pub enum PostData {
Form(Vec<(String, String)>),
Json(Value),
Raw(String, Option<String>),
}
pub struct SessionPage {
backend: HttpBackend,
extra_headers: Vec<(String, String)>,
jar: CookieJar,
max_redirects: usize,
last_url: String,
last_status: u16,
last_headers: Vec<(String, String)>,
last_body: String,
}
impl SessionPage {
pub fn new_default() -> Result<Self> {
Self::new(SessionOptions::default())
}
pub fn new(opts: SessionOptions) -> Result<Self> {
let extra_headers = opts.headers.clone();
let backend = HttpBackend::build(&opts)?;
Ok(Self {
backend,
extra_headers,
jar: CookieJar::default(),
max_redirects: opts.max_redirects,
last_url: String::new(),
last_status: 0,
last_headers: Vec::new(),
last_body: String::new(),
})
}
pub async fn get(&mut self, url: &str) -> Result<bool> {
self.request("GET", url, None).await
}
pub async fn post(&mut self, url: &str, data: PostData) -> Result<bool> {
self.request("POST", url, Some(data)).await
}
async fn request(&mut self, method: &str, url: &str, body: Option<PostData>) -> Result<bool> {
let mut current =
reqwest::Url::parse(url).map_err(|e| Error::Other(format!("非法 URL {url}: {e}")))?;
let mut method = method.to_string();
let mut body: Option<(String, Option<String>)> = match body {
Some(PostData::Form(f)) => Some((
form_encode(&f),
Some("application/x-www-form-urlencoded".to_string()),
)),
Some(PostData::Json(j)) => Some((serde_json::to_string(&j)?, Some("application/json".to_string()))),
Some(PostData::Raw(s, ct)) => Some((s, ct)),
None => None,
};
let mut hops = 0usize;
loop {
let mut headers = self.extra_headers.clone();
if let Some(cookie) = self.jar.header_for(¤t) {
headers.push(("Cookie".to_string(), cookie));
}
if let Some((_, Some(ct))) = &body {
headers.push(("Content-Type".to_string(), ct.clone()));
}
let body_str = body.as_ref().map(|(b, _)| b.as_str());
let resp: RawResponse = self
.backend
.send_once(&method, current.as_str(), &headers, body_str)
.await?;
for (k, v) in &resp.headers {
if k.eq_ignore_ascii_case("set-cookie") {
self.jar.store(v, ¤t);
}
}
let code = resp.status;
if (300..=399).contains(&code)
&& hops < self.max_redirects
&& let Some(loc) = resp
.headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case("location"))
.map(|(_, v)| v.clone())
{
let next = current
.join(&loc)
.map_err(|e| Error::Other(format!("非法重定向 Location {loc}: {e}")))?;
hops += 1;
current = next;
if code != 307 && code != 308 {
method = "GET".to_string();
body = None;
}
continue;
}
self.last_status = code;
self.last_url = current.to_string();
self.last_headers = resp.headers;
self.last_body = resp.body;
return Ok((200..=299).contains(&code));
}
}
pub fn html(&self) -> &str {
&self.last_body
}
pub fn text(&self) -> &str {
&self.last_body
}
pub fn status(&self) -> u16 {
self.last_status
}
pub fn url(&self) -> &str {
&self.last_url
}
pub fn headers(&self) -> &[(String, String)] {
&self.last_headers
}
pub fn header(&self, name: &str) -> Option<&str> {
self.last_headers
.iter()
.find(|(k, _)| k.eq_ignore_ascii_case(name))
.map(|(_, v)| v.as_str())
}
pub fn json(&self) -> Result<Value> {
Ok(serde_json::from_str(&self.last_body)?)
}
pub fn s_root(&self) -> Result<StaticElement> {
StaticElement::parse(&self.last_body)
}
pub fn s_ele(&self, selector: &str) -> Result<StaticElement> {
self.s_root()?.ele(selector)
}
pub fn s_eles(&self, selector: &str) -> Result<Vec<StaticElement>> {
self.s_root()?.eles(selector)
}
pub fn title(&self) -> Result<String> {
match self.s_root()?.ele("tag:title") {
Ok(t) => t.text(),
Err(_) => Ok(String::new()),
}
}
pub fn cookies(&self) -> Vec<CookieParam> {
self.jar.to_params()
}
pub fn set_cookies(&mut self, cookies: Vec<CookieParam>) {
for c in cookies {
self.jar.upsert_param(c);
}
}
pub fn clear_cookies(&mut self) {
self.jar.cookies.clear();
}
pub async fn load_cookies_from_tab(&mut self, tab: &Tab) -> Result<()> {
for c in tab.cookies().await? {
self.jar.upsert_cookie(c);
}
Ok(())
}
pub async fn apply_cookies_to_tab(&self, tab: &Tab) -> Result<()> {
tab.set_cookies(self.jar.to_params()).await
}
#[cfg(feature = "cdp")]
pub async fn load_cookies_from_cdp_tab(&mut self, tab: &crate::cdp::ChromiumTab) -> Result<()> {
let raw = tab.get_cookies().await?;
let mut params = Vec::with_capacity(raw.len());
for c in raw {
let name = c["name"].as_str().unwrap_or_default().to_string();
if name.is_empty() {
continue;
}
params.push(CookieParam {
name,
value: c["value"].as_str().unwrap_or_default().to_string(),
url: None,
domain: c["domain"].as_str().map(str::to_string),
path: c["path"].as_str().map(str::to_string),
secure: c["secure"].as_bool(),
http_only: c["httpOnly"].as_bool(),
expires: c["expires"].as_f64(),
});
}
self.set_cookies(params);
Ok(())
}
pub fn save_cookies(&self, path: &str) -> Result<()> {
std::fs::write(path, serde_json::to_string_pretty(&self.jar.cookies)?)?;
Ok(())
}
pub fn load_cookies_file(&mut self, path: &str) -> Result<()> {
let s = std::fs::read_to_string(path)?;
let list: Vec<StoredCookie> = serde_json::from_str(&s)?;
for c in list {
self.jar.upsert(c);
}
Ok(())
}
}
#[derive(Default)]
struct CookieJar {
cookies: Vec<StoredCookie>,
}
#[derive(Clone, Serialize, Deserialize)]
struct StoredCookie {
name: String,
value: String,
domain: String,
path: String,
secure: bool,
http_only: bool,
expires: Option<f64>,
host_only: bool,
}
impl CookieJar {
fn upsert(&mut self, c: StoredCookie) {
if let Some(e) = self
.cookies
.iter_mut()
.find(|x| x.name == c.name && x.domain == c.domain && x.path == c.path)
{
*e = c;
} else {
self.cookies.push(c);
}
}
fn upsert_cookie(&mut self, c: Cookie) {
let host_only = !c.domain.starts_with('.');
self.upsert(StoredCookie {
name: c.name,
value: c.value,
domain: c.domain.trim_start_matches('.').to_ascii_lowercase(),
path: if c.path.is_empty() {
"/".into()
} else {
c.path
},
secure: c.secure,
http_only: c.http_only,
expires: if c.expires > 0.0 {
Some(c.expires)
} else {
None
},
host_only,
});
}
fn upsert_param(&mut self, c: CookieParam) {
let raw_domain = c.domain.clone().unwrap_or_default();
let host_only = c
.domain
.as_deref()
.map(|d| !d.starts_with('.'))
.unwrap_or(true);
let domain = if raw_domain.is_empty() {
c.url
.as_deref()
.and_then(|u| reqwest::Url::parse(u).ok())
.and_then(|u| u.host_str().map(|h| h.to_string()))
.unwrap_or_default()
} else {
raw_domain.trim_start_matches('.').to_string()
}
.to_ascii_lowercase();
self.upsert(StoredCookie {
name: c.name,
value: c.value,
domain,
path: c.path.unwrap_or_else(|| "/".into()),
secure: c.secure.unwrap_or(false),
http_only: c.http_only.unwrap_or(false),
expires: c.expires.filter(|e| *e > 0.0),
host_only,
});
}
fn to_params(&self) -> Vec<CookieParam> {
self.cookies
.iter()
.map(|c| CookieParam {
name: c.name.clone(),
value: c.value.clone(),
url: None,
domain: Some(c.domain.clone()),
path: Some(c.path.clone()),
secure: Some(c.secure),
http_only: Some(c.http_only),
expires: c.expires,
})
.collect()
}
fn store(&mut self, set_cookie: &str, req: &reqwest::Url) {
let mut parts = set_cookie.split(';');
let nv = match parts.next() {
Some(s) => s.trim(),
None => return,
};
let (name, value) = match nv.split_once('=') {
Some((n, v)) => (n.trim().to_string(), v.trim().to_string()),
None => return,
};
if name.is_empty() {
return;
}
let mut domain = String::new();
let mut path = String::new();
let mut secure = false;
let mut http_only = false;
let expires: Option<f64> = None;
let mut max_age: Option<f64> = None;
for attr in parts {
let attr = attr.trim();
let (k, v) = match attr.split_once('=') {
Some((k, v)) => (k.trim().to_ascii_lowercase(), v.trim().to_string()),
None => (attr.to_ascii_lowercase(), String::new()),
};
match k.as_str() {
"domain" => domain = v.trim_start_matches('.').to_ascii_lowercase(),
"path" => path = v,
"secure" => secure = true,
"httponly" => http_only = true,
"max-age" => max_age = v.parse::<f64>().ok().map(|s| now_unix() + s),
_ => {} }
}
let host_only = domain.is_empty();
let domain = if domain.is_empty() {
req.host_str().unwrap_or_default().to_ascii_lowercase()
} else {
domain
};
let path = if path.starts_with('/') {
path
} else {
default_path(req)
};
let exp = max_age.or(expires);
if let Some(e) = exp
&& e <= now_unix()
{
self.cookies
.retain(|c| !(c.name == name && c.domain == domain && c.path == path));
return;
}
self.upsert(StoredCookie {
name,
value,
domain,
path,
secure,
http_only,
expires: exp,
host_only,
});
}
fn header_for(&self, url: &reqwest::Url) -> Option<String> {
let host = url.host_str()?.to_ascii_lowercase();
let path = url.path();
let secure_ctx = url.scheme() == "https";
let now = now_unix();
let mut matched: Vec<&StoredCookie> = self
.cookies
.iter()
.filter(|c| {
if let Some(e) = c.expires
&& e <= now
{
return false;
}
if c.secure && !secure_ctx {
return false;
}
let domain_ok = if c.host_only {
host == c.domain
} else {
host == c.domain || host.ends_with(&format!(".{}", c.domain))
};
domain_ok && path_match(path, &c.path)
})
.collect();
if matched.is_empty() {
return None;
}
matched.sort_by_key(|c| std::cmp::Reverse(c.path.len()));
Some(
matched
.iter()
.map(|c| format!("{}={}", c.name, c.value))
.collect::<Vec<_>>()
.join("; "),
)
}
}
fn now_unix() -> f64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs_f64())
.unwrap_or(0.0)
}
fn form_encode(pairs: &[(String, String)]) -> String {
pairs
.iter()
.map(|(k, v)| format!("{}={}", percent_encode(k), percent_encode(v)))
.collect::<Vec<_>>()
.join("&")
}
fn percent_encode(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for b in s.bytes() {
match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
out.push(b as char)
}
b' ' => out.push('+'),
_ => out.push_str(&format!("%{b:02X}")),
}
}
out
}
fn default_path(url: &reqwest::Url) -> String {
let p = url.path();
match p.rfind('/') {
None | Some(0) => "/".to_string(),
Some(i) => p[..i].to_string(),
}
}
fn path_match(req: &str, cookie: &str) -> bool {
if cookie == req {
return true;
}
if !req.starts_with(cookie) {
return false;
}
cookie.ends_with('/') || req[cookie.len()..].starts_with('/')
}
#[cfg(test)]
mod tests {
use super::*;
fn url(u: &str) -> reqwest::Url {
reqwest::Url::parse(u).unwrap()
}
#[test]
fn store_and_send_basic() {
let mut jar = CookieJar::default();
jar.store("sid=abc; Path=/; HttpOnly", &url("https://x.com/login"));
let h = jar.header_for(&url("https://x.com/anything")).unwrap();
assert_eq!(h, "sid=abc");
assert!(jar.header_for(&url("https://y.com/")).is_none());
}
#[test]
fn secure_only_on_https() {
let mut jar = CookieJar::default();
jar.store("s=1; Secure", &url("https://x.com/"));
assert!(jar.header_for(&url("http://x.com/")).is_none());
assert_eq!(
jar.header_for(&url("https://x.com/")).as_deref(),
Some("s=1")
);
}
#[test]
fn domain_attr_includes_subdomains() {
let mut jar = CookieJar::default();
jar.store("t=2; Domain=x.com", &url("https://www.x.com/"));
assert_eq!(
jar.header_for(&url("https://api.x.com/")).as_deref(),
Some("t=2")
);
assert_eq!(
jar.header_for(&url("https://x.com/")).as_deref(),
Some("t=2")
);
}
#[test]
fn max_age_zero_deletes() {
let mut jar = CookieJar::default();
jar.store("k=v; Path=/", &url("https://x.com/"));
assert!(jar.header_for(&url("https://x.com/")).is_some());
jar.store("k=; Path=/; Max-Age=0", &url("https://x.com/"));
assert!(jar.header_for(&url("https://x.com/")).is_none());
}
#[test]
fn path_scoping() {
let mut jar = CookieJar::default();
jar.store("a=1; Path=/admin", &url("https://x.com/admin/x"));
assert!(jar.header_for(&url("https://x.com/")).is_none());
assert_eq!(
jar.header_for(&url("https://x.com/admin/y")).as_deref(),
Some("a=1")
);
}
#[test]
fn interop_param_roundtrip() {
let mut jar = CookieJar::default();
jar.upsert_param(CookieParam {
name: "u".into(),
value: "1".into(),
url: None,
domain: Some(".x.com".into()),
path: Some("/".into()),
secure: Some(false),
http_only: Some(true),
expires: None,
});
let params = jar.to_params();
assert_eq!(params.len(), 1);
assert_eq!(params[0].name, "u");
assert_eq!(
jar.header_for(&url("https://a.x.com/")).as_deref(),
Some("u=1")
);
}
}