diff --git a/crates/transcription/Cargo.toml b/crates/transcription/Cargo.toml index 6500e48..f053e45 100644 --- a/crates/transcription/Cargo.toml +++ b/crates/transcription/Cargo.toml @@ -29,3 +29,8 @@ thiserror = "2" # Structured logging at backend boundaries (observability for initial_prompt flow). tracing = "0.1" + +[dev-dependencies] +# TcpListener fixture for the download resume tests (mirrors kon-llm). +tokio = { version = "1", features = ["rt", "sync", "net", "io-util", "macros"] } +tempfile = "3" diff --git a/crates/transcription/src/model_manager.rs b/crates/transcription/src/model_manager.rs index 0d42c5d..4c59fb9 100644 --- a/crates/transcription/src/model_manager.rs +++ b/crates/transcription/src/model_manager.rs @@ -52,6 +52,11 @@ pub fn list_downloaded() -> Vec { /// 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, @@ -64,7 +69,29 @@ pub async fn download( for file in &entry.files { let dest = dir.join(file.filename); if dest.exists() { - continue; + 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?; } @@ -72,6 +99,24 @@ pub async fn download( 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, @@ -118,8 +163,28 @@ async fn download_file( .await .map_err(|e| KonError::DownloadFailed(e.to_string()))?; - // Check if server supports range (206 Partial Content) or gave full file (200) - let actually_resuming = resuming && response.status().as_u16() == 206; + // 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 @@ -202,6 +267,10 @@ async fn download_file( #[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() { @@ -223,4 +292,124 @@ mod tests { // 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"); + } }