defter 0.5.2

//! 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 its library's files ([`FileScope`]):
//! `files/<item key>` for a library from before there were several,
//! `users/<id>/files/<item key>` or `groups/<id>/files/<item key>`. While
//! backups run, a deleted or replaced file moves to
//! `deleted/<date>/<time>-<n>/<its key>` until they purge it, so a restored
//! library finds its files (see `backup`).

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,
}

/// Where a library's attachment files are.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum FileScope {
    /// `files/`: the one library of a server from before there were several.
    Legacy,
    User(u64),
    #[allow(dead_code, reason = "group libraries come in a later release")]
    Group(u64),
}

impl FileScope {
    /// The key prefix of the library's files, ending in `/`.
    pub fn prefix(&self) -> String {
        match self {
            FileScope::Legacy => "files/".to_string(),
            FileScope::User(id) => format!("users/{id}/files/"),
            FileScope::Group(id) => format!("groups/{id}/files/"),
        }
    }

    /// The key of the file of `item_key`, without the storage prefix.
    pub fn key(&self, item_key: &str) -> String {
        format!("{}{item_key}", self.prefix())
    }
}

/// An object under the storage prefix, as listed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Listed {
    /// Its key, without the prefix.
    pub key: String,
    pub size: u64,
    /// When it was last written, in seconds since the epoch.
    pub modified: 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,
    /// Move deleted files aside for backups to find instead of removing them.
    keep_deleted: std::sync::atomic::AtomicBool,
    /// Tells apart files set aside in the same second.
    set_aside: std::sync::atomic::AtomicU64,
}

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,
            keep_deleted: false.into(),
            set_aside: 0.into(),
        };
        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, scope: &FileScope, item_key: &str) -> Path {
        self.raw_path(&scope.key(item_key))
    }

    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`. While
    /// backups run, a file it replaces is set aside first.
    pub async fn commit(
        &self,
        token: &str,
        scope: &FileScope,
        item_key: &str,
    ) -> Result<(), AppError> {
        if self.keeps_deleted() {
            self.set_aside(&scope.key(item_key)).await?;
        }
        self.store
            .rename(&self.pending_path(token), &self.file_path(scope, 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");
        }
    }

    /// Keep deleted files under `deleted/` for backups (see [`Storage::remove`]).
    pub fn keep_deleted(&self, keep: bool) {
        self.keep_deleted
            .store(keep, std::sync::atomic::Ordering::Relaxed);
    }

    fn keeps_deleted(&self) -> bool {
        self.keep_deleted.load(std::sync::atomic::Ordering::Relaxed)
    }

    /// Move `key`, if it exists, to `deleted/<date>/<time>-<n>/<key>`, for
    /// backups to find and later purge.
    pub async fn set_aside(&self, key: &str) -> Result<(), AppError> {
        let now = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map_or(0, |d| d.as_secs());
        let n = self
            .set_aside
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        let aside = format!(
            "deleted/{}/{}-{n:06}/{key}",
            crate::backup::date(now),
            crate::backup::timestamp(now)
        );
        match self
            .store
            .rename(&self.raw_path(key), &self.raw_path(&aside))
            .await
        {
            Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()),
            Err(e) => Err(e.into()),
        }
    }

    /// Remove the file of a deleted item, or, while backups run, set it aside
    /// for them.
    pub async fn remove(&self, scope: &FileScope, item_key: &str) {
        let result = if self.keeps_deleted() {
            self.set_aside(&scope.key(item_key)).await
        } else {
            match self.store.delete(&self.file_path(scope, item_key)).await {
                Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()),
                Err(e) => Err(e.into()),
            }
        };
        if let Err(e) = result {
            tracing::warn!(error = %e, item_key, "removing attachment file");
        }
    }

    fn raw_path(&self, key: &str) -> Path {
        Path::from(format!("{}{key}", self.prefix))
    }

    /// Every object whose key starts with `prefix`.
    pub async fn list(&self, prefix: &str) -> Result<Vec<Listed>, AppError> {
        let path = self.raw_path(prefix.trim_end_matches('/'));
        let mut objects = Vec::new();
        let mut stream = self.store.list(Some(&path));
        while let Some(meta) = stream.next().await {
            let meta = meta?;
            let location = meta.location.as_ref();
            let key = location.strip_prefix(&self.prefix).unwrap_or(location);
            objects.push(Listed {
                key: key.to_string(),
                size: meta.size,
                modified: u64::try_from(meta.last_modified.timestamp()).unwrap_or(0),
            });
        }
        Ok(objects)
    }

    pub async fn put_bytes(&self, key: &str, bytes: Vec<u8>) -> Result<(), AppError> {
        self.store
            .put(&self.raw_path(key), Bytes::from(bytes).into())
            .await?;
        Ok(())
    }

    /// Upload the file at `file` as `key`, in parts.
    pub async fn put_file(&self, key: &str, file: &std::path::Path) -> Result<(), AppError> {
        use tokio::io::AsyncReadExt;
        let mut input = tokio::fs::File::open(file)
            .await
            .map_err(|e| AppError::Internal(format!("reading {}: {e}", file.display())))?;
        let mut writer = WriteMultipart::new(self.store.put_multipart(&self.raw_path(key)).await?);
        let failure = loop {
            let mut chunk = vec![0u8; 8 << 20];
            let n = match input.read(&mut chunk).await {
                Ok(n) => n,
                Err(e) => {
                    break Some(AppError::Internal(format!(
                        "reading {}: {e}",
                        file.display()
                    )));
                }
            };
            if n == 0 {
                break None;
            }
            chunk.truncate(n);
            if let Err(e) = writer.wait_for_capacity(4).await {
                break Some(e.into());
            }
            writer.put(Bytes::from(chunk));
        };
        if let Some(error) = failure {
            // Leave no unfinished upload behind to be billed for.
            let _ = writer.abort().await;
            return Err(error);
        }
        writer.finish().await?;
        Ok(())
    }

    pub async fn get_bytes(&self, key: &str) -> Result<Vec<u8>, AppError> {
        let result = self.store.get(&self.raw_path(key)).await?;
        Ok(result.bytes().await?.to_vec())
    }

    /// Download `key` into the file at `file`.
    pub async fn get_to_file(&self, key: &str, file: &std::path::Path) -> Result<(), AppError> {
        use tokio::io::AsyncWriteExt;
        let io = |e: std::io::Error| AppError::Internal(format!("writing {}: {e}", file.display()));
        let mut out = tokio::fs::File::create(file).await.map_err(io)?;
        let mut stream = self.store.get(&self.raw_path(key)).await?.into_stream();
        while let Some(chunk) = stream.next().await {
            out.write_all(&chunk?).await.map_err(io)?;
        }
        out.sync_all().await.map_err(io)?;
        Ok(())
    }

    pub async fn exists(&self, key: &str) -> Result<bool, AppError> {
        match self.store.head(&self.raw_path(key)).await {
            Ok(_) => Ok(true),
            Err(object_store::Error::NotFound { .. }) => Ok(false),
            Err(e) => Err(e.into()),
        }
    }

    pub async fn copy(&self, from: &str, to: &str) -> Result<(), AppError> {
        self.store
            .copy(&self.raw_path(from), &self.raw_path(to))
            .await?;
        Ok(())
    }

    pub async fn delete(&self, key: &str) -> Result<(), AppError> {
        match self.store.delete(&self.raw_path(key)).await {
            Ok(()) | Err(object_store::Error::NotFound { .. }) => Ok(()),
            Err(e) => Err(e.into()),
        }
    }

    /// A pre-signed GET URL for the file of `item_key`, when the backend signs.
    pub async fn signed_download(
        &self,
        scope: &FileScope,
        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(scope, 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,
        scope: &FileScope,
        item_key: &str,
    ) -> Result<(Body, u64), AppError> {
        match self.store.get(&self.file_path(scope, 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")));
        let scope = FileScope::User(1);
        storage.commit("tok", &scope, "AAAAAAAA").await.unwrap();
        assert!(storage.exists("users/1/files/AAAAAAAA").await.unwrap());
        let (body, size) = storage.download(&scope, "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(&scope, "AAAAAAAA")
                .await
                .unwrap()
                .is_none()
        );
        // Another library's files are elsewhere.
        assert!(
            storage
                .download(&FileScope::User(2), "AAAAAAAA")
                .await
                .is_err()
        );
        storage.remove(&scope, "AAAAAAAA").await;
        assert!(matches!(
            storage.download(&scope, "AAAAAAAA").await,
            Err(AppError::NotFound(_))
        ));
    }

    #[tokio::test]
    async fn keeps_deleted_files_aside_for_backups() {
        let storage = Storage::new(&Backend::Memory, Duration::from_secs(60), false).unwrap();
        storage
            .put_bytes("files/AAAAAAAA", b"one".to_vec())
            .await
            .unwrap();
        storage.keep_deleted(true);
        storage.remove(&FileScope::Legacy, "AAAAAAAA").await;
        assert!(!storage.exists("files/AAAAAAAA").await.unwrap());
        let aside = storage.list("deleted/").await.unwrap();
        assert_eq!(aside.len(), 1);
        assert!(aside[0].key.ends_with("/files/AAAAAAAA"), "{aside:?}");
        assert_eq!(storage.get_bytes(&aside[0].key).await.unwrap(), b"one");
        // Without backups, deleting deletes.
        storage.keep_deleted(false);
        storage
            .put_bytes("files/BBBBBBBB", b"two".to_vec())
            .await
            .unwrap();
        storage.remove(&FileScope::Legacy, "BBBBBBBB").await;
        assert_eq!(storage.list("deleted/").await.unwrap().len(), 1);
        // Removing what is not there is fine either way.
        storage.remove(&FileScope::Legacy, "CCCCCCCC").await;
        // While backups run, an upload replacing a file sets the old one aside.
        storage.keep_deleted(true);
        for (token, bytes) in [("t1", b"old" as &'static [u8]), ("t2", b"new")] {
            storage
                .receive(token, chunks(&[bytes]), 1024)
                .await
                .unwrap();
            storage
                .commit(token, &FileScope::Legacy, "DDDDDDDD")
                .await
                .unwrap();
        }
        assert_eq!(storage.get_bytes("files/DDDDDDDD").await.unwrap(), b"new");
        let aside = storage.list("deleted/").await.unwrap();
        let old = aside
            .iter()
            .find(|o| o.key.ends_with("/files/DDDDDDDD"))
            .unwrap();
        assert_eq!(storage.get_bytes(&old.key).await.unwrap(), b"old");
        assert_eq!(aside.len(), 2);
    }

    #[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(&FileScope::Legacy, "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(&FileScope::Legacy, "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(&FileScope::Legacy, "AAAAAAAA")
            .await
            .unwrap()
            .unwrap();
        assert!(
            url.starts_with("http://127.0.0.1:9000/zotero/files/AAAAAAAA?"),
            "{url}"
        );
        assert!(url.contains("X-Amz-Signature="));
    }
}