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,
/// 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!(),
},

View File

@@ -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!")?;
}

View File

@@ -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<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
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)
})
}

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
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 dst_file = OpenOptions::new()
.write(true)
@@ -223,6 +227,12 @@ pub fn copy_sparse_file<P: AsRef<Path>, Q: AsRef<Path>>(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<P: AsRef<Path>, Q: AsRef<Path>>(src: P, dst: Q) -> Resul
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!()
}