1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

use std::collections::HashMap;

use crate::CacheManager;

use http_cache_semantics::CachePolicy;
use reqwest::{
    header::{HeaderName, HeaderValue},
    Request, Response,
};
use serde::{Deserialize, Serialize};

type Result<T> = std::result::Result<T, anyhow::Error>;

/// Implements [`CacheManager`] with [`cacache`](https://github.com/zkat/cacache-rs) as the backend.
#[derive(Debug, Clone)]
pub struct CACacheManager {
    /// Directory where the cache will be stored.
    pub path: String,
}

impl Default for CACacheManager {
    fn default() -> Self {
        CACacheManager {
            path: "./reqwest-cacache".into(),
        }
    }
}

#[derive(Debug, Deserialize, Serialize)]
struct Store {
    response: StoredResponse,
    policy: CachePolicy,
}

#[derive(Debug, Deserialize, Serialize)]
struct StoredResponse {
    body: Vec<u8>,
    headers: HashMap<String, String>,
}

async fn to_store(res: Response, policy: CachePolicy) -> Result<Store> {
    let mut headers = HashMap::new();
    for header in res.headers() {
        headers.insert(header.0.as_str().to_owned(), header.1.to_str()?.to_owned());
    }
    let body: Vec<u8> = res.bytes().await?.to_vec();
    Ok(Store {
        response: StoredResponse { body, headers },
        policy,
    })
}

fn from_store(store: &Store) -> Result<Response> {
    let mut res = http::Response::builder().body(store.response.body.clone())?;
    for header in &store.response.headers {
        res.headers_mut().insert(
            HeaderName::from_lowercase(header.0.clone().as_str().to_lowercase().as_bytes())?,
            HeaderValue::from_str(header.1.clone().as_str())?,
        );
    }
    Ok(Response::from(res))
}

fn req_key(req: &Request) -> String {
    format!("{}:{}", req.method(), req.url())
}

#[allow(dead_code)]
impl CACacheManager {
    /// Clears out the entire cache.
    pub async fn clear(&self) -> Result<()> {
        cacache::clear(&self.path).await?;
        Ok(())
    }
}

#[async_trait::async_trait]
impl CacheManager for CACacheManager {
    async fn get(&self, req: &Request) -> Result<Option<(Response, CachePolicy)>> {
        let store: Store = match cacache::read(&self.path, &req_key(req)).await {
            Ok(d) => bincode::deserialize(&d)?,
            Err(_e) => {
                return Ok(None);
            }
        };
        Ok(Some((from_store(&store)?, store.policy)))
    }

    // TODO - This needs some reviewing.
    async fn put(&self, req: &Request, res: Response, policy: CachePolicy) -> Result<Response> {
        let status = res.status();
        let version = res.version();
        let headers = res.headers().clone();
        let data = to_store(res, policy).await?;
        let bytes = bincode::serialize(&data)?;
        cacache::write(&self.path, &req_key(req), bytes).await?;
        let mut ret_res = http::Response::builder()
            .status(status)
            .body(data.response.body)?;
        for header in headers {
            ret_res
                .headers_mut()
                .insert(header.0.unwrap(), header.1.clone());
        }
        *ret_res.version_mut() = version;
        Ok(Response::from(ret_res))
    }

    async fn delete(&self, req: &Request) -> Result<()> {
        cacache::remove(&self.path, &req_key(req)).await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use anyhow::Result;
    use http::{Method, Response};
    use reqwest::Request;
    use std::str::FromStr;

    #[tokio::test]
    async fn can_cache_response() -> Result<()> {
        let url = reqwest::Url::from_str("https://example.com")?;
        let res = Response::new("test");
        let res = reqwest::Response::from(res);
        let req = Request::new(Method::GET, url);
        let policy = CachePolicy::new(&req, &res);
        let manager = CACacheManager::default();
        manager.put(&req, res, policy).await?;
        let data = manager.get(&req).await?;
        let body = match data {
            Some(d) => d.0.text().await?,
            None => String::new(),
        };
        assert_eq!(&body, "test");
        manager.delete(&req).await?;
        let data = manager.get(&req).await?;
        assert!(data.is_none());
        manager.clear().await?;
        Ok(())
    }
}