Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Store response cookies #863

Merged
merged 1 commit into from
Oct 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 12 additions & 4 deletions src/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -133,9 +133,11 @@ impl Agent {
}
}

/// Access the cookie jar.
/// Access the shared cookie jar.
///
/// Used to persist and manipulate the cookies.
/// Used to persist and manipulate the cookies. The jar is shared between
/// all clones of the same [`Agent`], meaning you must drop the CookieJar
/// before using the agent, or end up with a deadlock.
///
/// ```no_run
/// use std::io::Write;
Expand All @@ -148,11 +150,17 @@ impl Agent {
///
/// // Saves (persistent) cookies
/// let mut file = File::create("cookies.json")?;
/// agent.cookie_jar().save_json(&mut file)?;
/// let jar = agent.cookie_jar_lock();
///
/// jar.save_json(&mut file)?;
///
/// // Release the cookie jar to use agents again.
/// jar.release();
///
/// # Ok::<_, ureq::Error>(())
/// ```
#[cfg(feature = "cookies")]
pub fn cookie_jar(&self) -> crate::cookies::CookieJar<'_> {
pub fn cookie_jar_lock(&self) -> crate::cookies::CookieJar<'_> {
self.jar.lock()
}

Expand Down
24 changes: 23 additions & 1 deletion src/cookies.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ pub(crate) struct SharedCookieJar {

/// Collection of cookies.
///
/// The jar is accessed using [`Agent::cookie_jar`][crate::Agent::cookie_jar].
/// The jar is accessed using [`Agent::cookie_jar_lock`][crate::Agent::cookie_jar_lock].
/// It can be saved and loaded.
pub struct CookieJar<'a>(MutexGuard<'a, CookieStore>);

Expand Down Expand Up @@ -45,6 +45,15 @@ enum CookieInner<'a> {
Owned(cookie_store::Cookie<'a>),
}

impl<'a> CookieInner<'a> {
fn into_static(self) -> cookie_store::Cookie<'static> {
match self {
CookieInner::Borrowed(v) => v.clone().into_owned(),
CookieInner::Owned(v) => v.into_owned(),
}
}
}

impl<'a> Cookie<'a> {
/// Parses a new [`Cookie`] from a string
pub fn parse<S>(cookie_str: S, uri: &Uri) -> Result<Cookie<'a>, Error>
Expand Down Expand Up @@ -139,6 +148,19 @@ impl<'a> CookieJar<'a> {
*self.0 = store;
Ok(())
}

pub(crate) fn store_response_cookies<'b>(
&mut self,
iter: impl Iterator<Item = Cookie<'b>>,
uri: &Uri,
) {
let url = uri.try_into_url().expect("uri to be a url");
let raw_cookies = iter.map(|c| c.0.into_static().into());
self.0.store_response_cookies(raw_cookies, &url);
}

/// Release the cookie jar.
pub fn release(self) {}
}

impl SharedCookieJar {
Expand Down
39 changes: 39 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -687,6 +687,45 @@ pub(crate) mod test {
assert_eq!(err.to_string(), "http: invalid uri character");
}

#[test]
#[cfg(all(feature = "cookies", feature = "_test"))]
fn store_response_cookies() {
let agent = Agent::new_with_defaults();
let _ = agent.get("https://www.google.com").call().unwrap();

let mut all: Vec<_> = agent
.cookie_jar_lock()
.iter()
.map(|c| c.name().to_string())
.collect();

all.sort();

assert_eq!(all, ["AEC", "__Secure-ENID"])
}

#[test]
#[cfg(all(feature = "cookies", feature = "_test"))]
fn send_request_cookies() {
init_test_log();

let agent = Agent::new_with_defaults();
let uri = Uri::from_static("http://cookie.test/cookie-test");
let uri2 = Uri::from_static("http://cookie2.test/cookie-test");

let mut jar = agent.cookie_jar_lock();
jar.insert(Cookie::parse("a=1", &uri).unwrap(), &uri)
.unwrap();
jar.insert(Cookie::parse("b=2", &uri).unwrap(), &uri)
.unwrap();
jar.insert(Cookie::parse("c=3", &uri2).unwrap(), &uri2)
.unwrap();

jar.release();

let _ = agent.get("http://cookie.test/cookie-test").call().unwrap();
}

// This doesn't need to run, just compile.
fn _ensure_send_sync() {
fn is_send(_t: impl Send) {}
Expand Down
14 changes: 14 additions & 0 deletions src/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,20 @@ fn flow_run(

info!("{:?}", DebugResponse(&response));

#[cfg(feature = "cookies")]
{
let mut jar = agent.cookie_jar_lock();

let iter = response
.headers()
.get_all(http::header::SET_COOKIE)
.iter()
.filter_map(|h| h.to_str().ok())
.filter_map(|s| crate::Cookie::parse(s, &uri).ok());

jar.store_response_cookies(iter, &uri);
}

let ret = match response_result {
RecvResponseResult::RecvBody(flow) => {
let timings = mem::take(timings);
Expand Down
32 changes: 32 additions & 0 deletions src/transport/test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,14 @@ fn setup_default_handlers(handlers: &mut Vec<TestHandler>) {
w,
"HTTP/1.1 200 OK\r\n\
Content-Type: text/html;charset=ISO-8859-1\r\n\
set-cookie: AEC=AVYB7cpadYFS8ZgaioQ17NnxHl1QcSQ_2aH2WEIg1KGDXD5kjk2HhpGVhfk; \
expires=Mon, 14-Apr-2025 17:23:39 GMT; path=/; domain=.google.com; \
Secure; HttpOnly; SameSite=lax\r\n\
set-cookie: __Secure-ENID=23.SE=WaDe-mOBoV2nk-IwHr73boNt6dYcjzQh1X_k8zv2UmUXBL\
m80a3pzLJyx1N1NOqBxDDOR8OJyvuNYw5phFf0VnbqzVtcKPijo2FY8O_vymzyc7x2VwFhGlgU\
WXSWYinjWL7Zvz_EOcA4kfnEXweW5ZDzLrvaLuBIrz5CA_-454AMIXpDiZAVPChCawbkzMptAr\
lMTikkon2EQVXsicqq1XnrMEMPZR5Ld2JC6lpBM8A; expires=Sun, 16-Nov-2025 09:41:57 \
GMT; path=/; domain=.google.com; Secure; HttpOnly; SameSite=lax\r\n\
\r\n\
ureq test server here"
)
Expand Down Expand Up @@ -291,6 +299,30 @@ fn setup_default_handlers(handlers: &mut Vec<TestHandler>) {
handlers,
);

maybe_add(
TestHandler::new("/cookie-test", |_uri, req, w| {
let mut all: Vec<_> = req
.headers()
.get_all("cookie")
.iter()
.map(|c| c.to_str().unwrap())
.collect();

all.sort();

assert_eq!(all, ["a=1;b=2"]);

write!(
w,
"HTTP/1.1 200 OK\r\n\
content-length: 2\r\n\
\r\n\
ok",
)
}),
handlers,
);

#[cfg(feature = "charset")]
{
let (cow, _, _) =
Expand Down
Loading