feat(A.1 #8): harden transcription model downloads with sha + resume tests
Ports the kon-llm model_manager resume pattern the rest of the way
into kon-transcription::model_manager:
- download() now validates an existing complete file against its
sha256 before skipping; a hash mismatch removes the file and
re-fetches, instead of serving a corrupt file to whisper.cpp.
- download_file() now distinguishes 206 Partial Content, 200 OK
(resume silently ignored by server), and other statuses, rather
than treating any non-206 as 'just use it as a fresh start'.
200-on-Range is handled by discarding the partial and starting
over cleanly.
- New tests: download_file_resumes_from_partial_and_verifies_sha
(TcpListener fixture, same shape as kon-llm's), and
download_file_fails_on_sha_mismatch_and_cleans_part_file.
- sha256_of_file helper + unit test for the existing-file guard.
Dev-deps: tempfile + tokio(net/io-util/macros). Total workspace
lib-test count: 116 → 123.
Co-authored-by: jars <jakejars@users.noreply.github.com>
This commit is contained in:
@@ -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"
|
||||
|
||||
@@ -52,6 +52,11 @@ pub fn list_downloaded() -> Vec<ModelId> {
|
||||
|
||||
/// 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,14 +69,54 @@ pub async fn download(
|
||||
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<String> {
|
||||
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<u8>) -> 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::<usize>().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");
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user