diff --git a/.gitignore b/.gitignore index 36616af..55e3e47 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ /target config.yml -.env \ No newline at end of file +.env +.vscode \ No newline at end of file diff --git a/Cargo.lock b/Cargo.lock index 9261c58..3af5a8e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -38,54 +38,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "anstream" -version = "0.6.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ab91ebe16eb252986481c5b62f6098f3b698a45e34b5b98200cf20dd2484a44" -dependencies = [ - "anstyle", - "anstyle-parse", - "anstyle-query", - "anstyle-wincon", - "colorchoice", - "utf8parse", -] - -[[package]] -name = "anstyle" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7079075b41f533b8c61d2a4d073c4676e1f8b249ff94a393b0595db304e0dd87" - -[[package]] -name = "anstyle-parse" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "317b9a89c1868f5ea6ff1d9539a69f45dffc21ce321ac1fd1160dfa48c8e2140" -dependencies = [ - "utf8parse", -] - -[[package]] -name = "anstyle-query" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ca11d4be1bab0c8bc8734a9aa7bf4ee8316d462a08c6ac5052f888fef5b494b" -dependencies = [ - "windows-sys", -] - -[[package]] -name = "anstyle-wincon" -version = "3.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f0699d10d2f4d628a98ee7b57b289abbc98ff3bad977cb3152709d4bf2330628" -dependencies = [ - "anstyle", - "windows-sys", -] - [[package]] name = "anyhow" version = "1.0.75" @@ -275,52 +227,6 @@ dependencies = [ "generic-array", ] -[[package]] -name = "clap" -version = "4.4.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41fffed7514f420abec6d183b1d3acfd9099c79c3a10a06ade4f8203f1411272" -dependencies = [ - "clap_builder", - "clap_derive", -] - -[[package]] -name = "clap_builder" -version = "4.4.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "63361bae7eef3771745f02d8d892bec2fee5f6e34af316ba556e7f97a7069ff1" -dependencies = [ - "anstream", - "anstyle", - "clap_lex", - "strsim", -] - -[[package]] -name = "clap_derive" -version = "4.4.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf9804afaaf59a91e75b022a30fb7229a7901f60c755489cc61c9b423b836442" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "clap_lex" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "702fc72eb24e5a1e48ce58027a675bc24edd52096d5397d4aea7c6dd9eca0bd1" - -[[package]] -name = "colorchoice" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7" - [[package]] name = "core-foundation" version = "0.9.3" @@ -479,6 +385,21 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "futures" +version = "0.3.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da0290714b38af9b4a7b094b8a37086d1b4e61f2df9122c3cad2577669145335" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.29" @@ -486,6 +407,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ff4dd66668b557604244583e3e1e1eada8c5c2e96a6d0d6653ede395b78bbacb" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -494,6 +416,17 @@ version = "0.3.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eb1d22c66e66d9d72e1758f0bd7d4fd0bee04cad842ee34587d68c07e45d088c" +[[package]] +name = "futures-executor" +version = "0.3.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f4fb8693db0cf099eadcca0efe2a5a22e4550f98ed16aba6c48700da29597bc" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.29" @@ -529,6 +462,7 @@ version = "0.3.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a19526d624e703a3179b3d322efec918b6246ea0fa51d41124525f00f1cc8104" dependencies = [ + "futures-channel", "futures-core", "futures-io", "futures-macro", @@ -808,17 +742,16 @@ dependencies = [ "async-trait", "axum", "axum-macros", - "clap", + "base64", "dotenvy", + "futures", "magic-crypt", - "queryst", "regex", "reqwest", "sentry", "serde", "serde_regex", "serde_yaml", - "thiserror", "tokio", "tokio-stream", "tokio-util", @@ -1181,19 +1114,6 @@ dependencies = [ "unicode-ident", ] -[[package]] -name = "queryst" -version = "3.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1cbeb75ac695daf201ca2d66d9c684f873b135f28af4f2c79952478cab3b9d9" -dependencies = [ - "lazy_static", - "percent-encoding", - "regex", - "serde", - "serde_json", -] - [[package]] name = "quote" version = "1.0.33" @@ -1645,12 +1565,6 @@ dependencies = [ "windows-sys", ] -[[package]] -name = "strsim" -version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" - [[package]] name = "syn" version = "2.0.38" @@ -2012,12 +1926,6 @@ dependencies = [ "serde", ] -[[package]] -name = "utf8parse" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a" - [[package]] name = "uuid" version = "1.5.0" diff --git a/Cargo.toml b/Cargo.toml index c7912c4..2046fa9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +futures = "0.3" tokio = { version = "1.35", features = ["full", "tracing"] } tracing = "0.1" tracing-subscriber = "0.3" @@ -16,13 +17,11 @@ reqwest = { version = "0.11", features = ["stream"] } anyhow = "1.0" sentry = "0.32" dotenvy = "0.15" -thiserror = "1.0" axum-macros = "0.4" magic-crypt = "3.1" async-trait = "0.1" -queryst = "3" -clap = { version = "4.4.10", features = ["derive"] } serde = "1.0" serde_yaml = "0.9" +base64 = "0.21" regex = "1.10" -serde_regex = "1.1" \ No newline at end of file +serde_regex = "1.1" diff --git a/src/config.rs b/src/config.rs index a3a1913..d34002b 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,25 +1,22 @@ -use std::str::FromStr; - use anyhow::anyhow; -use serde::{Deserialize, de::DeserializeOwned}; - +use serde::Deserialize; #[derive(Deserialize)] pub struct MemoryStorageConfig { pub max_size: usize, - pub ttl: usize + pub ttl: usize, } #[derive(Deserialize)] pub struct DiskStorageConfig { pub path: String, - pub max_size: usize + pub max_size: usize, } #[derive(Deserialize)] pub enum StorageStrategyConfig { Memory(MemoryStorageConfig), - Disk(DiskStorageConfig) + Disk(DiskStorageConfig), } #[derive(Deserialize)] @@ -27,19 +24,21 @@ pub struct StorageConfig { pub strategy: StorageStrategyConfig, #[serde(with = "serde_regex")] - pub regex: regex::Regex + pub regex: regex::Regex, } #[derive(Deserialize, Default)] pub struct YAMLConfig { pub storages: Vec, - pub secret_key: Option + pub secret_key: Option, } pub async fn load() -> YAMLConfig { tokio::fs::read("config.yml") .await .map_err(|err| anyhow!(err)) - .and_then(|content| serde_yaml::from_slice::(&content).map_err(|err| anyhow!(err))) + .and_then(|content| { + serde_yaml::from_slice::(&content).map_err(|err| anyhow!(err)) + }) .unwrap_or_default() -} \ No newline at end of file +} diff --git a/src/crypt.rs b/src/crypt.rs new file mode 100644 index 0000000..e6d3acf --- /dev/null +++ b/src/crypt.rs @@ -0,0 +1,19 @@ +use std::io::Write; + +use base64::engine::general_purpose; +use magic_crypt::{new_magic_crypt, MagicCryptTrait}; + +pub fn decryt(value: I, secret_key: String) -> anyhow::Result +where + I: AsRef, +{ + let mc = new_magic_crypt!(&secret_key, 256); + Ok(mc.decrypt_base64_to_string(value)?) +} + +pub fn compute_key(from: String) -> String { + let mut enc = base64::write::EncoderStringWriter::new(&general_purpose::URL_SAFE_NO_PAD); + enc.write_all(from.as_bytes()).unwrap(); + + enc.into_inner() +} diff --git a/src/main.rs b/src/main.rs index 56b93cf..be55c00 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,28 +1,15 @@ -use anyhow::anyhow; -use axum::{ - http::StatusCode, - routing::get, Router, extract::{Path, State}, response::IntoResponse, -}; +use axum::{routing::get, Router}; +use tokio::sync::Mutex; -use magic_crypt::{new_magic_crypt, MagicCryptTrait}; -use tokio_stream::{self as stream, StreamExt}; -use tokio_util::{bytes::Bytes}; -use tracing::{warn}; -use clap::Parser; +use std::sync::Arc; -use std::{sync::Arc, path::PathBuf}; - -use crate::storages::StoragePool; +use crate::{state::AppState, storages::StoragePool}; mod config; -mod storages; +mod crypt; mod routes; mod state; - -struct AppState { - secret_key: String, - authorized_commands: Vec -} +mod storages; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -32,59 +19,20 @@ async fn main() -> anyhow::Result<()> { let secret_key = config.secret_key.clone(); let storage_pool = StoragePool::from_config(config.storages); - let app_state = Arc::new(state::AppState::new(storage_pool, secret_key)); - - println!("{}", config.secret_key.clone().unwrap_or_default()); + let app_state = AppState::new(storage_pool, secret_key); let app = Router::new() - .route("/unsecure/*src", get(routes::handle_unsecure)) - .route("/secure/*data", get(routes::handle_secure)) - .with_state(app_state); + .route( + "/*src", + match app_state.secret_key { + Some(_) => get(routes::handle_secure), + None => get(routes::handle_unsecure), + }, + ) + .with_state(Arc::new(Mutex::new(app_state))); let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await?; axum::serve(listener, app).await?; Ok(()) } - -/*async fn handle_secure(Path(data): Path, State(state): State>) -> Result { - let mc = new_magic_crypt!(&state.secret_key, 256); - let decrypt = mc.decrypt_base64_to_string(data)?; - let split = decrypt.splitn(2, "/").map(String::from).collect::>(); - - let opts = split.get(0).ok_or_else(|| anyhow::anyhow!("opts not found"))?; - let src = split.get(1).ok_or_else(|| anyhow::anyhow!("src not found"))?; - - let bytes = handle(opts.clone(), src.clone(), state).await?; - Ok(axum::body::Body::from(bytes)) -} - -async fn handle_unsecure(Path((opts, src)): Path<(String, String)>, State(state): State>) -> Result { - let bytes = handle(opts, src, state).await?; - - Ok(axum::body::Body::from(bytes)) -}*/ - -/*async fn handle(_opts: String, src: String) -> anyhow::Result>> { - let client = reqwest::Client::new(); - let res = client.get(src).send().await?; - let stream = res.bytes_stream(); - - Ok(stream) -}*/ - -/*async fn handle(opts: String, src: String, app_state: Arc) -> anyhow::Result { - let qs_opts = queryst::parse(&opts).map_err(|err| anyhow::anyhow!("{}", err.message))?; - - let client = reqwest::Client::new(); - let res = client.get(src.clone()).send().await?; - let bytes = res.bytes().await?; - - if let Some(storage) = &app_state.storage { - - } - - - Ok(bytes) -}*/ - diff --git a/src/routes.rs b/src/routes.rs index e6af8f0..f55351f 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -1,37 +1,66 @@ -use std::{sync::Arc}; +use std::{pin::Pin, sync::Arc}; -use axum::{extract::{Path, State}, response::IntoResponse, http::StatusCode}; -use magic_crypt::{new_magic_crypt, MagicCryptTrait}; +use axum::{ + extract::{Path, State}, + http::StatusCode, + response::IntoResponse, +}; +use futures::TryStreamExt; +use tokio::sync::Mutex; use tokio_stream::Stream; use tokio_util::bytes::Bytes; -use crate::state::AppState; +use crate::{crypt, state::AppState}; +pub async fn handle_secure( + Path(data): Path, + state: State>>, +) -> Result { + let app_state = state.lock().await; + let decrypt = crypt::decryt(data, app_state.secret_key.clone().unwrap())?; + std::mem::drop(app_state); -pub async fn handle_secure(Path(data): Path, State(state): State>) -> Result { - let mc = new_magic_crypt!(&state.secret_key.clone().unwrap(), 256); - let decrypt = mc.decrypt_base64_to_string(data)?; - let split = decrypt.splitn(2, "/").map(String::from).collect::>(); - - let opts = split.get(0).ok_or_else(|| anyhow::anyhow!("opts not found"))?; - let src = split.get(1).ok_or_else(|| anyhow::anyhow!("src not found"))?; - - let bytes = handle(opts.clone(), src.clone()).await?; + let bytes = handle(decrypt, state).await?; Ok(axum::body::Body::from_stream(bytes)) } -pub async fn handle_unsecure(Path((opts, src)): Path<(String, String)>, State(state): State>) -> Result { - let bytes = handle(opts, src).await?; +pub async fn handle_unsecure( + Path(src): Path, + state: State>>, +) -> Result { + let bytes = handle(src, state).await?; Ok(axum::body::Body::from_stream(bytes)) } -async fn handle(_opts: String, src: String) -> anyhow::Result>> { +async fn handle( + src: String, + state: State>>, +) -> anyhow::Result> + Send> { + let mut app_state = state.lock().await; + + match app_state.storage_pool.retrieve(src.clone()).await { + Some(stream) => Ok(stream), + None => { + let save_stream = fetch(src.clone()).await?; + app_state + .storage_pool + .save(src.clone(), save_stream) + .await?; + + Ok(app_state.storage_pool.retrieve(src).await.unwrap()) + } + } +} + +async fn fetch( + src: String, +) -> anyhow::Result> + Send>>> { let client = reqwest::Client::new(); let res = client.get(src).send().await?; - let stream = res.bytes_stream(); + let stream = res.bytes_stream().map_err(|err| anyhow::anyhow!(err)); - Ok(stream) + Ok(Box::pin(stream)) } pub struct MyAnyhow(anyhow::Error); @@ -42,8 +71,11 @@ impl IntoResponse for MyAnyhow { } } -impl From for MyAnyhow where T: Into { +impl From for MyAnyhow +where + T: Into, +{ fn from(value: T) -> Self { Self(value.into()) } -} \ No newline at end of file +} diff --git a/src/state.rs b/src/state.rs index 7196038..e34b89e 100644 --- a/src/state.rs +++ b/src/state.rs @@ -2,11 +2,14 @@ use crate::storages::StoragePool; pub struct AppState { pub secret_key: Option, - pub storage_pool: StoragePool + pub storage_pool: StoragePool, } impl AppState { pub fn new(storage_pool: StoragePool, secret_key: Option) -> Self { - Self { secret_key, storage_pool } + Self { + secret_key, + storage_pool, + } } -} \ No newline at end of file +} diff --git a/src/storages/disk.rs b/src/storages/disk.rs index 3caa5c3..c369270 100644 --- a/src/storages/disk.rs +++ b/src/storages/disk.rs @@ -1,9 +1,14 @@ +use std::pin::Pin; + use crate::config::DiskStorageConfig; +use async_trait::async_trait; +use tokio_stream::Stream; +use tokio_util::bytes::Bytes; use super::Storage; pub struct DiskStorage { - config: DiskStorageConfig + config: DiskStorageConfig, } impl DiskStorage { @@ -12,4 +17,28 @@ impl DiskStorage { } } -impl Storage for DiskStorage {} \ No newline at end of file +#[async_trait] +impl Storage for DiskStorage { + async fn eligible(&self, src: String) -> bool { + false + } + + async fn delete(&mut self, key: String) -> anyhow::Result<()> { + Ok(()) + } + + async fn retrieve( + &self, + key: String, + ) -> Option> + Send>>> { + None + } + + async fn save( + &mut self, + key: String, + stream: Pin> + Send>>, + ) -> anyhow::Result<()> { + todo!() + } +} diff --git a/src/storages/memory.rs b/src/storages/memory.rs index f7ed83c..e2ccf4e 100644 --- a/src/storages/memory.rs +++ b/src/storages/memory.rs @@ -1,15 +1,74 @@ +use std::{collections::HashMap, pin::Pin}; + +use anyhow::Ok; +use async_trait::async_trait; +use tokio_stream::{Stream, StreamExt}; +use tokio_util::bytes::Bytes; + use crate::config::MemoryStorageConfig; use super::Storage; pub struct MemoryStorage { - config: MemoryStorageConfig + config: MemoryStorageConfig, + items: HashMap>, } impl MemoryStorage { pub fn new(config: MemoryStorageConfig) -> Self { - Self { config } + Self { + config, + items: HashMap::new(), + } } } -impl Storage for MemoryStorage {} \ No newline at end of file +#[async_trait] +impl Storage for MemoryStorage { + async fn eligible(&self, src: String) -> bool { + true + } + + async fn delete(&mut self, key: String) -> anyhow::Result<()> { + self.items + .remove(&key) + .ok_or_else(|| anyhow::anyhow!("")) + .map(|_| ()) + } + + async fn retrieve( + &self, + key: String, + ) -> Option> + Send>>> { + match self.items.get(&key) { + Some(item) => { + let bytes = tokio_util::bytes::Bytes::from(item.clone()); + let stream = futures::stream::iter([Ok(bytes)]); + + Some(Box::pin(stream)) + } + None => None, + } + } + + async fn save( + &mut self, + key: String, + mut stream: Pin> + Send>>, + ) -> anyhow::Result<()> { + if !self.items.contains_key(&key) { + self.items.insert(key.clone(), Vec::new()); + } + + while let Some(chunk) = stream.next().await { + self.items + .get_mut(&key) + .unwrap() + .append(&mut chunk?.to_vec()); + } + + println!("{} saved", key); + + Ok(()) + } +} diff --git a/src/storages/mod.rs b/src/storages/mod.rs index a11e09c..c6ec435 100644 --- a/src/storages/mod.rs +++ b/src/storages/mod.rs @@ -1,28 +1,83 @@ -pub mod memory; pub mod disk; +pub mod memory; + +use std::pin::Pin; use async_trait::async_trait; +use tokio_stream::Stream; +use tokio_util::bytes::Bytes; -use crate::config::StorageConfig; +use crate::{ + config::{StorageConfig, StorageStrategyConfig}, + crypt, +}; -use self::{memory::MemoryStorage, disk::DiskStorage}; +use self::{disk::DiskStorage, memory::MemoryStorage}; #[async_trait] -pub trait Storage {} +pub trait Storage { + async fn eligible(&self, src: String) -> bool; + async fn delete(&mut self, key: String) -> anyhow::Result<()>; + async fn retrieve( + &self, + key: String, + ) -> Option> + Send>>>; + async fn save( + &mut self, + key: String, + stream: Pin> + Send>>, + ) -> anyhow::Result<()>; +} pub struct StoragePool { - storages: Vec> + storages: Vec>, } impl StoragePool { pub fn from_config(items: Vec) -> Self { - Self { storages: items.into_iter().map(create_storage).collect() } + Self { + storages: items.into_iter().map(create_storage).collect(), + } + } + + pub async fn retrieve( + &self, + src: String, + ) -> Option> + Send>>> { + let mut stream = None; + let key = crypt::compute_key(src); + + for item in &self.storages { + if let Some(storage_stream) = item.retrieve(key.clone()).await { + stream = Some(storage_stream); + break; + } + } + + stream + } + + pub async fn save( + &mut self, + src: String, + stream: Pin> + Send>>, + ) -> anyhow::Result<()> { + let key = crypt::compute_key(src); + + for item in self.storages.iter_mut() { + if item.eligible(key.clone()).await { + item.save(key, stream).await?; + break; + } + } + + Ok(()) } } fn create_storage(item: StorageConfig) -> Box { match item.strategy { - crate::config::StorageMethodConfig::Memory(config) => Box::new(MemoryStorage::new(config)), - crate::config::StorageMethodConfig::Disk(config) => Box::new(DiskStorage::new(config)), + StorageStrategyConfig::Memory(config) => Box::new(MemoryStorage::new(config)), + StorageStrategyConfig::Disk(config) => Box::new(DiskStorage::new(config)), } -} \ No newline at end of file +}