From b670db2d222b0fef3bbc5f4d450fdcb7b003cc0a Mon Sep 17 00:00:00 2001 From: weishu Date: Sun, 25 Feb 2024 20:08:25 +0800 Subject: [PATCH] ksud: fix punch hole --- userspace/ksud/src/cli.rs | 18 +++--- userspace/ksud/src/event.rs | 2 +- userspace/ksud/src/module.rs | 30 ++++------ userspace/ksud/src/utils.rs | 103 ++++------------------------------- 4 files changed, 31 insertions(+), 122 deletions(-) diff --git a/userspace/ksud/src/cli.rs b/userspace/ksud/src/cli.rs index 86483398..f8892b67 100644 --- a/userspace/ksud/src/cli.rs +++ b/userspace/ksud/src/cli.rs @@ -122,12 +122,9 @@ enum Debug { src: String, /// destination file dst: String, - }, - - /// Punch hole file - PunchHole { - /// file path - file: String, + /// punch hole + #[arg(short, long, default_value = "false")] + punch_hole: bool, }, /// For testing @@ -299,11 +296,14 @@ pub fn run() -> Result<()> { } Debug::Su { global_mnt } => crate::ksu::grant_root(global_mnt), Debug::Mount => event::mount_systemlessly(defs::MODULE_DIR), - Debug::Xcp { src, dst } => { - utils::copy_sparse_file(src, dst)?; + Debug::Xcp { + src, + dst, + punch_hole, + } => { + utils::copy_sparse_file(src, dst, punch_hole)?; Ok(()) } - Debug::PunchHole { file } => utils::punch_hole(file), Debug::Test => todo!(), }, diff --git a/userspace/ksud/src/event.rs b/userspace/ksud/src/event.rs index 84b709d9..a7a16b11 100644 --- a/userspace/ksud/src/event.rs +++ b/userspace/ksud/src/event.rs @@ -255,7 +255,7 @@ pub fn on_boot_completed() -> Result<()> { // this is a update and we successfully booted if std::fs::rename(module_update_img, module_img).is_err() { warn!("Failed to rename images, copy it now.",); - utils::copy_sparse_file(module_update_img, module_img) + utils::copy_sparse_file(module_update_img, module_img, false) .with_context(|| "Failed to copy images")?; std::fs::remove_file(module_update_img).with_context(|| "Failed to remove image!")?; } diff --git a/userspace/ksud/src/module.rs b/userspace/ksud/src/module.rs index 70822f76..fa711e50 100644 --- a/userspace/ksud/src/module.rs +++ b/userspace/ksud/src/module.rs @@ -367,7 +367,7 @@ fn _install_module(zip: &str) -> Result<()> { } else if modules_update_img_exist { // modules_update.img exists, we should use it as tmp img info!("Using existing modules_update.img as tmp image"); - utils::copy_sparse_file(modules_update_img, tmp_module_img).with_context(|| { + utils::copy_sparse_file(modules_update_img, tmp_module_img, true).with_context(|| { format!( "Failed to copy {} to {}", modules_update_img.display(), @@ -377,7 +377,7 @@ fn _install_module(zip: &str) -> Result<()> { } else { // modules.img exists, we should use it as tmp img info!("Using existing modules.img as tmp image"); - utils::copy_sparse_file(modules_img, tmp_module_img).with_context(|| { + utils::copy_sparse_file(modules_img, tmp_module_img, true).with_context(|| { format!( "Failed to copy {} to {}", modules_img.display(), @@ -445,15 +445,11 @@ fn _install_module(zip: &str) -> Result<()> { exec_install_script(zip)?; - if let Err(e) = utils::punch_hole(tmp_module_img) { - warn!("Failed to punch hole: {}", e); - } - info!("rename {tmp_module_img} to {}", defs::MODULE_UPDATE_IMG); // all done, rename the tmp image to modules_update.img if std::fs::rename(tmp_module_img, defs::MODULE_UPDATE_IMG).is_err() { warn!("Rename image failed, try copy it."); - utils::copy_sparse_file(tmp_module_img, defs::MODULE_UPDATE_IMG) + utils::copy_sparse_file(tmp_module_img, defs::MODULE_UPDATE_IMG, true) .with_context(|| "Failed to copy image.".to_string())?; let _ = std::fs::remove_file(tmp_module_img); } @@ -476,7 +472,7 @@ pub fn install_module(zip: &str) -> Result<()> { result } -fn update_module(update_dir: &str, id: &str, punch_hole: bool, func: F) -> Result<()> +fn update_module(update_dir: &str, id: &str, func: F) -> Result<()> where F: Fn(&str, &str) -> Result<()>, { @@ -493,14 +489,14 @@ where modules_update_img.display(), modules_update_tmp_img.display() ); - utils::copy_sparse_file(modules_update_img, modules_update_tmp_img)?; + utils::copy_sparse_file(modules_update_img, modules_update_tmp_img, true)?; } else { info!( "copy {} to {}", modules_img.display(), modules_update_tmp_img.display() ); - utils::copy_sparse_file(modules_img, modules_update_tmp_img)?; + utils::copy_sparse_file(modules_img, modules_update_tmp_img, true)?; } // ensure modules_update dir exist @@ -512,15 +508,9 @@ where // call the operation func let result = func(id, update_dir); - if punch_hole { - if let Err(e) = utils::punch_hole(modules_update_tmp_img) { - warn!("Failed to punch hole: {}", e); - } - } - if let Err(e) = std::fs::rename(modules_update_tmp_img, defs::MODULE_UPDATE_IMG) { warn!("Rename image failed: {e}, try copy it."); - utils::copy_sparse_file(modules_update_tmp_img, defs::MODULE_UPDATE_IMG) + utils::copy_sparse_file(modules_update_tmp_img, defs::MODULE_UPDATE_IMG, true) .with_context(|| "Failed to copy image.".to_string())?; let _ = std::fs::remove_file(modules_update_tmp_img); } @@ -531,7 +521,7 @@ where } pub fn uninstall_module(id: &str) -> Result<()> { - update_module(defs::MODULE_UPDATE_TMP_DIR, id, true, |mid, update_dir| { + update_module(defs::MODULE_UPDATE_TMP_DIR, id, |mid, update_dir| { let dir = Path::new(update_dir); ensure!(dir.exists(), "No module installed"); @@ -597,13 +587,13 @@ fn _enable_module(module_dir: &str, mid: &str, enable: bool) -> Result<()> { } pub fn enable_module(id: &str) -> Result<()> { - update_module(defs::MODULE_UPDATE_TMP_DIR, id, false, |mid, update_dir| { + update_module(defs::MODULE_UPDATE_TMP_DIR, id, |mid, update_dir| { _enable_module(update_dir, mid, true) }) } pub fn disable_module(id: &str) -> Result<()> { - update_module(defs::MODULE_UPDATE_TMP_DIR, id, false, |mid, update_dir| { + update_module(defs::MODULE_UPDATE_TMP_DIR, id, |mid, update_dir| { _enable_module(update_dir, mid, false) }) } diff --git a/userspace/ksud/src/utils.rs b/userspace/ksud/src/utils.rs index 1505b5c9..6ab1ff09 100644 --- a/userspace/ksud/src/utils.rs +++ b/userspace/ksud/src/utils.rs @@ -192,7 +192,11 @@ pub fn get_tmp_path() -> &'static str { } // TODO: use libxcp to improve the speed if cross's MSRV is 1.70 -pub fn copy_sparse_file, Q: AsRef>(src: P, dst: Q) -> Result<()> { +pub fn copy_sparse_file, Q: AsRef>( + src: P, + dst: Q, + punch_hole: bool, +) -> Result<()> { let mut src_file = File::open(src.as_ref())?; let mut dst_file = OpenOptions::new() .write(true) @@ -223,6 +227,12 @@ pub fn copy_sparse_file, Q: AsRef>(src: P, dst: Q) -> Resul break; } + if punch_hole && buffer[..bytes_read].iter().all(|&x| x == 0) { + // all zero, don't copy it at all! + dst_file.seek(SeekFrom::Current(bytes_read as i64))?; + total_bytes_copied += bytes_read as u64; + continue; + } dst_file.write_all(&buffer[..bytes_read])?; total_bytes_copied += bytes_read as u64; } @@ -231,94 +241,3 @@ pub fn copy_sparse_file, Q: AsRef>(src: P, dst: Q) -> Resul Ok(()) } - -#[cfg(any(target_os = "linux", target_os = "android"))] -pub fn punch_hole(src: impl AsRef) -> Result<()> { - let mut src_file = OpenOptions::new().write(true).read(true).open(src)?; - - let st = rustix::fs::fstat(&src_file)?; - - let bufsz = st.st_blksize; - let mut buf = vec![0u8; bufsz as usize]; - - let mut ct = 0; - let mut hole_sz = 0; - let mut hole_start = 0; - - let segments = src_file.scan_chunks()?; - for segment in segments { - if segment.segment_type != SegmentType::Data { - continue; - } - - let mut off = segment.start; - let end = segment.end + 1; - - while off < end { - let mut rsz = rustix::io::pread(&src_file, &mut buf, off)? as u64; - if rsz > 0 && rsz > end - off { - // exceed the end of the boundary - rsz = end - off; - } - - if rsz == 0 { - break; - } - - if buf.iter().all(|&x| x == 0) { - // the whole buf is zero, mark it as a hole - if hole_sz == 0 { - hole_start = off; - } - // for continuous zero, we can merge them into a bigger hole - hole_sz += rsz; - } else if hole_sz > 0 { - if let Err(e) = rustix::fs::fallocate( - &src_file, - rustix::fs::FallocateFlags::PUNCH_HOLE | rustix::fs::FallocateFlags::KEEP_SIZE, - hole_start, - hole_sz, - ) { - log::warn!("Failed to punch hole: {:?}", e); - } - - ct += hole_sz; - - hole_sz = 0; - hole_start = 0; - } - - off += rsz; - } - - // if the last segment is a hole, we need to punch it - if hole_sz > 0 { - let mut alloc_sz = hole_sz; - if off >= end { - alloc_sz += st.st_blksize as u64; - } - if let Err(e) = rustix::fs::fallocate( - &src_file, - rustix::fs::FallocateFlags::PUNCH_HOLE | rustix::fs::FallocateFlags::KEEP_SIZE, - hole_start, - alloc_sz, - ) { - log::warn!("Failed to punch hole: {:?}", e); - } - - ct += hole_sz; - } - } - - log::info!( - "Punched {} of hole", - humansize::format_size(ct, humansize::DECIMAL) - ); - - Ok(()) -} - -#[cfg(not(any(target_os = "linux", target_os = "android")))] -pub fn punch_hole(src: impl AsRef) -> Result<()> { - unimplemented!() -}