use async_trait::async_trait;
use std::path::PathBuf;
#[derive(Debug)]
pub enum StorageError {
IoError(std::io::Error),
ConfigError(String),
DriverError(String),
}
impl std::fmt::Display for StorageError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
StorageError::IoError(err) => write!(f, "IO error: {}", err),
StorageError::ConfigError(err) => write!(f, "Configuration error: {}", err),
StorageError::DriverError(err) => write!(f, "Storage driver error: {}", err),
}
}
}
impl std::error::Error for StorageError {}
impl From<std::io::Error> for StorageError {
fn from(err: std::io::Error) -> Self {
StorageError::IoError(err)
}
}
#[async_trait]
pub trait StorageDriver: Send + Sync {
async fn put(&self, path: &str, bytes: &[u8]) -> Result<(), StorageError>;
async fn get(&self, path: &str) -> Result<Vec<u8>, StorageError>;
async fn exists(&self, path: &str) -> Result<bool, StorageError>;
async fn delete(&self, path: &str) -> Result<(), StorageError>;
async fn url(&self, path: &str) -> Result<String, StorageError>;
}
pub struct LocalDriver {
pub root: PathBuf,
}
impl LocalDriver {
pub fn new(root: impl Into<PathBuf>) -> Self {
LocalDriver { root: root.into() }
}
fn resolve_path(&self, path: &str) -> PathBuf {
self.root.join(path.trim_start_matches('/'))
}
}
#[async_trait]
impl StorageDriver for LocalDriver {
async fn put(&self, path: &str, bytes: &[u8]) -> Result<(), StorageError> {
let full_path = self.resolve_path(path);
if let Some(parent) = full_path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(full_path, bytes)?;
Ok(())
}
async fn get(&self, path: &str) -> Result<Vec<u8>, StorageError> {
let full_path = self.resolve_path(path);
let data = std::fs::read(full_path)?;
Ok(data)
}
async fn exists(&self, path: &str) -> Result<bool, StorageError> {
let full_path = self.resolve_path(path);
Ok(full_path.exists())
}
async fn delete(&self, path: &str) -> Result<(), StorageError> {
let full_path = self.resolve_path(path);
if full_path.exists() {
std::fs::remove_file(full_path)?;
}
Ok(())
}
async fn url(&self, path: &str) -> Result<String, StorageError> {
Ok(format!("/storage/{}", path.trim_start_matches('/')))
}
}
#[cfg(feature = "storage-s3")]
pub struct S3Driver {
pub client: aws_sdk_s3::Client,
pub bucket: String,
pub endpoint: Option<String>,
}
#[cfg(feature = "storage-s3")]
#[async_trait]
impl StorageDriver for S3Driver {
async fn put(&self, path: &str, bytes: &[u8]) -> Result<(), StorageError> {
let body = aws_sdk_s3::primitives::ByteStream::from(bytes.to_vec());
self.client
.put_object()
.bucket(&self.bucket)
.key(path)
.body(body)
.send()
.await
.map_err(|e| StorageError::DriverError(format!("{}", e)))?;
Ok(())
}
async fn get(&self, path: &str) -> Result<Vec<u8>, StorageError> {
let res = self.client
.get_object()
.bucket(&self.bucket)
.key(path)
.send()
.await
.map_err(|e| StorageError::DriverError(format!("{}", e)))?;
let bytes = res
.body
.collect()
.await
.map_err(|e| StorageError::DriverError(format!("{}", e)))?
.into_bytes()
.to_vec();
Ok(bytes)
}
async fn exists(&self, path: &str) -> Result<bool, StorageError> {
let res = self.client
.head_object()
.bucket(&self.bucket)
.key(path)
.send()
.await;
match res {
Ok(_) => Ok(true),
Err(e) => {
let service_err = e.into_service_error();
if service_err.is_no_such_key() {
Ok(false)
} else {
Err(StorageError::DriverError(format!("{}", service_err)))
}
}
}
}
async fn delete(&self, path: &str) -> Result<(), StorageError> {
self.client
.delete_object()
.bucket(&self.bucket)
.key(path)
.send()
.await
.map_err(|e| StorageError::DriverError(format!("{}", e)))?;
Ok(())
}
async fn url(&self, path: &str) -> Result<String, StorageError> {
if let Some(ref endpoint) = self.endpoint {
Ok(format!(
"{}/{}/{}",
endpoint.trim_end_matches('/'),
self.bucket,
path
))
} else {
Ok(format!("https://{}.s3.amazonaws.com/{}", self.bucket, path))
}
}
}
#[cfg(not(feature = "storage-s3"))]
pub struct S3Driver;
#[cfg(not(feature = "storage-s3"))]
#[async_trait]
impl StorageDriver for S3Driver {
async fn put(&self, _path: &str, _bytes: &[u8]) -> Result<(), StorageError> {
Err(StorageError::DriverError(
"S3 storage driver requires the 'storage-s3' Cargo feature to be enabled".to_string(),
))
}
async fn get(&self, _path: &str) -> Result<Vec<u8>, StorageError> {
Err(StorageError::DriverError(
"S3 storage driver requires the 'storage-s3' Cargo feature to be enabled".to_string(),
))
}
async fn exists(&self, _path: &str) -> Result<bool, StorageError> {
Err(StorageError::DriverError(
"S3 storage driver requires the 'storage-s3' Cargo feature to be enabled".to_string(),
))
}
async fn delete(&self, _path: &str) -> Result<(), StorageError> {
Err(StorageError::DriverError(
"S3 storage driver requires the 'storage-s3' Cargo feature to be enabled".to_string(),
))
}
async fn url(&self, _path: &str) -> Result<String, StorageError> {
Err(StorageError::DriverError(
"S3 storage driver requires the 'storage-s3' Cargo feature to be enabled".to_string(),
))
}
}
pub struct Storage;
impl Storage {
pub async fn put(path: &str, bytes: &[u8]) -> Result<(), StorageError> {
Self::disk("local").put(path, bytes).await
}
pub async fn get(path: &str) -> Result<Vec<u8>, StorageError> {
Self::disk("local").get(path).await
}
pub async fn exists(path: &str) -> Result<bool, StorageError> {
Self::disk("local").exists(path).await
}
pub async fn delete(path: &str) -> Result<(), StorageError> {
Self::disk("local").delete(path).await
}
pub async fn url(path: &str) -> Result<String, StorageError> {
Self::disk("local").url(path).await
}
pub fn disk(name: &str) -> Box<dyn StorageDriver> {
match name {
"local" => {
let root = std::env::var("STORAGE_ROOT").unwrap_or_else(|_| "storage/app".to_string());
Box::new(LocalDriver::new(root))
}
"s3" => {
#[cfg(feature = "storage-s3")]
{
let bucket = std::env::var("AWS_BUCKET").unwrap_or_default();
let endpoint = std::env::var("AWS_ENDPOINT").ok();
let config = tokio::task::block_in_place(|| {
tokio::runtime::Handle::current().block_on(aws_config::load_from_env())
});
let client = aws_sdk_s3::Client::new(&config);
Box::new(S3Driver {
client,
bucket,
endpoint,
})
}
#[cfg(not(feature = "storage-s3"))]
{
Box::new(S3Driver)
}
}
other => panic!("Unknown storage disk: {}", other),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_local_storage_lifecycle() {
let temp_dir = "storage/test_app";
let _ = std::fs::remove_dir_all(temp_dir);
let driver = LocalDriver::new(temp_dir);
driver.put("images/avatar.png", b"fake-png-bytes").await.unwrap();
assert!(driver.exists("images/avatar.png").await.unwrap());
let data = driver.get("images/avatar.png").await.unwrap();
assert_eq!(data, b"fake-png-bytes");
let url = driver.url("images/avatar.png").await.unwrap();
assert_eq!(url, "/storage/images/avatar.png");
driver.delete("images/avatar.png").await.unwrap();
assert!(!driver.exists("images/avatar.png").await.unwrap());
let _ = std::fs::remove_dir_all(temp_dir);
}
}