acick_util/service/
cookie.rs1use std::convert::TryFrom as _;
2use std::fs::File;
3use std::io::{BufReader, Seek as _, SeekFrom};
4
5use anyhow::Context as _;
6use cookie::Cookie as RawCookie;
7use cookie_store::CookieStore;
8use fs2::FileExt as _;
9use reqwest::blocking::{Request, Response};
10use reqwest::header::{HeaderValue, COOKIE, SET_COOKIE};
11
12use crate::abs_path::AbsPathBuf;
13use crate::{Error, Result};
14
15pub struct CookieStorage {
16 file: File,
17 store: CookieStore,
18}
19
20impl CookieStorage {
21 pub fn open(path: &AbsPathBuf) -> Result<Self> {
22 let file = path
23 .create_dir_all_and_open(true, true)
24 .context("Could not open cookies file")?;
25 file.try_lock_exclusive()
26 .context("Could not lock cookies file")?;
27 let reader = BufReader::new(&file);
28 let store = CookieStore::load_json(reader).map_err(Error::msg)?;
29 Ok(Self { file, store })
30 }
31
32 pub fn load_into(&self, request: &mut Request) -> Result<()> {
33 let url = request.url();
34 let cookies = self.store.get_request_cookies(url).map(|rc| rc.to_string());
35 for cookie in cookies {
36 request
37 .headers_mut()
38 .append(COOKIE, HeaderValue::try_from(cookie)?);
39 }
40 Ok(())
41 }
42
43 pub fn store_from(&mut self, response: &Response) -> Result<()> {
44 let cookies = response
45 .headers()
46 .get_all(SET_COOKIE)
47 .iter()
48 .filter_map(|val| {
49 val.to_str().ok().and_then(|cookie_str| {
50 match RawCookie::parse(cookie_str.to_owned()) {
51 Ok(raw_cookie) => Some(raw_cookie),
52 Err(_) => None,
53 }
54 })
55 });
56 let url = response.url();
57 self.store.store_response_cookies(cookies, url);
58 self.save().context("Could not save cookies to json file")
59 }
60
61 pub fn save(&mut self) -> Result<()> {
62 self.file.seek(SeekFrom::Start(0))?;
63 self.file.set_len(0)?;
64 self.store.save_json(&mut self.file).map_err(Error::msg)
65 }
66}
67
68impl Drop for CookieStorage {
69 fn drop(&mut self) {
70 self.file.unlock().expect("Could no unlock cookies file");
71 }
72}