Skip to content

Commit

Permalink
Store response cookies
Browse files Browse the repository at this point in the history
  • Loading branch information
algesten committed Oct 16, 2024
1 parent 35910fb commit 3994caf
Show file tree
Hide file tree
Showing 2 changed files with 33 additions and 0 deletions.
19 changes: 19 additions & 0 deletions src/cookies.rs
Original file line number Diff line number Diff line change
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,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 {
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();

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

0 comments on commit 3994caf

Please sign in to comment.