rustpad/src/storage.rs

use std::{path::PathBuf, sync::Arc};

use aws_config::Region;
use aws_credential_types::Credentials;
use aws_sdk_s3::{Client, config::Builder as S3ConfigBuilder, primitives::ByteStream};
use bytes::Bytes;

#[derive(Debug, Clone)]
pub enum StorageConfig {
    Local {
        root: PathBuf,
    },
    S3 {
        endpoint: Option<String>,
        region: String,
        bucket: String,
        access_key: String,
        secret_key: String,
        force_path_style: bool,
    },
}

#[derive(Clone)]
pub enum Storage {
    Local { root: PathBuf },
    S3 { client: Client, bucket: Arc<str> },
}

impl std::fmt::Debug for Storage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Local { root } => f.debug_struct("LocalStorage").field("root", root).finish(),
            Self::S3 { bucket, .. } => f.debug_struct("S3Storage").field("bucket", bucket).finish(),
        }
    }
}

impl Storage {
    pub async fn from_config(config: StorageConfig) -> Result<Self, Box<dyn std::error::Error>> {
        match config {
            StorageConfig::Local { root } => {
                tokio::fs::create_dir_all(&root).await?;
                Ok(Self::Local { root })
            }
            StorageConfig::S3 {
                endpoint,
                region,
                bucket,
                access_key,
                secret_key,
                force_path_style,
            } => {
                let credentials =
                    Credentials::new(access_key, secret_key, None, None, "rustpad-env");
                let shared = aws_config::defaults(aws_config::BehaviorVersion::latest())
                    .region(Region::new(region.clone()))
                    .credentials_provider(credentials)
                    .load()
                    .await;
                let mut builder = S3ConfigBuilder::from(&shared)
                    .region(Region::new(region))
                    .force_path_style(force_path_style);
                if let Some(endpoint) = endpoint.filter(|value| !value.trim().is_empty()) {
                    builder = builder.endpoint_url(endpoint);
                }
                Ok(Self::S3 {
                    client: Client::from_conf(builder.build()),
                    bucket: Arc::from(bucket),
                })
            }
        }
    }

    pub fn backend_name(&self) -> &'static str {
        match self {
            Self::Local { .. } => "local",
            Self::S3 { .. } => "s3",
        }
    }

    pub async fn exists(&self, key: &str) -> Result<bool, StorageError> {
        match self {
            Self::Local { root } => Ok(root.join(key).is_file()),
            Self::S3 { client, bucket } => match client
                .head_object()
                .bucket(bucket.as_ref())
                .key(key)
                .send()
                .await
            {
                Ok(_) => Ok(true),
                Err(error)
                    if error
                        .as_service_error()
                        .is_some_and(|service| service.is_not_found()) =>
                {
                    Ok(false)
                }
                Err(error) => Err(StorageError::Backend(error.to_string())),
            },
        }
    }

    pub async fn put(
        &self,
        key: &str,
        bytes: Bytes,
        content_type: &str,
        cache_control: &str,
    ) -> Result<(), StorageError> {
        match self {
            Self::Local { root } => {
                let path = root.join(key);
                if let Some(parent) = path.parent() {
                    tokio::fs::create_dir_all(parent).await?;
                }
                tokio::fs::write(path, bytes).await?;
                Ok(())
            }
            Self::S3 { client, bucket } => {
                client
                    .put_object()
                    .bucket(bucket.as_ref())
                    .key(key)
                    .content_type(content_type)
                    .cache_control(cache_control)
                    .body(ByteStream::from(bytes))
                    .send()
                    .await
                    .map_err(|error| StorageError::Backend(error.to_string()))?;
                Ok(())
            }
        }
    }

    pub async fn get(&self, key: &str) -> Result<Bytes, StorageError> {
        match self {
            Self::Local { root } => Ok(Bytes::from(tokio::fs::read(root.join(key)).await?)),
            Self::S3 { client, bucket } => {
                let output = client
                    .get_object()
                    .bucket(bucket.as_ref())
                    .key(key)
                    .send()
                    .await
                    .map_err(|error| StorageError::Backend(error.to_string()))?;
                let bytes = output
                    .body
                    .collect()
                    .await
                    .map_err(|error| StorageError::Backend(error.to_string()))?
                    .into_bytes();
                Ok(bytes)
            }
        }
    }

    pub async fn get_local_with_legacy(
        &self,
        key: &str,
        legacy_key: &str,
    ) -> Result<Bytes, StorageError> {
        match self {
            Self::Local { root } => {
                let canonical = root.join(key);
                let path = if canonical.is_file() {
                    canonical
                } else {
                    root.join(legacy_key)
                };
                Ok(Bytes::from(tokio::fs::read(path).await?))
            }
            Self::S3 { .. } => self.get(key).await,
        }
    }

    pub async fn delete(&self, key: &str) -> Result<(), StorageError> {
        match self {
            Self::Local { root } => match tokio::fs::remove_file(root.join(key)).await {
                Ok(()) => Ok(()),
                Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
                Err(error) => Err(error.into()),
            },
            Self::S3 { client, bucket } => {
                client
                    .delete_object()
                    .bucket(bucket.as_ref())
                    .key(key)
                    .send()
                    .await
                    .map_err(|error| StorageError::Backend(error.to_string()))?;
                Ok(())
            }
        }
    }
}

#[derive(Debug)]
pub enum StorageError {
    Io(std::io::Error),
    Backend(String),
}

impl std::fmt::Display for StorageError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Io(error) => write!(f, "{error}"),
            Self::Backend(error) => f.write_str(error),
        }
    }
}
impl std::error::Error for StorageError {}
impl From<std::io::Error> for StorageError {
    fn from(value: std::io::Error) -> Self {
        Self::Io(value)
    }
}

pub fn object_key(kind: &str, owner_id: i64, token: &str, filename: &str) -> String {
    format!("{kind}/{owner_id}_{token}/{filename}")
}

pub fn legacy_key(owner_id: i64, token: &str, filename: &str) -> String {
    format!("{owner_id}_{token}/{filename}")
}