very bad basic caching

1 files changed, 22 insertions, 6 deletions

diff --git a/src/modules/cache.rs b/src/modules/cache.rs
index d306ed1..03dc1be 100644
--- a/src/modules/cache.rs
+++ b/src/modules/cache.rs
@@ -14,11 +14,12 @@ use crate::{config::DynNode, error::ServiceError};
 use anyhow::Result;
 use bytes::Bytes;
 use http::Response;
-use http_body_util::BodyExt;
+use http_body_util::{BodyExt, Full};
 use serde::Deserialize;
 use serde_yaml::Value;
 use sha2::{Digest, Sha256};
 use std::{collections::HashMap, future::Future, pin::Pin, sync::Arc};
+use tokio::sync::RwLock;
 
 pub struct CacheKind;
 
@@ -28,7 +29,7 @@ struct CacheConfig {
 }
 
 struct Cache {
-    entries: HashMap<[u8; 32], Response<Bytes>>,
+    entries: RwLock<HashMap<[u8; 32], Response<Bytes>>>,
     config: CacheConfig,
 }
 
@@ -40,7 +41,7 @@ impl NodeKind for CacheKind {
         let config = serde_yaml::from_value::<CacheConfig>(config)?;
         Ok(Arc::new(Cache {
             config,
-            entries: HashMap::new(),
+            entries: HashMap::new().into(),
         }))
     }
 }
@@ -60,13 +61,28 @@ impl Node for Cache {
             }
             let key: [u8; 32] = hasher.finalize().try_into().unwrap();
 
-            if let Some(resp) = self.entries.get(&key) {
-                // Ok(resp.to_owned().map(|b| b.map_err(|e| match e {}).boxed()))
+            if let Some(resp) = self.entries.read().await.get(&key) {
+                return Ok(resp
+                    .to_owned()
+                    .map(|b| Full::new(b).map_err(|e| match e {}).boxed()));
             }
 
             let response = self.config.next.handle(context, request).await?;
 
-            Ok(response)
+            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.clone();
+
+            let mut r2 = Response::new(body);
+            *r2.headers_mut() = h;
+            *r2.status_mut() = s;
+            self.entries.write().await.insert(key, r2);
+
+            Ok(r1)
         })
     }
 }