|
| 1 | +// taken from https://github.com/pfernie/reqwest_cookie_store/blob/2ec4afabcd55e24d3afe3f0626ee6dc97bed938d/src/lib.rs |
| 2 | + |
| 3 | +use std::sync::{Mutex, MutexGuard, PoisonError}; |
| 4 | + |
| 5 | +use cookie_store::{CookieStore, RawCookie, RawCookieParseError}; |
| 6 | +use reqwest::header::HeaderValue; |
| 7 | +use serde::{Deserialize, Serialize}; |
| 8 | + |
| 9 | +fn set_cookies( |
| 10 | + cookie_store: &mut CookieStore, |
| 11 | + cookie_headers: &mut dyn Iterator<Item = &HeaderValue>, |
| 12 | + url: &url::Url, |
| 13 | +) { |
| 14 | + let cookies = cookie_headers.filter_map(|val| { |
| 15 | + std::str::from_utf8(val.as_bytes()) |
| 16 | + .map_err(RawCookieParseError::from) |
| 17 | + .and_then(RawCookie::parse) |
| 18 | + .map(|c| c.into_owned()) |
| 19 | + .ok() |
| 20 | + }); |
| 21 | + cookie_store.store_response_cookies(cookies, url); |
| 22 | +} |
| 23 | + |
| 24 | +fn cookies(cookie_store: &CookieStore, url: &url::Url) -> Option<HeaderValue> { |
| 25 | + let s = cookie_store |
| 26 | + .get_request_values(url) |
| 27 | + .map(|(name, value)| format!("{}={}", name, value)) |
| 28 | + .collect::<Vec<_>>() |
| 29 | + .join("; "); |
| 30 | + |
| 31 | + if s.is_empty() { |
| 32 | + return None; |
| 33 | + } |
| 34 | + |
| 35 | + HeaderValue::from_maybe_shared(bytes::Bytes::from(s)).ok() |
| 36 | +} |
| 37 | + |
| 38 | +/// A [`cookie_store::CookieStore`] wrapped internally by a [`std::sync::Mutex`], suitable for use in |
| 39 | +/// async/concurrent contexts. |
| 40 | +#[derive(Debug, Serialize, Deserialize)] |
| 41 | +pub struct CookieStoreMutex(Mutex<CookieStore>); |
| 42 | + |
| 43 | +impl Default for CookieStoreMutex { |
| 44 | + /// Create a new, empty [`CookieStoreMutex`] |
| 45 | + fn default() -> Self { |
| 46 | + CookieStoreMutex::new(CookieStore::default()) |
| 47 | + } |
| 48 | +} |
| 49 | + |
| 50 | +impl CookieStoreMutex { |
| 51 | + /// Create a new [`CookieStoreMutex`] from an existing [`cookie_store::CookieStore`]. |
| 52 | + pub const fn new(cookie_store: CookieStore) -> CookieStoreMutex { |
| 53 | + CookieStoreMutex(Mutex::new(cookie_store)) |
| 54 | + } |
| 55 | + |
| 56 | + /// Lock and get a handle to the contained [`cookie_store::CookieStore`]. |
| 57 | + pub fn lock( |
| 58 | + &self, |
| 59 | + ) -> Result<MutexGuard<'_, CookieStore>, PoisonError<MutexGuard<'_, CookieStore>>> { |
| 60 | + self.0.lock() |
| 61 | + } |
| 62 | + |
| 63 | + pub fn load<R: std::io::BufRead>(reader: R) -> cookie_store::Result<CookieStoreMutex> { |
| 64 | + cookie_store::serde::load(reader, |c| serde_json::from_str(c)).map(CookieStoreMutex::new) |
| 65 | + } |
| 66 | + |
| 67 | + pub fn save<W: std::io::Write>(&self, writer: &mut W) -> cookie_store::Result<()> { |
| 68 | + let store = self.lock().expect("poisoned cookie jar mutex"); |
| 69 | + cookie_store::serde::save(&store, writer, |c| serde_json::to_string(c)) |
| 70 | + } |
| 71 | +} |
| 72 | + |
| 73 | +impl reqwest::cookie::CookieStore for CookieStoreMutex { |
| 74 | + fn set_cookies(&self, cookie_headers: &mut dyn Iterator<Item = &HeaderValue>, url: &url::Url) { |
| 75 | + let mut store = self.0.lock().unwrap(); |
| 76 | + set_cookies(&mut store, cookie_headers, url); |
| 77 | + } |
| 78 | + |
| 79 | + fn cookies(&self, url: &url::Url) -> Option<HeaderValue> { |
| 80 | + let store = self.0.lock().unwrap(); |
| 81 | + cookies(&store, url) |
| 82 | + } |
| 83 | +} |
0 commit comments