acick_util/service/
cookie.rs

1use 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}