diff --git a/src/main.rs b/src/main.rs index 6ead4d7..c1f6fd2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -313,6 +313,11 @@ fn run() -> Result<()> { // Route subsequent INFO/WARN/DEBUG logs through the cliclack/indicatif area polyscribe::progress::set_global_progress_manager(&pm); + // Show a friendly intro banner (TTY-aware via cliclack). Ignore errors. + if !polyscribe::is_quiet() { + let _ = cliclack::intro("PolyScribe"); + } + // Determine formats let out_formats = if args.out_format.is_empty() { OutputFormats::all() @@ -487,6 +492,15 @@ fn run() -> Result<()> { if had_error { polyscribe::elog!("One or more inputs failed"); } } + // Outro banner summarizing result; ignore errors. + if !polyscribe::is_quiet() { + if had_error { + let _ = cliclack::outro("Completed with errors. Some inputs failed."); + } else { + let _ = cliclack::outro("All done. Outputs written."); + } + } + if had_error { std::process::exit(2); } let _elapsed = start_overall.elapsed(); Ok(()) diff --git a/src/models.rs b/src/models.rs index a355ac1..9b26704 100644 --- a/src/models.rs +++ b/src/models.rs @@ -616,6 +616,30 @@ fn download_one_model_with_progress(client: &Client, models_dir: &Path, entry: & } if final_path.exists() { let _ = std::fs::remove_file(&final_path); } std::fs::rename(&tmp_path, &final_path).with_context(|| format!("Failed to move into place: {}", final_path.display()))?; + // Hardened verification after save + if let Some(expected) = &entry.sha256 { + match compute_file_sha256_hex(&final_path) { + Ok(rehash) => { + if !rehash.eq_ignore_ascii_case(expected) { + let _ = std::fs::remove_file(&final_path); + return Err(anyhow!( + "Downloaded file failed SHA-256 verification after save for {}: expected {}, got {}. The file has been removed. Please try downloading again. If the problem persists, check your network connection and disk space, or report this issue.", + entry.name, + expected, + rehash + )); + } + } + Err(e) => { + let _ = std::fs::remove_file(&final_path); + return Err(anyhow!( + "Failed to verify downloaded file {}: {}. The file has been removed. Please try again.", + final_path.display(), + e + )); + } + } + } item.set_progress(1.0); item.finish_with("done"); return Ok(()); @@ -668,6 +692,32 @@ fn download_one_model_with_progress(client: &Client, models_dir: &Path, entry: & if final_path.exists() { let _ = std::fs::remove_file(&final_path); } std::fs::rename(&tmp_path, &final_path) .with_context(|| format!("Failed to move into place: {}", final_path.display()))?; + + // Hardened verification: recompute SHA-256 from the saved file and compare to expected. + if let Some(expected) = &entry.sha256 { + match compute_file_sha256_hex(&final_path) { + Ok(rehash) => { + if !rehash.eq_ignore_ascii_case(expected) { + let _ = std::fs::remove_file(&final_path); + return Err(anyhow!( + "Downloaded file failed SHA-256 verification after save for {}: expected {}, got {}. The file has been removed. Please try downloading again. If the problem persists, check your network connection and disk space, or report this issue.", + entry.name, + expected, + rehash + )); + } + } + Err(e) => { + let _ = std::fs::remove_file(&final_path); + return Err(anyhow!( + "Failed to verify downloaded file {}: {}. The file has been removed. Please try again.", + final_path.display(), + e + )); + } + } + } + item.finish_with("done"); Ok(()) } @@ -772,6 +822,30 @@ fn download_one_model(client: &Client, models_dir: &Path, entry: &ModelEntry) -> } std::fs::rename(&tmp_path, &final_path) .with_context(|| format!("Failed to move into place: {}", final_path.display()))?; + // Hardened verification after save + if let Some(expected) = &entry.sha256 { + match compute_file_sha256_hex(&final_path) { + Ok(rehash) => { + if !rehash.eq_ignore_ascii_case(expected) { + let _ = std::fs::remove_file(&final_path); + return Err(anyhow!( + "Downloaded file failed SHA-256 verification after save for {}: expected {}, got {}. The file has been removed. Please try downloading again. If the problem persists, check your network connection and disk space, or report this issue.", + entry.name, + expected, + rehash + )); + } + } + Err(e) => { + let _ = std::fs::remove_file(&final_path); + return Err(anyhow!( + "Failed to verify downloaded file {}: {}. The file has been removed. Please try again.", + final_path.display(), + e + )); + } + } + } qlog!("Saved: {}", final_path.display()); return Ok(()); } @@ -837,6 +911,32 @@ fn download_one_model(client: &Client, models_dir: &Path, entry: &ModelEntry) -> } std::fs::rename(&tmp_path, &final_path) .with_context(|| format!("Failed to move into place: {}", final_path.display()))?; + + // Hardened verification: recompute SHA-256 from the saved file and compare to expected. + if let Some(expected) = &entry.sha256 { + match compute_file_sha256_hex(&final_path) { + Ok(rehash) => { + if !rehash.eq_ignore_ascii_case(expected) { + let _ = std::fs::remove_file(&final_path); + return Err(anyhow!( + "Downloaded file failed SHA-256 verification after save for {}: expected {}, got {}. The file has been removed. Please try downloading again. If the problem persists, check your network connection and disk space, or report this issue.", + entry.name, + expected, + rehash + )); + } + } + Err(e) => { + let _ = std::fs::remove_file(&final_path); + return Err(anyhow!( + "Failed to verify downloaded file {}: {}. The file has been removed. Please try again.", + final_path.display(), + e + )); + } + } + } + qlog!("Saved: {}", final_path.display()); Ok(()) } @@ -1274,4 +1374,49 @@ mod tests { panic!("progress manager did not expose current state"); } } + + #[test] + fn test_wrong_hash_deletes_temp_and_errors() { + use std::sync::{Mutex, OnceLock}; + static ENV_LOCK: OnceLock> = OnceLock::new(); + let _guard = ENV_LOCK.get_or_init(|| Mutex::new(())).lock().unwrap(); + + let tmp_models = tempdir().unwrap(); + let tmp_base = tempdir().unwrap(); + + // Prepare source model file content and a pre-existing local file to trigger update + let model_name = "tiny.en-q5_1"; + let src_path = tmp_base.path().join(format!("ggml-{}.bin", model_name)); + let content = b"model data"; + fs::write(&src_path, content).unwrap(); + let wrong_sha = "0000000000000000000000000000000000000000000000000000000000000000".to_string(); + + let local_path = tmp_models.path().join(format!("ggml-{}.bin", model_name)); + let original = b"old local"; + fs::write(&local_path, original).unwrap(); + + unsafe { + std::env::set_var("POLYSCRIBE_MODELS_BASE_COPY_DIR", tmp_base.path()); + } + + // Construct a ModelEntry with wrong expected sha and call the downloader directly + let client = Client::builder().build().unwrap(); + let entry = ModelEntry { + name: model_name.to_string(), + base: "tiny".to_string(), + subtype: "en-q5_1".to_string(), + size: content.len() as u64, + sha256: Some(wrong_sha), + repo: "ggerganov/whisper.cpp".to_string(), + }; + let res = super::download_one_model(&client, tmp_models.path(), &entry); + assert!(res.is_err(), "expected error due to wrong hash"); + + let final_path = tmp_models.path().join(format!("ggml-{}.bin", model_name)); + let tmp_path = tmp_models.path().join(format!("ggml-{}.bin.part", model_name)); + assert!(final_path.exists(), "existing local file should remain when new download fails"); + let preserved = fs::read(&final_path).unwrap(); + assert_eq!(preserved, original, "existing local file must be preserved"); + assert!(!tmp_path.exists(), ".part file should be deleted on hash mismatch"); + } } diff --git a/src/progress.rs b/src/progress.rs index 56a2e06..d67e4aa 100644 --- a/src/progress.rs +++ b/src/progress.rs @@ -359,13 +359,58 @@ impl ProgressManager { /// Print a line above the bars safely (TTY-aware). Falls back to eprintln! when disabled. pub fn println_above_bars(&self, line: &str) { + // Try to interpret certain INFO lines as a stable title + dynamic message. + // Examples to match: + // - "INFO: Fetching online data: listing models from ggerganov/whisper.cpp..." + // -> header = "INFO: Fetching online data"; info = "listing models from ..." + // - "INFO: Downloading tiny.en-q5_1 (252 MiB | https://...)..." + // -> header = "INFO: Downloading"; info = rest + // - "INFO: Total 1/3" (defensive): header = "INFO: Total"; info = rest + let parsed: Option<(String, String)> = { + let s = line.trim(); + if let Some(rest) = s.strip_prefix("INFO: ") { + // Case A: explicit title followed by colon + if let Some((title, body)) = rest.split_once(':') { + let title_clean = format!("INFO: {}", title.trim()); + let body_clean = body.trim().to_string(); + Some((title_clean, body_clean)) + } else if let Some(rest2) = rest.strip_prefix("Downloading ") { + Some(("INFO: Downloading".to_string(), rest2.trim().to_string())) + } else if let Some(rest2) = rest.strip_prefix("Total") { + Some(("INFO: Total".to_string(), rest2.trim().to_string())) + } else { + // Fallback: use first word as title, remainder as body + let mut it = rest.splitn(2, ' '); + let first = it.next().unwrap_or("").trim(); + let remainder = it.next().unwrap_or("").trim(); + if !first.is_empty() { + Some((format!("INFO: {}", first), remainder.to_string())) + } else { + None + } + } + } else { + None + } + }; + match &self.inner { ProgressInner::Noop => eprintln!("{}", line), ProgressInner::Single(s) => { - let _ = s._mp.println(line); + if let Some((title, body)) = parsed.as_ref() { + s.header.set_message(title.clone()); + s.info.set_message(body.clone()); + } else { + let _ = s._mp.println(line); + } } ProgressInner::Multi(m) => { - let _ = m._mp.println(line); + if let Some((title, body)) = parsed.as_ref() { + m.header.set_message(title.clone()); + m.info.set_message(body.clone()); + } else { + let _ = m._mp.println(line); + } } } }