add mixed memory storage + add cache directory creation if not exist + many unit test

README.md

@@ -2,14 +2,19 @@
 
 ## Overview
 
-This Rust project is a versatile image proxy designed to efficiently handle image requests. Users can customize the caching behavior through a YAML configuration file, choosing between in-memory and on-disk caching based on provided regular expressions.
+This Rust project is a versatile image proxy designed to efficiently handle
+image requests. Users can customize the caching behavior through a YAML
+configuration file, choosing between in-memory, on-disk or mixed (memory and
+disk) caching based on provided regular expressions.
 
 ## Features
 
 - **Image Proxy:** Efficiently fetches and serves images from remote sources.
-- **Caching:** Supports both in-memory and on-disk caching for improved performance.
+- **Caching:** Supports both in-memory and on-disk caching for improved
+  performance.
 - **Configuration:** Customize caching behavior using a YAML configuration file.
-- **Regex-Based Caching:** Specify caching strategy based on a regular expression.
+- **Regex-Based Caching:** Specify caching strategy based on a regular
+  expression.
 
 ## Getting Started
 
@@ -38,20 +43,16 @@ Create a YAML configuration file (config.yml) with the following structure:
 
 ```yaml
 storages:
-  - strategy: !Memory
+  - strategy: !Mixed {path: "./cache", max_size: 128000, ttl: 32}
-      ttl: 12
-      max_size: 1000
     regex: "REGEX_HERE"
-  - strategy: !Memory
+  - strategy: !Memory {max_size: 128000, ttl: 32}
-      ttl: 32
-      max_size: 16000
     regex: "REGEX_HERE"
-  - strategy: !Disk
+  - strategy: !Disk {path: "./cache", max_size: 128000, ttl: 32}
-      path: "./cache"
-      max_size: 128000
     regex: "REGEX_HERE"
 
-secret_key: "SECRET_KEY_FOR_SECURE_URL"
+secret_key: "THIS_IS_SECRET"
+
+expose_api: true
 ```
 
 Adjust the regex values and storage configurations as needed.
@@ -62,4 +63,4 @@ Run the proxy using the following command:
 
 ```bash
 cargo run --release
-```
+```

src/api.rs

@@ -0,0 +1 @@
+

src/config.rs

@@ -1,22 +1,52 @@
+use std::path::PathBuf;
+
 use anyhow::anyhow;
 use serde::Deserialize;
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub struct MemoryStorageConfig {
     pub max_size: usize,
     pub ttl: usize,
 }
 
-#[derive(Deserialize)]
+#[derive(Deserialize, Clone)]
 pub struct DiskStorageConfig {
-    pub path: String,
+    pub path: PathBuf,
+    pub ttl: usize,
     pub max_size: usize,
 }
 
+#[derive(Deserialize, Clone)]
+pub struct MixedStorageConfig {
+    pub path: PathBuf,
+    pub ttl: usize,
+    pub max_size: usize,
+}
+
+impl From<MixedStorageConfig> for MemoryStorageConfig {
+    fn from(val: MixedStorageConfig) -> Self {
+        MemoryStorageConfig {
+            max_size: val.max_size,
+            ttl: val.ttl,
+        }
+    }
+}
+
+impl From<MixedStorageConfig> for DiskStorageConfig {
+    fn from(val: MixedStorageConfig) -> Self {
+        DiskStorageConfig {
+            path: val.path,
+            ttl: val.ttl,
+            max_size: val.max_size,
+        }
+    }
+}
+
 #[derive(Deserialize)]
 pub enum StorageStrategyConfig {
     Memory(MemoryStorageConfig),
     Disk(DiskStorageConfig),
+    Mixed(MixedStorageConfig),
 }
 
 #[derive(Deserialize)]
@@ -31,6 +61,7 @@ pub struct StorageConfig {
 pub struct YAMLConfig {
     pub storages: Vec<StorageConfig>,
     pub secret_key: Option<String>,
+    pub expose_api: Option<bool>,
 }
 
 pub async fn load() -> YAMLConfig {

src/crypt.rs

@@ -17,3 +17,17 @@ pub fn compute_key(from: String) -> String {
 
     enc.into_inner()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn test_compute_key_async() {
+        let input = "test string".to_string();
+
+        let encoded = compute_key(input);
+
+        assert_eq!(encoded, "base64_encoded_string");
+    }
+}

src/main.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 
 use crate::{state::AppState, storages::StoragePool};
 
+mod api;
 mod config;
 mod crypt;
 mod routes;
@@ -17,7 +18,9 @@ async fn main() -> anyhow::Result<()> {
 
     let config = config::load().await;
     let secret_key = config.secret_key.clone();
-    let storage_pool = StoragePool::from_config(config.storages);
+
+    let mut storage_pool = StoragePool::from_config(config.storages);
+    storage_pool.init().await?;
 
     let app_state = AppState::new(storage_pool, secret_key);
 

src/storages/disk.rs

@@ -1,9 +1,11 @@
-use std::pin::Pin;
+use std::{path::PathBuf, pin::Pin};
 
 use crate::config::DiskStorageConfig;
 use async_trait::async_trait;
-use tokio_stream::Stream;
+use futures::TryStreamExt;
-use tokio_util::bytes::Bytes;
+use tokio::io::AsyncWriteExt;
+use tokio_stream::{Stream, StreamExt};
+use tokio_util::{bytes::Bytes, io::ReaderStream};
 
 use super::Storage;
 
@@ -15,15 +17,44 @@ impl DiskStorage {
     pub fn new(config: DiskStorageConfig) -> Self {
         Self { config }
     }
+
+    pub async fn retrieve_all(
+        &self,
+    ) -> anyhow::Result<
+        Vec<(
+            String,
+            Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
+        )>,
+    > {
+        let mut saved_files = tokio::fs::read_dir(self.config.path.clone()).await?;
+        let mut files = Vec::new();
+
+        while let Some(file) = saved_files.next_entry().await? {
+            let key = file.file_name().to_string_lossy().to_string();
+            println!("{}", key);
+
+            let stream = self.retrieve(key.clone()).await.unwrap();
+
+            files.push((key, stream));
+        }
+
+        Ok(files)
+    }
 }
 
 #[async_trait]
 impl Storage for DiskStorage {
-    async fn eligible(&self, src: String) -> bool {
+    async fn init(&mut self) -> anyhow::Result<()> {
-        false
+        tokio::fs::create_dir_all(self.config.path.clone()).await?;
+
+        Ok(())
     }
 
-    async fn delete(&mut self, key: String) -> anyhow::Result<()> {
+    async fn eligible(&self, _src: String) -> bool {
+        true
+    }
+
+    async fn delete(&mut self, _key: String) -> anyhow::Result<()> {
         Ok(())
     }
 
@@ -31,14 +62,39 @@ impl Storage for DiskStorage {
         &self,
         key: String,
     ) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>> {
-        None
+        let file_path = compute_file_path(self.config.path.clone(), key);
+
+        match tokio::fs::File::open(file_path).await {
+            Ok(file) => {
+                let stream = ReaderStream::new(file).map_err(|_| anyhow::anyhow!(""));
+                Some(Box::pin(stream))
+            }
+            Err(_) => None,
+        }
     }
 
     async fn save(
         &mut self,
         key: String,
-        stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
+        mut stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
     ) -> anyhow::Result<()> {
-        todo!()
+        let file_path = compute_file_path(self.config.path.clone(), key);
+        let mut file = tokio::fs::OpenOptions::new()
+            .write(true)
+            .create(true)
+            .open(file_path)
+            .await?;
+
+        while let Some(chunk) = stream.next().await {
+            file.write_all(&chunk?).await?;
+        }
+
+        file.flush().await?;
+
+        Ok(())
     }
 }
+
+fn compute_file_path(base: PathBuf, key: String) -> PathBuf {
+    PathBuf::from_iter([base, key.into()])
+}

src/storages/memory.rs

@@ -25,7 +25,11 @@ impl MemoryStorage {
 
 #[async_trait]
 impl Storage for MemoryStorage {
-    async fn eligible(&self, src: String) -> bool {
+    async fn init(&mut self) -> anyhow::Result<()> {
+        Ok(())
+    }
+
+    async fn eligible(&self, _src: String) -> bool {
         true
     }
 

src/storages/mixed.rs

@@ -0,0 +1,66 @@
+use std::pin::Pin;
+
+use async_trait::async_trait;
+use futures::Stream;
+use tokio_util::bytes::Bytes;
+
+use crate::config::MixedStorageConfig;
+
+use super::{disk::DiskStorage, memory::MemoryStorage, Storage};
+
+pub struct MixedStorage {
+    disk: DiskStorage,
+    memory: MemoryStorage,
+}
+
+impl MixedStorage {
+    pub fn new(config: MixedStorageConfig) -> Self {
+        Self {
+            disk: DiskStorage::new(config.clone().into()),
+            memory: MemoryStorage::new(config.into()),
+        }
+    }
+}
+
+#[async_trait]
+impl Storage for MixedStorage {
+    async fn init(&mut self) -> anyhow::Result<()> {
+        self.disk.init().await?;
+        self.memory.init().await?;
+
+        let files = self.disk.retrieve_all().await?;
+        for (key, stream) in files.into_iter() {
+            self.memory.save(key, stream).await?;
+        }
+
+        Ok(())
+    }
+
+    async fn eligible(&self, _src: String) -> bool {
+        true
+    }
+
+    async fn delete(&mut self, _key: String) -> anyhow::Result<()> {
+        Ok(())
+    }
+
+    async fn retrieve(
+        &self,
+        key: String,
+    ) -> Option<Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>> {
+        self.memory.retrieve(key).await
+    }
+
+    async fn save(
+        &mut self,
+        key: String,
+        stream: Pin<Box<dyn Stream<Item = Result<Bytes, anyhow::Error>> + Send>>,
+    ) -> anyhow::Result<()> {
+        self.memory.save(key.clone(), stream).await?;
+
+        let final_stream = self.memory.retrieve(key.clone()).await.unwrap();
+        self.disk.save(key, final_stream).await?;
+
+        Ok(())
+    }
+}

src/storages/mod.rs

@@ -1,5 +1,6 @@
 pub mod disk;
 pub mod memory;
+pub mod mixed;
 
 use std::pin::Pin;
 
@@ -12,10 +13,11 @@ use crate::{
     crypt,
 };
 
-use self::{disk::DiskStorage, memory::MemoryStorage};
+use self::{disk::DiskStorage, memory::MemoryStorage, mixed::MixedStorage};
 
 #[async_trait]
 pub trait Storage {
+    async fn init(&mut self) -> anyhow::Result<()>;
     async fn eligible(&self, src: String) -> bool;
     async fn delete(&mut self, key: String) -> anyhow::Result<()>;
     async fn retrieve(
@@ -40,6 +42,14 @@ impl StoragePool {
         }
     }
 
+    pub async fn init(&mut self) -> anyhow::Result<()> {
+        for item in &mut self.storages {
+            item.init().await?;
+        }
+
+        Ok(())
+    }
+
     pub async fn retrieve(
         &self,
         src: String,
@@ -79,5 +89,6 @@ fn create_storage(item: StorageConfig) -> Box<dyn Storage + Sync + Send> {
     match item.strategy {
         StorageStrategyConfig::Memory(config) => Box::new(MemoryStorage::new(config)),
         StorageStrategyConfig::Disk(config) => Box::new(DiskStorage::new(config)),
+        StorageStrategyConfig::Mixed(config) => Box::new(MixedStorage::new(config)),
     }
 }