defter 0.2.0

//! Attachment bytes in object storage, through the `object_store` crate.
//!
//! In production this is an S3-compatible bucket (AWS S3, Cloudflare R2,
//! MinIO, RustFS, …). Uploads stream *through* the server — the client POSTs
//! the bytes, the server verifies their md5 while writing them to the bucket —
//! and downloads redirect the client to a short-lived pre-signed GET URL, so
//! the bytes are read straight from the bucket.
//!
//! Without a bucket (development) the files go to a local directory, and
//! downloads are streamed by the server under an expiring token instead.
//!
//! Layout: a verified upload waits under `pending/<token>` until the client
//! registers it, then moves to `files/<item key>`.

use std::sync::Arc;
use std::time::Duration;

use axum::body::{Body, Bytes};
use futures::{Stream, StreamExt};
use object_store::aws::AmazonS3Builder;
use object_store::local::LocalFileSystem;
#[cfg(test)]
use object_store::memory::InMemory;
use object_store::path::Path;
use object_store::signer::{Method, Signer};
use object_store::{ObjectStore, ObjectStoreExt, WriteMultipart};

use crate::error::AppError;

/// Where attachment bytes go.
#[derive(Debug, Clone)]
pub enum Backend {
    S3(S3Config),
    Local(std::path::PathBuf),
    #[cfg(test)]
    Memory,
}

/// An S3-compatible bucket. Unset fields fall back to the standard `AWS_*`
/// environment variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`,
/// `AWS_REGION`, `AWS_ENDPOINT`, …) that `object_store` reads.
#[derive(Debug, Clone, Default)]
pub struct S3Config {
    pub bucket: String,
    pub endpoint: Option<String>,
    pub region: Option<String>,
    pub access_key_id: Option<String>,
    pub secret_access_key: Option<String>,
    pub virtual_hosted_style: bool,
    /// Prepended to every object key, so one bucket can hold several things.
    pub prefix: String,
}

/// What an upload stream turned out to be.
#[derive(Debug, PartialEq, Eq)]
pub struct Received {
    pub md5: String,
    pub size: u64,
}

pub struct Storage {
    store: Arc<dyn ObjectStore>,
    /// Signs download URLs; `None` when downloads go through the server.
    signer: Option<Arc<dyn Signer>>,
    prefix: String,
    presign_ttl: Duration,
    description: String,
}

impl Storage {
    /// Connect to `backend`. With `proxy_downloads`, even an S3 backend serves
    /// downloads through the server (for a bucket clients cannot reach).
    pub fn new(
        backend: &Backend,
        presign_ttl: Duration,
        proxy_downloads: bool,
    ) -> Result<Self, AppError> {
        let storage = |store: Arc<dyn ObjectStore>, signer, prefix, description| Storage {
            store,
            signer,
            prefix,
            presign_ttl,
            description,
        };
        Ok(match backend {
            Backend::S3(cfg) => {
                let mut builder = AmazonS3Builder::from_env()
                    .with_bucket_name(&cfg.bucket)
                    .with_virtual_hosted_style_request(cfg.virtual_hosted_style);
                if let Some(endpoint) = &cfg.endpoint {
                    builder = builder
                        .with_endpoint(endpoint)
                        .with_allow_http(endpoint.starts_with("http://"));
                }
                if let Some(region) = &cfg.region {
                    builder = builder.with_region(region);
                }
                if let Some(key) = &cfg.access_key_id {
                    builder = builder.with_access_key_id(key);
                }
                if let Some(secret) = &cfg.secret_access_key {
                    builder = builder.with_secret_access_key(secret);
                }
                let s3 = Arc::new(builder.build()?);
                let signer = (!proxy_downloads).then(|| s3.clone() as Arc<dyn Signer>);
                let prefix = normalize_prefix(&cfg.prefix);
                let description = format!("s3://{}/{}", cfg.bucket, prefix);
                storage(s3, signer, prefix, description)
            }
            Backend::Local(dir) => {
                std::fs::create_dir_all(dir)
                    .map_err(|e| AppError::Internal(format!("create {}: {e}", dir.display())))?;
                storage(
                    Arc::new(LocalFileSystem::new_with_prefix(dir)?),
                    None,
                    String::new(),
                    format!("file://{}", dir.display()),
                )
            }
            #[cfg(test)]
            Backend::Memory => storage(
                Arc::new(InMemory::new()),
                None,
                String::new(),
                "memory://".into(),
            ),
        })
    }

    /// Where the bytes go, for logs and the dashboard.
    pub fn description(&self) -> &str {
        &self.description
    }

    /// Whether downloads redirect to pre-signed bucket URLs.
    pub fn signs_downloads(&self) -> bool {
        self.signer.is_some()
    }

    fn file_path(&self, item_key: &str) -> Path {
        Path::from(format!("{}files/{item_key}", self.prefix))
    }

    fn pending_path(&self, token: &str) -> Path {
        Path::from(format!("{}pending/{token}", self.prefix))
    }

    /// Stream an upload into `pending/<token>`, hashing it on the way. Fails
    /// with [`AppError::TooLarge`] past `max_size` bytes, removing what was
    /// written.
    pub async fn receive<S, E>(
        &self,
        token: &str,
        mut body: S,
        max_size: u64,
    ) -> Result<Received, AppError>
    where
        S: Stream<Item = Result<Bytes, E>> + Unpin,
        E: std::fmt::Display,
    {
        let path = self.pending_path(token);
        let mut writer = WriteMultipart::new(self.store.put_multipart(&path).await?);
        let mut hash = md5::Context::new();
        let mut size = 0u64;
        let failure = loop {
            match body.next().await {
                None => break None,
                Some(Err(e)) => {
                    break Some(AppError::BadRequest(format!("upload interrupted: {e}")));
                }
                Some(Ok(chunk)) => {
                    size += chunk.len() as u64;
                    if size > max_size {
                        break Some(AppError::TooLarge);
                    }
                    hash.consume(&chunk);
                    if let Err(e) = writer.wait_for_capacity(4).await {
                        break Some(e.into());
                    }
                    writer.put(chunk);
                }
            }
        };
        if let Some(error) = failure {
            let _ = writer.abort().await;
            return Err(error);
        }
        writer.finish().await?;
        Ok(Received {
            md5: format!("{:x}", hash.finalize()),
            size,
        })
    }

    /// Move a verified upload into place as the file of `item_key`.
    pub async fn commit(&self, token: &str, item_key: &str) -> Result<(), AppError> {
        self.store
            .rename(&self.pending_path(token), &self.file_path(item_key))
            .await?;
        Ok(())
    }

    /// Drop an upload that will not be registered.
    pub async fn discard(&self, token: &str) {
        if let Err(e) = self.store.delete(&self.pending_path(token)).await {
            tracing::debug!(error = %e, token, "discarding pending upload");
        }
    }

    /// Remove the file of a deleted item.
    pub async fn remove(&self, item_key: &str) {
        match self.store.delete(&self.file_path(item_key)).await {
            Ok(()) | Err(object_store::Error::NotFound { .. }) => {}
            Err(e) => tracing::warn!(error = %e, item_key, "removing attachment file"),
        }
    }

    /// A pre-signed GET URL for the file of `item_key`, when the backend signs.
    pub async fn signed_download(&self, item_key: &str) -> Result<Option<String>, AppError> {
        let Some(signer) = &self.signer else {
            return Ok(None);
        };
        let url = signer
            .signed_url(Method::GET, &self.file_path(item_key), self.presign_ttl)
            .await?;
        Ok(Some(url.to_string()))
    }

    /// The file of `item_key` as a response body, with its size.
    pub async fn download(&self, item_key: &str) -> Result<(Body, u64), AppError> {
        match self.store.get(&self.file_path(item_key)).await {
            Ok(result) => {
                let size = result.meta.size;
                Ok((Body::from_stream(result.into_stream()), size))
            }
            Err(object_store::Error::NotFound { .. }) => {
                Err(AppError::NotFound(format!("file of {item_key}")))
            }
            Err(e) => Err(e.into()),
        }
    }
}

fn normalize_prefix(prefix: &str) -> String {
    let trimmed = prefix.trim_matches('/');
    if trimmed.is_empty() {
        String::new()
    } else {
        format!("{trimmed}/")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn chunks(parts: &[&'static [u8]]) -> impl Stream<Item = Result<Bytes, String>> + Unpin {
        futures::stream::iter(
            parts
                .iter()
                .map(|p| Ok(Bytes::from_static(p)))
                .collect::<Vec<_>>(),
        )
    }

    #[tokio::test]
    async fn receives_hashes_commits_and_downloads() {
        let storage = Storage::new(&Backend::Memory, Duration::from_secs(60), false).unwrap();
        let received = storage
            .receive("tok", chunks(&[b"hello ", b"world"]), 1024)
            .await
            .unwrap();
        assert_eq!(received.size, 11);
        assert_eq!(received.md5, format!("{:x}", md5::compute(b"hello world")));
        storage.commit("tok", "AAAAAAAA").await.unwrap();
        let (body, size) = storage.download("AAAAAAAA").await.unwrap();
        assert_eq!(size, 11);
        let bytes = axum::body::to_bytes(body, 1024).await.unwrap();
        assert_eq!(&bytes[..], b"hello world");
        assert!(storage.signed_download("AAAAAAAA").await.unwrap().is_none());
        storage.remove("AAAAAAAA").await;
        assert!(matches!(
            storage.download("AAAAAAAA").await,
            Err(AppError::NotFound(_))
        ));
    }

    #[tokio::test]
    async fn rejects_oversized_uploads() {
        let storage = Storage::new(&Backend::Memory, Duration::from_secs(60), false).unwrap();
        let result = storage.receive("tok", chunks(&[b"0123456789"]), 5).await;
        assert!(matches!(result, Err(AppError::TooLarge)));
    }

    #[test]
    fn s3_signs_downloads_unless_proxied() {
        let backend = Backend::S3(S3Config {
            bucket: "zotero".into(),
            endpoint: Some("http://127.0.0.1:9000".into()),
            region: Some("auto".into()),
            access_key_id: Some("key".into()),
            secret_access_key: Some("secret".into()),
            prefix: "/defter/".into(),
            ..Default::default()
        });
        let storage = Storage::new(&backend, Duration::from_secs(60), false).unwrap();
        assert!(storage.signs_downloads());
        assert_eq!(storage.description(), "s3://zotero/defter/");
        assert_eq!(storage.file_path("K").as_ref(), "defter/files/K");
        let proxied = Storage::new(&backend, Duration::from_secs(60), true).unwrap();
        assert!(!proxied.signs_downloads());
    }

    #[tokio::test]
    async fn presigns_s3_urls() {
        let backend = Backend::S3(S3Config {
            bucket: "zotero".into(),
            endpoint: Some("http://127.0.0.1:9000".into()),
            region: Some("auto".into()),
            access_key_id: Some("key".into()),
            secret_access_key: Some("secret".into()),
            ..Default::default()
        });
        let storage = Storage::new(&backend, Duration::from_secs(60), false).unwrap();
        let url = storage.signed_download("AAAAAAAA").await.unwrap().unwrap();
        assert!(
            url.starts_with("http://127.0.0.1:9000/zotero/files/AAAAAAAA?"),
            "{url}"
        );
        assert!(url.contains("X-Amz-Signature="));
    }
}