more cache code

1 files changed, 38 insertions, 17 deletions

diff --git a/src/modules/cache.rs b/src/modules/cache.rs
index f5cd72d..34db902 100644
--- a/src/modules/cache.rs
+++ b/src/modules/cache.rs
@@ -18,12 +18,14 @@ use super::{Node, NodeContext, NodeKind, NodeRequest, NodeResponse};
 use crate::{config::DynNode, error::ServiceError};
 use anyhow::Result;
 use bytes::Bytes;
+use headers::{CacheControl, HeaderMapExt};
 use http::Response;
 use http_body_util::{BodyExt, Full};
+use log::debug;
 use serde::Deserialize;
 use serde_yml::Value;
 use sha2::{Digest, Sha256};
-use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
+use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc, time::Duration};
 use tokio::sync::RwLock;
 
 pub struct CacheKind;
@@ -57,21 +59,24 @@ impl Node for Cache {
         request: NodeRequest,
     ) -> Pin<Box<dyn Future<Output = Result<NodeResponse, ServiceError>> + Send + Sync + 'a>> {
         Box::pin(async move {
+            let allow_cache = request.method().is_safe();
+
+            if !allow_cache {
+                return self.config.next.handle(context, request).await;
+            }
+
             // not very fast
             let mut hasher = Sha256::new();
             hasher.update(request.method().as_str().len().to_be_bytes());
             hasher.update(request.method().as_str());
             hasher.update(request.uri().path().len().to_be_bytes());
             hasher.update(request.uri().path());
-            hasher.update(if request.uri().query().is_some() {
-                [1]
-            } else {
-                [0]
-            });
+            hasher.update([request.uri().query().is_some() as u8]);
             if let Some(q) = request.uri().query() {
                 hasher.update(q.len().to_be_bytes());
                 hasher.update(q);
             }
+            // TODO which headers are important to caching?
             hasher.update(request.headers().len().to_be_bytes());
             for (k, v) in request.headers() {
                 hasher.update(k.as_str().len().to_be_bytes());
@@ -82,27 +87,43 @@ impl Node for Cache {
             let key: [u8; 32] = hasher.finalize().into();
 
             if let Some(resp) = self.entries.read().await.get(&key) {
+                debug!("hit");
                 return Ok(resp
                     .to_owned()
                     .map(|b| Full::new(b).map_err(|e| match e {}).boxed()));
             }
+            debug!("miss");
 
             let response = self.config.next.handle(context, request).await?;
 
-            let h = response.headers().to_owned();
-            let s = response.status().to_owned();
-            let body = response.collect().await?.to_bytes();
+            let cache_control = response
+                .headers()
+                .typed_get::<CacheControl>()
+                .unwrap_or(CacheControl::new().with_max_age(Duration::from_secs(120))); // TODO what is the correct default?
 
-            let mut r1 = Response::new(Full::new(body.clone()).map_err(|e| match e {}).boxed());
-            *r1.headers_mut() = h.clone();
-            *r1.status_mut() = s;
+            let allow_store = !cache_control.no_store() && !cache_control.no_cache();
 
-            let mut r2 = Response::new(body);
-            *r2.headers_mut() = h;
-            *r2.status_mut() = s;
-            self.entries.write().await.insert(key, r2);
+            let response = if allow_store {
+                debug!("store");
+                let h = response.headers().to_owned();
+                let s = response.status().to_owned();
+                let body = response.collect().await?.to_bytes();
+
+                let mut r1 = Response::new(Full::new(body.clone()).map_err(|e| match e {}).boxed());
+                *r1.headers_mut() = h.clone();
+                *r1.status_mut() = s;
+
+                let mut r2 = Response::new(body);
+                *r2.headers_mut() = h;
+                *r2.status_mut() = s;
+                self.entries.write().await.insert(key, r2);
+                r1
+            } else {
+                debug!("no store");
+                response
+            };
 
-            Ok(r1)
+            Ok(response)
         })
     }
 }