starred/yaak-mountain-loop
1
0
Fork 0
mirror of https://github.com/mountain-loop/yaak.git synced 2026-05-18 05:37:09 +02:00
Code Issues 8 Packages Projects Releases 10 Wiki Activity

Add cookie editing and inherited request settings

Browse Source
This commit is contained in:
Gregory Schier
2026-05-17 07:58:12 -07:00
parent dcfdf077e7
commit dc47b54b1c
42 changed files with 3789 additions and 932 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/yaak-http/src/cookies.rs
+33 -29
View File
@@ -7,7 +7,7 @@ use log::debug;
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use url::Url;
use yaak_models::models::{Cookie, CookieDomain, CookieExpires};
use yaak_models::models::{Cookie, CookieDomain, CookieExpires, CookieSameSite};
/// A thread-safe cookie store that can be shared across requests
#[derive(Debug, Clone)]
@@ -45,10 +45,7 @@ impl CookieStore {
let matching_cookies: Vec<_> = cookies
.iter()
.filter(|cookie| self.cookie_matches(cookie, url, &now))
.filter_map(|cookie| {
// Parse the raw cookie to get name=value
parse_cookie_name_value(&cookie.raw_cookie)
})
.map(|cookie| (cookie.name.clone(), cookie.value.clone()))
.collect();
if matching_cookies.is_empty() {
@@ -72,13 +69,7 @@ impl CookieStore {
if let Some(cookie) = parse_set_cookie(header_value, url) {
// Remove any existing cookie with the same name and domain
cookies.retain(|existing| !cookies_match(existing, &cookie));
debug!(
"Storing cookie: {} for domain {:?}",
parse_cookie_name_value(&cookie.raw_cookie)
.map(|(n, _)| n)
.unwrap_or_else(|| "unknown".to_string()),
cookie.domain
);
debug!("Storing cookie: {} for domain {:?}", cookie.name, cookie.domain);
cookies.push(cookie);
}
}
@@ -117,10 +108,9 @@ impl CookieStore {
}
// Check path
let (cookie_path, _) = &cookie.path;
let url_path = url.path();
path_matches(url_path, cookie_path)
path_matches(url_path, &cookie.path)
}
}
@@ -133,8 +123,7 @@ pub fn get_cookie_value_from_jar(
let domain = domain.and_then(normalize_cookie_domain_filter);
cookies.into_iter().find_map(|cookie| {
let (cookie_name, value) = parse_cookie_name_value(&cookie.raw_cookie)?;
if cookie_name != name {
if cookie.name != name {
return None;
}
@@ -144,11 +133,12 @@ pub fn get_cookie_value_from_jar(
}
}
Some(value)
Some(cookie.value)
})
}
/// Parse name=value from a cookie string (raw_cookie format)
#[cfg(test)]
fn parse_cookie_name_value(raw_cookie: &str) -> Option<(String, String)> {
// The raw_cookie typically looks like "name=value" or "name=value; attr1; attr2=..."
let first_part = raw_cookie.split(';').next()?;
@@ -177,8 +167,6 @@ fn cookie_domain_matches_filter(cookie_domain: &CookieDomain, domain: &str) -> b
fn parse_set_cookie(header_value: &str, request_url: &Url) -> Option<Cookie> {
let parsed = cookie::Cookie::parse(header_value).ok()?;
let raw_cookie = format!("{}={}", parsed.name(), parsed.value());
// Determine domain
let domain = if let Some(domain_attr) = parsed.domain() {
// Domain attribute present - this is a suffix match
@@ -216,14 +204,28 @@ fn parse_set_cookie(header_value: &str, request_url: &Url) -> Option<Cookie> {
// Determine path
let path = if let Some(path_attr) = parsed.path() {
(path_attr.to_string(), true)
path_attr.to_string()
} else {
// Default path is the directory of the request URI
let default_path = default_cookie_path(request_url.path());
(default_path, false)
default_cookie_path(request_url.path())
};
Some(Cookie { raw_cookie, domain, expires, path })
let same_site = parsed.same_site().map(|same_site| match same_site {
cookie::SameSite::Strict => CookieSameSite::Strict,
cookie::SameSite::Lax => CookieSameSite::Lax,
cookie::SameSite::None => CookieSameSite::None,
});
Some(Cookie {
name: parsed.name().to_string(),
value: parsed.value().to_string(),
domain,
expires,
path,
secure: parsed.secure().unwrap_or(false),
http_only: parsed.http_only().unwrap_or(false),
same_site,
})
}
/// Get the default cookie path from a request path (RFC 6265 Section 5.1.4)
@@ -261,10 +263,7 @@ fn path_matches(request_path: &str, cookie_path: &str) -> bool {
/// Check if two cookies match (same name and domain)
fn cookies_match(a: &Cookie, b: &Cookie) -> bool {
let name_a = parse_cookie_name_value(&a.raw_cookie).map(|(n, _)| n);
let name_b = parse_cookie_name_value(&b.raw_cookie).map(|(n, _)| n);
if name_a != name_b {
if a.name != b.name {
return false;
}
@@ -317,11 +316,16 @@ mod tests {
use super::*;
fn cookie(raw_cookie: &str, domain: CookieDomain) -> Cookie {
let (name, value) = parse_cookie_name_value(raw_cookie).unwrap();
Cookie {
raw_cookie: raw_cookie.to_string(),
name,
value,
domain,
expires: CookieExpires::SessionEnd,
path: ("/".to_string(), false),
path: "/".to_string(),
secure: false,
http_only: false,
same_site: None,
}
}
crates/yaak-http/src/transaction.rs
+182 -18
View File
@@ -12,22 +12,58 @@ pub struct HttpTransaction<S: HttpSender> {
sender: S,
max_redirects: usize,
cookie_store: Option<CookieStore>,
send_cookies: bool,
store_cookies: bool,
}
impl<S: HttpSender> HttpTransaction<S> {
/// Create a new transaction with default settings
pub fn new(sender: S) -> Self {
Self { sender, max_redirects: 10, cookie_store: None }
Self {
sender,
max_redirects: 10,
cookie_store: None,
send_cookies: false,
store_cookies: false,
}
}
/// Create a new transaction with custom max redirects
pub fn with_max_redirects(sender: S, max_redirects: usize) -> Self {
Self { sender, max_redirects, cookie_store: None }
Self {
sender,
max_redirects,
cookie_store: None,
send_cookies: false,
store_cookies: false,
}
}
/// Create a new transaction with a cookie store
pub fn with_cookie_store(sender: S, cookie_store: CookieStore) -> Self {
Self { sender, max_redirects: 10, cookie_store: Some(cookie_store) }
Self {
sender,
max_redirects: 10,
cookie_store: Some(cookie_store),
send_cookies: true,
store_cookies: true,
}
}
/// Create a new transaction with a cookie store and explicit send/store behavior
pub fn with_cookie_behavior(
sender: S,
cookie_store: CookieStore,
send_cookies: bool,
store_cookies: bool,
) -> Self {
Self {
sender,
max_redirects: 10,
cookie_store: Some(cookie_store),
send_cookies,
store_cookies,
}
}
/// Create a new transaction with custom max redirects and a cookie store
@@ -36,7 +72,13 @@ impl<S: HttpSender> HttpTransaction<S> {
max_redirects: usize,
cookie_store: Option<CookieStore>,
) -> Self {
Self { sender, max_redirects, cookie_store }
Self {
sender,
max_redirects,
send_cookies: cookie_store.is_some(),
store_cookies: cookie_store.is_some(),
cookie_store,
}
}
/// Execute the request with cancellation support.
@@ -66,9 +108,11 @@ impl<S: HttpSender> HttpTransaction<S> {
}
// Inject cookies into headers if we have a cookie store
let headers_with_cookies = if let Some(cookie_store) = &self.cookie_store {
let headers_with_cookies = if self.send_cookies {
let mut headers = current_headers.clone();
if let Ok(url) = Url::parse(&current_url) {
if let (Some(cookie_store), Ok(url)) =
(&self.cookie_store, Url::parse(&current_url))
{
if let Some(cookie_header) = cookie_store.get_cookie_header(&url) {
debug!("Injecting Cookie header: {}", cookie_header);
// Check if there's already a Cookie header and merge if so
@@ -115,8 +159,10 @@ impl<S: HttpSender> HttpTransaction<S> {
};
// Parse Set-Cookie headers and store cookies
if let Some(cookie_store) = &self.cookie_store {
if let Ok(url) = Url::parse(&current_url) {
if self.store_cookies {
if let (Some(cookie_store), Ok(url)) =
(&self.cookie_store, Url::parse(&current_url))
{
let set_cookie_headers: Vec<String> = response
.headers
.iter()
@@ -579,10 +625,14 @@ mod tests {
// Create a cookie store with a test cookie
let cookie = Cookie {
raw_cookie: "session=abc123".to_string(),
name: "session".to_string(),
value: "abc123".to_string(),
domain: CookieDomain::HostOnly("example.com".to_string()),
expires: CookieExpires::SessionEnd,
path: ("/".to_string(), false),
path: "/".to_string(),
secure: false,
http_only: false,
same_site: None,
};
let cookie_store = CookieStore::from_cookies(vec![cookie]);
@@ -602,6 +652,67 @@ mod tests {
assert!(result.is_ok());
}
#[tokio::test]
async fn test_cookie_injection_can_be_disabled() {
struct CookieRejectingSender;
#[async_trait]
impl HttpSender for CookieRejectingSender {
async fn send(
&self,
request: SendableHttpRequest,
_event_tx: mpsc::Sender<HttpResponseEvent>,
) -> Result<HttpResponse> {
let cookie_header =
request.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case("cookie"));
assert!(cookie_header.is_none(), "Cookie header should not be present");
let body_stream: Pin<Box<dyn AsyncRead + Send>> =
Box::pin(std::io::Cursor::new(vec![]));
Ok(HttpResponse::new(
200,
None,
Vec::new(),
Vec::new(),
None,
"https://example.com".to_string(),
None,
Some("HTTP/1.1".to_string()),
body_stream,
ContentEncoding::Identity,
))
}
}
use yaak_models::models::{Cookie, CookieDomain, CookieExpires};
let cookie = Cookie {
name: "session".to_string(),
value: "abc123".to_string(),
domain: CookieDomain::HostOnly("example.com".to_string()),
expires: CookieExpires::SessionEnd,
path: "/".to_string(),
secure: false,
http_only: false,
same_site: None,
};
let cookie_store = CookieStore::from_cookies(vec![cookie]);
let transaction =
HttpTransaction::with_cookie_behavior(CookieRejectingSender, cookie_store, false, true);
let request = SendableHttpRequest {
url: "https://example.com/api".to_string(),
method: "GET".to_string(),
headers: vec![],
..Default::default()
};
let (_tx, rx) = tokio::sync::watch::channel(false);
let (event_tx, _event_rx) = mpsc::channel(100);
let result = transaction.execute_with_cancellation(request, rx, event_tx).await;
assert!(result.is_ok());
}
#[tokio::test]
async fn test_set_cookie_parsing() {
// Create a cookie store
@@ -655,7 +766,62 @@ mod tests {
// Verify the cookie was stored
let cookies = cookie_store.get_all_cookies();
assert_eq!(cookies.len(), 1);
assert!(cookies[0].raw_cookie.contains("session=xyz789"));
assert_eq!(cookies[0].name, "session");
assert_eq!(cookies[0].value, "xyz789");
}
#[tokio::test]
async fn test_set_cookie_storage_can_be_disabled() {
let cookie_store = CookieStore::new();
struct SetCookieSender;
#[async_trait]
impl HttpSender for SetCookieSender {
async fn send(
&self,
_request: SendableHttpRequest,
_event_tx: mpsc::Sender<HttpResponseEvent>,
) -> Result<HttpResponse> {
let headers =
vec![("set-cookie".to_string(), "session=xyz789; Path=/".to_string())];
let body_stream: Pin<Box<dyn AsyncRead + Send>> =
Box::pin(std::io::Cursor::new(vec![]));
Ok(HttpResponse::new(
200,
None,
headers,
Vec::new(),
None,
"https://example.com".to_string(),
None,
Some("HTTP/1.1".to_string()),
body_stream,
ContentEncoding::Identity,
))
}
}
let transaction = HttpTransaction::with_cookie_behavior(
SetCookieSender,
cookie_store.clone(),
true,
false,
);
let request = SendableHttpRequest {
url: "https://example.com/login".to_string(),
method: "POST".to_string(),
headers: vec![],
..Default::default()
};
let (_tx, rx) = tokio::sync::watch::channel(false);
let (event_tx, _event_rx) = mpsc::channel(100);
let result = transaction.execute_with_cancellation(request, rx, event_tx).await;
assert!(result.is_ok());
assert!(cookie_store.get_all_cookies().is_empty());
}
#[tokio::test]
@@ -719,17 +885,15 @@ mod tests {
let cookies = cookie_store.get_all_cookies();
assert_eq!(cookies.len(), 3, "All three Set-Cookie headers should be parsed and stored");
let cookie_values: Vec<&str> = cookies.iter().map(|c| c.raw_cookie.as_str()).collect();
let cookie_values: Vec<_> =
cookies.iter().map(|c| format!("{}={}", c.name, c.value)).collect();
assert!(
cookie_values.iter().any(|c| c.contains("session=abc123")),
cookie_values.iter().any(|c| c == "session=abc123"),
"session cookie should be stored"
);
assert!(cookie_values.iter().any(|c| c == "user_id=42"), "user_id cookie should be stored");
assert!(
cookie_values.iter().any(|c| c.contains("user_id=42")),
"user_id cookie should be stored"
);
assert!(
cookie_values.iter().any(|c| c.contains("preferences=dark")),
cookie_values.iter().any(|c| c == "preferences=dark"),
"preferences cookie should be stored"
);
}