Ramdisk: Enable detaching bootconfig from initrd

This is the reverse of appending bootconfigs to initrd.
initrd_bootconfig detach-bootconfig takes initrd_with_bootconfig as
input & outputs the containing ramdisk & bootconfigs.

This is required because our signing infrastructure requires changing
vbmeta related configs (when images are resigned using a different key).
This can be achieved by detaching the configs, modifying them &
re-attaching them.

Bug: 245277660
Test: Attach & then detach bootconfigs to initrd & the diff the
artifacts.

Change-Id: Ia8803496549d887dae84cf7c4f545213591c6a78
diff --git a/microdroid/initrd/Android.bp b/microdroid/initrd/Android.bp
index 7a95ce6..ff6314b 100644
--- a/microdroid/initrd/Android.bp
+++ b/microdroid/initrd/Android.bp
@@ -68,7 +68,7 @@
         ":microdroid_bootconfig_debuggable_src",
     ] + bootconfigs_arm64,
     out: ["microdroid_initrd_debuggable_arm64"],
-    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
 genrule {
@@ -79,7 +79,7 @@
         ":microdroid_bootconfig_debuggable_src",
     ] + bootconfigs_x86_64,
     out: ["microdroid_initrd_debuggable_x86_64"],
-    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
 genrule {
@@ -90,7 +90,7 @@
         ":microdroid_bootconfig_normal_src",
     ] + bootconfigs_arm64,
     out: ["microdroid_initrd_normal_arm64"],
-    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
 genrule {
@@ -101,7 +101,7 @@
         ":microdroid_bootconfig_normal_src",
     ] + bootconfigs_x86_64,
     out: ["microdroid_initrd_normal_x86_64"],
-    cmd: "$(location initrd_bootconfig) --output $(out) $(in)",
+    cmd: "$(location initrd_bootconfig) attach --output $(out) $(in)",
 }
 
 prebuilt_etc {
diff --git a/microdroid/initrd/src/main.rs b/microdroid/initrd/src/main.rs
index 74e4ba6..bf4f3a2 100644
--- a/microdroid/initrd/src/main.rs
+++ b/microdroid/initrd/src/main.rs
@@ -12,31 +12,90 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! Append bootconfig to initrd image
-use anyhow::Result;
+//! Attach/Detach bootconfigs to initrd image
+use anyhow::{bail, Result};
 use clap::Parser;
+use std::cmp::min;
 use std::fs::File;
-use std::io::{Read, Write};
+use std::io::{Read, Seek, SeekFrom, Write};
+use std::mem::size_of;
 use std::path::PathBuf;
 
 const FOOTER_ALIGNMENT: usize = 4;
 const ZEROS: [u8; 4] = [0u8; 4_usize];
+const BOOTCONFIG_MAGIC: &str = "#BOOTCONFIG\n";
+// Footer includes [size(le32)][checksum(le32)][#BOOTCONFIG\n] at the end of bootconfigs.
+const INITRD_FOOTER_LEN: usize = 2 * std::mem::size_of::<u32>() + BOOTCONFIG_MAGIC.len();
 
 #[derive(Parser, Debug)]
-struct Args {
-    /// Initrd (without bootconfig)
-    initrd: PathBuf,
-    /// Bootconfig
-    bootconfigs: Vec<PathBuf>,
-    /// Output
-    #[clap(long = "output")]
-    output: PathBuf,
+enum Opt {
+    /// Append bootconfig(s) to initrd image
+    Attach {
+        /// Initrd (without bootconfigs) <- Input
+        initrd: PathBuf,
+        /// Bootconfigs <- Input
+        bootconfigs: Vec<PathBuf>,
+        /// Initrd (with bootconfigs) <- Output
+        #[clap(long = "output")]
+        output: PathBuf,
+    },
+
+    /// Detach the initrd & bootconfigs - this is required for cases when we update
+    /// bootconfigs in sign_virt_apex
+    Detach {
+        /// Initrd (with bootconfigs) <- Input
+        initrd_with_bootconfigs: PathBuf,
+        /// Initrd (without bootconfigs) <- Output
+        initrd: PathBuf,
+        /// Bootconfigs <- Output
+        bootconfigs: PathBuf,
+    },
 }
 
 fn get_checksum(file_path: &PathBuf) -> Result<u32> {
     File::open(file_path)?.bytes().map(|x| Ok(x? as u32)).sum()
 }
 
+// Copy n bytes of file_in to file_out. Note: copying starts from the current cursors of files.
+// On successful return, the files' cursors would have moved forward by k bytes.
+fn copyfile2file(file_in: &mut File, file_out: &mut File, n: usize) -> Result<()> {
+    let mut buf = vec![0; 1024];
+    let mut copied: usize = 0;
+    while copied < n {
+        let k = min(n - copied, buf.len());
+        file_in.read_exact(&mut buf[..k])?;
+        file_out.write_all(&buf[..k])?;
+        copied += k;
+    }
+    Ok(())
+}
+
+// Note: attaching & then detaching bootconfigs can lead to extra padding in bootconfigs
+fn detach_bootconfig(initrd_bc: PathBuf, initrd: PathBuf, bootconfig: PathBuf) -> Result<()> {
+    let mut initrd_bc = File::open(&initrd_bc)?;
+    let mut bootconfig = File::create(&bootconfig)?;
+    let mut initrd = File::create(&initrd)?;
+    let initrd_bc_size: usize = initrd_bc.metadata()?.len().try_into()?;
+
+    initrd_bc.seek(SeekFrom::End(-(BOOTCONFIG_MAGIC.len() as i64)))?;
+    let mut magic_buf = [0; BOOTCONFIG_MAGIC.len()];
+    initrd_bc.read_exact(&mut magic_buf)?;
+    if magic_buf != BOOTCONFIG_MAGIC.as_bytes() {
+        bail!("BOOTCONFIG_MAGIC not found in initrd. Bootconfigs might not be attached correctly");
+    }
+    let mut size_buf = [0; size_of::<u32>()];
+    initrd_bc.seek(SeekFrom::End(-(INITRD_FOOTER_LEN as i64)))?;
+    initrd_bc.read_exact(&mut size_buf)?;
+    let bc_size: usize = u32::from_le_bytes(size_buf) as usize;
+
+    let initrd_size: usize = initrd_bc_size - bc_size - INITRD_FOOTER_LEN;
+
+    initrd_bc.seek(SeekFrom::Start(0))?;
+    copyfile2file(&mut initrd_bc, &mut initrd, initrd_size)?;
+    copyfile2file(&mut initrd_bc, &mut bootconfig, bc_size)?;
+    Ok(())
+}
+
 // Bootconfig is attached to the initrd in the following way:
 // [initrd][bootconfig][padding][size(le32)][checksum(le32)][#BOOTCONFIG\n]
 fn attach_bootconfig(initrd: PathBuf, bootconfigs: Vec<PathBuf>, output: PathBuf) -> Result<()> {
@@ -59,14 +118,21 @@
     output_file.write_all(&ZEROS[..padding_size])?;
     output_file.write_all(&((padding_size + bootconfig_size) as u32).to_le_bytes())?;
     output_file.write_all(&checksum.to_le_bytes())?;
-    output_file.write_all(b"#BOOTCONFIG\n")?;
+    output_file.write_all(BOOTCONFIG_MAGIC.as_bytes())?;
     output_file.flush()?;
     Ok(())
 }
 
 fn try_main() -> Result<()> {
-    let args = Args::parse();
-    attach_bootconfig(args.initrd, args.bootconfigs, args.output)?;
+    let args = Opt::parse();
+    match args {
+        Opt::Attach { initrd, bootconfigs, output } => {
+            attach_bootconfig(initrd, bootconfigs, output)?
+        }
+        Opt::Detach { initrd_with_bootconfigs, initrd, bootconfigs } => {
+            detach_bootconfig(initrd_with_bootconfigs, initrd, bootconfigs)?
+        }
+    };
     Ok(())
 }
 
@@ -82,6 +148,6 @@
     #[test]
     fn verify_args() {
         // Check that the command parsing has been configured in a valid way.
-        Args::command().debug_assert();
+        Opt::command().debug_assert();
     }
 }