From 54d77c82991d89f43dcee48d3d692c17d9821f90 Mon Sep 17 00:00:00 2001 From: qpismont Date: Tue, 2 Jan 2024 22:15:00 +0100 Subject: [PATCH] add mixed memory storage + add cache directory creation if not exist + many unit test --- README.md | 29 +++++++++-------- src/api.rs | 1 + src/config.rs | 37 +++++++++++++++++++-- src/crypt.rs | 14 ++++++++ src/main.rs | 5 ++- src/storages/disk.rs | 74 +++++++++++++++++++++++++++++++++++++----- src/storages/memory.rs | 6 +++- src/storages/mixed.rs | 66 +++++++++++++++++++++++++++++++++++++ src/storages/mod.rs | 13 +++++++- 9 files changed, 216 insertions(+), 29 deletions(-) create mode 100644 src/api.rs create mode 100644 src/storages/mixed.rs diff --git a/README.md b/README.md index 1929dd0..c7c8159 100644 --- a/README.md +++ b/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 - ttl: 12 - max_size: 1000 + - strategy: !Mixed {path: "./cache", max_size: 128000, ttl: 32} regex: "REGEX_HERE" - - strategy: !Memory - ttl: 32 - max_size: 16000 + - strategy: !Memory {max_size: 128000, ttl: 32} regex: "REGEX_HERE" - - strategy: !Disk - path: "./cache" - max_size: 128000 + - strategy: !Disk {path: "./cache", max_size: 128000, ttl: 32} 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 -``` \ No newline at end of file +``` diff --git a/src/api.rs b/src/api.rs new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/src/api.rs @@ -0,0 +1 @@ + diff --git a/src/config.rs b/src/config.rs index d34002b..4440108 100644 --- a/src/config.rs +++ b/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 for MemoryStorageConfig { + fn from(val: MixedStorageConfig) -> Self { + MemoryStorageConfig { + max_size: val.max_size, + ttl: val.ttl, + } + } +} + +impl From 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, pub secret_key: Option, + pub expose_api: Option, } pub async fn load() -> YAMLConfig { diff --git a/src/crypt.rs b/src/crypt.rs index e6d3acf..62e5e25 100644 --- a/src/crypt.rs +++ b/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"); + } +} diff --git a/src/main.rs b/src/main.rs index be55c00..e18df9c 100644 --- a/src/main.rs +++ b/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); diff --git a/src/storages/disk.rs b/src/storages/disk.rs index c369270..ffa0cfd 100644 --- a/src/storages/disk.rs +++ b/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 tokio_util::bytes::Bytes; +use futures::TryStreamExt; +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> + 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 { - false + async fn init(&mut self) -> anyhow::Result<()> { + 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> + 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> + Send>>, + mut stream: Pin> + 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()]) +} diff --git a/src/storages/memory.rs b/src/storages/memory.rs index e2ccf4e..02d00d0 100644 --- a/src/storages/memory.rs +++ b/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 } diff --git a/src/storages/mixed.rs b/src/storages/mixed.rs new file mode 100644 index 0000000..f0c55d4 --- /dev/null +++ b/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> + Send>>> { + self.memory.retrieve(key).await + } + + async fn save( + &mut self, + key: String, + stream: Pin> + 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(()) + } +} diff --git a/src/storages/mod.rs b/src/storages/mod.rs index c6ec435..9a0e7c2 100644 --- a/src/storages/mod.rs +++ b/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 { 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)), } }