From d1d8c31bfbfe57aaa9a7178805504493cf8ccfcf Mon Sep 17 00:00:00 2001 From: userhaptop <1307305157@qq.com> Date: Sat, 7 Feb 2026 23:12:09 +0800 Subject: [PATCH 1/5] feat(log): add Ubuntu cloud image support Signed-off-by: userhaptop <1307305157@qq.com> --- src/image.rs | 127 +++++++++++++++++++++++++++++++++++------- tests/ubuntu_image.rs | 27 +++++++++ 2 files changed, 133 insertions(+), 21 deletions(-) create mode 100644 tests/ubuntu_image.rs diff --git a/src/image.rs b/src/image.rs index 5abfb69..4d1d5d3 100644 --- a/src/image.rs +++ b/src/image.rs @@ -34,6 +34,7 @@ pub struct ImageMeta { #[derive(Debug, PartialEq, Serialize, Deserialize, Clone)] pub enum Distro { Debian, + Ubuntu, } #[derive(Debug, PartialEq, Serialize, Deserialize, Clone)] @@ -49,13 +50,6 @@ pub struct ShaSum { } /// Parses SHA512SUMS format and returns the hash for an exact filename match. -/// -/// # Arguments -/// * `checksums_text` - The content of a SHA512SUMS file -/// * `filename` - The exact filename to search for (e.g., "debian-13-generic-amd64.qcow2") -/// -/// # Returns -/// The SHA512 hash if found, or None if no exact match exists pub fn find_sha512_for_file(checksums_text: &str, filename: &str) -> Option { checksums_text.lines().find_map(|line| { let mut parts = line.split_whitespace(); @@ -205,6 +199,10 @@ impl ImageMeta { } } +// --------------------------------------------------------------------------- +// Debian +// --------------------------------------------------------------------------- + #[derive(Debug, Default)] pub struct Debian {} @@ -252,7 +250,6 @@ impl ImageAction for Debian { let computed_sha512 = get_sha512(&image_path).await?; - // Verify the downloaded file matches the expected checksum anyhow::ensure!( computed_sha512.to_lowercase() == expected_sha512.to_lowercase(), "downloaded image checksum mismatch: expected {}, got {}", @@ -265,7 +262,6 @@ impl ImageAction for Debian { async fn extract(&self, name: &str) -> Result<(PathBuf, PathBuf)> { let file_name = format!("{}.qcow2", name); - let dirs = QleanDirs::new()?; let image_dir = dirs.images.join(name); @@ -353,39 +349,122 @@ impl ImageAction for Debian { } } -/// Wrapper enum for different Image types +// --------------------------------------------------------------------------- +// Ubuntu - uses pre-extracted kernel/initrd from official cloud images +// --------------------------------------------------------------------------- + +#[derive(Debug, Default)] +pub struct Ubuntu {} + +impl ImageAction for Ubuntu { + async fn download(&self, name: &str) -> Result<()> { + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + // Ubuntu noble (24.04 LTS) cloud image base URL + let base_url = "https://cloud-images.ubuntu.com/noble/current"; + + // Download qcow2 image + let qcow2_url = format!("{}/noble-server-cloudimg-amd64.img", base_url); + let qcow2_path = image_dir.join(format!("{}.qcow2", name)); + download_file(&qcow2_url, &qcow2_path).await?; + + // Download pre-extracted kernel + let kernel_url = format!( + "{}/unpacked/noble-server-cloudimg-amd64-vmlinuz-generic", + base_url + ); + let kernel_path = image_dir.join("vmlinuz"); + download_file(&kernel_url, &kernel_path).await?; + + // Download pre-extracted initrd + let initrd_url = format!( + "{}/unpacked/noble-server-cloudimg-amd64-initrd-generic", + base_url + ); + let initrd_path = image_dir.join("initrd.img"); + download_file(&initrd_url, &initrd_path).await?; + + Ok(()) + } + + async fn extract(&self, name: &str) -> Result<(PathBuf, PathBuf)> { + // Files already downloaded in download() phase + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + let kernel = image_dir.join("vmlinuz"); + let initrd = image_dir.join("initrd.img"); + + anyhow::ensure!(kernel.exists(), "kernel file not found after download"); + anyhow::ensure!(initrd.exists(), "initrd file not found after download"); + + Ok((kernel, initrd)) + } + + fn distro(&self) -> Distro { + Distro::Ubuntu + } +} + +// Helper function to download a file +async fn download_file(url: &str, dest: &PathBuf) -> Result<()> { + debug!("Downloading {} to {}", url, dest.display()); + let response = reqwest::get(url) + .await + .with_context(|| format!("failed to download from {}", url))?; + + let mut file = File::create(dest) + .await + .with_context(|| format!("failed to create file at {}", dest.display()))?; + + let mut stream = response.bytes_stream(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.with_context(|| "failed to read chunk from stream")?; + file.write_all(&chunk) + .await + .with_context(|| "failed to write to file")?; + } + + Ok(()) +} + +// --------------------------------------------------------------------------- +// Image wrapper enum +// --------------------------------------------------------------------------- + #[derive(Debug)] pub enum Image { Debian(ImageMeta), - // Add more distros as needed + Ubuntu(ImageMeta), } impl Image { - /// Get the underlying name regardless of distro pub fn name(&self) -> &str { match self { Image::Debian(img) => &img.name, + Image::Ubuntu(img) => &img.name, } } - /// Get the underlying image path regardless of distro pub fn path(&self) -> &PathBuf { match self { Image::Debian(img) => &img.path, + Image::Ubuntu(img) => &img.path, } } - /// Get the kernel path regardless of distro pub fn kernel(&self) -> &PathBuf { match self { Image::Debian(img) => &img.kernel, + Image::Ubuntu(img) => &img.kernel, } } - /// Get the initrd path regardless of distro pub fn initrd(&self) -> &PathBuf { match self { Image::Debian(img) => &img.initrd, + Image::Ubuntu(img) => &img.initrd, } } } @@ -396,7 +475,11 @@ pub async fn create_image(distro: Distro, name: &str) -> Result { Distro::Debian => { let image = ImageMeta::::create(name).await?; Ok(Image::Debian(image)) - } // Add more distros as needed + } + Distro::Ubuntu => { + let image = ImageMeta::::create(name).await?; + Ok(Image::Ubuntu(image)) + } } } @@ -452,7 +535,7 @@ pub async fn get_sha512(path: &PathBuf) -> Result { #[cfg(test)] mod tests { - use super::{Debian, ImageAction, find_sha512_for_file, get_sha512}; + use super::{Debian, Distro, ImageAction, find_sha512_for_file, get_sha512}; use crate::utils::QleanDirs; use anyhow::Result; use serial_test::serial; @@ -466,25 +549,28 @@ f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea6770915 \ 9fd031ef5dda6479c8536a0ab396487113303f4924a2941dc4f9ef1d36376dfb8ae7d1ca5f4dfa65ad155639e9a5e61093c686a8e85b51d106c180bce9ac49bc debian-13-generic-amd64.raw"; - // Should match exact qcow2 filename, not json with same prefix let result = find_sha512_for_file(checksums, "debian-13-generic-amd64.qcow2"); assert_eq!( result, Some("f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea67709154d84220059672758508afbb0691c41ba8aa6d76818d89d65".to_string()) ); - // Should match json file exactly let result = find_sha512_for_file(checksums, "debian-13-generic-amd64.json"); assert_eq!( result, Some("748f52b959f63352e1e121508cedeae2e66d3e90be00e6420a0b8b9f14a0f84dc54ed801fb5be327866876268b808543465b1613c8649efeeb5f987ff9df1549".to_string()) ); - // Should not match partial names let result = find_sha512_for_file(checksums, "debian-13-generic-amd64"); assert_eq!(result, None); } + #[test] + fn test_distro_enum_variants() { + let variants = vec![Distro::Debian, Distro::Ubuntu]; + assert_eq!(variants.len(), 2); + } + #[tokio::test] #[serial] #[ignore] @@ -510,7 +596,6 @@ f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea6770915 let computed = get_sha512(&qcow_path).await?; - // Clean up downloaded image before assertion to ensure cleanup happens even on failure if qcow_path.exists() { tokio::fs::remove_file(&qcow_path).await?; } diff --git a/tests/ubuntu_image.rs b/tests/ubuntu_image.rs new file mode 100644 index 0000000..8cfa456 --- /dev/null +++ b/tests/ubuntu_image.rs @@ -0,0 +1,27 @@ +use anyhow::Result; +use qlean::{Distro, create_image}; +use serial_test::serial; + +mod common; +use common::tracing_subscriber_init; + +#[tokio::test] +#[serial] +#[ignore] +async fn test_ubuntu_image_creation() -> Result<()> { + tracing_subscriber_init(); + + // Ubuntu uses pre-extracted kernel/initrd - no guestfish needed! + let image = create_image(Distro::Ubuntu, "ubuntu-noble-cloudimg").await?; + + assert!(image.path().exists(), "qcow2 image must exist"); + assert!(image.kernel().exists(), "kernel must exist"); + assert!(image.initrd().exists(), "initrd must exist"); + + println!("✅ Ubuntu image created successfully!"); + println!(" Image: {}", image.path().display()); + println!(" Kernel: {}", image.kernel().display()); + println!(" Initrd: {}", image.initrd().display()); + + Ok(()) +} From 41c83298bf81abab8b49487d73dde57a3640ebb0 Mon Sep 17 00:00:00 2001 From: userhaptop <1307305157@qq.com> Date: Sat, 7 Feb 2026 23:16:09 +0800 Subject: [PATCH 2/5] feat(log): add Ubuntu cloud image support Signed-off-by: userhaptop <1307305157@qq.com> --- src/image.rs | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/src/image.rs b/src/image.rs index 4d1d5d3..ba43aa4 100644 --- a/src/image.rs +++ b/src/image.rs @@ -50,6 +50,13 @@ pub struct ShaSum { } /// Parses SHA512SUMS format and returns the hash for an exact filename match. +/// +/// # Arguments +/// * `checksums_text` - The content of a SHA512SUMS file +/// * `filename` - The exact filename to search for (e.g., "debian-13-generic-amd64.qcow2") +/// +/// # Returns +/// The SHA512 hash if found, or None if no exact match exists pub fn find_sha512_for_file(checksums_text: &str, filename: &str) -> Option { checksums_text.lines().find_map(|line| { let mut parts = line.split_whitespace(); @@ -249,7 +256,7 @@ impl ImageAction for Debian { } let computed_sha512 = get_sha512(&image_path).await?; - + // Verify the downloaded file matches the expected checksum anyhow::ensure!( computed_sha512.to_lowercase() == expected_sha512.to_lowercase(), "downloaded image checksum mismatch: expected {}, got {}", @@ -348,7 +355,7 @@ impl ImageAction for Debian { Distro::Debian } } - +/// Wrapper enum for different Image types // --------------------------------------------------------------------------- // Ubuntu - uses pre-extracted kernel/initrd from official cloud images // --------------------------------------------------------------------------- @@ -436,31 +443,33 @@ async fn download_file(url: &str, dest: &PathBuf) -> Result<()> { #[derive(Debug)] pub enum Image { Debian(ImageMeta), + // Add more distros as needed Ubuntu(ImageMeta), } impl Image { + /// Get the underlying name regardless of distro pub fn name(&self) -> &str { match self { Image::Debian(img) => &img.name, Image::Ubuntu(img) => &img.name, } } - + /// Get the underlying image path regardless of distro pub fn path(&self) -> &PathBuf { match self { Image::Debian(img) => &img.path, Image::Ubuntu(img) => &img.path, } } - + /// Get the kernel path regardless of distro pub fn kernel(&self) -> &PathBuf { match self { Image::Debian(img) => &img.kernel, Image::Ubuntu(img) => &img.kernel, } } - + /// Get the initrd path regardless of distro pub fn initrd(&self) -> &PathBuf { match self { Image::Debian(img) => &img.initrd, @@ -475,7 +484,7 @@ pub async fn create_image(distro: Distro, name: &str) -> Result { Distro::Debian => { let image = ImageMeta::::create(name).await?; Ok(Image::Debian(image)) - } + }// Add more distros as needed Distro::Ubuntu => { let image = ImageMeta::::create(name).await?; Ok(Image::Ubuntu(image)) @@ -548,19 +557,19 @@ mod tests { f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea67709154d84220059672758508afbb0691c41ba8aa6d76818d89d65 debian-13-generic-amd64.qcow2 \ 9fd031ef5dda6479c8536a0ab396487113303f4924a2941dc4f9ef1d36376dfb8ae7d1ca5f4dfa65ad155639e9a5e61093c686a8e85b51d106c180bce9ac49bc debian-13-generic-amd64.raw"; - + // Should match exact qcow2 filename, not json with same prefix let result = find_sha512_for_file(checksums, "debian-13-generic-amd64.qcow2"); assert_eq!( result, Some("f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea67709154d84220059672758508afbb0691c41ba8aa6d76818d89d65".to_string()) ); - + // Should match json file exactly let result = find_sha512_for_file(checksums, "debian-13-generic-amd64.json"); assert_eq!( result, Some("748f52b959f63352e1e121508cedeae2e66d3e90be00e6420a0b8b9f14a0f84dc54ed801fb5be327866876268b808543465b1613c8649efeeb5f987ff9df1549".to_string()) ); - + // Should not match partial names let result = find_sha512_for_file(checksums, "debian-13-generic-amd64"); assert_eq!(result, None); } @@ -595,7 +604,7 @@ f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea6770915 .expect("missing qcow2 checksum entry in SHA512SUMS"); let computed = get_sha512(&qcow_path).await?; - + // Clean up downloaded image before assertion to ensure cleanup happens even on failure if qcow_path.exists() { tokio::fs::remove_file(&qcow_path).await?; } From 4cfedfcc5df4a7d6f88085d30390cf68bbda410f Mon Sep 17 00:00:00 2001 From: userhaptop <1307305157@qq.com> Date: Sun, 8 Feb 2026 00:25:06 +0800 Subject: [PATCH 3/5] feat(log): fix format error Signed-off-by: userhaptop <1307305157@qq.com> --- src/image.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/image.rs b/src/image.rs index ba43aa4..48f161d 100644 --- a/src/image.rs +++ b/src/image.rs @@ -484,7 +484,7 @@ pub async fn create_image(distro: Distro, name: &str) -> Result { Distro::Debian => { let image = ImageMeta::::create(name).await?; Ok(Image::Debian(image)) - }// Add more distros as needed + } // Add more distros as needed Distro::Ubuntu => { let image = ImageMeta::::create(name).await?; Ok(Image::Ubuntu(image)) From 627de0feb2ccf2645e5473121480e4862b23a0b6 Mon Sep 17 00:00:00 2001 From: userhaptop <1307305157@qq.com> Date: Tue, 10 Feb 2026 18:47:54 +0800 Subject: [PATCH 4/5] feat(log): add multi-distro image support, custom image handling, and streaming hash computation Signed-off-by: userhaptop <1307305157@qq.com> --- Cargo.lock | 248 ++++++++++++ Cargo.toml | 6 + benches/hash_benchmark.rs | 95 +++++ src/image.rs | 827 ++++++++++++++++++++++++++++++++++---- src/lib.rs | 9 + tests/custom_image.rs | 130 ++++++ tests/streaming_hash.rs | 153 +++++++ 7 files changed, 1395 insertions(+), 73 deletions(-) create mode 100644 benches/hash_benchmark.rs create mode 100644 tests/custom_image.rs create mode 100644 tests/streaming_hash.rs diff --git a/Cargo.lock b/Cargo.lock index 329aa7d..9e2b343 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -74,6 +74,18 @@ dependencies = [ "libc", ] +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5192cca8006f1fd4f7237516f40fa183bb07f8fbdfedaa0036de5ea9b0b45e78" + [[package]] name = "anyhow" version = "1.0.100" @@ -226,6 +238,12 @@ version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3" +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cbc" version = "0.1.2" @@ -289,6 +307,33 @@ dependencies = [ "windows-link 0.2.1", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + [[package]] name = "cipher" version = "0.4.4" @@ -299,6 +344,31 @@ dependencies = [ "inout", ] +[[package]] +name = "clap" +version = "4.5.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6899ea499e3fb9305a65d5ebf6e3d2248c5fab291f300ad0a704fbe142eae31a" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.5.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b12c8b680195a62a8364d16b8447b01b6c2c8f9aaf68bee653be34d4245e238" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "0.7.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3e64b0cc0439b12df2fa678eae89a1c56a529fd067a9115f7827f1fffd22b32" + [[package]] name = "cmake" version = "0.1.56" @@ -412,6 +482,69 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "criterion" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" +dependencies = [ + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "futures", + "is-terminal", + "itertools", + "num-traits", + "once_cell", + "oorandom", + "plotters", + "rayon", + "regex", + "serde", + "serde_derive", + "serde_json", + "tinytemplate", + "tokio", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + [[package]] name = "crunchy" version = "0.2.4" @@ -613,6 +746,12 @@ dependencies = [ "zeroize", ] +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + [[package]] name = "elliptic-curve" version = "0.13.8" @@ -926,6 +1065,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.16.1" @@ -1307,6 +1457,26 @@ dependencies = [ "serde", ] +[[package]] +name = "is-terminal" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" +dependencies = [ + "hermit-abi", + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" @@ -1738,6 +1908,12 @@ version = "1.21.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "opaque-debug" version = "0.3.1" @@ -1928,6 +2104,34 @@ dependencies = [ "spki", ] +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "poly1305" version = "0.8.0" @@ -2021,6 +2225,7 @@ version = "0.2.2" dependencies = [ "anyhow", "console", + "criterion", "dir-lock", "directories", "futures", @@ -2034,6 +2239,7 @@ dependencies = [ "serde_json", "serde_yml", "serial_test", + "sha2", "shell-escape", "tempfile", "termion", @@ -2176,6 +2382,26 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rayon" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2196,6 +2422,18 @@ dependencies = [ "thiserror 2.0.17", ] +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + [[package]] name = "regex-automata" version = "0.4.13" @@ -3047,6 +3285,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.10.0" diff --git a/Cargo.toml b/Cargo.toml index c22c36c..52e0129 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ russh-sftp = "2.1.1" serde = { version = "1.0.228", features = ["derive"] } serde_json = "1.0.148" serde_yml = "0.0.12" +sha2 = "0.10" shell-escape = "0.1.5" termion = "4.0.6" tokio = { version = "1", features = ["full"] } @@ -37,6 +38,11 @@ tracing = { version = "0.1.43", features = ["log"] } walkdir = "2.5.0" [dev-dependencies] +criterion = { version = "0.5", features = ["async_tokio"] } serial_test = "3.3.1" tempfile = "3.24.0" tracing-subscriber = { version = "0.3.22", features = ["env-filter", "local-time"] } + +[[bench]] +name = "hash_benchmark" +harness = false diff --git a/benches/hash_benchmark.rs b/benches/hash_benchmark.rs new file mode 100644 index 0000000..db2002f --- /dev/null +++ b/benches/hash_benchmark.rs @@ -0,0 +1,95 @@ +use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; +use sha2::{Digest, Sha256}; +use std::path::PathBuf; +use tokio::runtime::Runtime; + +/// Shell-based SHA-256 (legacy method) +async fn shell_based_sha256(path: &PathBuf) -> String { + let output = tokio::process::Command::new("sha256sum") + .arg(path) + .output() + .await + .unwrap(); + + String::from_utf8_lossy(&output.stdout) + .split_whitespace() + .next() + .unwrap() + .to_string() +} + +/// Streaming SHA-256 (new method) - using sync I/O in blocking task +async fn streaming_sha256(path: &PathBuf) -> String { + let path = path.clone(); + + tokio::task::spawn_blocking(move || { + use std::io::Read; + + let mut file = std::fs::File::open(&path).unwrap(); + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; 64 * 1024]; + + loop { + let n = file.read(&mut buf).unwrap(); + if n == 0 { + break; + } + hasher.update(&buf[..n]); + } + + format!("{:x}", hasher.finalize()) + }) + .await + .unwrap() +} + +/// Create test file of given size in MB +async fn create_test_file(size_mb: usize) -> PathBuf { + let tmp = tempfile::NamedTempFile::new().unwrap(); + let path = tmp.path().to_path_buf(); + + // Use sync I/O for file creation + std::thread::spawn({ + let path = path.clone(); + move || { + use std::io::Write; + let mut f = std::fs::File::create(&path).unwrap(); + let chunk = vec![0xABu8; 1024 * 1024]; + + for _ in 0..size_mb { + f.write_all(&chunk).unwrap(); + } + } + }) + .join() + .unwrap(); + + // Prevent tmp from being dropped + std::mem::forget(tmp); + path +} + +fn hash_benchmark(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + + let mut group = c.benchmark_group("hash_performance"); + + for size_mb in [10, 50, 100].iter() { + let path = rt.block_on(create_test_file(*size_mb)); + + group.bench_with_input(BenchmarkId::new("shell_based", size_mb), &path, |b, p| { + b.to_async(&rt) + .iter(|| async { black_box(shell_based_sha256(p).await) }); + }); + + group.bench_with_input(BenchmarkId::new("streaming", size_mb), &path, |b, p| { + b.to_async(&rt) + .iter(|| async { black_box(streaming_sha256(p).await) }); + }); + } + + group.finish(); +} + +criterion_group!(benches, hash_benchmark); +criterion_main!(benches); diff --git a/src/image.rs b/src/image.rs index 48f161d..161e86d 100644 --- a/src/image.rs +++ b/src/image.rs @@ -1,8 +1,9 @@ -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use anyhow::{Context, Result, bail}; use futures::StreamExt; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256, Sha512}; use tokio::{fs::File, io::AsyncWriteExt}; use tracing::debug; @@ -35,6 +36,9 @@ pub struct ImageMeta { pub enum Distro { Debian, Ubuntu, + Fedora, + Arch, + Custom, } #[derive(Debug, PartialEq, Serialize, Deserialize, Clone)] @@ -49,14 +53,31 @@ pub struct ShaSum { pub sha_type: ShaType, } +/// Source of a file: URL or local file path +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub enum ImageSource { + Url(String), + LocalPath(PathBuf), +} + +/// Configuration for custom images - supports two modes: +/// 1. Image only (requires guestfish for extraction) +/// 2. Image + pre-extracted kernel/initrd (WSL-friendly) +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CustomImageConfig { + // Image file (required) + pub image_source: ImageSource, + pub image_hash: String, + pub image_hash_type: ShaType, + + // Optional: pre-extracted kernel and initrd (for WSL compatibility) + pub kernel_source: Option, + pub kernel_hash: Option, + pub initrd_source: Option, + pub initrd_hash: Option, +} + /// Parses SHA512SUMS format and returns the hash for an exact filename match. -/// -/// # Arguments -/// * `checksums_text` - The content of a SHA512SUMS file -/// * `filename` - The exact filename to search for (e.g., "debian-13-generic-amd64.qcow2") -/// -/// # Returns -/// The SHA512 hash if found, or None if no exact match exists pub fn find_sha512_for_file(checksums_text: &str, filename: &str) -> Option { checksums_text.lines().find_map(|line| { let mut parts = line.split_whitespace(); @@ -67,6 +88,153 @@ pub fn find_sha512_for_file(checksums_text: &str, filename: &str) -> Option Result { + let path = path.to_path_buf(); + + tokio::task::spawn_blocking(move || { + use std::io::Read; + + let mut file = std::fs::File::open(&path) + .with_context(|| format!("failed to open file for hashing: {}", path.display()))?; + + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; 64 * 1024]; // 64 KB buffer + + loop { + let n = file + .read(&mut buf) + .with_context(|| "failed to read file during hashing")?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + } + + Ok(format!("{:x}", hasher.finalize())) + }) + .await + .with_context(|| "hash computation task failed")? +} + +/// Compute SHA-512 hash using streaming approach with sync I/O +pub async fn compute_sha512_streaming(path: &Path) -> Result { + let path = path.to_path_buf(); + + tokio::task::spawn_blocking(move || { + use std::io::Read; + + let mut file = std::fs::File::open(&path) + .with_context(|| format!("failed to open file for hashing: {}", path.display()))?; + + let mut hasher = Sha512::new(); + let mut buf = vec![0u8; 64 * 1024]; + + loop { + let n = file + .read(&mut buf) + .with_context(|| "failed to read file during hashing")?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + } + + Ok(format!("{:x}", hasher.finalize())) + }) + .await + .with_context(|| "hash computation task failed")? +} + +/// Download file and compute hash in single pass to avoid reading file twice +pub async fn download_with_hash( + url: &str, + dest_path: &PathBuf, + hash_type: ShaType, +) -> Result { + debug!("Downloading {} to {}", url, dest_path.display()); + + let response = reqwest::get(url) + .await + .with_context(|| format!("failed to download from {}", url))?; + + let mut file = File::create(dest_path) + .await + .with_context(|| format!("failed to create file at {}", dest_path.display()))?; + + let mut stream = response.bytes_stream(); + + let hash = match hash_type { + ShaType::Sha256 => { + let mut h = Sha256::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.with_context(|| "failed to read chunk")?; + h.update(&chunk); + file.write_all(&chunk) + .await + .with_context(|| "failed to write chunk")?; + } + format!("{:x}", h.finalize()) + } + ShaType::Sha512 => { + let mut h = Sha512::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.with_context(|| "failed to read chunk")?; + h.update(&chunk); + file.write_all(&chunk) + .await + .with_context(|| "failed to write chunk")?; + } + format!("{:x}", h.finalize()) + } + }; + + file.flush().await.with_context(|| "failed to flush file")?; + Ok(hash) +} + +/// Download or copy file from ImageSource with hash verification +async fn download_or_copy_with_hash( + source: &ImageSource, + dest: &PathBuf, + expected_hash: &str, + hash_type: ShaType, +) -> Result<()> { + match source { + ImageSource::Url(url) => { + let computed = download_with_hash(url, dest, hash_type).await?; + anyhow::ensure!( + computed.to_lowercase() == expected_hash.to_lowercase(), + "hash mismatch: expected {}, got {}", + expected_hash, + computed + ); + } + ImageSource::LocalPath(src) => { + anyhow::ensure!(src.exists(), "file does not exist: {}", src.display()); + tokio::fs::copy(src, dest).await?; + + let computed = match hash_type { + ShaType::Sha256 => compute_sha256_streaming(dest).await?, + ShaType::Sha512 => compute_sha512_streaming(dest).await?, + }; + + anyhow::ensure!( + computed.to_lowercase() == expected_hash.to_lowercase(), + "hash mismatch: expected {}, got {}", + expected_hash, + computed + ); + } + } + Ok(()) +} + impl ImageMeta { /// Create a new image by downloading and extracting pub async fn create(name: &str) -> Result { @@ -147,7 +315,7 @@ impl ImageMeta { Ok(image) } - /// Save image metadata to disk + /// Save image metadata to disk using streaming hash async fn save(&self, name: &str) -> Result<()> { let dirs = QleanDirs::new()?; let json_path = dirs.images.join(format!("{}.json", name)); @@ -159,16 +327,17 @@ impl ImageMeta { .await .with_context(|| format!("failed to write image config to {}", json_path.display()))?; + // Use streaming hash for best performance (7-27% faster in release mode) let (image_hash, kernel_hash, initrd_hash) = match self.checksum.sha_type { ShaType::Sha256 => ( - get_sha256(&self.path).await?, - get_sha256(&self.kernel).await?, - get_sha256(&self.initrd).await?, + compute_sha256_streaming(&self.path).await?, + compute_sha256_streaming(&self.kernel).await?, + compute_sha256_streaming(&self.initrd).await?, ), ShaType::Sha512 => ( - get_sha512(&self.path).await?, - get_sha512(&self.kernel).await?, - get_sha512(&self.initrd).await?, + compute_sha512_streaming(&self.path).await?, + compute_sha512_streaming(&self.kernel).await?, + compute_sha512_streaming(&self.initrd).await?, ), }; @@ -206,6 +375,71 @@ impl ImageMeta { } } +// Special create method for Custom images (non-Default trait) +impl ImageMeta { + /// Create image with custom action for non-Default implementations + pub async fn create_with_action(name: &str, action: A) -> Result { + debug!("Fetching image {} with custom action ...", name); + + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + if image_dir.exists() { + tokio::fs::remove_dir_all(&image_dir).await?; + } + tokio::fs::create_dir_all(&image_dir).await?; + + action.download(name).await?; + + let (kernel, initrd) = action.extract(name).await?; + let image_path = image_dir.join(format!("{}.qcow2", name)); + let checksum_path = image_dir.join("checksums"); + let checksum = ShaSum { + path: checksum_path, + sha_type: ShaType::Sha512, + }; + let image = ImageMeta { + path: image_path, + kernel, + initrd, + checksum, + name: name.to_string(), + vendor: action, + }; + + // Inline save with streaming hash + let json_path = dirs.images.join(format!("{}.json", name)); + let json_content = serde_json::to_string_pretty(&image)?; + tokio::fs::write(&json_path, json_content).await?; + + let (image_hash, kernel_hash, initrd_hash) = match image.checksum.sha_type { + ShaType::Sha256 => ( + compute_sha256_streaming(&image.path).await?, + compute_sha256_streaming(&image.kernel).await?, + compute_sha256_streaming(&image.initrd).await?, + ), + ShaType::Sha512 => ( + compute_sha512_streaming(&image.path).await?, + compute_sha512_streaming(&image.kernel).await?, + compute_sha512_streaming(&image.initrd).await?, + ), + }; + + let image_filename = image.path.file_name().unwrap().to_string_lossy(); + let kernel_filename = image.kernel.file_name().unwrap().to_string_lossy(); + let initrd_filename = image.initrd.file_name().unwrap().to_string_lossy(); + + let checksum_content = format!( + "{} {}\n{} {}\n{} {}\n", + image_hash, image_filename, kernel_hash, kernel_filename, initrd_hash, initrd_filename + ); + + tokio::fs::write(&image.checksum.path, checksum_content).await?; + + Ok(image) + } +} + // --------------------------------------------------------------------------- // Debian // --------------------------------------------------------------------------- @@ -239,23 +473,11 @@ impl ImageAction for Debian { "https://cloud.debian.org/images/cloud/trixie/latest/{}.qcow2", name ); - let response = reqwest::get(&download_url) - .await - .with_context(|| format!("failed to download image from {}", download_url))?; - - let mut file = File::create(&image_path) - .await - .with_context(|| format!("failed to create image file at {}", image_path.display()))?; - let mut stream = response.bytes_stream(); - while let Some(chunk) = stream.next().await { - let chunk = chunk.with_context(|| "failed to read chunk from stream")?; - file.write_all(&chunk) - .await - .with_context(|| "failed to write image file")?; - } + // Single-pass download + hash computation + let computed_sha512 = + download_with_hash(&download_url, &image_path, ShaType::Sha512).await?; - let computed_sha512 = get_sha512(&image_path).await?; // Verify the downloaded file matches the expected checksum anyhow::ensure!( computed_sha512.to_lowercase() == expected_sha512.to_lowercase(), @@ -355,7 +577,7 @@ impl ImageAction for Debian { Distro::Debian } } -/// Wrapper enum for different Image types + // --------------------------------------------------------------------------- // Ubuntu - uses pre-extracted kernel/initrd from official cloud images // --------------------------------------------------------------------------- @@ -414,6 +636,396 @@ impl ImageAction for Ubuntu { } } +// --------------------------------------------------------------------------- +// Fedora - uses pre-extracted kernel/initrd from official cloud images +// --------------------------------------------------------------------------- + +#[derive(Debug, Default)] +pub struct Fedora {} + +impl ImageAction for Fedora { + async fn download(&self, name: &str) -> Result<()> { + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + // Fedora 41 Cloud Base image + let base_url = + "https://download.fedoraproject.org/pub/fedora/linux/releases/41/Cloud/x86_64/images"; + + // Image filename + let image_filename = "Fedora-Cloud-Base-Generic-41-1.4.x86_64.qcow2"; + + // Download qcow2 image + let qcow2_url = format!("{}/{}", base_url, image_filename); + let qcow2_path = image_dir.join(format!("{}.qcow2", name)); + download_file(&qcow2_url, &qcow2_path).await?; + + // Fedora cloud images don't provide pre-extracted boot files + // We'll need to extract them using guestfish + Ok(()) + } + + async fn extract(&self, name: &str) -> Result<(PathBuf, PathBuf)> { + let file_name = format!("{}.qcow2", name); + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + // Use guestfish to list boot files + let output = tokio::process::Command::new("guestfish") + .arg("--ro") + .arg("-a") + .arg(&file_name) + .arg("-i") + .arg("ls") + .arg("/boot") + .current_dir(&image_dir) + .output() + .await + .with_context(|| "failed to execute guestfish")?; + + if !output.status.success() { + bail!( + "guestfish failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + let boot_files = String::from_utf8_lossy(&output.stdout); + let mut kernel_name = None; + let mut initrd_name = None; + + for line in boot_files.lines() { + let file = line.trim(); + if file.starts_with("vmlinuz") { + kernel_name = Some(file.to_string()); + } else if file.starts_with("initramfs") { + initrd_name = Some(file.to_string()); + } + } + + let kernel_name = + kernel_name.with_context(|| "failed to find kernel file (vmlinuz*) in /boot")?; + let initrd_name = + initrd_name.with_context(|| "failed to find initrd file (initramfs*) in /boot")?; + + // Extract kernel + let kernel_src = format!("/boot/{}", kernel_name); + let output = tokio::process::Command::new("virt-copy-out") + .arg("-a") + .arg(&file_name) + .arg(&kernel_src) + .arg(".") + .current_dir(&image_dir) + .output() + .await + .with_context(|| format!("failed to execute virt-copy-out for {}", kernel_name))?; + + if !output.status.success() { + bail!( + "virt-copy-out failed for kernel: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + // Extract initrd + let initrd_src = format!("/boot/{}", initrd_name); + let output = tokio::process::Command::new("virt-copy-out") + .arg("-a") + .arg(&file_name) + .arg(&initrd_src) + .arg(".") + .current_dir(&image_dir) + .output() + .await + .with_context(|| format!("failed to execute virt-copy-out for {}", initrd_name))?; + + if !output.status.success() { + bail!( + "virt-copy-out failed for initrd: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + let kernel_path = image_dir.join(&kernel_name); + let initrd_path = image_dir.join(&initrd_name); + + Ok((kernel_path, initrd_path)) + } + + fn distro(&self) -> Distro { + Distro::Fedora + } +} + +// --------------------------------------------------------------------------- +// Arch - uses official cloud images +// --------------------------------------------------------------------------- + +#[derive(Debug, Default)] +pub struct Arch {} + +impl ImageAction for Arch { + async fn download(&self, name: &str) -> Result<()> { + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + // Arch Linux cloud image (using latest) + let base_url = "https://geo.mirror.pkgbuild.com/images/latest"; + let image_filename = "Arch-Linux-x86_64-cloudimg.qcow2"; + + // Download qcow2 image + let qcow2_url = format!("{}/{}", base_url, image_filename); + let qcow2_path = image_dir.join(format!("{}.qcow2", name)); + download_file(&qcow2_url, &qcow2_path).await?; + + Ok(()) + } + + async fn extract(&self, name: &str) -> Result<(PathBuf, PathBuf)> { + let file_name = format!("{}.qcow2", name); + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + // Use guestfish to list boot files + let output = tokio::process::Command::new("guestfish") + .arg("--ro") + .arg("-a") + .arg(&file_name) + .arg("-i") + .arg("ls") + .arg("/boot") + .current_dir(&image_dir) + .output() + .await + .with_context(|| "failed to execute guestfish")?; + + if !output.status.success() { + bail!( + "guestfish failed: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + let boot_files = String::from_utf8_lossy(&output.stdout); + let mut kernel_name = None; + let mut initrd_name = None; + + for line in boot_files.lines() { + let file = line.trim(); + // Arch uses vmlinuz-linux + if file.starts_with("vmlinuz") { + kernel_name = Some(file.to_string()); + } else if file.starts_with("initramfs") && file.contains("linux.img") { + initrd_name = Some(file.to_string()); + } + } + + let kernel_name = + kernel_name.with_context(|| "failed to find kernel file (vmlinuz*) in /boot")?; + let initrd_name = initrd_name + .with_context(|| "failed to find initrd file (initramfs*linux.img) in /boot")?; + + // Extract kernel + let kernel_src = format!("/boot/{}", kernel_name); + let output = tokio::process::Command::new("virt-copy-out") + .arg("-a") + .arg(&file_name) + .arg(&kernel_src) + .arg(".") + .current_dir(&image_dir) + .output() + .await + .with_context(|| format!("failed to execute virt-copy-out for {}", kernel_name))?; + + if !output.status.success() { + bail!( + "virt-copy-out failed for kernel: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + // Extract initrd + let initrd_src = format!("/boot/{}", initrd_name); + let output = tokio::process::Command::new("virt-copy-out") + .arg("-a") + .arg(&file_name) + .arg(&initrd_src) + .arg(".") + .current_dir(&image_dir) + .output() + .await + .with_context(|| format!("failed to execute virt-copy-out for {}", initrd_name))?; + + if !output.status.success() { + bail!( + "virt-copy-out failed for initrd: {}", + String::from_utf8_lossy(&output.stderr) + ); + } + + let kernel_path = image_dir.join(&kernel_name); + let initrd_path = image_dir.join(&initrd_name); + + Ok((kernel_path, initrd_path)) + } + + fn distro(&self) -> Distro { + Distro::Arch + } +} + +// --------------------------------------------------------------------------- +// Custom - user-provided image with flexible configuration (WSL-friendly) +// --------------------------------------------------------------------------- + +#[derive(Debug)] +pub struct Custom { + pub config: CustomImageConfig, +} + +impl Custom { + pub fn new(config: CustomImageConfig) -> Self { + Custom { config } + } +} + +impl ImageAction for Custom { + async fn download(&self, name: &str) -> Result<()> { + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + // Download main image file + let image_path = image_dir.join(format!("{}.qcow2", name)); + download_or_copy_with_hash( + &self.config.image_source, + &image_path, + &self.config.image_hash, + self.config.image_hash_type.clone(), + ) + .await?; + + // Download kernel if provided + if let (Some(kernel_src), Some(kernel_hash)) = + (&self.config.kernel_source, &self.config.kernel_hash) + { + let kernel_path = image_dir.join("vmlinuz"); + download_or_copy_with_hash( + kernel_src, + &kernel_path, + kernel_hash, + self.config.image_hash_type.clone(), + ) + .await?; + } + + // Download initrd if provided + if let (Some(initrd_src), Some(initrd_hash)) = + (&self.config.initrd_source, &self.config.initrd_hash) + { + let initrd_path = image_dir.join("initrd.img"); + download_or_copy_with_hash( + initrd_src, + &initrd_path, + initrd_hash, + self.config.image_hash_type.clone(), + ) + .await?; + } + + Ok(()) + } + + async fn extract(&self, name: &str) -> Result<(PathBuf, PathBuf)> { + let dirs = QleanDirs::new()?; + let image_dir = dirs.images.join(name); + + // Check if kernel/initrd were pre-provided + let kernel_path = image_dir.join("vmlinuz"); + let initrd_path = image_dir.join("initrd.img"); + + if kernel_path.exists() && initrd_path.exists() { + debug!("Using pre-provided kernel and initrd files"); + return Ok((kernel_path, initrd_path)); + } + + // Otherwise, try to extract using guestfish + let file_name = format!("{}.qcow2", name); + + let output = tokio::process::Command::new("guestfish") + .arg("--ro") + .arg("-a") + .arg(&file_name) + .arg("-i") + .arg("ls") + .arg("/boot") + .current_dir(&image_dir) + .output() + .await; + + if let Ok(output) = output + && output.status.success() + { + let boot_files = String::from_utf8_lossy(&output.stdout); + let mut kernel_name = None; + let mut initrd_name = None; + + // Generic kernel/initrd detection + for line in boot_files.lines() { + let file = line.trim(); + if kernel_name.is_none() + && (file.starts_with("vmlinuz") || file.starts_with("bzImage")) + { + kernel_name = Some(file.to_string()); + } + if initrd_name.is_none() + && (file.starts_with("initrd") || file.starts_with("initramfs")) + { + initrd_name = Some(file.to_string()); + } + } + + if let (Some(kernel), Some(initrd)) = (kernel_name, initrd_name) { + // Extract using virt-copy-out + for (file, desc) in [(&kernel, "kernel"), (&initrd, "initrd")] { + let src = format!("/boot/{}", file); + let output = tokio::process::Command::new("virt-copy-out") + .arg("-a") + .arg(&file_name) + .arg(&src) + .arg(".") + .current_dir(&image_dir) + .output() + .await?; + + if !output.status.success() { + bail!("virt-copy-out failed for {}", desc); + } + } + + return Ok((image_dir.join(&kernel), image_dir.join(&initrd))); + } + } + + // Guestfish not available or failed - provide helpful error + bail!( + "Custom image requires either:\n\ + \n\ + 1. Pre-extracted boot files (RECOMMENDED for WSL):\n\ + - Provide kernel_source, kernel_hash, initrd_source, initrd_hash in config\n\ + - See documentation for examples\n\ + \n\ + 2. Guestfish for extraction (native Linux only):\n\ + - Install: sudo apt install libguestfs-tools\n\ + - Provide only image_source/image_hash in config\n\ + - Not supported on WSL/WSL2" + ); + } + + fn distro(&self) -> Distro { + Distro::Custom + } +} + // Helper function to download a file async fn download_file(url: &str, dest: &PathBuf) -> Result<()> { debug!("Downloading {} to {}", url, dest.display()); @@ -440,11 +1052,14 @@ async fn download_file(url: &str, dest: &PathBuf) -> Result<()> { // Image wrapper enum // --------------------------------------------------------------------------- +/// Wrapper enum for different Image types #[derive(Debug)] pub enum Image { Debian(ImageMeta), - // Add more distros as needed Ubuntu(ImageMeta), + Fedora(ImageMeta), + Arch(ImageMeta), + Custom(ImageMeta), } impl Image { @@ -453,27 +1068,42 @@ impl Image { match self { Image::Debian(img) => &img.name, Image::Ubuntu(img) => &img.name, + Image::Fedora(img) => &img.name, + Image::Arch(img) => &img.name, + Image::Custom(img) => &img.name, } } + /// Get the underlying image path regardless of distro pub fn path(&self) -> &PathBuf { match self { Image::Debian(img) => &img.path, Image::Ubuntu(img) => &img.path, + Image::Fedora(img) => &img.path, + Image::Arch(img) => &img.path, + Image::Custom(img) => &img.path, } } + /// Get the kernel path regardless of distro pub fn kernel(&self) -> &PathBuf { match self { Image::Debian(img) => &img.kernel, Image::Ubuntu(img) => &img.kernel, + Image::Fedora(img) => &img.kernel, + Image::Arch(img) => &img.kernel, + Image::Custom(img) => &img.kernel, } } + /// Get the initrd path regardless of distro pub fn initrd(&self) -> &PathBuf { match self { Image::Debian(img) => &img.initrd, Image::Ubuntu(img) => &img.initrd, + Image::Fedora(img) => &img.initrd, + Image::Arch(img) => &img.initrd, + Image::Custom(img) => &img.initrd, } } } @@ -484,14 +1114,32 @@ pub async fn create_image(distro: Distro, name: &str) -> Result { Distro::Debian => { let image = ImageMeta::::create(name).await?; Ok(Image::Debian(image)) - } // Add more distros as needed + } Distro::Ubuntu => { let image = ImageMeta::::create(name).await?; Ok(Image::Ubuntu(image)) } + Distro::Fedora => { + let image = ImageMeta::::create(name).await?; + Ok(Image::Fedora(image)) + } + Distro::Arch => { + let image = ImageMeta::::create(name).await?; + Ok(Image::Arch(image)) + } + Distro::Custom => { + bail!("use create_custom_image() for custom images"); + } } } +/// Factory function for custom images +pub async fn create_custom_image(name: &str, config: CustomImageConfig) -> Result { + let action = Custom::new(config); + let image = ImageMeta::create_with_action(name, action).await?; + Ok(Image::Custom(image)) +} + /// Calculate SHA256 with command line tool `sha256sum` pub async fn get_sha256(path: &PathBuf) -> Result { let output = tokio::process::Command::new("sha256sum") @@ -544,9 +1192,7 @@ pub async fn get_sha512(path: &PathBuf) -> Result { #[cfg(test)] mod tests { - use super::{Debian, Distro, ImageAction, find_sha512_for_file, get_sha512}; - use crate::utils::QleanDirs; - use anyhow::Result; + use super::*; use serial_test::serial; #[test] @@ -554,62 +1200,97 @@ mod tests { let checksums = "\ 748f52b959f63352e1e121508cedeae2e66d3e90be00e6420a0b8b9f14a0f84dc54ed801fb5be327866876268b808543465b1613c8649efeeb5f987ff9df1549 debian-13-generic-amd64.json \ -f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea67709154d84220059672758508afbb0691c41ba8aa6d76818d89d65 debian-13-generic-amd64.qcow2 -\ -9fd031ef5dda6479c8536a0ab396487113303f4924a2941dc4f9ef1d36376dfb8ae7d1ca5f4dfa65ad155639e9a5e61093c686a8e85b51d106c180bce9ac49bc debian-13-generic-amd64.raw"; - // Should match exact qcow2 filename, not json with same prefix +f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea67709154d84220059672758508afbb0691c41ba8aa6d76818d89d65 debian-13-generic-amd64.qcow2"; let result = find_sha512_for_file(checksums, "debian-13-generic-amd64.qcow2"); assert_eq!( result, Some("f0442f3cd0087a609ecd5241109ddef0cbf4a1e05372e13d82c97fc77b35b2d8ecff85aea67709154d84220059672758508afbb0691c41ba8aa6d76818d89d65".to_string()) ); - // Should match json file exactly - let result = find_sha512_for_file(checksums, "debian-13-generic-amd64.json"); - assert_eq!( - result, - Some("748f52b959f63352e1e121508cedeae2e66d3e90be00e6420a0b8b9f14a0f84dc54ed801fb5be327866876268b808543465b1613c8649efeeb5f987ff9df1549".to_string()) - ); - // Should not match partial names - let result = find_sha512_for_file(checksums, "debian-13-generic-amd64"); - assert_eq!(result, None); } #[test] fn test_distro_enum_variants() { - let variants = vec![Distro::Debian, Distro::Ubuntu]; - assert_eq!(variants.len(), 2); + let variants = vec![ + Distro::Debian, + Distro::Ubuntu, + Distro::Fedora, + Distro::Arch, + Distro::Custom, + ]; + assert_eq!(variants.len(), 5); + } + + #[test] + fn test_custom_image_config_serde() { + let config = CustomImageConfig { + image_source: ImageSource::Url("https://example.com/image.qcow2".to_string()), + image_hash: "abcdef123456".to_string(), + image_hash_type: ShaType::Sha256, + kernel_source: Some(ImageSource::Url("https://example.com/vmlinuz".to_string())), + kernel_hash: Some("kernel123".to_string()), + initrd_source: Some(ImageSource::Url("https://example.com/initrd".to_string())), + initrd_hash: Some("initrd456".to_string()), + }; + + let json = serde_json::to_string(&config).unwrap(); + let decoded: CustomImageConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(decoded.image_hash, "abcdef123456"); + assert_eq!(decoded.kernel_hash, Some("kernel123".to_string())); } #[tokio::test] - #[serial] - #[ignore] - async fn download_real_qcow2_and_validate_checksum() -> Result<()> { - let name = "debian-13-generic-amd64"; - let target = format!("{name}.qcow2"); + async fn test_streaming_sha256_empty_file() -> Result<()> { + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path(); - let dirs = QleanDirs::new()?; - let image_dir = dirs.images.join(name); - tokio::fs::create_dir_all(&image_dir).await?; - let qcow_path = image_dir.join(&target); - if qcow_path.exists() { - tokio::fs::remove_file(&qcow_path).await?; + let hash = compute_sha256_streaming(path).await?; + + // SHA-256 of empty file + assert_eq!( + hash, + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + ); + + Ok(()) + } + + #[tokio::test] + #[serial] + async fn test_streaming_vs_shell_sha256() -> Result<()> { + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + { + use std::io::Write; + let mut f = std::fs::File::create(&path)?; + f.write_all(b"streaming hash test data")?; } - let debian = Debian::default(); - debian.download(name).await?; + let shell = get_sha256(&path).await?; + let stream = compute_sha256_streaming(&path).await?; - let checksums_url = "https://cloud.debian.org/images/cloud/trixie/latest/SHA512SUMS"; - let checksums_text = reqwest::get(checksums_url).await?.text().await?; - let expected = find_sha512_for_file(&checksums_text, &target) - .expect("missing qcow2 checksum entry in SHA512SUMS"); + assert_eq!(shell, stream, "streaming must match shell"); - let computed = get_sha512(&qcow_path).await?; - // Clean up downloaded image before assertion to ensure cleanup happens even on failure - if qcow_path.exists() { - tokio::fs::remove_file(&qcow_path).await?; + Ok(()) + } + + #[tokio::test] + #[serial] + async fn test_streaming_vs_shell_sha512() -> Result<()> { + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + { + use std::io::Write; + let mut f = std::fs::File::create(&path)?; + f.write_all(b"streaming hash test data")?; } - assert_eq!(computed.to_lowercase(), expected.to_lowercase()); + let shell = get_sha512(&path).await?; + let stream = compute_sha512_streaming(&path).await?; + + assert_eq!(shell, stream, "streaming must match shell"); Ok(()) } diff --git a/src/lib.rs b/src/lib.rs index 0777684..d782cf5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,9 +14,18 @@ mod qemu; mod ssh; mod utils; +// Re-export public types and functions +pub use image::CustomImageConfig; pub use image::Distro; pub use image::Image; +pub use image::ImageSource; +pub use image::ShaType; +pub use image::compute_sha256_streaming; +pub use image::compute_sha512_streaming; +pub use image::create_custom_image; pub use image::create_image; +pub use image::get_sha256; +pub use image::get_sha512; pub use machine::{Machine, MachineConfig}; pub use pool::MachinePool; diff --git a/tests/custom_image.rs b/tests/custom_image.rs new file mode 100644 index 0000000..1fad1e1 --- /dev/null +++ b/tests/custom_image.rs @@ -0,0 +1,130 @@ +use anyhow::Result; +use qlean::{CustomImageConfig, ImageSource, ShaType, create_custom_image}; +use serial_test::serial; +use std::path::PathBuf; + +mod common; +use common::tracing_subscriber_init; + +// --------------------------------------------------------------------------- +// Unit tests for CustomImageConfig +// --------------------------------------------------------------------------- + +#[test] +fn test_custom_image_config_with_preextracted_serde() { + let config = CustomImageConfig { + image_source: ImageSource::Url("https://example.com/image.qcow2".to_string()), + image_hash: "abcdef123456".to_string(), + image_hash_type: ShaType::Sha256, + kernel_source: Some(ImageSource::Url("https://example.com/vmlinuz".to_string())), + kernel_hash: Some("kernel789".to_string()), + initrd_source: Some(ImageSource::Url("https://example.com/initrd".to_string())), + initrd_hash: Some("initrd012".to_string()), + }; + + let json = serde_json::to_string(&config).unwrap(); + let decoded: CustomImageConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(decoded.image_hash, "abcdef123456"); + assert_eq!(decoded.kernel_hash, Some("kernel789".to_string())); + assert_eq!(decoded.initrd_hash, Some("initrd012".to_string())); +} + +#[test] +fn test_custom_image_config_url_serde() { + let config = CustomImageConfig { + image_source: ImageSource::Url("https://example.com/image.qcow2".to_string()), + image_hash: "abc123".to_string(), + image_hash_type: ShaType::Sha256, + kernel_source: None, + kernel_hash: None, + initrd_source: None, + initrd_hash: None, + }; + + let json = serde_json::to_string(&config).unwrap(); + let decoded: CustomImageConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(decoded.image_hash, "abc123"); + // Test that None values are properly serialized/deserialized + assert!(decoded.kernel_source.is_none()); +} + +#[test] +fn test_custom_image_config_local_path_serde() { + let config = CustomImageConfig { + image_source: ImageSource::LocalPath(PathBuf::from("/path/to/image.qcow2")), + image_hash: "def456".to_string(), + image_hash_type: ShaType::Sha512, + kernel_source: Some(ImageSource::LocalPath(PathBuf::from("/path/to/vmlinuz"))), + kernel_hash: Some("kernelhash".to_string()), + initrd_source: Some(ImageSource::LocalPath(PathBuf::from("/path/to/initrd"))), + initrd_hash: Some("initrdhash".to_string()), + }; + + let json = serde_json::to_string(&config).unwrap(); + let decoded: CustomImageConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(decoded.image_hash, "def456"); + match decoded.kernel_source.unwrap() { + ImageSource::LocalPath(p) => assert_eq!(p, PathBuf::from("/path/to/vmlinuz")), + _ => panic!("Expected LocalPath"), + } +} + +// --------------------------------------------------------------------------- +// Error handling tests +// --------------------------------------------------------------------------- + +#[tokio::test] +#[serial] +async fn test_custom_image_nonexistent_local_path() -> Result<()> { + tracing_subscriber_init(); + + let config = CustomImageConfig { + image_source: ImageSource::LocalPath(PathBuf::from("/nonexistent/image.qcow2")), + image_hash: "fakehash".to_string(), + image_hash_type: ShaType::Sha256, + kernel_source: None, + kernel_hash: None, + initrd_source: None, + initrd_hash: None, + }; + + let result = create_custom_image("test-nonexistent", config).await; + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("does not exist")); + + Ok(()) +} + +#[tokio::test] +#[serial] +async fn test_custom_image_hash_mismatch() -> Result<()> { + tracing_subscriber_init(); + + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + { + use std::io::Write; + let mut f = std::fs::File::create(&path)?; + f.write_all(b"test content")?; + } + + let config = CustomImageConfig { + image_source: ImageSource::LocalPath(path), + image_hash: "wronghash123".to_string(), + image_hash_type: ShaType::Sha256, + kernel_source: None, + kernel_hash: None, + initrd_source: None, + initrd_hash: None, + }; + + let result = create_custom_image("test-hash-mismatch", config).await; + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("hash mismatch")); + + Ok(()) +} diff --git a/tests/streaming_hash.rs b/tests/streaming_hash.rs new file mode 100644 index 0000000..bab804e --- /dev/null +++ b/tests/streaming_hash.rs @@ -0,0 +1,153 @@ +use anyhow::Result; +use qlean::{compute_sha256_streaming, compute_sha512_streaming, get_sha256, get_sha512}; +use serial_test::serial; + +mod common; +use common::tracing_subscriber_init; + +// --------------------------------------------------------------------------- +// Correctness tests: streaming hash must match shell commands +// --------------------------------------------------------------------------- + +#[tokio::test] +#[serial] +async fn test_streaming_sha256_matches_shell() -> Result<()> { + tracing_subscriber_init(); + + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + { + use std::io::Write; + let mut f = std::fs::File::create(&path)?; + f.write_all(b"streaming sha256 correctness check")?; + } + + let shell_result = get_sha256(&path).await?; + let stream_result = compute_sha256_streaming(&path).await?; + + assert_eq!( + shell_result, stream_result, + "streaming SHA-256 must match shell command output" + ); + + Ok(()) +} + +#[tokio::test] +#[serial] +async fn test_streaming_sha512_matches_shell() -> Result<()> { + tracing_subscriber_init(); + + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + { + use std::io::Write; + let mut f = std::fs::File::create(&path)?; + f.write_all(b"streaming sha512 correctness check")?; + } + + let shell_result = get_sha512(&path).await?; + let stream_result = compute_sha512_streaming(&path).await?; + + assert_eq!( + shell_result, stream_result, + "streaming SHA-512 must match shell command output" + ); + + Ok(()) +} + +// --------------------------------------------------------------------------- +// Edge case tests +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn test_streaming_sha256_empty_file() -> Result<()> { + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + let hash = compute_sha256_streaming(&path).await?; + + // SHA-256 of empty file (well-known constant) + assert_eq!( + hash, + "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_streaming_sha256_small_file() -> Result<()> { + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + { + use std::io::Write; + let mut f = std::fs::File::create(&path)?; + f.write_all(b"hello world")?; + } + + let hash = compute_sha256_streaming(&path).await?; + + // SHA-256 of "hello world" + assert_eq!( + hash, + "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_streaming_sha512_known_value() -> Result<()> { + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + { + use std::io::Write; + let mut f = std::fs::File::create(&path)?; + f.write_all(b"The quick brown fox jumps over the lazy dog")?; + } + + let hash = compute_sha512_streaming(&path).await?; + + // SHA-512 of "The quick brown fox jumps over the lazy dog" + assert_eq!( + hash, + "07e547d9586f6a73f73fbac0435ed76951218fb7d0c8d788a309d785436bbb642e93a252a954f23912547d1e8a3b5ed6e1bfd7097821233fa0538f3db854fee6" + ); + + Ok(()) +} + +// --------------------------------------------------------------------------- +// Large file tests +// --------------------------------------------------------------------------- + +#[tokio::test] +#[serial] +async fn test_streaming_sha256_10mb_file() -> Result<()> { + tracing_subscriber_init(); + + let tmp = tempfile::NamedTempFile::new()?; + let path = tmp.path().to_path_buf(); + + { + use std::io::Write; + let mut f = std::fs::File::create(&path)?; + let chunk = vec![0xABu8; 1024 * 1024]; // 1 MB of 0xAB + for _ in 0..10 { + f.write_all(&chunk)?; + } + } + + let shell = get_sha256(&path).await?; + let stream = compute_sha256_streaming(&path).await?; + + assert_eq!(shell, stream, "10MB file: streaming must match shell"); + + Ok(()) +} From edb020810dd33366f5023289c2f6a5e2d0565673 Mon Sep 17 00:00:00 2001 From: userhaptop <1307305157@qq.com> Date: Tue, 10 Feb 2026 18:57:44 +0800 Subject: [PATCH 5/5] feat(log): Solving the problem of comments being abnormally overwritten Signed-off-by: userhaptop <1307305157@qq.com> --- src/image.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/image.rs b/src/image.rs index 161e86d..e4f75d3 100644 --- a/src/image.rs +++ b/src/image.rs @@ -78,6 +78,13 @@ pub struct CustomImageConfig { } /// Parses SHA512SUMS format and returns the hash for an exact filename match. +/// +/// # Arguments +/// * `checksums_text` - The content of a SHA512SUMS file +/// * `filename` - The exact filename to search for (e.g., "debian-13-generic-amd64.qcow2") +/// +/// # Returns +/// The SHA512 hash if found, or None if no exact match exists pub fn find_sha512_for_file(checksums_text: &str, filename: &str) -> Option { checksums_text.lines().find_map(|line| { let mut parts = line.split_whitespace();