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

use std::collections::HashMap;

use crate::CacheManager;

use http_cache_semantics::CachePolicy;
use serde::{Deserialize, Serialize};
use surf::{Request, Response};

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

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

impl Default for CACacheManager {
    fn default() -> Self {
        CACacheManager {
            path: "./surf-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: &mut Response, policy: CachePolicy) -> Result<Store> {
    let mut headers = HashMap::new();
    for header in res.iter() {
        headers.insert(header.0.as_str().to_owned(), header.1.as_str().to_owned());
    }
    let body: Vec<u8> = res.body_bytes().await?;
    Ok(Store {
        response: StoredResponse { body, headers },
        policy,
    })
}

fn from_store(store: &Store) -> Response {
    let mut res = http_types::Response::new(http_types::StatusCode::Ok);
    for header in &store.response.headers {
        let val =
            http_types::headers::HeaderValue::from_bytes(header.1.as_bytes().to_vec()).unwrap();
        res.insert_header(header.0.as_str(), val);
    }
    res.set_body(store.response.body.clone());
    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(())
    }
}

#[surf::utils::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: &mut Response,
        policy: CachePolicy,
    ) -> Result<Response> {
        let data = to_store(res, policy).await?;
        let bytes = bincode::serialize(&data).unwrap();
        cacache::write(&self.path, &req_key(req), bytes).await?;
        let mut ret_res = http_types::Response::new(res.status());
        ret_res.set_body(data.response.body);
        for header in res.iter() {
            ret_res.insert_header(header.0, header.1);
        }
        ret_res.set_version(res.version());
        Ok(Response::from(ret_res))
    }

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{get_request_parts, get_response_parts};
    use http_types::{Method, Response, StatusCode};
    use std::str::FromStr;
    use surf::{Request, Result};

    #[async_std::test]
    async fn can_cache_response() -> Result<()> {
        let url = surf::http::Url::from_str("https://example.com")?;
        let mut res = Response::new(StatusCode::Ok);
        res.set_body("test");
        let mut res = surf::Response::from(res);
        let req = Request::new(Method::Get, url);
        let policy = CachePolicy::new(&get_request_parts(&req), &get_response_parts(&res));
        let manager = CACacheManager::default();
        manager.put(&req, &mut res, policy).await?;
        let data = manager.get(&req).await?;
        let body = match data {
            Some(mut d) => d.0.body_string().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(())
    }
}