Mangalord/backend/src/storage/local.rs

use std::path::{Path, PathBuf};

use async_trait::async_trait;
use futures_util::StreamExt as _;
use tokio::fs;
use tokio::io::AsyncWriteExt as _;
use tokio_util::io::ReaderStream;

use super::{PutByteStream, Storage, StorageError, StreamingFile};

pub struct LocalStorage {
    root: PathBuf,
}

impl LocalStorage {
    pub fn new(root: impl Into<PathBuf>) -> Self {
        Self { root: root.into() }
    }

    fn resolve(&self, key: &str) -> Result<PathBuf, StorageError> {
        // NUL bytes are rejected by the Linux syscall layer, but the
        // error surfaces as an opaque IO failure rather than the
        // explicit `BadKey` the rest of the contract uses. Catch it
        // here so the error path is consistent.
        if key.contains('\0') {
            return Err(StorageError::BadKey);
        }
        let key = key.trim_start_matches('/');
        if key.is_empty() {
            return Err(StorageError::BadKey);
        }
        if key.split('/').any(|seg| seg.is_empty() || seg == "." || seg == "..") {
            return Err(StorageError::BadKey);
        }
        Ok(self.root.join(key))
    }
}

#[async_trait]
impl Storage for LocalStorage {
    async fn put(&self, key: &str, bytes: &[u8]) -> Result<(), StorageError> {
        let path = self.resolve(key)?;
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).await?;
        }
        fs::write(path, bytes).await?;
        Ok(())
    }

    async fn put_stream(
        &self,
        key: &str,
        mut stream: PutByteStream<'_>,
    ) -> Result<u64, StorageError> {
        let path = self.resolve(key)?;
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent).await?;
        }
        // Atomic install via temp + rename. A failure mid-stream
        // removes the temp so nothing is visible at `key`. The temp
        // name uses a UUID suffix so concurrent puts of the same key
        // (e.g. two workers racing a retry) don't clobber each
        // other's in-progress file before the rename.
        let tmp = path.with_extension(format!(
            "{}.tmp.{}",
            path.extension().and_then(|e| e.to_str()).unwrap_or(""),
            uuid::Uuid::new_v4().simple()
        ));
        let mut written: u64 = 0;
        let result: Result<(), StorageError> = async {
            let mut f = fs::File::create(&tmp).await?;
            while let Some(chunk) = stream.next().await {
                let chunk = chunk?;
                f.write_all(&chunk).await?;
                written = written.saturating_add(chunk.len() as u64);
            }
            // fsync before rename so a power-loss can't leave a
            // zero-byte file at the destination.
            f.sync_all().await?;
            fs::rename(&tmp, &path).await?;
            Ok(())
        }
        .await;
        if let Err(e) = result {
            // Best-effort cleanup; ignore "no such file" if the temp
            // was never created.
            let _ = fs::remove_file(&tmp).await;
            return Err(e);
        }
        Ok(written)
    }

    async fn get(&self, key: &str) -> Result<Vec<u8>, StorageError> {
        let path = self.resolve(key)?;
        match fs::read(&path).await {
            Ok(b) => Ok(b),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Err(StorageError::NotFound),
            Err(e) => Err(e.into()),
        }
    }

    async fn get_stream(&self, key: &str) -> Result<StreamingFile, StorageError> {
        let path = self.resolve(key)?;
        let file = match fs::File::open(&path).await {
            Ok(f) => f,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
                return Err(StorageError::NotFound)
            }
            Err(e) => return Err(e.into()),
        };
        let size_bytes = file.metadata().await?.len();
        // 64 KiB chunks: small enough that a few-MB page emits many frames
        // (so streaming is observable), large enough to keep syscalls cheap.
        let stream = ReaderStream::with_capacity(file, 64 * 1024);
        Ok(StreamingFile {
            stream: Box::pin(stream),
            size_bytes,
        })
    }

    async fn delete(&self, key: &str) -> Result<(), StorageError> {
        let path = self.resolve(key)?;
        match fs::remove_file(&path).await {
            Ok(()) => Ok(()),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Err(StorageError::NotFound),
            Err(e) => Err(e.into()),
        }
    }

    async fn exists(&self, key: &str) -> Result<bool, StorageError> {
        let path: &Path = &self.resolve(key)?;
        Ok(fs::try_exists(path).await?)
    }

    fn local_root(&self) -> Option<&Path> {
        Some(&self.root)
    }
}

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

    #[tokio::test]
    async fn put_get_delete_roundtrip() {
        let dir = tempdir().unwrap();
        let s = LocalStorage::new(dir.path());

        s.put("mangas/abc/cover.jpg", b"hello").await.unwrap();
        assert!(s.exists("mangas/abc/cover.jpg").await.unwrap());
        assert_eq!(s.get("mangas/abc/cover.jpg").await.unwrap(), b"hello");
        s.delete("mangas/abc/cover.jpg").await.unwrap();
        assert!(!s.exists("mangas/abc/cover.jpg").await.unwrap());
    }

    #[tokio::test]
    async fn rejects_path_traversal() {
        let dir = tempdir().unwrap();
        let s = LocalStorage::new(dir.path());
        // Parent-dir reference at the start.
        assert!(matches!(s.put("../escape", b"x").await, Err(StorageError::BadKey)));
        // Parent-dir reference mid-path.
        assert!(matches!(s.get("a/../../b").await, Err(StorageError::BadKey)));
        // Empty key.
        assert!(matches!(s.exists("").await, Err(StorageError::BadKey)));
        // Current-dir reference (the implementation rejects `.` segments
        // alongside `..`; this exercises that arm).
        assert!(matches!(s.get("a/./b").await, Err(StorageError::BadKey)));
        assert!(matches!(s.get(".").await, Err(StorageError::BadKey)));
        // Empty segment via doubled slash.
        assert!(matches!(s.get("a//b").await, Err(StorageError::BadKey)));
        // NUL byte (rejected explicitly so callers see BadKey rather
        // than an opaque IO error from the kernel).
        assert!(matches!(s.put("a\0b", b"x").await, Err(StorageError::BadKey)));
    }

    #[tokio::test]
    async fn missing_key_is_not_found() {
        let dir = tempdir().unwrap();
        let s = LocalStorage::new(dir.path());
        assert!(matches!(s.get("nope").await, Err(StorageError::NotFound)));
        assert!(matches!(s.delete("nope").await, Err(StorageError::NotFound)));
        assert!(matches!(
            s.get_stream("nope").await.err(),
            Some(StorageError::NotFound)
        ));
    }

    #[tokio::test]
    async fn put_stream_writes_full_body_and_removes_temp_on_error() {
        use bytes::Bytes;
        use futures_util::stream;

        let dir = tempdir().unwrap();
        let s = LocalStorage::new(dir.path());

        // Success: a 3-chunk stream writes the concatenation and
        // returns the right byte count.
        let chunks: Vec<Result<Bytes, StorageError>> = vec![
            Ok(Bytes::from_static(b"alpha-")),
            Ok(Bytes::from_static(b"beta-")),
            Ok(Bytes::from_static(b"gamma")),
        ];
        let bytes_written = s
            .put_stream("streamed/ok.bin", Box::pin(stream::iter(chunks)))
            .await
            .unwrap();
        assert_eq!(bytes_written, b"alpha-beta-gamma".len() as u64);
        assert_eq!(s.get("streamed/ok.bin").await.unwrap(), b"alpha-beta-gamma");

        // Failure mid-stream: nothing is visible at the destination
        // and no .tmp file is left behind.
        let chunks: Vec<Result<Bytes, StorageError>> = vec![
            Ok(Bytes::from_static(b"good")),
            Err(StorageError::Io(std::io::Error::other("boom"))),
        ];
        let err = s
            .put_stream("streamed/bad.bin", Box::pin(stream::iter(chunks)))
            .await
            .unwrap_err();
        assert!(matches!(err, StorageError::Io(_)));
        assert!(matches!(
            s.get("streamed/bad.bin").await,
            Err(StorageError::NotFound)
        ));
        // No stray temp file in the streamed/ directory.
        let entries: Vec<_> = std::fs::read_dir(dir.path().join("streamed"))
            .unwrap()
            .filter_map(|e| e.ok())
            .map(|e| e.file_name().to_string_lossy().to_string())
            .collect();
        assert_eq!(entries, vec!["ok.bin"]);
    }

    #[tokio::test]
    async fn get_stream_emits_multiple_chunks_for_large_files() {
        use futures_util::StreamExt as _;

        let dir = tempdir().unwrap();
        let s = LocalStorage::new(dir.path());
        // 256 KiB blob → at 64 KiB chunks should emit ~4 chunks.
        let big = vec![7u8; 256 * 1024];
        s.put("big.bin", &big).await.unwrap();

        let StreamingFile { mut stream, size_bytes } = s.get_stream("big.bin").await.unwrap();
        assert_eq!(size_bytes, big.len() as u64);

        let mut chunks = 0usize;
        let mut total = 0usize;
        while let Some(frame) = stream.next().await {
            let bytes = frame.unwrap();
            chunks += 1;
            total += bytes.len();
        }
        assert_eq!(total, big.len());
        assert!(chunks > 1, "expected >1 chunk, got {chunks}");
    }
}