use std::collections::HashSet; use std::path::{Path, PathBuf}; use std::sync::{LazyLock, Mutex}; use lumotia_core::error::{Error, Result}; use lumotia_core::model_registry::{find_model, ModelFile}; use lumotia_core::types::{DownloadProgress, ModelId}; static ACTIVE_DOWNLOADS: LazyLock>> = LazyLock::new(|| Mutex::new(HashSet::new())); struct DownloadReservation { id: String, } impl DownloadReservation { fn acquire(id: &ModelId) -> Result { let id = id.as_str().to_string(); let mut active = ACTIVE_DOWNLOADS .lock() .map_err(|_| Error::DownloadFailed("download lock poisoned".into()))?; if !active.insert(id.clone()) { return Err(Error::DownloadFailed(format!( "download already in progress for {id}" ))); } Ok(Self { id }) } } impl Drop for DownloadReservation { fn drop(&mut self) { if let Ok(mut active) = ACTIVE_DOWNLOADS.lock() { active.remove(&self.id); } } } /// Resolve the models storage directory. /// Windows: %LOCALAPPDATA%/magnotia/models /// Unix: ~/.magnotia/models pub fn models_dir() -> PathBuf { lumotia_core::paths::app_paths().models_dir() } /// Get the directory path where a specific model's files are stored. pub fn model_dir(id: &ModelId) -> PathBuf { lumotia_core::paths::app_paths().speech_model_dir(id) } /// Check whether all files for a model have been downloaded. pub fn is_downloaded(id: &ModelId) -> bool { let entry = match find_model(id) { Some(e) => e, None => return false, }; let dir = model_dir(id); entry.files.iter().all(|f| dir.join(f.filename).exists()) && verified_manifest_matches(entry, &dir) } /// List all downloaded model IDs. pub fn list_downloaded() -> Vec { lumotia_core::model_registry::all_models() .iter() .filter(|m| is_downloaded(&m.id)) .map(|m| m.id.clone()) .collect() } /// Download all files for a model, calling the progress callback per chunk. /// Files are downloaded to a .part suffix and atomically renamed on completion. /// /// For files that declare a `sha256` checksum we validate an existing /// complete file before skipping the download — a truncated or /// tampered file gets redownloaded automatically (pattern ported from /// `lumotia-llm`'s model_manager, item #8 in the Whisper ecosystem brief). pub async fn download( id: &ModelId, progress: impl Fn(DownloadProgress) + Send + 'static, ) -> Result<()> { let _reservation = DownloadReservation::acquire(id)?; let entry = find_model(id).ok_or_else(|| Error::ModelNotFound(id.clone()))?; let dir = model_dir(id); std::fs::create_dir_all(&dir)?; for file in &entry.files { let dest = dir.join(file.filename); if dest.exists() { // Validate the existing file. If the hash doesn't match, // the file is corrupt (partial download, tampering, bit // rot) and we must re-fetch it to avoid crashing on // model load later. match sha256_of_file(&dest) { Ok(actual) if actual.eq_ignore_ascii_case(file.sha256) => continue, Ok(_actual) => { let _ = std::fs::remove_file(&dest); } Err(e) => { return Err(Error::DownloadFailed(format!( "failed to verify existing {}: {e}", file.filename ))); } } } download_file(file, &dest, id, &progress).await?; } write_verified_manifest(entry, &dir)?; Ok(()) } fn verified_manifest_path(dir: &Path) -> PathBuf { dir.join(".lumotia-verified") } fn verified_manifest_matches( entry: &lumotia_core::model_registry::ModelEntry, dir: &Path, ) -> bool { let manifest = match std::fs::read_to_string(verified_manifest_path(dir)) { Ok(contents) => contents, Err(_) => return false, }; for file in &entry.files { let path = dir.join(file.filename); let size = match std::fs::metadata(&path) { Ok(metadata) => metadata.len(), Err(_) => return false, }; let expected_line = format!("{}\t{}\t{}", file.filename, file.sha256, size); if !manifest.lines().any(|line| line == expected_line) { return false; } } true } fn write_verified_manifest( entry: &lumotia_core::model_registry::ModelEntry, dir: &Path, ) -> std::io::Result<()> { let mut lines = Vec::with_capacity(entry.files.len() + 1); lines.push("version\t1".to_string()); for file in &entry.files { let size = std::fs::metadata(dir.join(file.filename))?.len(); lines.push(format!("{}\t{}\t{}", file.filename, file.sha256, size)); } std::fs::write( verified_manifest_path(dir), format!("{}\n", lines.join("\n")), ) } /// Non-streaming SHA256 of a file on disk. Used by `download()` to /// validate an existing complete file before trusting it. fn sha256_of_file(path: &Path) -> std::io::Result { use sha2::{Digest, Sha256}; let mut hasher = Sha256::new(); let mut file = std::fs::File::open(path)?; let mut buffer = [0u8; 8192]; loop { let n = std::io::Read::read(&mut file, &mut buffer)?; if n == 0 { break; } hasher.update(&buffer[..n]); } Ok(format!("{:x}", hasher.finalize())) } /// Download a single file with HTTP Range resume and optional SHA256 verification. /// /// Resume pattern from Buzz (chidiwilliams/buzz): if a .part file exists, /// send a Range header to resume from where we left off. SHA256 is checked /// incrementally during download — no second pass over the file. async fn download_file( file: &ModelFile, dest: &Path, model_id: &ModelId, progress: &(impl Fn(DownloadProgress) + Send), ) -> Result<()> { use futures_util::StreamExt; use sha2::{Digest, Sha256}; let part_path = dest.with_extension( dest.extension() .map(|e| format!("{}.part", e.to_string_lossy())) .unwrap_or_else(|| "part".to_string()), ); let client = reqwest::Client::builder() .connect_timeout(std::time::Duration::from_secs(30)) .build() .map_err(|e| Error::DownloadFailed(e.to_string()))?; // Check for existing partial download (resume support) let existing_bytes = if part_path.exists() { std::fs::metadata(&part_path).map(|m| m.len()).unwrap_or(0) } else { 0 }; let mut request = client.get(file.url); let resuming = existing_bytes > 0; if resuming { request = request.header("Range", format!("bytes={existing_bytes}-")); } let response = request .send() .await .map_err(|e| Error::DownloadFailed(e.to_string()))?; // If we requested Range but the server returned 200 (full file), the // server does not support resume. Rather than blindly appending a // full file on top of our partial bytes (which would produce a // corrupt result), restart cleanly. This mirrors the lumotia-llm // ResumeUnsupported branch — item #8 of the brief. // // For the non-resume path, we still have to validate the status: // reqwest does not error on 4xx/5xx by default, so without this // check a 404 or 500 would be streamed into `.part` and renamed // over the destination as if the download succeeded // (2026-04-22 review MAJOR). let actually_resuming = if resuming { match response.status().as_u16() { 206 => true, 200 => { // Server ignored our Range header — treat as fresh start. // The old .part bytes are discarded below. false } other => { return Err(Error::DownloadFailed(format!( "resume request returned unexpected status {other}" ))); } } } else { if !response.status().is_success() { return Err(Error::DownloadFailed(format!( "download returned HTTP {} for {}", response.status(), file.filename ))); } false }; let total_bytes = if actually_resuming { // Content-Range: bytes START-END/TOTAL — extract TOTAL response .headers() .get("content-range") .and_then(|v| v.to_str().ok()) .and_then(|s| s.rsplit('/').next()) .and_then(|s| s.parse::().ok()) .unwrap_or(0) } else { response.content_length().unwrap_or(0) }; let mut stream = response.bytes_stream(); let mut downloaded: u64 = if actually_resuming { existing_bytes } else { 0 }; let mut last_percent: u8 = 0; // Open file for append (resume) or create (fresh start) let mut out = if actually_resuming { std::fs::OpenOptions::new().append(true).open(&part_path)? } else { std::fs::File::create(&part_path)? }; let mut hasher = Sha256::new(); if actually_resuming { let mut partial = std::fs::File::open(&part_path)?; let mut buffer = [0u8; 8192]; loop { let n = std::io::Read::read(&mut partial, &mut buffer)?; if n == 0 { break; } hasher.update(&buffer[..n]); } } while let Some(chunk) = stream.next().await { let chunk = chunk.map_err(|e| Error::DownloadFailed(e.to_string()))?; std::io::Write::write_all(&mut out, &chunk)?; hasher.update(&chunk); downloaded += chunk.len() as u64; let percent = if total_bytes > 0 { ((downloaded as f64 / total_bytes as f64) * 100.0) as u8 } else { 0 }; if percent != last_percent { last_percent = percent; progress(DownloadProgress { model_id: model_id.clone(), file_name: file.filename.to_string(), bytes_downloaded: downloaded, total_bytes, percent, }); } } drop(out); let actual = format!("{:x}", hasher.finalize()); if actual != file.sha256 { let _ = std::fs::remove_file(&part_path); return Err(Error::DownloadFailed(format!( "SHA256 mismatch for {}: expected {}, got {}", file.filename, file.sha256, actual ))); } // Atomic rename — file is complete and verified std::fs::rename(&part_path, dest)?; Ok(()) } #[cfg(test)] mod tests { use super::*; use sha2::Digest; use tempfile::tempdir; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpListener; #[test] fn model_dir_returns_correct_path() { let id = ModelId::new("whisper-tiny-en"); let path = model_dir(&id); assert!(path.to_string_lossy().contains("whisper-tiny-en")); } #[test] fn is_downloaded_returns_false_for_missing_model() { let id = ModelId::new("nonexistent-model"); assert!(!is_downloaded(&id)); } #[test] fn list_downloaded_returns_empty_when_no_models() { let list = list_downloaded(); // In test environment, no models are downloaded // This just verifies the function doesn't panic assert!(list.len() <= lumotia_core::model_registry::all_models().len()); } #[test] fn sha256_of_file_matches_sha2() { let dir = tempdir().unwrap(); let path = dir.path().join("f.bin"); std::fs::write(&path, b"hello world").unwrap(); let expected = format!("{:x}", sha2::Sha256::digest(b"hello world")); assert_eq!(sha256_of_file(&path).unwrap(), expected); } /// A minimal HTTP server that sends a Range response when a Range /// header is present and otherwise sends the full body. Ported from /// crates/llm/src/model_manager.rs to give the transcription /// download stack the same fixture-backed coverage. async fn spawn_range_server(content: Vec) -> std::net::SocketAddr { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { let (mut socket, _) = listener.accept().await.unwrap(); let mut buf = vec![0u8; 2048]; let size = socket.read(&mut buf).await.unwrap(); let request = String::from_utf8_lossy(&buf[..size]).to_lowercase(); let range_start = request .lines() .find_map(|line| line.strip_prefix("range: bytes=")) .and_then(|line| line.strip_suffix('-')) .and_then(|line| line.trim().parse::().ok()); if let Some(start) = range_start { let body = &content[start..]; let response = format!( "HTTP/1.1 206 Partial Content\r\n\ Content-Length: {}\r\n\ Content-Range: bytes {}-{}/{}\r\n\ Accept-Ranges: bytes\r\n\r\n", body.len(), start, content.len() - 1, content.len(), ); socket.write_all(response.as_bytes()).await.unwrap(); socket.write_all(body).await.unwrap(); } else { let response = format!( "HTTP/1.1 200 OK\r\n\ Content-Length: {}\r\n\ Accept-Ranges: bytes\r\n\r\n", content.len(), ); socket.write_all(response.as_bytes()).await.unwrap(); socket.write_all(&content).await.unwrap(); } }); addr } /// A minimal HTTP server that responds with 200 + full body **iff** /// the request actually carries a `Range` header, and 400 otherwise. /// This models a mirror / proxy that accepts Range requests but /// refuses to honour them (returning a fresh full body), which is /// exactly the ResumeUnsupported branch `download_file` needs to /// handle. The 400-on-missing-Range behaviour is load-bearing for /// the test: it turns "client never sent Range" into a download /// failure, so deleting the resume-detection logic causes the test /// to fail rather than pass coincidentally through File::create's /// truncation semantics. async fn spawn_no_range_server(content: Vec) -> std::net::SocketAddr { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { let (mut socket, _) = listener.accept().await.unwrap(); let mut buf = vec![0u8; 2048]; let size = socket.read(&mut buf).await.unwrap(); let request = String::from_utf8_lossy(&buf[..size]).to_lowercase(); let saw_range_header = request .lines() .any(|line| line.trim_start().starts_with("range:")); if !saw_range_header { let response = "HTTP/1.1 400 Bad Request\r\n\ Content-Length: 0\r\n\r\n"; socket.write_all(response.as_bytes()).await.unwrap(); return; } let response = format!( "HTTP/1.1 200 OK\r\n\ Content-Length: {}\r\n\r\n", content.len(), ); socket.write_all(response.as_bytes()).await.unwrap(); socket.write_all(&content).await.unwrap(); }); addr } /// ModelFile stores `&'static str` fields, so we leak the strings /// once per test — tests are one-shot, so the cost is noise. fn leak(s: String) -> &'static str { Box::leak(s.into_boxed_str()) } #[tokio::test] async fn download_file_resumes_from_partial_and_verifies_sha() { let body = b"resumable transcription payload".to_vec(); let expected_sha = format!("{:x}", sha2::Sha256::digest(&body)); let addr = spawn_range_server(body.clone()).await; let dir = tempdir().unwrap(); let dest = dir.path().join("fixture.bin"); let part = dest.with_extension("bin.part"); // Pretend we already downloaded the first 7 bytes. std::fs::write(&part, &body[..7]).unwrap(); let file = ModelFile { filename: leak(dest.file_name().unwrap().to_string_lossy().into_owned()), url: leak(format!("http://{addr}/fixture.bin")), size: lumotia_core::types::Megabytes(0), sha256: leak(expected_sha.clone()), }; let id = ModelId::new("test-fixture"); download_file(&file, &dest, &id, &|_| ()).await.unwrap(); let bytes = std::fs::read(&dest).unwrap(); assert_eq!(bytes, body); assert!(!part.exists()); assert_eq!(sha256_of_file(&dest).unwrap(), expected_sha); } #[tokio::test] async fn download_file_restarts_when_server_ignores_range() { // Covers the ResumeUnsupported branch documented in `download_file`: // when a partial `.part` file exists and the server returns 200 // (full body) to our Range request, we must discard the stale // partial bytes and write the fresh body from offset zero rather // than appending on top. let body = b"fresh transcription payload that replaces any stale partial".to_vec(); let expected_sha = format!("{:x}", sha2::Sha256::digest(&body)); let addr = spawn_no_range_server(body.clone()).await; let dir = tempdir().unwrap(); let dest = dir.path().join("fixture.bin"); let part = dest.with_extension("bin.part"); // Pretend a previous attempt downloaded 12 bytes of something // entirely unrelated. If the client naively appended the 200 // body, the final file would start with these bytes. std::fs::write(&part, b"STALE_BYTES1").unwrap(); let file = ModelFile { filename: leak(dest.file_name().unwrap().to_string_lossy().into_owned()), url: leak(format!("http://{addr}/fixture.bin")), size: lumotia_core::types::Megabytes(0), sha256: leak(expected_sha), }; let id = ModelId::new("test-fixture"); download_file(&file, &dest, &id, &|_| ()).await.unwrap(); let bytes = std::fs::read(&dest).unwrap(); assert_eq!( bytes, body, "server returned 200 to Range — downloader must discard stale .part and rewrite from scratch" ); assert!(!part.exists(), ".part → dest rename must run after restart"); } /// Always returns HTTP 500 with a short error body. Used to verify /// the non-resume download path validates status codes rather than /// writing error bodies into `.part` and renaming them over the /// destination. async fn spawn_500_server() -> std::net::SocketAddr { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); tokio::spawn(async move { let (mut socket, _) = listener.accept().await.unwrap(); let mut buf = vec![0u8; 2048]; let _ = socket.read(&mut buf).await.unwrap(); let body = b"internal error"; let response = format!( "HTTP/1.1 500 Internal Server Error\r\n\ Content-Length: {}\r\n\r\n", body.len() ); socket.write_all(response.as_bytes()).await.unwrap(); socket.write_all(body).await.unwrap(); }); addr } #[tokio::test] async fn download_file_rejects_5xx_on_non_resume_path() { // Regression for the 2026-04-22 review: reqwest does not // auto-error on 4xx/5xx, and the non-resume branch previously // streamed any status' body into `.part` and renamed it over // the destination. let addr = spawn_500_server().await; let dir = tempdir().unwrap(); let dest = dir.path().join("fixture.bin"); let part = dest.with_extension("bin.part"); let file = ModelFile { filename: leak(dest.file_name().unwrap().to_string_lossy().into_owned()), url: leak(format!("http://{addr}/fixture.bin")), size: lumotia_core::types::Megabytes(0), sha256: leak("0".repeat(64)), }; let id = ModelId::new("test-fixture"); let err = download_file(&file, &dest, &id, &|_| ()) .await .expect_err("5xx must fail"); let msg = err.to_string(); assert!( msg.contains("HTTP 500"), "error should name the HTTP status, got: {msg}" ); assert!(!dest.exists(), "5xx must not leave a destination file"); assert!(!part.exists(), "5xx must not leave a .part file"); } #[tokio::test] async fn download_file_fails_on_sha_mismatch_and_cleans_part_file() { let body = b"speech-to-text fixture body".to_vec(); let addr = spawn_range_server(body.clone()).await; let dir = tempdir().unwrap(); let dest = dir.path().join("fixture.bin"); let file = ModelFile { filename: leak(dest.file_name().unwrap().to_string_lossy().into_owned()), url: leak(format!("http://{addr}/fixture.bin")), size: lumotia_core::types::Megabytes(0), sha256: leak("deadbeef".repeat(8)), }; let id = ModelId::new("test-fixture"); let err = download_file(&file, &dest, &id, &|_| ()) .await .expect_err("mismatched sha must fail"); let msg = err.to_string(); assert!(msg.contains("SHA256 mismatch"), "unexpected error: {msg}"); assert!( !dest.exists(), ".part → dest rename must not run on mismatch" ); let part = dest.with_extension("bin.part"); assert!(!part.exists(), "failed hash must clean up the .part file"); } }