From 6aa0bf4c82ef4512a88e1c720d62a7b12ce12746 Mon Sep 17 00:00:00 2001 From: Ariel Ebersberger Date: Thu, 3 Sep 2026 21:44:19 +0200 Subject: [PATCH 1/2] Deduplicate progress plumbing and unify writer structure across platforms --- crates/hai-core/src/disk_writer.rs | 29 ++ crates/hai-core/src/disk_writer/linux.rs | 492 ++++++++++----------- crates/hai-core/src/disk_writer/macos.rs | 229 ++++------ crates/hai-core/src/disk_writer/windows.rs | 269 +++++------ crates/hai-core/src/types.rs | 18 + 5 files changed, 474 insertions(+), 563 deletions(-) diff --git a/crates/hai-core/src/disk_writer.rs b/crates/hai-core/src/disk_writer.rs index 8dba8cb..b52aa6f 100644 --- a/crates/hai-core/src/disk_writer.rs +++ b/crates/hai-core/src/disk_writer.rs @@ -47,6 +47,35 @@ fn is_drive_disconnected(io_err: &std::io::Error) -> bool { }) } +/// Drive a blocking task while forwarding its progress updates to the +/// callback, then drain updates buffered after the task finished (e.g. the +/// final "Write complete" / "Verification complete") so they aren't lost. +async fn run_with_progress( + handle: tokio::task::JoinHandle>, + progress_rx: std::sync::mpsc::Receiver, + progress_callback: &P, +) -> Result<()> { + loop { + match progress_rx.recv_timeout(std::time::Duration::from_millis(100)) { + Ok(update) => progress_callback.on_progress(update), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { + if handle.is_finished() { + break; + } + } + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break, + } + } + + while let Ok(update) = progress_rx.try_recv() { + progress_callback.on_progress(update); + } + + handle + .await + .map_err(|e| Error::Io(std::io::Error::other(e)))? +} + /// Write an image file to a block device with progress updates pub async fn write_image( image_path: &PathBuf, diff --git a/crates/hai-core/src/disk_writer/linux.rs b/crates/hai-core/src/disk_writer/linux.rs index c386859..9a0af6b 100644 --- a/crates/hai-core/src/disk_writer/linux.rs +++ b/crates/hai-core/src/disk_writer/linux.rs @@ -10,14 +10,6 @@ use zbus::proxy::CacheProperties; use zbus::zvariant::{OwnedFd, OwnedObjectPath, Value}; use zbus::Connection; -/// Progress update sent from blocking task -struct ProgressUpdate { - stage: FlashStage, - bytes_processed: u64, - total_bytes: u64, - message: String, -} - pub async fn write_image( image_path: &PathBuf, device_id: &str, @@ -29,22 +21,21 @@ pub async fn write_image( .map_err(|e| map_udisks_error(e, "connecting to the system bus"))?; let block_path = resolve_block_path(&connection, device_id).await?; - unmount_device(&connection, &block_path).await?; + unmount_disk(&connection, &block_path).await?; let image_size = std::fs::metadata(image_path)?.len(); let device = open_device_rw(&connection, &block_path).await?; - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Writing, - progress: 0, - bytes_processed: 0, - total_bytes: image_size, - message: "Writing image to device...".to_string(), - }); + progress_callback.on_progress(FlashProgress::new( + FlashStage::Writing, + 0, + image_size, + "Writing image to device...", + )); - // Create channel for progress updates from the blocking task. - let (progress_tx, progress_rx) = mpsc::channel::(); + // Send progress updates from the blocking task through a channel. + let (progress_tx, progress_rx) = mpsc::channel::(); let image_path_clone = image_path.clone(); @@ -52,196 +43,33 @@ pub async fn write_image( write_and_verify(&image_path_clone, device, image_size, verify, progress_tx) }); - let forward = |update: ProgressUpdate| { - let progress = if update.total_bytes > 0 { - ((update.bytes_processed as f64 / update.total_bytes as f64) * 100.0) as u8 - } else { - 0 - }; - progress_callback.on_progress(FlashProgress { - stage: update.stage, - progress, - bytes_processed: update.bytes_processed, - total_bytes: update.total_bytes, - message: update.message, - }); - }; - - loop { - match progress_rx.recv_timeout(std::time::Duration::from_millis(100)) { - Ok(update) => forward(update), - Err(mpsc::RecvTimeoutError::Timeout) => { - if write_handle.is_finished() { - break; - } - } - Err(mpsc::RecvTimeoutError::Disconnected) => break, - } - } + run_with_progress(write_handle, progress_rx, progress_callback).await?; - // Drain updates the task buffered after it finished (e.g. the final - // "Write complete" / "Verification complete") so they aren't lost. - while let Ok(update) = progress_rx.try_recv() { - forward(update); - } - - write_handle - .await - .map_err(|e| Error::Io(std::io::Error::other(e)))??; - - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Finalizing, - progress: 0, - bytes_processed: 0, - total_bytes: 0, - message: "Syncing data...".to_string(), - }); + progress_callback.on_progress(FlashProgress::new( + FlashStage::Finalizing, + 0, + 0, + "Syncing data...", + )); let _ = Command::new("sync").output(); - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Complete, - progress: 100, - bytes_processed: image_size, - total_bytes: image_size, - message: "Complete".to_string(), - }); + progress_callback.on_progress(FlashProgress::new( + FlashStage::Complete, + image_size, + image_size, + "Complete", + )); Ok(()) } -fn map_udisks_error(err: zbus::Error, context: &str) -> Error { - // Remedy only; the DiskServiceUnavailable variant supplies the - // "Disk service unavailable:" prefix. - const UNAVAILABLE: &str = "install and enable udisks2 to flash drives."; - - if let zbus::Error::MethodError(name, message, _) = &err { - let name = name.as_str(); - // udisks2 isn't installed / not activatable on the bus. - if name.contains("ServiceUnknown") || name.contains("NameHasNoOwner") { - return Error::DiskServiceUnavailable(UNAVAILABLE.to_string()); - } - // User dismissed the polkit dialog. - if name.contains("NotAuthorizedDismissed") { - return Error::PermissionDenied("Authorization was canceled".to_string()); - } - if name.contains("NotAuthorized") { - return Error::PermissionDenied( - message - .clone() - .unwrap_or_else(|| "Not authorized to access the device".to_string()), - ); - } - // udisks may use the dedicated DeviceBusy name, but often reports a - // busy device as Error.Failed with "Device or resource busy" in the - // message instead, so check both. - if name.contains("DeviceBusy") - || message - .as_deref() - .is_some_and(|m| m.contains("Device or resource busy")) - { - return Error::DeviceBusy(context.to_string()); - } - return Error::Io(std::io::Error::other(format!( - "udisks2 error while {context}: {err}" - ))); - } - - // Not a method error → couldn't reach the bus/service at all. - Error::DiskServiceUnavailable(format!("{UNAVAILABLE} ({err})")) -} - -#[zbus::proxy( - interface = "org.freedesktop.UDisks2.Manager", - default_service = "org.freedesktop.UDisks2", - default_path = "/org/freedesktop/UDisks2/Manager", - gen_blocking = false -)] -trait UDisks2Manager { - /// Resolve a device spec like `{"path": "/dev/sdb"}` to block-object paths. - fn resolve_device( - &self, - devspec: HashMap<&str, Value<'_>>, - options: HashMap<&str, Value<'_>>, - ) -> zbus::Result>; -} - -#[zbus::proxy( - interface = "org.freedesktop.UDisks2.Block", - default_service = "org.freedesktop.UDisks2", - gen_blocking = false -)] -trait UDisks2Block { - /// Open the whole device; `mode` is `"r"`, `"w"`, or `"rw"`. - fn open_device(&self, mode: &str, options: HashMap<&str, Value<'_>>) -> zbus::Result; -} - -#[zbus::proxy( - interface = "org.freedesktop.UDisks2.Filesystem", - default_service = "org.freedesktop.UDisks2", - gen_blocking = false -)] -trait UDisks2Filesystem { - fn unmount(&self, options: HashMap<&str, Value<'_>>) -> zbus::Result<()>; -} - -#[zbus::proxy( - interface = "org.freedesktop.UDisks2.PartitionTable", - default_service = "org.freedesktop.UDisks2", - gen_blocking = false -)] -trait UDisks2PartitionTable { - #[zbus(property)] - fn partitions(&self) -> zbus::Result>; -} - -async fn resolve_block_path(conn: &Connection, device_id: &str) -> Result { - let manager = UDisks2ManagerProxy::new(conn) - .await - .map_err(|e| map_udisks_error(e, "connecting to udisks2"))?; - - let mut devspec = HashMap::new(); - devspec.insert("path", Value::from(device_id)); - - let paths = manager - .resolve_device(devspec, HashMap::new()) - .await - .map_err(|e| map_udisks_error(e, "resolving the device"))?; - - paths - .into_iter() - .next() - .ok_or_else(|| Error::DeviceNotFound(device_id.to_string())) -} - -/// Open the device read-write via udisks2 (raises the polkit prompt). -async fn open_device_rw(conn: &Connection, path: &OwnedObjectPath) -> Result { - let block = UDisks2BlockProxy::builder(conn) - .path(path.to_string()) - .map_err(|e| map_udisks_error(e, "addressing the device"))? - .cache_properties(CacheProperties::No) - .build() - .await - .map_err(|e| map_udisks_error(e, "addressing the device"))?; - - // O_EXCL: to make sure we have exclusive access to the disk and error if not - // O_SYNC: so each write reaches the card before returning to keep the - // progress bar in sync - let mut options = HashMap::new(); - options.insert("flags", Value::from(libc::O_EXCL | libc::O_SYNC)); - let fd = block - .open_device("rw", options) - .await - .map_err(|e| map_udisks_error(e, "opening the device"))?; - Ok(File::from(std::os::fd::OwnedFd::from(fd))) -} - fn write_and_verify( image_path: &PathBuf, mut device: File, total_size: u64, verify: bool, - progress_tx: mpsc::Sender, + progress_tx: mpsc::Sender, ) -> Result<()> { write_to_device(image_path, &mut device, total_size, &progress_tx)?; @@ -272,59 +100,11 @@ fn write_and_verify( Ok(()) } -/// Best-effort: drop the kernel page cache for the whole device so a -/// subsequent read-back hits the physical medium instead of the copy we -/// just wrote. Failures are advisory and ignored. -fn drop_device_cache(device: &File) { - use std::os::fd::AsRawFd; - // SAFETY: `device` owns the fd and keeps it open for this call. A `len` - // of 0 means "to end of file"; the return value is purely advisory. - unsafe { - libc::posix_fadvise(device.as_raw_fd(), 0, 0, libc::POSIX_FADV_DONTNEED); - } -} - -/// Unmount everything on the device before writing. -async fn unmount_device(conn: &Connection, block_path: &OwnedObjectPath) -> Result<()> { - let mut targets = vec![block_path.clone()]; - // Unpartitioned media has no PartitionTable interface; the property - // read fails and only the whole-disk filesystem is unmounted. - if let Ok(builder) = UDisks2PartitionTableProxy::builder(conn).path(block_path.clone()) { - if let Ok(table) = builder.cache_properties(CacheProperties::No).build().await { - if let Ok(partitions) = table.partitions().await { - targets.extend(partitions); - } - } - } - - for path in targets { - let builder = match UDisks2FilesystemProxy::builder(conn).path(path) { - Ok(builder) => builder, - Err(_) => continue, - }; - let proxy = match builder.cache_properties(CacheProperties::No).build().await { - Ok(proxy) => proxy, - Err(_) => continue, - }; - let mut options = HashMap::new(); - options.insert("force", Value::from(true)); - if let Err(e) = proxy.unmount(options).await { - if let err @ (Error::PermissionDenied(_) | Error::DiskServiceUnavailable(_)) = - map_udisks_error(e, "unmounting a volume") - { - return Err(err); - } - } - } - - Ok(()) -} - fn write_to_device( image_path: &PathBuf, dest: &mut File, total_size: u64, - progress_tx: &mpsc::Sender, + progress_tx: &mpsc::Sender, ) -> Result<()> { let mut source = File::open(image_path)?; @@ -351,12 +131,12 @@ fn write_to_device( // Update progress periodically if bytes_written - last_progress_bytes >= PROGRESS_UPDATE_INTERVAL { last_progress_bytes = bytes_written; - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Writing, - bytes_processed: bytes_written, - total_bytes: total_size, - message: "Writing image to device...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Writing, + bytes_written, + total_size, + "Writing image to device...", + )); } } @@ -369,12 +149,12 @@ fn write_to_device( })?; // Send final progress - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Writing, - bytes_processed: bytes_written, - total_bytes: total_size, - message: "Write complete".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Writing, + bytes_written, + total_size, + "Write complete", + )); Ok(()) } @@ -383,7 +163,7 @@ fn verify_write( image_path: &PathBuf, dest: &mut File, total_size: u64, - progress_tx: &mpsc::Sender, + progress_tx: &mpsc::Sender, ) -> Result<()> { let mut source = File::open(image_path)?; @@ -418,26 +198,200 @@ fn verify_write( // Update progress periodically if bytes_verified - last_progress_bytes >= PROGRESS_UPDATE_INTERVAL { last_progress_bytes = bytes_verified; - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Verifying, - bytes_processed: bytes_verified, - total_bytes: total_size, - message: "Verifying written data...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Verifying, + bytes_verified, + total_size, + "Verifying written data...", + )); } } // Send final progress - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Verifying, - bytes_processed: bytes_verified, - total_bytes: total_size, - message: "Verification complete".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Verifying, + bytes_verified, + total_size, + "Verification complete", + )); + + Ok(()) +} + +/// Unmount everything on the disk before writing. +async fn unmount_disk(conn: &Connection, block_path: &OwnedObjectPath) -> Result<()> { + let mut targets = vec![block_path.clone()]; + // Unpartitioned media has no PartitionTable interface; the property + // read fails and only the whole-disk filesystem is unmounted. + if let Ok(builder) = UDisks2PartitionTableProxy::builder(conn).path(block_path.clone()) { + if let Ok(table) = builder.cache_properties(CacheProperties::No).build().await { + if let Ok(partitions) = table.partitions().await { + targets.extend(partitions); + } + } + } + + for path in targets { + let builder = match UDisks2FilesystemProxy::builder(conn).path(path) { + Ok(builder) => builder, + Err(_) => continue, + }; + let proxy = match builder.cache_properties(CacheProperties::No).build().await { + Ok(proxy) => proxy, + Err(_) => continue, + }; + let mut options = HashMap::new(); + options.insert("force", Value::from(true)); + if let Err(e) = proxy.unmount(options).await { + if let err @ (Error::PermissionDenied(_) | Error::DiskServiceUnavailable(_)) = + map_udisks_error(e, "unmounting a volume") + { + return Err(err); + } + } + } Ok(()) } +async fn resolve_block_path(conn: &Connection, device_id: &str) -> Result { + let manager = UDisks2ManagerProxy::new(conn) + .await + .map_err(|e| map_udisks_error(e, "connecting to udisks2"))?; + + let mut devspec = HashMap::new(); + devspec.insert("path", Value::from(device_id)); + + let paths = manager + .resolve_device(devspec, HashMap::new()) + .await + .map_err(|e| map_udisks_error(e, "resolving the device"))?; + + paths + .into_iter() + .next() + .ok_or_else(|| Error::DeviceNotFound(device_id.to_string())) +} + +/// Open the device read-write via udisks2 (raises the polkit prompt). +async fn open_device_rw(conn: &Connection, path: &OwnedObjectPath) -> Result { + let block = UDisks2BlockProxy::builder(conn) + .path(path.to_string()) + .map_err(|e| map_udisks_error(e, "addressing the device"))? + .cache_properties(CacheProperties::No) + .build() + .await + .map_err(|e| map_udisks_error(e, "addressing the device"))?; + + // O_EXCL: to make sure we have exclusive access to the disk and error if not + // O_SYNC: so each write reaches the card before returning to keep the + // progress bar in sync + let mut options = HashMap::new(); + options.insert("flags", Value::from(libc::O_EXCL | libc::O_SYNC)); + let fd = block + .open_device("rw", options) + .await + .map_err(|e| map_udisks_error(e, "opening the device"))?; + Ok(File::from(std::os::fd::OwnedFd::from(fd))) +} + +/// Best-effort: drop the kernel page cache for the whole device so a +/// subsequent read-back hits the physical medium instead of the copy we +/// just wrote. Failures are advisory and ignored. +fn drop_device_cache(device: &File) { + use std::os::fd::AsRawFd; + // SAFETY: `device` owns the fd and keeps it open for this call. A `len` + // of 0 means "to end of file"; the return value is purely advisory. + unsafe { + libc::posix_fadvise(device.as_raw_fd(), 0, 0, libc::POSIX_FADV_DONTNEED); + } +} + +fn map_udisks_error(err: zbus::Error, context: &str) -> Error { + // Remedy only; the DiskServiceUnavailable variant supplies the + // "Disk service unavailable:" prefix. + const UNAVAILABLE: &str = "install and enable udisks2 to flash drives."; + + if let zbus::Error::MethodError(name, message, _) = &err { + let name = name.as_str(); + // udisks2 isn't installed / not activatable on the bus. + if name.contains("ServiceUnknown") || name.contains("NameHasNoOwner") { + return Error::DiskServiceUnavailable(UNAVAILABLE.to_string()); + } + // User dismissed the polkit dialog. + if name.contains("NotAuthorizedDismissed") { + return Error::PermissionDenied("Authorization was canceled".to_string()); + } + if name.contains("NotAuthorized") { + return Error::PermissionDenied( + message + .clone() + .unwrap_or_else(|| "Not authorized to access the device".to_string()), + ); + } + // udisks may use the dedicated DeviceBusy name, but often reports a + // busy device as Error.Failed with "Device or resource busy" in the + // message instead, so check both. + if name.contains("DeviceBusy") + || message + .as_deref() + .is_some_and(|m| m.contains("Device or resource busy")) + { + return Error::DeviceBusy(context.to_string()); + } + return Error::Io(std::io::Error::other(format!( + "udisks2 error while {context}: {err}" + ))); + } + + // Not a method error → couldn't reach the bus/service at all. + Error::DiskServiceUnavailable(format!("{UNAVAILABLE} ({err})")) +} + +#[zbus::proxy( + interface = "org.freedesktop.UDisks2.Manager", + default_service = "org.freedesktop.UDisks2", + default_path = "/org/freedesktop/UDisks2/Manager", + gen_blocking = false +)] +trait UDisks2Manager { + /// Resolve a device spec like `{"path": "/dev/sdb"}` to block-object paths. + fn resolve_device( + &self, + devspec: HashMap<&str, Value<'_>>, + options: HashMap<&str, Value<'_>>, + ) -> zbus::Result>; +} + +#[zbus::proxy( + interface = "org.freedesktop.UDisks2.Block", + default_service = "org.freedesktop.UDisks2", + gen_blocking = false +)] +trait UDisks2Block { + /// Open the whole device; `mode` is `"r"`, `"w"`, or `"rw"`. + fn open_device(&self, mode: &str, options: HashMap<&str, Value<'_>>) -> zbus::Result; +} + +#[zbus::proxy( + interface = "org.freedesktop.UDisks2.Filesystem", + default_service = "org.freedesktop.UDisks2", + gen_blocking = false +)] +trait UDisks2Filesystem { + fn unmount(&self, options: HashMap<&str, Value<'_>>) -> zbus::Result<()>; +} + +#[zbus::proxy( + interface = "org.freedesktop.UDisks2.PartitionTable", + default_service = "org.freedesktop.UDisks2", + gen_blocking = false +)] +trait UDisks2PartitionTable { + #[zbus(property)] + fn partitions(&self) -> zbus::Result>; +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/hai-core/src/disk_writer/macos.rs b/crates/hai-core/src/disk_writer/macos.rs index e13bad7..8b594ec 100644 --- a/crates/hai-core/src/disk_writer/macos.rs +++ b/crates/hai-core/src/disk_writer/macos.rs @@ -7,14 +7,6 @@ use std::path::Path; use std::process::Command; use std::sync::mpsc; -/// Progress update sent from blocking task -struct ProgressUpdate { - stage: FlashStage, - bytes_processed: u64, - total_bytes: u64, - message: String, -} - pub async fn write_image( image_path: &PathBuf, device_id: &str, @@ -34,16 +26,15 @@ pub async fn write_image( let image_size = std::fs::metadata(image_path)?.len(); // Send initial progress - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Writing, - progress: 0, - bytes_processed: 0, - total_bytes: image_size, - message: "Requesting administrator access...".to_string(), - }); + progress_callback.on_progress(FlashProgress::new( + FlashStage::Writing, + 0, + image_size, + "Requesting administrator access...", + )); - // Create channel for progress updates from blocking task - let (progress_tx, progress_rx) = mpsc::channel::(); + // Send progress updates from the blocking task through a channel. + let (progress_tx, progress_rx) = mpsc::channel::(); // Perform write and optional verify in a blocking task let image_path_clone = image_path.clone(); @@ -51,7 +42,7 @@ pub async fn write_image( let disk_id_clone = disk_id.to_string(); let write_handle = tokio::task::spawn_blocking(move || { - write_and_verify_blocking( + write_and_verify( &image_path_clone, &raw_device_clone, &disk_id_clone, @@ -61,68 +52,31 @@ pub async fn write_image( ) }); - // Forward progress updates while waiting for write to complete - loop { - // Check for progress updates (non-blocking with timeout) - match progress_rx.recv_timeout(std::time::Duration::from_millis(100)) { - Ok(update) => { - let progress = if update.total_bytes > 0 { - ((update.bytes_processed as f64 / update.total_bytes as f64) * 100.0) as u8 - } else { - 0 - }; - progress_callback.on_progress(FlashProgress { - stage: update.stage, - progress, - bytes_processed: update.bytes_processed, - total_bytes: update.total_bytes, - message: update.message, - }); - } - Err(mpsc::RecvTimeoutError::Timeout) => { - // Check if the blocking task is done - if write_handle.is_finished() { - break; - } - } - Err(mpsc::RecvTimeoutError::Disconnected) => { - // Sender dropped, task is done - break; - } - } - } - - // Wait for the result - let result = write_handle - .await - .map_err(|e| Error::Io(std::io::Error::other(e)))?; + run_with_progress(write_handle, progress_rx, progress_callback).await?; - result?; - - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Complete, - progress: 100, - bytes_processed: image_size, - total_bytes: image_size, - message: "Complete".to_string(), - }); + progress_callback.on_progress(FlashProgress::new( + FlashStage::Complete, + image_size, + image_size, + "Complete", + )); Ok(()) } -fn write_and_verify_blocking( +fn write_and_verify( image_path: &PathBuf, device_path: &str, disk_id: &str, total_size: u64, verify: bool, - progress_tx: mpsc::Sender, + progress_tx: mpsc::Sender, ) -> Result<()> { // Request authorization let auth = request_authorization()?; // Write the image and compute checksum if verification is requested - let source_checksum = write_with_auth( + let source_checksum = write_to_device( &auth, image_path, device_path, @@ -135,7 +89,7 @@ fn write_and_verify_blocking( if verify { let checksum = source_checksum.expect("Checksum should have been computed when verify=true"); - verify_with_auth(&auth, &checksum, device_path, total_size, &progress_tx)?; + verify_write(&auth, &checksum, device_path, total_size, &progress_tx)?; } // Finalize - eject @@ -144,35 +98,13 @@ fn write_and_verify_blocking( Ok(()) } -fn request_authorization() -> Result { - let rights = AuthorizationItemSetBuilder::new() - .add_right("system.privilege.admin") - .map_err(|e| Error::PermissionDenied(format!("Failed to create rights: {}", e)))? - .build(); - - Authorization::new( - Some(rights), - None, - Flags::INTERACTION_ALLOWED | Flags::EXTEND_RIGHTS | Flags::PREAUTHORIZE, - ) - .map_err(|e| { - if e.code() == -60006 { - Error::PermissionDenied("Administrator access was denied by user".to_string()) - } else if e.code() == -60005 { - Error::PermissionDenied("Authorization was canceled".to_string()) - } else { - Error::PermissionDenied(format!("Authorization failed: {}", e)) - } - }) -} - -fn write_with_auth( +fn write_to_device( auth: &Authorization, image_path: &PathBuf, device_path: &str, total_size: u64, compute_checksum: bool, - progress_tx: &mpsc::Sender, + progress_tx: &mpsc::Sender, ) -> Result> { use sha2::{Digest, Sha256}; use std::io::Write; @@ -184,12 +116,12 @@ fn write_with_auth( let bs_arg = "bs=64m".to_string(); // Send progress update before requesting privilege - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Writing, - bytes_processed: 0, - total_bytes: total_size, - message: "Starting write...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Writing, + 0, + total_size, + "Starting write...", + )); let mut pipe = auth .execute_with_privileges_piped(dd_path, [&of_arg, &bs_arg], Flags::empty()) @@ -226,23 +158,23 @@ fn write_with_auth( // Send progress update every PROGRESS_UPDATE_INTERVAL bytes if bytes_written - last_progress_update >= PROGRESS_UPDATE_INTERVAL { - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Writing, - bytes_processed: bytes_written, - total_bytes: total_size, - message: "Writing image to drive...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Writing, + bytes_written, + total_size, + "Writing image to drive...", + )); last_progress_update = bytes_written; } } // Send final write progress - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Writing, - bytes_processed: bytes_written, - total_bytes: total_size, - message: "Syncing data to drive...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Writing, + bytes_written, + total_size, + "Syncing data to drive...", + )); drop(pipe); let _ = Command::new("sync").output(); @@ -251,22 +183,22 @@ fn write_with_auth( Ok(checksum) } -fn verify_with_auth( +fn verify_write( auth: &Authorization, source_checksum: &str, device_path: &str, total_size: u64, - progress_tx: &mpsc::Sender, + progress_tx: &mpsc::Sender, ) -> Result<()> { use sha2::{Digest, Sha256}; // Send initial verify progress - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Verifying, - bytes_processed: 0, - total_bytes: total_size, - message: "Starting verification...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Verifying, + 0, + total_size, + "Starting verification...", + )); let block_count = total_size.div_ceil(FAST_DRIVE_BUFFER_SIZE as u64); let dd_path = Path::new("/bin/dd"); @@ -292,12 +224,12 @@ fn verify_with_auth( // Send progress update every PROGRESS_UPDATE_INTERVAL bytes if bytes_read_total - last_progress_update >= PROGRESS_UPDATE_INTERVAL { - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Verifying, - bytes_processed: bytes_read_total, - total_bytes: total_size, - message: "Verifying written data...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Verifying, + bytes_read_total, + total_size, + "Verifying written data...", + )); last_progress_update = bytes_read_total; } } @@ -321,12 +253,12 @@ fn verify_with_auth( } // Send final verify progress - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Verifying, - bytes_processed: total_size, - total_bytes: total_size, - message: "Verification complete".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Verifying, + total_size, + total_size, + "Verification complete", + )); Ok(()) } @@ -365,6 +297,28 @@ fn unmount_disk(disk_id: &str) -> Result<()> { Ok(()) } +fn request_authorization() -> Result { + let rights = AuthorizationItemSetBuilder::new() + .add_right("system.privilege.admin") + .map_err(|e| Error::PermissionDenied(format!("Failed to create rights: {}", e)))? + .build(); + + Authorization::new( + Some(rights), + None, + Flags::INTERACTION_ALLOWED | Flags::EXTEND_RIGHTS | Flags::PREAUTHORIZE, + ) + .map_err(|e| { + if e.code() == -60006 { + Error::PermissionDenied("Administrator access was denied by user".to_string()) + } else if e.code() == -60005 { + Error::PermissionDenied("Authorization was canceled".to_string()) + } else { + Error::PermissionDenied(format!("Authorization failed: {}", e)) + } + }) +} + fn eject_disk(disk_id: &str) -> Result<()> { let output = Command::new("diskutil").args(["eject", disk_id]).output()?; @@ -379,13 +333,20 @@ fn eject_disk(disk_id: &str) -> Result<()> { #[cfg(test)] mod tests { use super::*; + use serial_test::serial; - #[test] - fn test_unmount_disk_nonexistent() { - // Test unmounting a disk that doesn't exist - let result = unmount_disk("disk999"); - // This should either succeed (if disk not mounted) or fail with I/O error - assert!(result.is_ok() || result.is_err()); + #[tokio::test] + #[serial] + async fn test_write_image_with_invalid_device() { + let temp_file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(temp_file.path(), b"test data").unwrap(); + let image_path = temp_file.path().to_path_buf(); + + // Fails at unmount_disk: the device does not exist. + let device_id = "/dev/nonexistent_disk999"; + + let result = write_image(&image_path, device_id, false, &crate::NoOpProgress).await; + assert!(result.is_err()); } #[test] diff --git a/crates/hai-core/src/disk_writer/windows.rs b/crates/hai-core/src/disk_writer/windows.rs index 6b21dbd..04e723a 100644 --- a/crates/hai-core/src/disk_writer/windows.rs +++ b/crates/hai-core/src/disk_writer/windows.rs @@ -6,14 +6,6 @@ use std::io::{Read, Write}; use std::process::Command; use std::sync::mpsc; -/// Progress update sent from blocking task -struct ProgressUpdate { - stage: FlashStage, - bytes_processed: u64, - total_bytes: u64, - message: String, -} - pub async fn write_image( image_path: &PathBuf, device_id: &str, @@ -28,141 +20,71 @@ pub async fn write_image( let image_size = std::fs::metadata(image_path)?.len(); - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Writing, - progress: 0, - bytes_processed: 0, - total_bytes: image_size, - message: "Writing image to device...".to_string(), - }); + progress_callback.on_progress(FlashProgress::new( + FlashStage::Writing, + 0, + image_size, + "Writing image to device...", + )); - // Create channel for progress updates from blocking task - let (progress_tx, progress_rx) = mpsc::channel::(); + // Send progress updates from the blocking task through a channel. + let (progress_tx, progress_rx) = mpsc::channel::(); let image_path_clone = image_path.clone(); let device_id_clone = device_id.to_string(); let write_handle = tokio::task::spawn_blocking(move || { - write_to_device(&image_path_clone, &device_id_clone, image_size, progress_tx) + write_and_verify( + &image_path_clone, + &device_id_clone, + image_size, + verify, + progress_tx, + ) }); - // Forward progress updates while waiting for write to complete - loop { - match progress_rx.recv_timeout(std::time::Duration::from_millis(100)) { - Ok(update) => { - let progress = if update.total_bytes > 0 { - ((update.bytes_processed as f64 / update.total_bytes as f64) * 100.0) as u8 - } else { - 0 - }; - progress_callback.on_progress(FlashProgress { - stage: update.stage, - progress, - bytes_processed: update.bytes_processed, - total_bytes: update.total_bytes, - message: update.message, - }); - } - Err(mpsc::RecvTimeoutError::Timeout) => { - if write_handle.is_finished() { - break; - } - } - Err(mpsc::RecvTimeoutError::Disconnected) => { - break; - } - } - } + run_with_progress(write_handle, progress_rx, progress_callback).await?; - write_handle - .await - .map_err(|e| Error::Io(std::io::Error::new(std::io::ErrorKind::Other, e)))??; + progress_callback.on_progress(FlashProgress::new( + FlashStage::Finalizing, + 0, + 0, + "Finalizing...", + )); - if verify { - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Verifying, - progress: 0, - bytes_processed: 0, - total_bytes: image_size, - message: "Verifying written data...".to_string(), - }); - - let (verify_tx, verify_rx) = mpsc::channel::(); - - let image_path_clone = image_path.clone(); - let device_id_clone = device_id.to_string(); - - let verify_handle = tokio::task::spawn_blocking(move || { - verify_write(&image_path_clone, &device_id_clone, image_size, verify_tx) - }); - - // Forward verify progress updates - loop { - match verify_rx.recv_timeout(std::time::Duration::from_millis(100)) { - Ok(update) => { - let progress = if update.total_bytes > 0 { - ((update.bytes_processed as f64 / update.total_bytes as f64) * 100.0) as u8 - } else { - 0 - }; - progress_callback.on_progress(FlashProgress { - stage: update.stage, - progress, - bytes_processed: update.bytes_processed, - total_bytes: update.total_bytes, - message: update.message, - }); - } - Err(mpsc::RecvTimeoutError::Timeout) => { - if verify_handle.is_finished() { - break; - } - } - Err(mpsc::RecvTimeoutError::Disconnected) => { - break; - } - } - } - - verify_handle - .await - .map_err(|e| Error::Io(std::io::Error::new(std::io::ErrorKind::Other, e)))??; - } - - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Finalizing, - progress: 0, - bytes_processed: 0, - total_bytes: 0, - message: "Finalizing...".to_string(), - }); - - progress_callback.on_progress(FlashProgress { - stage: FlashStage::Complete, - progress: 100, - bytes_processed: image_size, - total_bytes: image_size, - message: "Complete".to_string(), - }); + progress_callback.on_progress(FlashProgress::new( + FlashStage::Complete, + image_size, + image_size, + "Complete", + )); Ok(()) } -fn clean_disk(disk_number: &str) -> Result<()> { - let ps_script = format!( - "Clear-Disk -Number {} -RemoveData -RemoveOEM -Confirm:$false", - disk_number - ); - - let output = Command::new("powershell") - .args(["-NoProfile", "-Command", &ps_script]) - .output()?; +fn write_and_verify( + image_path: &PathBuf, + device_path: &str, + total_size: u64, + verify: bool, + progress_tx: mpsc::Sender, +) -> Result<()> { + write_to_device(image_path, device_path, total_size, &progress_tx)?; - if !output.status.success() { - let stderr = String::from_utf8_lossy(&output.stderr); - if !stderr.contains("not found") && !stderr.contains("no media") { - return Err(Error::DeviceBusy(stderr.to_string())); - } + if verify { + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Verifying, + 0, + total_size, + "Verifying written data...", + )); + + // Tag verify-phase failures as VerificationFailed so the caller can + // label them "Verification failed" rather than "Write failed". + verify_write(image_path, device_path, total_size, &progress_tx).map_err(|e| match e { + Error::VerificationFailed(_) | Error::DriveDisconnected => e, + other => Error::VerificationFailed(other.to_string()), + })?; } Ok(()) @@ -172,7 +94,7 @@ fn write_to_device( image_path: &PathBuf, device_path: &str, total_size: u64, - progress_tx: mpsc::Sender, + progress_tx: &mpsc::Sender, ) -> Result<()> { let mut source = File::open(image_path)?; @@ -214,12 +136,12 @@ fn write_to_device( // Update progress periodically if bytes_written - last_progress_bytes >= PROGRESS_UPDATE_INTERVAL { last_progress_bytes = bytes_written; - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Writing, - bytes_processed: bytes_written, - total_bytes: total_size, - message: "Writing image to device...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Writing, + bytes_written, + total_size, + "Writing image to device...", + )); } } @@ -232,12 +154,12 @@ fn write_to_device( })?; // Send final progress - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Writing, - bytes_processed: bytes_written, - total_bytes: total_size, - message: "Write complete".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Writing, + bytes_written, + total_size, + "Write complete", + )); Ok(()) } @@ -246,7 +168,7 @@ fn verify_write( image_path: &PathBuf, device_path: &str, total_size: u64, - progress_tx: mpsc::Sender, + progress_tx: &mpsc::Sender, ) -> Result<()> { let mut source = File::open(image_path)?; let mut dest = File::open(device_path).map_err(|e| { @@ -288,22 +210,42 @@ fn verify_write( // Update progress periodically if bytes_verified - last_progress_bytes >= PROGRESS_UPDATE_INTERVAL { last_progress_bytes = bytes_verified; - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Verifying, - bytes_processed: bytes_verified, - total_bytes: total_size, - message: "Verifying written data...".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Verifying, + bytes_verified, + total_size, + "Verifying written data...", + )); } } // Send final progress - let _ = progress_tx.send(ProgressUpdate { - stage: FlashStage::Verifying, - bytes_processed: bytes_verified, - total_bytes: total_size, - message: "Verification complete".to_string(), - }); + let _ = progress_tx.send(FlashProgress::new( + FlashStage::Verifying, + bytes_verified, + total_size, + "Verification complete", + )); + + Ok(()) +} + +fn clean_disk(disk_number: &str) -> Result<()> { + let ps_script = format!( + "Clear-Disk -Number {} -RemoveData -RemoveOEM -Confirm:$false", + disk_number + ); + + let output = Command::new("powershell") + .args(["-NoProfile", "-Command", &ps_script]) + .output()?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + if !stderr.contains("not found") && !stderr.contains("no media") { + return Err(Error::DeviceBusy(stderr.to_string())); + } + } Ok(()) } @@ -311,12 +253,19 @@ fn verify_write( #[cfg(test)] mod tests { use super::*; + use serial_test::serial; + + #[tokio::test] + #[serial] + async fn test_write_image_invalid_device() { + let temp_file = tempfile::NamedTempFile::new().unwrap(); + std::fs::write(temp_file.path(), b"test data").unwrap(); + let image_path = temp_file.path().to_path_buf(); + + // Fails at clean_disk: the device does not exist. + let device_id = "\\\\.\\PhysicalDrive999"; - #[test] - fn test_clean_disk_nonexistent() { - // Test cleaning a disk that doesn't exist - let result = clean_disk("999"); - // Should either succeed or fail, but not panic - assert!(result.is_ok() || result.is_err()); + let result = write_image(&image_path, device_id, false, &crate::NoOpProgress).await; + assert!(result.is_err()); } } diff --git a/crates/hai-core/src/types.rs b/crates/hai-core/src/types.rs index f945db6..fce597a 100644 --- a/crates/hai-core/src/types.rs +++ b/crates/hai-core/src/types.rs @@ -51,6 +51,24 @@ pub struct FlashProgress { pub message: String, } +impl FlashProgress { + /// Build a progress event, deriving the percentage from the byte counts. + pub fn new(stage: FlashStage, bytes_processed: u64, total_bytes: u64, message: &str) -> Self { + let progress = if total_bytes > 0 { + ((bytes_processed as f64 / total_bytes as f64) * 100.0) as u8 + } else { + 0 + }; + Self { + stage, + progress, + bytes_processed, + total_bytes, + message: message.to_string(), + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(rename_all = "snake_case")] pub enum FlashStage { From 9982bf34bebef1d8746934ce0d47abdbfdfd12c5 Mon Sep 17 00:00:00 2001 From: Ariel Ebersberger Date: Thu, 3 Sep 2026 21:44:28 +0200 Subject: [PATCH 2/2] Merge devices and disk_writer into disk/{os}/{device,writer} and deduplicate tests --- crates/hai-core/src/devices.rs | 75 --- crates/hai-core/src/disk.rs | 181 ++++++++ .../linux.rs => disk/linux/device.rs} | 34 +- crates/hai-core/src/disk/linux/mod.rs | 5 + .../linux.rs => disk/linux/writer.rs} | 2 +- .../macos.rs => disk/macos/device.rs} | 2 +- crates/hai-core/src/disk/macos/mod.rs | 5 + .../macos.rs => disk/macos/writer.rs} | 2 +- .../windows.rs => disk/windows/device.rs} | 2 +- crates/hai-core/src/disk/windows/mod.rs | 5 + .../windows.rs => disk/windows/writer.rs} | 2 +- crates/hai-core/src/disk_writer.rs | 432 ------------------ crates/hai-core/src/lib.rs | 3 +- crates/hai-desktop/src/commands.rs | 12 +- 14 files changed, 209 insertions(+), 553 deletions(-) delete mode 100644 crates/hai-core/src/devices.rs create mode 100644 crates/hai-core/src/disk.rs rename crates/hai-core/src/{devices/linux.rs => disk/linux/device.rs} (90%) create mode 100644 crates/hai-core/src/disk/linux/mod.rs rename crates/hai-core/src/{disk_writer/linux.rs => disk/linux/writer.rs} (99%) rename crates/hai-core/src/{devices/macos.rs => disk/macos/device.rs} (99%) create mode 100644 crates/hai-core/src/disk/macos/mod.rs rename crates/hai-core/src/{disk_writer/macos.rs => disk/macos/writer.rs} (99%) rename crates/hai-core/src/{devices/windows.rs => disk/windows/device.rs} (99%) create mode 100644 crates/hai-core/src/disk/windows/mod.rs rename crates/hai-core/src/{disk_writer/windows.rs => disk/windows/writer.rs} (99%) delete mode 100644 crates/hai-core/src/disk_writer.rs diff --git a/crates/hai-core/src/devices.rs b/crates/hai-core/src/devices.rs deleted file mode 100644 index 965cb9c..0000000 --- a/crates/hai-core/src/devices.rs +++ /dev/null @@ -1,75 +0,0 @@ -//! Block device enumeration for different platforms -//! -//! This module provides platform-specific implementations for listing -//! block devices (SD cards, USB drives, etc.) that can be used as -//! installation targets. - -use crate::error::Result; -use crate::types::{BlockDevice, DeviceType}; - -#[cfg(target_os = "linux")] -#[path = "devices/linux.rs"] -mod imp; -#[cfg(target_os = "macos")] -#[path = "devices/macos.rs"] -mod imp; -#[cfg(target_os = "windows")] -#[path = "devices/windows.rs"] -mod imp; - -#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))] -compile_error!("hai-core supports only Linux, macOS and Windows"); - -/// List all available block devices on the system -/// -/// Returns removable devices suitable for flashing (SD cards, USB drives, etc.) -/// Filters out internal and system drives for safety. -pub async fn list_devices() -> Result> { - imp::list_devices().await -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_device_type_values() { - // Ensure all device types can be created - let types = vec![ - DeviceType::SdCard, - DeviceType::UsbDrive, - DeviceType::Ssd, - DeviceType::Hdd, - DeviceType::NvMe, - DeviceType::Unknown, - ]; - - for device_type in types { - let json = serde_json::to_string(&device_type).unwrap(); - assert!(!json.is_empty()); - } - } - - #[test] - fn test_block_device_creation() { - let device = BlockDevice { - id: "/dev/sdb".to_string(), - name: "Test Device".to_string(), - size: 32_000_000_000, - device_type: DeviceType::UsbDrive, - removable: true, - model: Some("Test Model".to_string()), - vendor: Some("Test Vendor".to_string()), - }; - - assert_eq!(device.id, "/dev/sdb"); - assert_eq!(device.size, 32_000_000_000); - assert!(device.removable); - } - - #[tokio::test] - async fn test_list_devices_succeeds() { - // Smoke test: the real platform enumeration runs and succeeds. - assert!(list_devices().await.is_ok()); - } -} diff --git a/crates/hai-core/src/disk.rs b/crates/hai-core/src/disk.rs new file mode 100644 index 0000000..bec3f99 --- /dev/null +++ b/crates/hai-core/src/disk.rs @@ -0,0 +1,181 @@ +//! Block device enumeration and raw disk writing. +//! +//! This module provides platform-specific implementations for listing +//! block devices (SD cards, USB drives, etc.) and writing raw disk +//! images to them. + +use crate::error::{Error, Result}; +use crate::types::{BlockDevice, DeviceType, FlashProgress, FlashStage}; +use crate::ProgressCallback; +use std::path::PathBuf; + +#[cfg(target_os = "linux")] +#[path = "disk/linux/mod.rs"] +mod imp; +#[cfg(target_os = "macos")] +#[path = "disk/macos/mod.rs"] +mod imp; +#[cfg(target_os = "windows")] +#[path = "disk/windows/mod.rs"] +mod imp; + +#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))] +compile_error!("hai-core supports only Linux, macOS and Windows"); + +/// Buffer size for disk writes (4 MB for SD cards) +#[allow(dead_code)] +const WRITE_BUFFER_SIZE: usize = 4 * 1024 * 1024; + +/// Buffer size for fast drives like NVMe/SSDs (64 MB) +#[allow(dead_code)] +const FAST_DRIVE_BUFFER_SIZE: usize = 64 * 1024 * 1024; + +/// How often to send progress updates (every N bytes) +#[allow(dead_code)] +const PROGRESS_UPDATE_INTERVAL: u64 = 10 * 1024 * 1024; // 10 MB + +/// Check if an I/O error indicates the drive was disconnected +fn is_drive_disconnected(io_err: &std::io::Error) -> bool { + matches!( + io_err.kind(), + std::io::ErrorKind::NotFound + | std::io::ErrorKind::BrokenPipe + | std::io::ErrorKind::UnexpectedEof + ) || io_err.raw_os_error().is_some_and(|code| { + // macOS: ENXIO (6) = "Device not configured" + // Linux: ENODEV (19) = "No such device", ENXIO (6) + matches!(code, 6 | 19) + }) +} + +/// Drive a blocking task while forwarding its progress updates to the +/// callback, then drain updates buffered after the task finished (e.g. the +/// final "Write complete" / "Verification complete") so they aren't lost. +async fn run_with_progress( + handle: tokio::task::JoinHandle>, + progress_rx: std::sync::mpsc::Receiver, + progress_callback: &P, +) -> Result<()> { + loop { + match progress_rx.recv_timeout(std::time::Duration::from_millis(100)) { + Ok(update) => progress_callback.on_progress(update), + Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { + if handle.is_finished() { + break; + } + } + Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break, + } + } + + while let Ok(update) = progress_rx.try_recv() { + progress_callback.on_progress(update); + } + + handle + .await + .map_err(|e| Error::Io(std::io::Error::other(e)))? +} + +/// List all available block devices on the system +/// +/// Returns removable devices suitable for flashing (SD cards, USB drives, etc.) +/// Filters out internal and system drives for safety. +pub async fn list_devices() -> Result> { + imp::list_devices().await +} + +/// Write an image file to a block device with progress updates +pub async fn write_image( + image_path: &PathBuf, + device_id: &str, + verify: bool, + progress_callback: &P, +) -> Result<()> { + std::fs::metadata(image_path)?; + + imp::write_image(image_path, device_id, verify, progress_callback).await +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_is_drive_disconnected_all_matching_kinds() { + use std::io::ErrorKind; + + for kind in [ + ErrorKind::NotFound, + ErrorKind::BrokenPipe, + ErrorKind::UnexpectedEof, + ] { + assert!( + is_drive_disconnected(&std::io::Error::new(kind, "test")), + "{kind:?} should be detected as disconnected" + ); + } + + // Raw OS error codes: ENXIO (6) and ENODEV (19) + for code in [6, 19] { + assert!( + is_drive_disconnected(&std::io::Error::from_raw_os_error(code)), + "os error {code} should be detected as disconnected" + ); + } + } + + #[test] + fn test_is_drive_disconnected_non_matching_kinds() { + use std::io::ErrorKind; + + for kind in [ + ErrorKind::PermissionDenied, + ErrorKind::ConnectionRefused, + ErrorKind::ConnectionReset, + ErrorKind::ConnectionAborted, + ErrorKind::AddrInUse, + ErrorKind::AddrNotAvailable, + ErrorKind::InvalidInput, + ErrorKind::InvalidData, + ErrorKind::TimedOut, + ErrorKind::WriteZero, + ErrorKind::Interrupted, + ErrorKind::Other, + ErrorKind::WouldBlock, + ] { + assert!( + !is_drive_disconnected(&std::io::Error::new(kind, "test")), + "{kind:?} should NOT be detected as disconnected" + ); + } + + // Non-matching raw OS error codes: EPERM (1) and EACCES (13) + for code in [1, 13] { + assert!( + !is_drive_disconnected(&std::io::Error::from_raw_os_error(code)), + "os error {code} should NOT be detected as disconnected" + ); + } + + // Error without a raw OS error code + let err = std::io::Error::new(ErrorKind::Other, "generic error"); + assert!(!is_drive_disconnected(&err)); + } + + #[tokio::test] + async fn test_list_devices_succeeds() { + // Smoke test: the real platform enumeration runs and succeeds. + assert!(list_devices().await.is_ok()); + } + + #[tokio::test] + async fn test_write_image_nonexistent_image() { + // The image metadata check runs before any platform code, so a + // missing image surfaces as Io and the device id is never touched. + let image_path = PathBuf::from("/nonexistent/image/file.img"); + let result = + write_image(&image_path, "unused-device-id", false, &crate::NoOpProgress).await; + assert!(matches!(result.unwrap_err(), Error::Io(_))); + } +} diff --git a/crates/hai-core/src/devices/linux.rs b/crates/hai-core/src/disk/linux/device.rs similarity index 90% rename from crates/hai-core/src/devices/linux.rs rename to crates/hai-core/src/disk/linux/device.rs index 10327ef..2fe0d55 100644 --- a/crates/hai-core/src/devices/linux.rs +++ b/crates/hai-core/src/disk/linux/device.rs @@ -1,6 +1,6 @@ //! Linux block device enumeration via `lsblk`. -use super::*; +use super::super::*; use crate::error::Error; use serde::Deserialize; use std::process::Command; @@ -282,38 +282,6 @@ mod tests { assert_eq!(determine_device_type(&dev), DeviceType::SdCard); } - #[test] - fn test_build_device_name_with_vendor_and_model() { - let vendor = Some("SanDisk".to_string()); - let model = Some("Ultra".to_string()); - let result = build_device_name("sdb", &vendor, &model); - assert_eq!(result, "SanDisk Ultra"); - } - - #[test] - fn test_build_device_name_with_vendor_only() { - let vendor = Some("Samsung".to_string()); - let model = None; - let result = build_device_name("sdb", &vendor, &model); - assert_eq!(result, "Samsung"); - } - - #[test] - fn test_build_device_name_with_model_only() { - let vendor = None; - let model = Some("Generic USB Drive".to_string()); - let result = build_device_name("sdb", &vendor, &model); - assert_eq!(result, "Generic USB Drive"); - } - - #[test] - fn test_build_device_name_with_neither() { - let vendor = None; - let model = None; - let result = build_device_name("sdb", &vendor, &model); - assert_eq!(result, "sdb"); - } - #[test] fn test_determine_device_type_sd_in_model_with_space() { let dev = LsblkDevice { diff --git a/crates/hai-core/src/disk/linux/mod.rs b/crates/hai-core/src/disk/linux/mod.rs new file mode 100644 index 0000000..60a2412 --- /dev/null +++ b/crates/hai-core/src/disk/linux/mod.rs @@ -0,0 +1,5 @@ +mod device; +mod writer; + +pub use device::list_devices; +pub use writer::write_image; diff --git a/crates/hai-core/src/disk_writer/linux.rs b/crates/hai-core/src/disk/linux/writer.rs similarity index 99% rename from crates/hai-core/src/disk_writer/linux.rs rename to crates/hai-core/src/disk/linux/writer.rs index 9a0af6b..b66877a 100644 --- a/crates/hai-core/src/disk_writer/linux.rs +++ b/crates/hai-core/src/disk/linux/writer.rs @@ -1,6 +1,6 @@ //! Linux disk writing via udisks2 over D-Bus (polkit handles authorization). -use super::*; +use super::super::*; use std::collections::HashMap; use std::fs::File; use std::io::{Read, Seek, SeekFrom, Write}; diff --git a/crates/hai-core/src/devices/macos.rs b/crates/hai-core/src/disk/macos/device.rs similarity index 99% rename from crates/hai-core/src/devices/macos.rs rename to crates/hai-core/src/disk/macos/device.rs index 26ffd30..6c5dd36 100644 --- a/crates/hai-core/src/devices/macos.rs +++ b/crates/hai-core/src/disk/macos/device.rs @@ -1,6 +1,6 @@ //! macOS block device enumeration via `diskutil`. -use super::*; +use super::super::*; use crate::error::Error; use serde::Deserialize; use std::process::Command; diff --git a/crates/hai-core/src/disk/macos/mod.rs b/crates/hai-core/src/disk/macos/mod.rs new file mode 100644 index 0000000..60a2412 --- /dev/null +++ b/crates/hai-core/src/disk/macos/mod.rs @@ -0,0 +1,5 @@ +mod device; +mod writer; + +pub use device::list_devices; +pub use writer::write_image; diff --git a/crates/hai-core/src/disk_writer/macos.rs b/crates/hai-core/src/disk/macos/writer.rs similarity index 99% rename from crates/hai-core/src/disk_writer/macos.rs rename to crates/hai-core/src/disk/macos/writer.rs index 8b594ec..e2312a9 100644 --- a/crates/hai-core/src/disk_writer/macos.rs +++ b/crates/hai-core/src/disk/macos/writer.rs @@ -1,6 +1,6 @@ //! macOS disk writing via privileged `dd` (Authorization Services) and `diskutil`. -use super::*; +use super::super::*; use security_framework::authorization::{Authorization, AuthorizationItemSetBuilder, Flags}; use std::io::Read; use std::path::Path; diff --git a/crates/hai-core/src/devices/windows.rs b/crates/hai-core/src/disk/windows/device.rs similarity index 99% rename from crates/hai-core/src/devices/windows.rs rename to crates/hai-core/src/disk/windows/device.rs index dfc8978..891572b 100644 --- a/crates/hai-core/src/devices/windows.rs +++ b/crates/hai-core/src/disk/windows/device.rs @@ -1,6 +1,6 @@ //! Windows block device enumeration via PowerShell `Get-Disk`. -use super::*; +use super::super::*; use crate::error::Error; use serde::Deserialize; use std::process::Command; diff --git a/crates/hai-core/src/disk/windows/mod.rs b/crates/hai-core/src/disk/windows/mod.rs new file mode 100644 index 0000000..60a2412 --- /dev/null +++ b/crates/hai-core/src/disk/windows/mod.rs @@ -0,0 +1,5 @@ +mod device; +mod writer; + +pub use device::list_devices; +pub use writer::write_image; diff --git a/crates/hai-core/src/disk_writer/windows.rs b/crates/hai-core/src/disk/windows/writer.rs similarity index 99% rename from crates/hai-core/src/disk_writer/windows.rs rename to crates/hai-core/src/disk/windows/writer.rs index 04e723a..733f57f 100644 --- a/crates/hai-core/src/disk_writer/windows.rs +++ b/crates/hai-core/src/disk/windows/writer.rs @@ -1,6 +1,6 @@ //! Windows disk writing via direct `\\.\PhysicalDrive` access (requires Administrator). -use super::*; +use super::super::*; use std::fs::File; use std::io::{Read, Write}; use std::process::Command; diff --git a/crates/hai-core/src/disk_writer.rs b/crates/hai-core/src/disk_writer.rs deleted file mode 100644 index b52aa6f..0000000 --- a/crates/hai-core/src/disk_writer.rs +++ /dev/null @@ -1,432 +0,0 @@ -//! Raw disk writing functionality for flashing images to devices. -//! -//! This module provides platform-specific implementations for writing -//! raw disk images to block devices (SD cards, USB drives, etc.). - -use crate::error::{Error, Result}; -use crate::types::{FlashProgress, FlashStage}; -use crate::ProgressCallback; -use std::path::PathBuf; - -#[cfg(target_os = "linux")] -#[path = "disk_writer/linux.rs"] -mod imp; -#[cfg(target_os = "macos")] -#[path = "disk_writer/macos.rs"] -mod imp; -#[cfg(target_os = "windows")] -#[path = "disk_writer/windows.rs"] -mod imp; - -#[cfg(not(any(target_os = "linux", target_os = "macos", target_os = "windows")))] -compile_error!("hai-core supports only Linux, macOS and Windows"); - -/// Buffer size for disk writes (4 MB for SD cards) -#[allow(dead_code)] -const WRITE_BUFFER_SIZE: usize = 4 * 1024 * 1024; - -/// Buffer size for fast drives like NVMe/SSDs (64 MB) -#[allow(dead_code)] -const FAST_DRIVE_BUFFER_SIZE: usize = 64 * 1024 * 1024; - -/// How often to send progress updates (every N bytes) -#[allow(dead_code)] -const PROGRESS_UPDATE_INTERVAL: u64 = 10 * 1024 * 1024; // 10 MB - -/// Check if an I/O error indicates the drive was disconnected -fn is_drive_disconnected(io_err: &std::io::Error) -> bool { - matches!( - io_err.kind(), - std::io::ErrorKind::NotFound - | std::io::ErrorKind::BrokenPipe - | std::io::ErrorKind::UnexpectedEof - ) || io_err.raw_os_error().is_some_and(|code| { - // macOS: ENXIO (6) = "Device not configured" - // Linux: ENODEV (19) = "No such device", ENXIO (6) - matches!(code, 6 | 19) - }) -} - -/// Drive a blocking task while forwarding its progress updates to the -/// callback, then drain updates buffered after the task finished (e.g. the -/// final "Write complete" / "Verification complete") so they aren't lost. -async fn run_with_progress( - handle: tokio::task::JoinHandle>, - progress_rx: std::sync::mpsc::Receiver, - progress_callback: &P, -) -> Result<()> { - loop { - match progress_rx.recv_timeout(std::time::Duration::from_millis(100)) { - Ok(update) => progress_callback.on_progress(update), - Err(std::sync::mpsc::RecvTimeoutError::Timeout) => { - if handle.is_finished() { - break; - } - } - Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => break, - } - } - - while let Ok(update) = progress_rx.try_recv() { - progress_callback.on_progress(update); - } - - handle - .await - .map_err(|e| Error::Io(std::io::Error::other(e)))? -} - -/// Write an image file to a block device with progress updates -pub async fn write_image( - image_path: &PathBuf, - device_id: &str, - verify: bool, - progress_callback: &P, -) -> Result<()> { - std::fs::metadata(image_path)?; - - imp::write_image(image_path, device_id, verify, progress_callback).await -} - -#[cfg(test)] -mod tests { - use super::*; - use serial_test::serial; - - #[test] - fn test_is_drive_disconnected_not_found() { - let err = std::io::Error::new(std::io::ErrorKind::NotFound, "device not found"); - assert!(is_drive_disconnected(&err)); - } - - #[test] - fn test_is_drive_disconnected_broken_pipe() { - let err = std::io::Error::new(std::io::ErrorKind::BrokenPipe, "pipe broken"); - assert!(is_drive_disconnected(&err)); - } - - #[test] - fn test_is_drive_disconnected_permission_denied_is_not_disconnect() { - let err = std::io::Error::new(std::io::ErrorKind::PermissionDenied, "access denied"); - assert!(!is_drive_disconnected(&err)); - } - - #[test] - fn test_is_drive_disconnected_unexpected_eof() { - let err = std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "unexpected eof"); - assert!(is_drive_disconnected(&err)); - } - - #[test] - fn test_is_drive_disconnected_other_error_kinds() { - let err = std::io::Error::new(std::io::ErrorKind::PermissionDenied, "permission denied"); - assert!(!is_drive_disconnected(&err)); - - let err = std::io::Error::new(std::io::ErrorKind::Other, "other error"); - assert!(!is_drive_disconnected(&err)); - } - - #[test] - fn test_is_drive_disconnected_with_os_error_code_6() { - // ENXIO = 6 on macOS and Linux - let err = std::io::Error::from_raw_os_error(6); - assert!(is_drive_disconnected(&err)); - } - - #[test] - #[cfg(target_os = "linux")] - fn test_is_drive_disconnected_with_os_error_code_19() { - // ENODEV = 19 on Linux - let err = std::io::Error::from_raw_os_error(19); - assert!(is_drive_disconnected(&err)); - } - - #[test] - fn test_is_drive_disconnected_all_matching_kinds() { - // All these should return true - assert!(is_drive_disconnected(&std::io::Error::new( - std::io::ErrorKind::NotFound, - "" - ))); - assert!(is_drive_disconnected(&std::io::Error::new( - std::io::ErrorKind::BrokenPipe, - "" - ))); - assert!(is_drive_disconnected(&std::io::Error::new( - std::io::ErrorKind::UnexpectedEof, - "" - ))); - } - - #[test] - fn test_is_drive_disconnected_non_matching_kinds() { - // All these should return false - assert!(!is_drive_disconnected(&std::io::Error::new( - std::io::ErrorKind::PermissionDenied, - "" - ))); - assert!(!is_drive_disconnected(&std::io::Error::new( - std::io::ErrorKind::InvalidInput, - "" - ))); - assert!(!is_drive_disconnected(&std::io::Error::new( - std::io::ErrorKind::TimedOut, - "" - ))); - assert!(!is_drive_disconnected(&std::io::Error::new( - std::io::ErrorKind::WriteZero, - "" - ))); - assert!(!is_drive_disconnected(&std::io::Error::new( - std::io::ErrorKind::Interrupted, - "" - ))); - } - - #[test] - fn test_is_drive_disconnected_with_other_os_error_codes() { - // Test non-matching OS error codes - let err = std::io::Error::from_raw_os_error(1); // EPERM - assert!(!is_drive_disconnected(&err)); - - let err = std::io::Error::from_raw_os_error(13); // EACCES - assert!(!is_drive_disconnected(&err)); - } - - #[test] - fn test_is_drive_disconnected_no_os_error() { - // Test error without raw OS error code - let err = std::io::Error::new(std::io::ErrorKind::Other, "generic error"); - assert!(!is_drive_disconnected(&err)); - } - - // Helper struct for testing progress callbacks - struct TestProgressCallback { - updates: std::sync::Arc>>, - } - - impl TestProgressCallback { - fn new() -> Self { - Self { - updates: std::sync::Arc::new(std::sync::Mutex::new(Vec::new())), - } - } - } - - impl ProgressCallback for TestProgressCallback { - fn on_progress(&self, progress: FlashProgress) { - self.updates.lock().unwrap().push(progress); - } - } - - // macOS-specific tests - #[cfg(target_os = "macos")] - mod macos_tests { - use super::*; - - #[tokio::test] - #[serial] - async fn test_write_image_nonexistent_file() { - let callback = TestProgressCallback::new(); - let image_path = PathBuf::from("/tmp/nonexistent_image_file.img"); - let device_id = "/dev/disk99"; - - let result = write_image(&image_path, device_id, false, &callback).await; - assert!(result.is_err()); - } - - #[tokio::test] - #[serial] - async fn test_write_image_with_invalid_device() { - let callback = TestProgressCallback::new(); - let temp_file = tempfile::NamedTempFile::new().unwrap(); - - // Write some test data - std::fs::write(temp_file.path(), b"test data").unwrap(); - let image_path = temp_file.path().to_path_buf(); - - // Use an invalid device path - let device_id = "/dev/nonexistent_disk999"; - - let result = write_image(&image_path, device_id, false, &callback).await; - // Should fail when trying to unmount or access the device - assert!(result.is_err()); - } - } - - // Linux-specific tests - #[cfg(target_os = "linux")] - mod linux_tests { - use super::*; - - #[tokio::test] - #[serial] - async fn test_write_image_nonexistent_file() { - let callback = TestProgressCallback::new(); - let image_path = PathBuf::from("/tmp/nonexistent_image_file.img"); - let device_id = "/dev/hai-test-nonexistent"; - - // The image check runs before device validation and any D-Bus - // call, so a missing image surfaces as Io even with a bad device. - let result = write_image(&image_path, device_id, false, &callback).await; - assert!(matches!(result.unwrap_err(), Error::Io(_))); - } - } - - // Windows-specific tests - #[cfg(target_os = "windows")] - mod windows_tests { - use super::*; - - #[tokio::test] - #[serial] - async fn test_write_image_nonexistent_file() { - let callback = TestProgressCallback::new(); - let image_path = PathBuf::from("C:\\nonexistent_image_file.img"); - let device_id = "\\\\.\\PhysicalDrive1"; - - let result = write_image(&image_path, device_id, false, &callback).await; - assert!(result.is_err()); - } - - #[tokio::test] - #[serial] - async fn test_write_image_invalid_device() { - let callback = TestProgressCallback::new(); - let temp_file = tempfile::NamedTempFile::new().unwrap(); - std::fs::write(temp_file.path(), b"test data").unwrap(); - let image_path = temp_file.path().to_path_buf(); - - let device_id = "\\\\.\\PhysicalDrive999"; - - let result = write_image(&image_path, device_id, false, &callback).await; - assert!(result.is_err()); - } - } - - // Test with Mock mode enabled to exercise more code paths - #[tokio::test] - #[serial] - async fn test_write_image_mock_mode_not_implemented() { - // Enable mock mode - std::env::set_var("HA_INSTALLER_MOCK", "1"); - - let callback = TestProgressCallback::new(); - let temp_file = tempfile::NamedTempFile::new().unwrap(); - std::fs::write(temp_file.path(), b"test data").unwrap(); - let image_path = temp_file.path().to_path_buf(); - - #[cfg(target_os = "macos")] - let device_id = "/dev/disk2"; - // Non-resolvable device so write_image fails before any privileged - // udisks2 call (mock mode isn't honoured by write_image itself). - #[cfg(target_os = "linux")] - let device_id = "/dev/hai-test-nonexistent"; - #[cfg(target_os = "windows")] - let device_id = "\\\\.\\PhysicalDrive1"; - - // The write_image function doesn't have mock support yet, - // so it will try to execute real commands which will fail - let result = write_image(&image_path, device_id, false, &callback).await; - - // Clean up - std::env::remove_var("HA_INSTALLER_MOCK"); - - // Expected to fail as there's no mock implementation for write_image - assert!(result.is_err()); - } - - #[tokio::test] - #[serial] - async fn test_write_image_with_verification_flag() { - let callback = TestProgressCallback::new(); - let temp_file = tempfile::NamedTempFile::new().unwrap(); - std::fs::write(temp_file.path(), b"test data").unwrap(); - let image_path = temp_file.path().to_path_buf(); - - #[cfg(target_os = "macos")] - let device_id = "/dev/disk99"; - #[cfg(target_os = "linux")] - let device_id = "/dev/hai-test-nonexistent"; - #[cfg(target_os = "windows")] - let device_id = "\\\\.\\PhysicalDrive99"; - - // Test with verify=true flag - let result = write_image(&image_path, device_id, true, &callback).await; - - // Should fail because device doesn't exist, but this exercises the verify code path - assert!(result.is_err()); - } - - // Test creating actual temp file and trying to write (will fail safely) - #[tokio::test] - #[serial] - async fn test_write_image_with_real_temp_file() { - let callback = TestProgressCallback::new(); - - // Create a temp file with some content - let temp_file = tempfile::NamedTempFile::new().unwrap(); - std::fs::write(temp_file.path(), b"test image data").unwrap(); - let image_path = temp_file.path().to_path_buf(); - - // Try to write to a nonexistent device - #[cfg(target_os = "macos")] - let device_id = "/dev/disk999"; - #[cfg(target_os = "linux")] - let device_id = "/dev/hai-test-nonexistent"; - #[cfg(target_os = "windows")] - let device_id = "\\\\.\\PhysicalDrive999"; - - let result = write_image(&image_path, device_id, false, &callback).await; - - // Should fail because device doesn't exist - assert!(result.is_err()); - - // Progress updates may or may not have been sent depending on when the failure occurred - // Just verify the function was called and failed appropriately - } - - // Test all ErrorKind variants for is_drive_disconnected - #[test] - fn test_is_drive_disconnected_comprehensive() { - use std::io::ErrorKind; - - // These should return true - let disconnect_kinds = [ - ErrorKind::NotFound, - ErrorKind::BrokenPipe, - ErrorKind::UnexpectedEof, - ]; - for kind in &disconnect_kinds { - assert!( - is_drive_disconnected(&std::io::Error::new(*kind, "test")), - "{:?} should be detected as disconnected", - kind - ); - } - - // These should return false - let other_kinds = [ - ErrorKind::PermissionDenied, - ErrorKind::ConnectionRefused, - ErrorKind::ConnectionReset, - ErrorKind::ConnectionAborted, - ErrorKind::AddrInUse, - ErrorKind::AddrNotAvailable, - ErrorKind::InvalidInput, - ErrorKind::InvalidData, - ErrorKind::TimedOut, - ErrorKind::WriteZero, - ErrorKind::Interrupted, - ErrorKind::Other, - ErrorKind::WouldBlock, - ]; - for kind in &other_kinds { - assert!( - !is_drive_disconnected(&std::io::Error::new(*kind, "test")), - "{:?} should NOT be detected as disconnected", - kind - ); - } - } -} diff --git a/crates/hai-core/src/lib.rs b/crates/hai-core/src/lib.rs index 3816223..6e1a3b7 100644 --- a/crates/hai-core/src/lib.rs +++ b/crates/hai-core/src/lib.rs @@ -7,8 +7,7 @@ //! The library is designed to be frontend-agnostic, supporting both the //! Tauri desktop application and potential TUI implementations. -pub mod devices; -pub mod disk_writer; +pub mod disk; pub mod download; pub mod error; pub mod types; diff --git a/crates/hai-desktop/src/commands.rs b/crates/hai-desktop/src/commands.rs index d4f0c1d..59d4890 100644 --- a/crates/hai-desktop/src/commands.rs +++ b/crates/hai-desktop/src/commands.rs @@ -4,9 +4,9 @@ //! It handles the bridge between Tauri's Channel and hai-core's ProgressCallback trait. use hai_core::{ - devices, download, is_mock_enabled, mock, BlockDevice, DeviceManifest, FlashProgress, - FlashStage, HaosRelease, ProgressCallback, ProxmoxCredentials, ProxmoxNode, ProxmoxSession, - ProxmoxStorage, ProxmoxVmConfig, ProxmoxVmResult, UpdateInfo, + disk, download, is_mock_enabled, mock, BlockDevice, DeviceManifest, FlashProgress, FlashStage, + HaosRelease, ProgressCallback, ProxmoxCredentials, ProxmoxNode, ProxmoxSession, ProxmoxStorage, + ProxmoxVmConfig, ProxmoxVmResult, UpdateInfo, }; use std::time::Duration; use tauri::ipc::Channel; @@ -86,7 +86,7 @@ pub async fn list_block_devices() -> Result, String> { if is_mock_enabled() { Ok(mock::get_mock_block_devices()) } else { - devices::list_devices().await.map_err(|e| e.to_string()) + disk::list_devices().await.map_err(|e| e.to_string()) } } @@ -195,7 +195,7 @@ pub async fn flash_image( .map_err(|e| format!("Failed to get image size: {}", e))? .len(); - let device_list = devices::list_devices() + let device_list = disk::list_devices() .await .map_err(|e| format!("Failed to list devices: {}", e))?; @@ -210,7 +210,7 @@ pub async fn flash_image( } // Write to device - hai_core::disk_writer::write_image( + disk::write_image( &extracted_path, &request.device_id, request.verify,