defter 0.3.8

//! 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>,
    /// Address the bucket as `<bucket>.<endpoint host>`. Always on for Hetzner
    /// Object Storage, whose pre-signed URLs need it.
    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 virtual_hosted =
                    cfg.virtual_hosted_style || cfg.endpoint.as_deref().is_some_and(is_hetzner);
                let mut builder = AmazonS3Builder::from_env()
                    .with_bucket_name(&cfg.bucket)
                    .with_virtual_hosted_style_request(virtual_hosted);
                if let Some(endpoint) = &cfg.endpoint {
                    // object_store expects a virtual-hosted endpoint to name
                    // the bucket already.
                    let endpoint = if virtual_hosted {
                        bucket_endpoint(endpoint, &cfg.bucket)
                    } else {
                        endpoint.trim_end_matches('/').to_string()
                    };
                    builder = builder
                        .with_allow_http(endpoint.starts_with("http://"))
                        .with_endpoint(endpoint);
                }
                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()),
        }
    }
}

/// Whether `endpoint` is Hetzner Object Storage (`https://<location>.your-objectstorage.com`).
fn is_hetzner(endpoint: &str) -> bool {
    host_of(endpoint).is_some_and(|host| host.ends_with(".your-objectstorage.com"))
}

fn host_of(endpoint: &str) -> Option<&str> {
    let rest = endpoint
        .split_once("://")
        .map_or(endpoint, |(_, rest)| rest);
    let host = rest.split(['/', ':']).next()?;
    (!host.is_empty()).then_some(host)
}

/// The endpoint with the bucket as the first label of its host:
/// `https://fsn1.your-objectstorage.com` →
/// `https://zotero.fsn1.your-objectstorage.com`. Unchanged if it names the
/// bucket already.
fn bucket_endpoint(endpoint: &str, bucket: &str) -> String {
    let endpoint = endpoint.trim_end_matches('/');
    let (scheme, rest) = endpoint
        .split_once("://")
        .map_or(("https", endpoint), |(scheme, rest)| (scheme, rest));
    if rest.starts_with(&format!("{bucket}.")) {
        return format!("{scheme}://{rest}");
    }
    format!("{scheme}://{bucket}.{rest}")
}

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());
    }

    #[test]
    fn hetzner_endpoints_are_addressed_by_bucket() {
        assert!(is_hetzner("https://fsn1.your-objectstorage.com"));
        assert!(is_hetzner("https://zotero.nbg1.your-objectstorage.com/"));
        assert!(!is_hetzner("http://127.0.0.1:9000"));
        assert!(!is_hetzner("https://your-objectstorage.com.example.org"));
        assert_eq!(
            bucket_endpoint("https://fsn1.your-objectstorage.com/", "zotero"),
            "https://zotero.fsn1.your-objectstorage.com"
        );
        assert_eq!(
            bucket_endpoint("https://zotero.fsn1.your-objectstorage.com", "zotero"),
            "https://zotero.fsn1.your-objectstorage.com"
        );
        assert_eq!(
            bucket_endpoint("http://localhost:9000", "zotero"),
            "http://zotero.localhost:9000"
        );
    }

    #[tokio::test]
    async fn presigns_hetzner_urls_with_the_bucket_in_the_host() {
        let backend = Backend::S3(S3Config {
            bucket: "zotero".into(),
            endpoint: Some("https://fsn1.your-objectstorage.com".into()),
            region: Some("fsn1".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();
        let url = storage.signed_download("AAAAAAAA").await.unwrap().unwrap();
        assert!(
            url.starts_with("https://zotero.fsn1.your-objectstorage.com/defter/files/AAAAAAAA?"),
            "{url}"
        );
        assert!(url.contains("X-Amz-Credential=key%2F"), "{url}");
        assert!(url.contains("%2Ffsn1%2Fs3%2Faws4_request"), "{url}");
    }

    #[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="));
    }
}