gitdataai/lib/storage/s3.rs

use std::time::Duration;

use async_trait::async_trait;
use aws_config::BehaviorVersion;
use aws_sdk_s3::{
    Client,
    config::{Credentials, Region},
    presigning::PresigningConfig,
    primitives::ByteStream,
};

use crate::{
    ObjectStorage, PutObjectOptions, StorageError, StorageObject,
    StorageObjectStream, StorageResult, StoredObject,
};

#[derive(Clone, Debug)]
pub struct S3StorageConfig {
    pub bucket: String,
    pub region: String,
    pub endpoint_url: Option<String>,
    pub access_key_id: Option<String>,
    pub secret_access_key: Option<String>,
    pub session_token: Option<String>,
    pub force_path_style: bool,
    pub public_url: Option<String>,
    pub presigned_url_ttl: Duration,
}

#[derive(Clone)]
pub struct S3Storage {
    client: Client,
    config: S3StorageConfig,
}

impl S3Storage {
    pub async fn connect(config: S3StorageConfig) -> StorageResult<Self> {
        let mut sdk_config = aws_config::defaults(BehaviorVersion::latest())
            .region(Region::new(config.region.clone()));

        match (&config.access_key_id, &config.secret_access_key) {
            (Some(access_key_id), Some(secret_access_key)) => {
                sdk_config = sdk_config.credentials_provider(Credentials::new(
                    access_key_id,
                    secret_access_key,
                    config.session_token.clone(),
                    None,
                    "app-storage-config",
                ));
            }
            (None, None) => {}
            _ => {
                return Err(StorageError::Config(
                    "APP_STORAGE_S3_ACCESS_KEY_ID and APP_STORAGE_S3_SECRET_ACCESS_KEY must be set together"
                        .to_string(),
                ));
            }
        }

        let sdk_config = sdk_config.load().await;
        let mut s3_config = aws_sdk_s3::config::Builder::from(&sdk_config)
            .force_path_style(config.force_path_style);

        if let Some(endpoint_url) = &config.endpoint_url {
            s3_config = s3_config.endpoint_url(endpoint_url);
        }

        Ok(Self {
            client: Client::from_conf(s3_config.build()),
            config,
        })
    }

    pub async fn put_bytes(
        &self,
        key: &str,
        bytes: Vec<u8>,
        options: PutObjectOptions,
    ) -> StorageResult<StoredObject> {
        self.put_stream(key, ByteStream::from(bytes), options).await
    }

    pub async fn put_file(
        &self,
        key: &str,
        path: impl AsRef<std::path::Path>,
        options: PutObjectOptions,
    ) -> StorageResult<StoredObject> {
        let body = ByteStream::read_from()
            .path(path.as_ref())
            .build()
            .await
            .map_err(|error| StorageError::Stream(error.to_string()))?;
        self.put_stream(key, body, options).await
    }

    fn normalize_key(key: &str) -> StorageResult<String> {
        let key = key.trim().trim_start_matches('/');
        if key.is_empty() || key.contains("..") {
            return Err(StorageError::InvalidKey(key.to_string()));
        }
        Ok(key.to_string())
    }

    fn public_url_for_config(
        config: &S3StorageConfig,
        key: &str,
    ) -> Option<String> {
        let base_url = config.public_url.as_ref()?.trim_end_matches('/');
        Some(format!("{base_url}/{key}"))
    }
}

#[async_trait]
impl ObjectStorage for S3Storage {
    async fn put_stream(
        &self,
        key: &str,
        body: ByteStream,
        options: PutObjectOptions,
    ) -> StorageResult<StoredObject> {
        let key = Self::normalize_key(key)?;

        let mut request = self
            .client
            .put_object()
            .bucket(&self.config.bucket)
            .key(&key)
            .body(body);

        if let Some(content_type) = options.content_type {
            request = request.content_type(content_type);
        }
        if let Some(content_length) = options.content_length {
            request = request.content_length(content_length);
        }
        if let Some(cache_control) = options.cache_control {
            request = request.cache_control(cache_control);
        }

        let output = request
            .send()
            .await
            .map_err(|error| StorageError::S3(error.to_string()))?;
        let url = match self.public_url(&key)? {
            Some(url) => url,
            None => {
                self.presigned_get_url(&key, self.config.presigned_url_ttl)
                    .await?
            }
        };

        Ok(StoredObject {
            key,
            url,
            e_tag: output.e_tag,
            version_id: output.version_id,
        })
    }

    async fn put_bytes(
        &self,
        key: &str,
        bytes: Vec<u8>,
        options: PutObjectOptions,
    ) -> StorageResult<StoredObject> {
        S3Storage::put_bytes(self, key, bytes, options).await
    }

    async fn get_stream(
        &self,
        key: &str,
    ) -> StorageResult<StorageObjectStream> {
        let key = Self::normalize_key(key)?;
        let output = self
            .client
            .get_object()
            .bucket(&self.config.bucket)
            .key(&key)
            .send()
            .await
            .map_err(|error| StorageError::S3(error.to_string()))?;

        Ok(StorageObjectStream {
            body: output.body,
            content_length: output.content_length,
            content_type: output.content_type,
            e_tag: output.e_tag,
        })
    }

    async fn get_bytes(&self, key: &str) -> StorageResult<StorageObject> {
        let stream = self.get_stream(key).await?;
        let bytes = crate::collect_byte_stream(stream.body)
            .await
            .map_err(|error| StorageError::Stream(error.to_string()))?;

        Ok(StorageObject {
            bytes,
            content_length: stream.content_length,
            content_type: stream.content_type,
            e_tag: stream.e_tag,
        })
    }

    async fn delete(&self, key: &str) -> StorageResult<()> {
        let key = Self::normalize_key(key)?;
        self.client
            .delete_object()
            .bucket(&self.config.bucket)
            .key(key)
            .send()
            .await
            .map_err(|error| StorageError::S3(error.to_string()))?;
        Ok(())
    }

    fn public_url(&self, key: &str) -> StorageResult<Option<String>> {
        let key = Self::normalize_key(key)?;
        Ok(Self::public_url_for_config(&self.config, &key))
    }

    async fn presigned_get_url(
        &self,
        key: &str,
        expires_in: Duration,
    ) -> StorageResult<String> {
        let key = Self::normalize_key(key)?;
        let config = PresigningConfig::expires_in(expires_in)
            .map_err(|error| StorageError::Config(error.to_string()))?;
        let request = self
            .client
            .get_object()
            .bucket(&self.config.bucket)
            .key(key)
            .presigned(config)
            .await
            .map_err(|error| StorageError::S3(error.to_string()))?;
        Ok(request.uri().to_string())
    }
}

impl TryFrom<&config::AppConfig> for S3StorageConfig {
    type Error = StorageError;

    fn try_from(config: &config::AppConfig) -> Result<Self, Self::Error> {
        Ok(Self {
            bucket: config
                .storage_s3_bucket()
                .map_err(|error| StorageError::Config(error.to_string()))?,
            region: config.storage_s3_region(),
            endpoint_url: config.storage_s3_endpoint_url(),
            access_key_id: config.storage_s3_access_key_id(),
            secret_access_key: config.storage_s3_secret_access_key(),
            session_token: config.storage_s3_session_token(),
            force_path_style: config
                .storage_s3_force_path_style()
                .map_err(|error| StorageError::Config(error.to_string()))?,
            public_url: config.storage_public_url_base(),
            presigned_url_ttl: config
                .storage_presigned_url_ttl()
                .map_err(|error| StorageError::Config(error.to_string()))?,
        })
    }
}