use std::{fs, io, path::PathBuf}; use crate::{ control::{ ControlMetadata, decode_control_metadata, encode_control_metadata, read_aria2_control_file, write_aria2_control_file, }, model::{PieceIndex, PieceMap, PieceState}, resume::{ResumeData, ResumeStore}, session::{Aria2MetadataState, SessionFile, load_session_file, save_session_file}, }; /// Persists control metadata by download gid. pub trait ControlStore { /// Concrete error type returned by the store implementation. type Error; /// Loads persisted control metadata for a download gid. /// /// # Errors /// /// Returns the store-specific error when control metadata cannot be read. fn load_control(&self, gid: &str) -> Result, Self::Error>; /// Persists control metadata for a download gid. /// /// # Errors /// /// Returns the store-specific error when control metadata cannot be written. fn save_control(&self, gid: &str, value: &ControlMetadata) -> Result<(), Self::Error>; /// Deletes control metadata for a download gid. /// /// # Errors /// /// Returns the store-specific error when control metadata cannot be removed. fn delete_control(&self, gid: &str) -> Result<(), Self::Error>; } /// Persists the session file. pub trait SessionStore { /// Concrete error type returned by the store implementation. type Error; /// Loads the persisted session file. /// /// # Errors /// /// Returns the store-specific error when the session file cannot be read. fn load_session(&self) -> Result; /// Persists the session file. /// /// # Errors /// /// Returns the store-specific error when the session file cannot be written. fn save_session(&self, session: &SessionFile) -> Result<(), Self::Error>; } /// Persists metadata and resume state keyed by gid. pub trait MetadataStore { /// Concrete error type returned by the store implementation. type Error; /// Loads persisted metadata state for a gid. /// /// # Errors /// /// Returns the store-specific error when metadata state cannot be read. fn load_metadata_state(&self, gid: &str) -> Result, Self::Error>; /// Persists metadata state for a gid. /// /// # Errors /// /// Returns the store-specific error when metadata state cannot be written. fn save_metadata_state(&self, state: &Aria2MetadataState) -> Result<(), Self::Error>; /// Persists resume data through the metadata-oriented store surface. /// /// # Errors /// /// Returns the store-specific error when resume data cannot be written. fn save_resume_data(&self, resume: &ResumeData) -> Result<(), Self::Error>; } /// Error type used by the local filesystem-backed store. #[derive(Debug)] pub enum LocalStoreError { /// Raw filesystem I/O failure. Io(io::Error), /// Format or decoding failure. Parse(String), } impl From for LocalStoreError { fn from(value: io::Error) -> Self { Self::Io(value) } } /// Local filesystem-backed implementation of the storage traits. #[derive(Debug)] pub struct LocalFileStore { /// Root directory that contains the store layout. root: PathBuf, } impl LocalFileStore { /// Creates a store rooted at the provided directory. #[must_use] pub const fn new(root: PathBuf) -> Self { Self { root } } /// Creates the on-disk directory layout used by the local store. /// /// # Errors /// /// Returns [`LocalStoreError`] when any required directory cannot be created. pub fn ensure_layout(&self) -> Result<(), LocalStoreError> { fs::create_dir_all(self.controls_dir())?; fs::create_dir_all(self.metadata_dir())?; fs::create_dir_all(self.resume_dir())?; if let Some(parent) = self.session_path().parent() { fs::create_dir_all(parent)?; } Ok(()) } /// Returns the control-file directory. fn controls_dir(&self) -> PathBuf { self.root.join("control") } /// Returns the metadata directory. fn metadata_dir(&self) -> PathBuf { self.root.join("metadata") } /// Returns the resume directory. fn resume_dir(&self) -> PathBuf { self.root.join("resume") } /// Returns the persisted session file path. fn session_path(&self) -> PathBuf { self.root.join("session").join("session.txt") } /// Returns the control-file path for one gid. fn control_path(&self, gid: &str) -> PathBuf { self.controls_dir().join(format!("{gid}.aria2")) } /// Returns the metadata path for one gid. fn metadata_path(&self, gid: &str) -> PathBuf { self.metadata_dir().join(format!("{gid}.meta")) } /// Returns the resume-file path for one gid. fn resume_path(&self, gid: &str) -> PathBuf { self.resume_dir().join(format!("{gid}.resume")) } } impl ControlStore for LocalFileStore { type Error = LocalStoreError; fn load_control(&self, gid: &str) -> Result, Self::Error> { let path = self.control_path(gid); if !path.exists() { return Ok(None); } read_aria2_control_file(&path) .map(Some) .map_err(|e| LocalStoreError::Parse(e.to_string())) } fn save_control(&self, gid: &str, value: &ControlMetadata) -> Result<(), Self::Error> { self.ensure_layout()?; write_aria2_control_file(&self.control_path(gid), value) .map_err(|e| LocalStoreError::Parse(e.to_string())) } fn delete_control(&self, gid: &str) -> Result<(), Self::Error> { let path = self.control_path(gid); if path.exists() { fs::remove_file(path)?; } Ok(()) } } impl SessionStore for LocalFileStore { type Error = LocalStoreError; fn load_session(&self) -> Result { let path = self.session_path(); if !path.exists() { return Ok(SessionFile::default()); } load_session_file(&path).map_err(LocalStoreError::Io) } fn save_session(&self, session: &SessionFile) -> Result<(), Self::Error> { self.ensure_layout()?; save_session_file(&self.session_path(), session).map_err(LocalStoreError::Io) } } impl MetadataStore for LocalFileStore { type Error = LocalStoreError; fn load_metadata_state(&self, gid: &str) -> Result, Self::Error> { let path = self.metadata_path(gid); if !path.exists() { return Ok(None); } let content = fs::read_to_string(path)?; let mut kv = std::collections::BTreeMap::new(); for line in content.lines() { let mut parts = line.splitn(2, '='); let k = parts.next().unwrap_or_default(); let v = parts.next().unwrap_or_default(); if !k.is_empty() { kv.insert(k.to_owned(), v.to_owned()); } } Ok(Some(Aria2MetadataState { gid: gid.to_owned(), kv, })) } fn save_metadata_state(&self, state: &Aria2MetadataState) -> Result<(), Self::Error> { self.ensure_layout()?; let body = state .kv .iter() .map(|(k, v)| format!("{k}={v}")) .collect::>() .join("\n"); fs::write(self.metadata_path(&state.gid), body)?; Ok(()) } fn save_resume_data(&self, resume: &ResumeData) -> Result<(), Self::Error> { ::save(self, resume) } } impl ResumeStore for LocalFileStore { type Error = LocalStoreError; fn load(&self, gid: &str) -> Result, Self::Error> { let path = self.resume_path(gid); if !path.exists() { return Ok(None); } let content = fs::read_to_string(path)?; let mut download_path = None; let mut verified = Vec::new(); let mut inflight = Vec::new(); let mut failed = Vec::new(); let mut metadata_blob = String::new(); let mut in_metadata = false; for line in content.lines() { if line == "metadata<<" { in_metadata = true; continue; } if line == ">>metadata" { in_metadata = false; continue; } if in_metadata { metadata_blob.push_str(line); metadata_blob.push('\n'); continue; } if let Some(v) = line.strip_prefix("download_path=") { download_path = Some(PathBuf::from(v)); } if let Some(v) = line.strip_prefix("verified=") { verified = parse_csv_u32(v)?; } if let Some(v) = line.strip_prefix("inflight=") { inflight = parse_csv_u32(v)?; } if let Some(v) = line.strip_prefix("failed=") { failed = parse_csv_u32(v)?; } } let mut piece_map = PieceMap::default(); for idx in verified { piece_map.set_state(PieceIndex(idx), PieceState::Verified); } for idx in inflight { piece_map.set_state(PieceIndex(idx), PieceState::InFlight); } for idx in failed { piece_map.set_state(PieceIndex(idx), PieceState::Failed); } let metadata = if metadata_blob.trim().is_empty() { None } else { Some( decode_control_metadata(metadata_blob.trim_end()) .map_err(|e| LocalStoreError::Parse(e.to_string()))?, ) }; Ok(Some(ResumeData { gid: gid.to_owned(), download_path: download_path.unwrap_or_default(), metadata, piece_map, })) } fn save(&self, resume: &ResumeData) -> Result<(), Self::Error> { self.ensure_layout()?; let mut lines = Vec::new(); lines.push(format!( "download_path={}", resume.download_path.to_string_lossy() )); lines.push(format!( "verified={}", join_piece_indexes(&resume.piece_map, PieceState::Verified) )); lines.push(format!( "inflight={}", join_piece_indexes(&resume.piece_map, PieceState::InFlight) )); lines.push(format!( "failed={}", join_piece_indexes(&resume.piece_map, PieceState::Failed) )); if let Some(metadata) = &resume.metadata { lines.push("metadata<<".to_owned()); lines.push(encode_control_metadata(metadata)); lines.push(">>metadata".to_owned()); } fs::write(self.resume_path(&resume.gid), lines.join("\n"))?; Ok(()) } fn remove(&self, gid: &str) -> Result<(), Self::Error> { let path = self.resume_path(gid); if path.exists() { fs::remove_file(path)?; } Ok(()) } } /// Parses a comma-separated list of piece indexes. fn parse_csv_u32(raw: &str) -> Result, LocalStoreError> { if raw.trim().is_empty() { return Ok(Vec::new()); } raw.split(',') .map(|s| { s.parse::() .map_err(|_| LocalStoreError::Parse(format!("invalid piece index: {s}"))) }) .collect() } /// Joins piece indexes in the provided state into a comma-separated string. fn join_piece_indexes(piece_map: &PieceMap, target: PieceState) -> String { piece_map .iter() .filter(|(_, state)| **state == target) .map(|(idx, _)| idx.0.to_string()) .collect::>() .join(",") } #[cfg(test)] mod tests { use std::{ collections::BTreeMap, time::{SystemTime, UNIX_EPOCH}, }; use super::*; use crate::session::SessionFileEntry; fn temp_root() -> PathBuf { let millis = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_millis(); std::env::temp_dir().join(format!("aria2-rust-pro-storage-test-{millis}")) } #[test] fn local_store_session_roundtrip() { let root = temp_root(); let store = LocalFileStore::new(root.clone()); let mut metadata = BTreeMap::new(); metadata.insert("bt.name".to_owned(), "ubuntu".to_owned()); let session = SessionFile { entries: vec![SessionFileEntry { gid: "gid-1".to_owned(), uri: "https://mirror-1.example/file.iso".to_owned(), uris: vec![ "https://mirror-1.example/file.iso".to_owned(), "https://mirror-2.example/file.iso".to_owned(), ], target_path: PathBuf::from("D:/downloads/file.iso"), metadata_path: Some(PathBuf::from("D:/downloads/file.meta")), metadata: Some(metadata), }], }; store.save_session(&session).unwrap(); let loaded = store.load_session().unwrap(); assert_eq!(loaded, session); let _ = fs::remove_dir_all(root); } #[test] fn local_store_resume_roundtrip() { let root = temp_root(); let store = LocalFileStore::new(root.clone()); let mut piece_map = PieceMap::default(); piece_map.set_state(PieceIndex(0), PieceState::Verified); piece_map.set_state(PieceIndex(1), PieceState::InFlight); piece_map.set_state(PieceIndex(2), PieceState::Failed); let resume = ResumeData { gid: "gid-resume-1".to_owned(), download_path: PathBuf::from("D:/downloads/file.iso"), metadata: None, piece_map, }; store.save(&resume).unwrap(); let loaded = store.load(&resume.gid).unwrap().unwrap(); assert_eq!(loaded, resume); let _ = fs::remove_dir_all(root); } }