Store response cookies

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,16 @@ 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);
+    }
 }
 
 impl SharedCookieJar {

src/run.rs

@@ -7,7 +7,7 @@ use hoot::client::flow::state::{Prepare, SendBody as SendBodyState};
 use hoot::client::flow::{Await100Result, RecvBodyResult, RecvResponseResult, SendRequestResult};
 use hoot::BodyMode;
 use http::uri::Scheme;
-use http::{HeaderValue, Request, Response, Uri};
+use http::{header, HeaderValue, Request, Response, Uri};
 
 use crate::body::ResponseInfo;
 use crate::config::{Config, RequestLevelConfig};
@@ -149,6 +149,20 @@ fn flow_run(
 
     info!("{:?}", DebugResponse(&response));
 
+    #[cfg(feature = "cookies")]
+    {
+        let mut jar = agent.cookie_jar();
+
+        let iter = response
+            .headers()
+            .get_all(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);