Store response cookies

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()
     }
 

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<S>(cookie_str: S, uri: &Uri) -> Result<Cookie<'a>, 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<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 {

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) {}

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);

src/transport/test.rs

@@ -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"
             )
@@ -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, _, _) =