forked/reqwest-impersonate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add CookieStore trait and expose default Jar (#1203)

Browse Source 
This adds a new trait, `CookieStore`, which allows for custom
implementations of storage for a client's cookies. After implementing,
you can call `ClientBuilder::cookie_provider` to use it over the
default storage.

The default store is exposed as `Jar`, to ease the most common use case
which is to add a few cookies to the store when creating a client.

Co-authored-by: Patrick Fernie <patrick.fernie@gmail.com>
This commit is contained in:
Sean McArthur
2021-03-08 15:40:58 -08:00
committed by GitHub
parent 2414042269
commit 12d7905520
3 changed files with 140 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
src/async_impl/client.rs
View File

@@ -3,8 +3,6 @@ use std::any::Any;
use std::convert::TryInto;
use std::net::IpAddr;
use std::sync::Arc;
#[cfg(feature = "cookies")]
use std::sync::RwLock;
use std::time::Duration;
use std::{fmt, str};
@@ -105,7 +103,7 @@ struct Config {
local_address: Option<IpAddr>,
nodelay: bool,
#[cfg(feature = "cookies")]
cookie_store: Option<cookie::CookieStore>,
cookie_store: Option<Arc<dyn cookie::CookieStore>>,
trust_dns: bool,
error: Option<crate::Error>,
https_only: bool,
@@ -350,7 +348,7 @@ impl ClientBuilder {
inner: Arc::new(ClientRef {
accepts: config.accepts,
#[cfg(feature = "cookies")]
cookie_store: config.cookie_store.map(RwLock::new),
cookie_store: config.cookie_store,
hyper: hyper_client,
headers: config.headers,
redirect_policy: config.redirect_policy,
@@ -464,11 +462,31 @@ impl ClientBuilder {
#[cfg(feature = "cookies")]
#[cfg_attr(docsrs, doc(cfg(feature = "cookies")))]
pub fn cookie_store(mut self, enable: bool) -> ClientBuilder {
self.config.cookie_store = if enable {
Some(cookie::CookieStore::default())
if enable {
self.cookie_provider(Arc::new(cookie::Jar::default()))
} else {
None
};
self.config.cookie_store = None;
self
}
}
/// Set the persistent cookie store for the client.
///
/// Cookies received in responses will be passed to this store, and
/// additional requests will query this store for cookies.
///
/// By default, no cookie store is used.
///
/// # Optional
///
/// This requires the optional `cookies` feature to be enabled.
#[cfg(feature = "cookies")]
#[cfg_attr(docsrs, doc(cfg(feature = "cookies")))]
pub fn cookie_provider<C: cookie::CookieStore + 'static>(
mut self,
cookie_store: Arc<C>,
) -> ClientBuilder {
self.config.cookie_store = Some(cookie_store as _);
self
}
@@ -1109,10 +1127,9 @@ impl Client {
// Add cookies from the cookie store.
#[cfg(feature = "cookies")]
{
if let Some(cookie_store_wrapper) = self.inner.cookie_store.as_ref() {
if let Some(cookie_store) = self.inner.cookie_store.as_ref() {
if headers.get(crate::header::COOKIE).is_none() {
let cookie_store = cookie_store_wrapper.read().unwrap();
add_cookie_header(&mut headers, &cookie_store, &url);
add_cookie_header(&mut headers, &**cookie_store, &url);
}
}
}
@@ -1289,7 +1306,7 @@ impl Config {
struct ClientRef {
accepts: Accepts,
#[cfg(feature = "cookies")]
cookie_store: Option<RwLock<cookie::CookieStore>>,
cookie_store: Option<Arc<dyn cookie::CookieStore>>,
headers: HeaderMap,
hyper: HyperClient,
redirect_policy: redirect::Policy,
@@ -1431,14 +1448,11 @@ impl Future for PendingRequest {
#[cfg(feature = "cookies")]
{
if let Some(store_wrapper) = self.client.cookie_store.as_ref() {
let mut cookies = cookie::extract_response_cookies(&res.headers())
.filter_map(|res| res.ok())
.map(|cookie| cookie.into_inner().into_owned())
.peekable();
if let Some(ref cookie_store) = self.client.cookie_store {
let mut cookies =
cookie::extract_response_cookie_headers(&res.headers()).peekable();
if cookies.peek().is_some() {
let mut store = store_wrapper.write().unwrap();
store.0.store_response_cookies(cookies, &self.url);
cookie_store.set_cookies(&mut cookies, &self.url);
}
}
}
@@ -1531,11 +1545,8 @@ impl Future for PendingRequest {
// Add cookies from the cookie store.
#[cfg(feature = "cookies")]
{
if let Some(cookie_store_wrapper) =
self.client.cookie_store.as_ref()
{
let cookie_store = cookie_store_wrapper.read().unwrap();
add_cookie_header(&mut headers, &cookie_store, &self.url);
if let Some(ref cookie_store) = self.client.cookie_store {
add_cookie_header(&mut headers, &**cookie_store, &self.url);
}
}
@@ -1592,18 +1603,9 @@ fn make_referer(next: &Url, previous: &Url) -> Option<HeaderValue> {
}
#[cfg(feature = "cookies")]
fn add_cookie_header(headers: &mut HeaderMap, cookie_store: &cookie::CookieStore, url: &Url) {
let header = cookie_store
.0
.get_request_cookies(url)
.map(|c| format!("{}={}", c.name(), c.value()))
.collect::<Vec<_>>()
.join("; ");
if !header.is_empty() {
headers.insert(
crate::header::COOKIE,
HeaderValue::from_bytes(header.as_bytes()).unwrap(),
);
fn add_cookie_header(headers: &mut HeaderMap, cookie_store: &dyn cookie::CookieStore, url: &Url) {
if let Some(header) = cookie_store.cookies(url) {
headers.insert(crate::header::COOKIE, header);
}
}