cfs/uki: Use buffered readers

Use buffered readers for reading UKI sections instead of loading the
entire thing in memory

Signed-off-by: Pragyan Poudyal <pragyanpoudyal41999@gmail.com>

Author: Johan-Liebert1

crates/lib/src/bootc_composefs/boot.rs

@@ -62,8 +62,7 @@
 //! 2. **Secondary**: Currently booted deployment (rollback option)
 
 use std::fs::create_dir_all;
-use std::io::Read;
-use std::io::Write;
+use std::io::{Read, Seek, SeekFrom, Write};
 use std::path::Path;
 use std::sync::Arc;
 
@@ -347,12 +346,13 @@ fn compute_boot_digest_type1(dir: &Dir) -> Result<String> {
 /// * entry - BootEntry containing VMlinuz and Initrd
 /// * repo - The composefs repository
 #[context("Computing boot digest")]
-pub(crate) fn compute_boot_digest_uki(uki: &[u8]) -> Result<String> {
-    let vmlinuz =
-        uki::get_section(uki, ".linux").ok_or_else(|| anyhow::anyhow!(".linux not present"))??;
-
-    let initramfs = uki::get_section(uki, ".initrd")
-        .ok_or_else(|| anyhow::anyhow!(".initrd not present"))??;
+pub(crate) fn compute_boot_digest_uki<R: Read + Seek>(uki_reader: &mut R) -> Result<String> {
+    let vmlinuz = uki::get_section_buffered(uki_reader, ".linux").context(".linux not present")?;
+    uki_reader
+        .seek(SeekFrom::Start(0))
+        .context("Moving seek to 0")?;
+    let initramfs =
+        uki::get_section_buffered(uki_reader, ".initrd").context(".initrd not present")?;
 
     let mut hasher = openssl::hash::Hasher::new(openssl::hash::MessageDigest::sha256())
         .context("Creating hasher")?;
@@ -796,17 +796,23 @@ fn write_pe_to_esp(
     missing_fsverity_allowed: bool,
     mounted_efi: impl AsRef<Path>,
 ) -> Result<Option<UKIInfo>> {
-    let efi_bin = read_file(file, &repo).context("Reading .efi binary")?;
+    let mut uki_reader = match file {
+        RegularFile::Inline(..) => {
+            // UKI/Addons would always be large enough to be an external object
+            anyhow::bail!("File too small to be UKI/Addon")
+        }
+        RegularFile::External(id, ..) => std::fs::File::from(repo.open_object(id)?),
+    };
 
     let mut boot_label: Option<UKIInfo> = None;
 
     // UKI Extension might not even have a cmdline
     // TODO: UKI Addon might also have a composefs= cmdline?
     if matches!(pe_type, PEType::Uki) {
-        let cmdline = uki::get_cmdline(&efi_bin).context("Getting UKI cmdline")?;
+        let cmdline = uki::get_cmdline_buffered(&mut uki_reader).context("Getting UKI cmdline")?;
 
         let (composefs_cmdline, missing_verity_allowed_cmdline) =
-            get_cmdline_composefs::<Sha512HashValue>(cmdline).context("Parsing composefs=")?;
+            get_cmdline_composefs::<Sha512HashValue>(&cmdline).context("Parsing composefs=")?;
 
         // If the UKI cmdline does not match what the user has passed as cmdline option
         // NOTE: This will only be checked for new installs and now upgrades/switches
@@ -830,14 +836,18 @@ fn write_pe_to_esp(
             );
         }
 
-        let osrel = uki::get_text_section(&efi_bin, ".osrel")?;
+        uki_reader.seek(SeekFrom::Start(0))?;
+        let osrel = uki::get_text_section_buffered(&mut uki_reader, ".osrel")?;
 
-        let parsed_osrel = OsReleaseInfo::parse(osrel);
+        let parsed_osrel = OsReleaseInfo::parse(&osrel);
 
-        let boot_digest = compute_boot_digest_uki(&efi_bin)?;
+        uki_reader.seek(SeekFrom::Start(0))?;
+        let boot_digest = compute_boot_digest_uki(&mut uki_reader)?;
 
+        uki_reader.seek(SeekFrom::Start(0))?;
         boot_label = Some(UKIInfo {
-            boot_label: uki::get_boot_label(&efi_bin).context("Getting UKI boot label")?,
+            boot_label: uki::get_boot_label_buffered(&mut uki_reader)
+                .context("Getting UKI boot label")?,
             version: parsed_osrel.get_version(),
             os_id: parsed_osrel.get_value(&["ID"]),
             boot_digest,
@@ -883,8 +893,9 @@ fn write_pe_to_esp(
             .as_str(),
     };
 
+    uki_reader.seek(SeekFrom::Start(0))?;
     pe_dir
-        .atomic_write(pe_name, efi_bin)
+        .atomic_replace_with(pe_name, |writer| std::io::copy(&mut uki_reader, writer))
         .context("Writing UKI")?;
 
     rustix::fs::fsync(

crates/lib/src/bootc_composefs/utils.rs

@@ -10,7 +10,7 @@ use bootc_kernel_cmdline::utf8::Cmdline;
 use composefs_ctl::composefs_boot;
 use fn_error_context::context;
 
-fn get_uki(storage: &Storage, deployment_verity: &str) -> Result<Vec<u8>> {
+fn get_uki(storage: &Storage, deployment_verity: &str) -> Result<cap_std_ext::cap_std::fs::File> {
     let uki_dir = storage.require_esp()?.fd.open_dir(BOOTC_UKI_DIR)?;
 
     let req_fname = get_uki_name(deployment_verity);
@@ -24,7 +24,7 @@ fn get_uki(storage: &Storage, deployment_verity: &str) -> Result<Vec<u8>> {
             continue;
         }
 
-        return Ok(uki_dir.read(filename)?);
+        return Ok(uki_dir.open(filename)?);
     }
 
     anyhow::bail!("UKI for deployment {deployment_verity} not found")
@@ -35,8 +35,8 @@ pub(crate) fn compute_store_boot_digest_for_uki(
     storage: &Storage,
     deployment_verity: &str,
 ) -> Result<String> {
-    let uki = get_uki(storage, deployment_verity)?;
-    let digest = compute_boot_digest_uki(&uki)?;
+    let mut uki = get_uki(storage, deployment_verity)?;
+    let digest = compute_boot_digest_uki(&mut uki)?;
 
     update_boot_digest_in_origin(storage, &deployment_verity, &digest)?;
     return Ok(digest);
@@ -47,8 +47,8 @@ pub(crate) fn get_uki_cmdline(
     storage: &Storage,
     deployment_verity: &str,
 ) -> Result<Cmdline<'static>> {
-    let uki = get_uki(storage, deployment_verity)?;
-    let cmdline = composefs_boot::uki::get_cmdline(&uki)?;
+    let mut uki = get_uki(storage, deployment_verity)?;
+    let cmdline = composefs_boot::uki::get_cmdline_buffered(&mut uki)?;
 
     return Ok(Cmdline::from(cmdline.to_owned()));
 }

crates/lib/src/kernel.rs

@@ -75,18 +75,18 @@ pub(crate) fn find_kernel(root: &Dir) -> Result<Option<KernelInternal>> {
     if let Some(uki_path) = find_uki_path(root)? {
         let version = uki_path.file_stem().unwrap_or(uki_path.as_str()).to_owned();
 
-        let uki = root.read(&uki_path).context("Reading UKI")?;
+        let mut uki = root.open(&uki_path).context("Opening UKI")?;
 
         // Best effort to check for composefs=?verity in the UKI cmdline
-        let cmdline = composefs_boot::uki::get_section(&uki, ".cmdline");
+        let cmdline = composefs_boot::uki::get_section_buffered(&mut uki, ".cmdline");
 
         let cmdline = match cmdline {
-            Some(Ok(cmdline)) => {
-                let cmdline_str = std::str::from_utf8(cmdline)?;
+            Ok(cmdline) => {
+                let cmdline_str = std::str::from_utf8(&cmdline)?;
                 Some(Cmdline::from(cmdline_str.to_owned()))
             }
 
-            Some(Err(uki_error)) => match uki_error {
+            Err(uki_error) => match uki_error {
                 composefs_boot::uki::UkiError::MissingSection(_) => {
                     // TODO(Johan-Liebert1): Check this when we have full UKI Addons support
                     // The cmdline might be in an addon, so don't allow missing verity
@@ -95,8 +95,6 @@ pub(crate) fn find_kernel(root: &Dir) -> Result<Option<KernelInternal>> {
 
                 e => anyhow::bail!("Failed to read UKI cmdline: {e:?}"),
             },
-
-            None => None,
         };
 
         return Ok(Some(KernelInternal {