ksud: fix punch hole

This commit is contained in:
weishu
2024-02-25 20:08:25 +08:00
parent 62a31b3dc2
commit b670db2d22
4 changed files with 31 additions and 122 deletions

View File

@@ -122,12 +122,9 @@ enum Debug {
src: String, src: String,
/// destination file /// destination file
dst: String, dst: String,
}, /// punch hole
#[arg(short, long, default_value = "false")]
/// Punch hole file punch_hole: bool,
PunchHole {
/// file path
file: String,
}, },
/// For testing /// For testing
@@ -299,11 +296,14 @@ pub fn run() -> Result<()> {
} }
Debug::Su { global_mnt } => crate::ksu::grant_root(global_mnt), Debug::Su { global_mnt } => crate::ksu::grant_root(global_mnt),
Debug::Mount => event::mount_systemlessly(defs::MODULE_DIR), Debug::Mount => event::mount_systemlessly(defs::MODULE_DIR),
Debug::Xcp { src, dst } => { Debug::Xcp {
utils::copy_sparse_file(src, dst)?; src,
dst,
punch_hole,
} => {
utils::copy_sparse_file(src, dst, punch_hole)?;
Ok(()) Ok(())
} }
Debug::PunchHole { file } => utils::punch_hole(file),
Debug::Test => todo!(), Debug::Test => todo!(),
}, },

View File

@@ -255,7 +255,7 @@ pub fn on_boot_completed() -> Result<()> {
// this is a update and we successfully booted // this is a update and we successfully booted
if std::fs::rename(module_update_img, module_img).is_err() { if std::fs::rename(module_update_img, module_img).is_err() {
warn!("Failed to rename images, copy it now.",); 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")?; .with_context(|| "Failed to copy images")?;
std::fs::remove_file(module_update_img).with_context(|| "Failed to remove image!")?; std::fs::remove_file(module_update_img).with_context(|| "Failed to remove image!")?;
} }

View File

@@ -367,7 +367,7 @@ fn _install_module(zip: &str) -> Result<()> {
} else if modules_update_img_exist { } else if modules_update_img_exist {
// modules_update.img exists, we should use it as tmp img // modules_update.img exists, we should use it as tmp img
info!("Using existing modules_update.img as tmp image"); 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!( format!(
"Failed to copy {} to {}", "Failed to copy {} to {}",
modules_update_img.display(), modules_update_img.display(),
@@ -377,7 +377,7 @@ fn _install_module(zip: &str) -> Result<()> {
} else { } else {
// modules.img exists, we should use it as tmp img // modules.img exists, we should use it as tmp img
info!("Using existing modules.img as tmp image"); 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!( format!(
"Failed to copy {} to {}", "Failed to copy {} to {}",
modules_img.display(), modules_img.display(),
@@ -445,15 +445,11 @@ fn _install_module(zip: &str) -> Result<()> {
exec_install_script(zip)?; 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); info!("rename {tmp_module_img} to {}", defs::MODULE_UPDATE_IMG);
// all done, rename the tmp image to modules_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() { if std::fs::rename(tmp_module_img, defs::MODULE_UPDATE_IMG).is_err() {
warn!("Rename image failed, try copy it."); 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())?; .with_context(|| "Failed to copy image.".to_string())?;
let _ = std::fs::remove_file(tmp_module_img); let _ = std::fs::remove_file(tmp_module_img);
} }
@@ -476,7 +472,7 @@ pub fn install_module(zip: &str) -> Result<()> {
result result
} }
fn update_module<F>(update_dir: &str, id: &str, punch_hole: bool, func: F) -> Result<()> fn update_module<F>(update_dir: &str, id: &str, func: F) -> Result<()>
where where
F: Fn(&str, &str) -> Result<()>, F: Fn(&str, &str) -> Result<()>,
{ {
@@ -493,14 +489,14 @@ where
modules_update_img.display(), modules_update_img.display(),
modules_update_tmp_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 { } else {
info!( info!(
"copy {} to {}", "copy {} to {}",
modules_img.display(), modules_img.display(),
modules_update_tmp_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 // ensure modules_update dir exist
@@ -512,15 +508,9 @@ where
// call the operation func // call the operation func
let result = func(id, update_dir); 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) { if let Err(e) = std::fs::rename(modules_update_tmp_img, defs::MODULE_UPDATE_IMG) {
warn!("Rename image failed: {e}, try copy it."); 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())?; .with_context(|| "Failed to copy image.".to_string())?;
let _ = std::fs::remove_file(modules_update_tmp_img); let _ = std::fs::remove_file(modules_update_tmp_img);
} }
@@ -531,7 +521,7 @@ where
} }
pub fn uninstall_module(id: &str) -> Result<()> { 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); let dir = Path::new(update_dir);
ensure!(dir.exists(), "No module installed"); 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<()> { 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) _enable_module(update_dir, mid, true)
}) })
} }
pub fn disable_module(id: &str) -> Result<()> { 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) _enable_module(update_dir, mid, false)
}) })
} }

View File

@@ -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 // TODO: use libxcp to improve the speed if cross's MSRV is 1.70
pub fn copy_sparse_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> Result<()> { pub fn copy_sparse_file<P: AsRef<Path>, Q: AsRef<Path>>(
src: P,
dst: Q,
punch_hole: bool,
) -> Result<()> {
let mut src_file = File::open(src.as_ref())?; let mut src_file = File::open(src.as_ref())?;
let mut dst_file = OpenOptions::new() let mut dst_file = OpenOptions::new()
.write(true) .write(true)
@@ -223,6 +227,12 @@ pub fn copy_sparse_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> Resul
break; 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])?; dst_file.write_all(&buffer[..bytes_read])?;
total_bytes_copied += bytes_read as u64; total_bytes_copied += bytes_read as u64;
} }
@@ -231,94 +241,3 @@ pub fn copy_sparse_file<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> Resul
Ok(()) Ok(())
} }
#[cfg(any(target_os = "linux", target_os = "android"))]
pub fn punch_hole(src: impl AsRef<Path>) -> 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<Path>) -> Result<()> {
unimplemented!()
}