From e0802b0d152bcd903458f0c0eaa70a18ec2a69ca Mon Sep 17 00:00:00 2001
From: weishu <twsxtd@gmail.com>
Date: Mon, 19 Feb 2024 12:31:13 +0800
Subject: [PATCH] ksud: check image before shrink

---
 userspace/ksud/src/cli.rs    |  2 +-
 userspace/ksud/src/module.rs | 13 +++++++++++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/userspace/ksud/src/cli.rs b/userspace/ksud/src/cli.rs
index 8c37821d987c..1f63583a82cf 100644
--- a/userspace/ksud/src/cli.rs
+++ b/userspace/ksud/src/cli.rs
@@ -254,7 +254,7 @@ pub fn run() -> Result<()> {
                 Module::Enable { id } => module::enable_module(&id),
                 Module::Disable { id } => module::disable_module(&id),
                 Module::List => module::list_modules(),
-                Module::Shrink => module::shrink_image(),
+                Module::Shrink => module::shrink_ksu_images(),
             }
         }
         Commands::Install => event::install(),
diff --git a/userspace/ksud/src/module.rs b/userspace/ksud/src/module.rs
index b5d89ac56219..e71e211361eb 100644
--- a/userspace/ksud/src/module.rs
+++ b/userspace/ksud/src/module.rs
@@ -674,11 +674,20 @@ pub fn list_modules() -> Result<()> {
     Ok(())
 }
 
-pub fn shrink_image() -> Result<()> {
+pub fn shrink_image(img: &str) -> Result<()> {
+    check_image(img)?;
     Command::new("resize2fs")
         .arg("-M")
-        .arg(defs::MODULE_IMG)
+        .arg(img)
         .stdout(Stdio::piped())
         .status()?;
     Ok(())
 }
+
+pub fn shrink_ksu_images() -> Result<()> {
+    shrink_image(defs::MODULE_IMG)?;
+    if Path::new(defs::MODULE_UPDATE_IMG).exists() {
+        shrink_image(defs::MODULE_UPDATE_IMG)?;
+    }
+    Ok(())
+}