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) -> Self { Self { root: root.into() } } fn resolve(&self, key: &str) -> Result { // 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 { 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, 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 { 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 { let path: &Path = &self.resolve(key)?; Ok(fs::try_exists(path).await?) } async fn size(&self, key: &str) -> Result { let path = self.resolve(key)?; match fs::metadata(&path).await { // A key resolving to a directory isn't a stored blob; treat it // as absent rather than returning the directory's inode size as // a bogus "measured" byte count. Ok(m) if !m.is_file() => Err(StorageError::NotFound), Ok(m) => Ok(m.len()), Err(e) if e.kind() == std::io::ErrorKind::NotFound => Err(StorageError::NotFound), Err(e) => Err(e.into()), } } 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 size_returns_byte_length() { let dir = tempdir().unwrap(); let s = LocalStorage::new(dir.path()); s.put("mangas/abc/cover.jpg", b"hello world").await.unwrap(); assert_eq!(s.size("mangas/abc/cover.jpg").await.unwrap(), 11); } #[tokio::test] async fn size_missing_is_not_found() { let dir = tempdir().unwrap(); let s = LocalStorage::new(dir.path()); assert!(matches!(s.size("nope").await, Err(StorageError::NotFound))); // Bad keys still surface as BadKey, consistent with the rest. assert!(matches!(s.size("../escape").await, Err(StorageError::BadKey))); } #[tokio::test] async fn size_on_directory_is_not_found() { // A key resolving to a directory must not report the dir inode size // as a "measured" blob size. let dir = tempdir().unwrap(); let s = LocalStorage::new(dir.path()); std::fs::create_dir(dir.path().join("adir")).unwrap(); assert!(matches!(s.size("adir").await, Err(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> = 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> = 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}"); } }