use std::path::{Path, PathBuf}; use kon_core::error::{KonError, Result}; use kon_core::model_registry::{find_model, ModelFile}; use kon_core::types::{DownloadProgress, ModelId}; /// Resolve the models storage directory. /// Windows: %LOCALAPPDATA%/kon/models /// Unix: ~/.kon/models pub fn models_dir() -> PathBuf { if cfg!(target_os = "windows") { let local_app_data = std::env::var("LOCALAPPDATA").unwrap_or_else(|_| ".".to_string()); PathBuf::from(local_app_data).join("kon").join("models") } else { dirs_path().join("models") } } fn dirs_path() -> PathBuf { if cfg!(target_os = "windows") { let local_app_data = std::env::var("LOCALAPPDATA").unwrap_or_else(|_| ".".to_string()); PathBuf::from(local_app_data).join("kon") } else { let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string()); PathBuf::from(home).join(".kon") } } /// Get the directory path where a specific model's files are stored. pub fn model_dir(id: &ModelId) -> PathBuf { models_dir().join(id.as_str()) } /// 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()) } /// List all downloaded model IDs. pub fn list_downloaded() -> Vec { kon_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 /// `kon-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 entry = find_model(id).ok_or_else(|| KonError::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() { if let Some(expected_sha) = file.sha256 { // 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(expected_sha) => continue, Ok(_actual) => { // Corrupt — remove + fall through to re-download. let _ = std::fs::remove_file(&dest); } Err(e) => { return Err(KonError::DownloadFailed(format!( "failed to verify existing {}: {e}", file.filename ))); } } } else { // No checksum — honour the existing file as-is; the // engine will barf on load if it's broken. continue; } } download_file(file, &dest, id, &progress).await?; } Ok(()) } /// 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| KonError::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); // If we have a partial file and no SHA256 to verify (can't verify partial), // request a range resume. If SHA256 is set, we restart to ensure integrity. let resuming = existing_bytes > 0 && file.sha256.is_none(); if resuming { request = request.header("Range", format!("bytes={existing_bytes}-")); } let response = request .send() .await .map_err(|e| KonError::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 kon-llm // ResumeUnsupported branch — item #8 of the brief. 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(KonError::DownloadFailed(format!( "resume request returned unexpected status {other}" ))); } } } else { 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)? }; // Incremental SHA256 — only when a checksum is provided let mut hasher = file.sha256.map(|_| Sha256::new()); // If resuming without SHA256, we can't hash the already-downloaded portion, // but we also don't need to — we only hash when sha256 is set, and we // restart from scratch in that case. while let Some(chunk) = stream.next().await { let chunk = chunk.map_err(|e| KonError::DownloadFailed(e.to_string()))?; std::io::Write::write_all(&mut out, &chunk)?; if let Some(ref mut h) = hasher { h.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); // Verify SHA256 if provided if let (Some(expected), Some(hasher)) = (file.sha256, hasher) { let actual = format!("{:x}", hasher.finalize()); if actual != expected { // Delete corrupt file so next attempt starts fresh let _ = std::fs::remove_file(&part_path); return Err(KonError::DownloadFailed(format!( "SHA256 mismatch for {}: expected {}, got {}", file.filename, expected, 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() <= kon_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 } /// 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: kon_core::types::Megabytes(0), sha256: None, // resume path only kicks in when sha256 is absent }; 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()); // Confirm the full file hash matches what we would have got via // a clean download — gives the resume path indirect integrity // coverage even when the ModelFile has no sha256 set. assert_eq!(sha256_of_file(&dest).unwrap(), expected_sha); } #[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: kon_core::types::Megabytes(0), sha256: Some(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"); } }