From ddd8a5395dfe6daf4b464003450465644594b158 Mon Sep 17 00:00:00 2001 From: Martin Algesten Date: Wed, 16 Oct 2024 19:01:03 +0200 Subject: [PATCH] Store response cookies --- src/agent.rs | 16 ++++++++++++---- src/cookies.rs | 22 ++++++++++++++++++++++ src/lib.rs | 39 +++++++++++++++++++++++++++++++++++++++ src/run.rs | 14 ++++++++++++++ src/transport/test.rs | 32 ++++++++++++++++++++++++++++++++ 5 files changed, 119 insertions(+), 4 deletions(-) diff --git a/src/agent.rs b/src/agent.rs index f94b82cc..06c682d5 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -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; @@ -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() } diff --git a/src/cookies.rs b/src/cookies.rs index a2adcf92..ac528257 100644 --- a/src/cookies.rs +++ b/src/cookies.rs @@ -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(cookie_str: S, uri: &Uri) -> Result, Error> @@ -139,6 +148,19 @@ impl<'a> CookieJar<'a> { *self.0 = store; Ok(()) } + + pub(crate) fn store_response_cookies<'b>( + &mut self, + iter: impl Iterator>, + 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 { diff --git a/src/lib.rs b/src/lib.rs index d8675daf..fa2db945 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -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) {} diff --git a/src/run.rs b/src/run.rs index 2d313f9a..ea6414c9 100644 --- a/src/run.rs +++ b/src/run.rs @@ -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); diff --git a/src/transport/test.rs b/src/transport/test.rs index 2ad48e44..fc6af76b 100644 --- a/src/transport/test.rs +++ b/src/transport/test.rs @@ -143,6 +143,14 @@ fn setup_default_handlers(handlers: &mut Vec) { 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" ) @@ -291,6 +299,30 @@ fn setup_default_handlers(handlers: &mut Vec) { 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, _, _) =