Implement buffered readers for UKI

contrib/packaging/lib.sh

@@ -33,9 +33,13 @@ pkg_install() {
             dnf clean all
             ;;
         debian|ubuntu)
+            export DEBIAN_FRONTEND=noninteractive
             debian_apt_init
             apt-get -o APT::Sandbox::User=root update
-            apt-get -o APT::Sandbox::User=root install -y --no-install-recommends "$@"
+            apt-get -o APT::Sandbox::User=root \
+                    -o Dpkg::Options::="--force-confold" \
+                    -o Dpkg::Options::="--force-confdef" \
+                install -y --no-install-recommends "$@"
             rm -rf /var/lib/apt/lists/*
             ;;
         *)

crates/composefs-boot/src/uki.rs

@@ -5,6 +5,7 @@
 //! Specification Type 2 requirements for UKI boot entries, including extraction of boot
 //! labels from os-release information embedded in the UKI binary.
 
+use std::io::{Read, Seek, SeekFrom};
 use thiserror::Error;
 use zerocopy::{
     FromBytes, Immutable, KnownLayout,
@@ -62,17 +63,20 @@ struct SectionHeader {
 }
 
 /// Errors that can occur when parsing UKI files.
-#[derive(Debug, Error, PartialEq)]
+#[derive(Debug, Error)]
 pub enum UkiError {
+    /// IO Error while reading or seeking
+    #[error("IO Error")]
+    Io(#[from] std::io::Error),
     /// The file is not a valid Portable Executable (PE/EFI) format
     #[error("UKI is not valid EFI executable")]
     PortableExecutableError,
     /// A required PE section is missing from the UKI
     #[error("UKI doesn't contain a '{0}' section")]
-    MissingSection(&'static str),
+    MissingSection(String),
     /// A PE section contains invalid UTF-8
     #[error("UKI section '{0}' is not UTF-8")]
-    UnicodeError(&'static str),
+    UnicodeError(String),
     /// The .osrel section lacks name information
     #[error("No name information found in .osrel section")]
     NoName,
@@ -97,7 +101,19 @@ pub fn get_text_section<'a>(
     section_name: &'static str,
 ) -> Result<&'a str, UkiError> {
     let bytes = get_section(image, section_name).ok_or(UkiError::PortableExecutableError)??;
-    std::str::from_utf8(bytes).or(Err(UkiError::UnicodeError(section_name)))
+    std::str::from_utf8(bytes).or(Err(UkiError::UnicodeError(section_name.into())))
+}
+
+/// Buffered version of [`get_text_section`].
+///
+/// See [`get_text_section`] for details. This version works with any [`Read`] + [`Seek`]
+/// source instead of requiring the entire image in memory.
+pub fn get_text_section_buffered<'a, R: Read + Seek>(
+    image: &'a mut R,
+    section_name: &'a str,
+) -> Result<String, UkiError> {
+    let bytes = get_section_buffered(image, section_name)?;
+    String::from_utf8(bytes).or(Err(UkiError::UnicodeError(section_name.into())))
 }
 
 /// Extracts a raw section from a UKI PE file by name.
@@ -158,7 +174,64 @@ pub fn get_section<'a>(
         }
     }
 
-    Some(Err(UkiError::MissingSection(section_name)))
+    Some(Err(UkiError::MissingSection(section_name.into())))
+}
+
+/// Buffered version of [`get_section`].
+///
+/// See [`get_section`] for details. This version works with any [`Read`] + [`Seek`]
+/// source and returns owned data instead of borrowed slices.
+pub fn get_section_buffered<R: Read + Seek>(
+    image: &mut R,
+    section_name: &str,
+) -> Result<Vec<u8>, UkiError> {
+    use std::io::Error as IOError;
+
+    // Turn the section_name ".osrel" into a section_key b".osrel\0\0".
+    // This will panic if section_name.len() > 8, which is what we want.
+    let mut section_key = [0u8; 8];
+    section_key[..section_name.len()].copy_from_slice(section_name.as_bytes());
+
+    // Skip the DOS stub
+    let mut buf: Vec<u8> = vec![0; std::mem::size_of::<DosStub>()];
+    image.read_exact(&mut buf)?;
+    let dos_stub =
+        DosStub::ref_from_bytes(&buf).map_err(|e| UkiError::Io(IOError::other(e.to_string())))?;
+    image.seek(SeekFrom::Start(dos_stub.pe_offset.get() as u64))?;
+
+    // Get the PE header
+    let mut buf: Vec<u8> = vec![0; std::mem::size_of::<PeHeader>()];
+    image.read_exact(&mut buf)?;
+    let pe_header =
+        PeHeader::ref_from_bytes(&buf).map_err(|e| UkiError::Io(IOError::other(e.to_string())))?;
+    if pe_header.pe_magic != PE_MAGIC {
+        return Err(UkiError::PortableExecutableError);
+    }
+
+    // Skip the optional header
+    image.seek(SeekFrom::Current(
+        pe_header.coff_file_header.size_of_optional_header.get() as i64,
+    ))?;
+
+    // Try to load the section headers
+    let n_sections = pe_header.coff_file_header.number_of_sections.get() as usize;
+    let mut sections = vec![0; std::mem::size_of::<SectionHeader>() * n_sections];
+    image.read_exact(&mut sections)?;
+    let sections = <[SectionHeader]>::ref_from_bytes_with_elems(&sections, n_sections)
+        .map_err(|e| UkiError::Io(IOError::other(e.to_string())))?;
+
+    for section in sections {
+        if section.name != section_key {
+            continue;
+        }
+
+        let mut buffer = vec![0; section.virtual_size.get() as usize];
+        image.seek(SeekFrom::Start(section.pointer_to_raw_data.get() as u64))?;
+        image.read_exact(&mut buffer)?;
+        return Ok(buffer);
+    }
+
+    Err(UkiError::MissingSection(section_name.to_string()))
 }
 
 /// Gets an appropriate label for display in the boot menu for the given UKI image, according to
@@ -189,11 +262,26 @@ pub fn get_boot_label(image: &[u8]) -> Result<String, UkiError> {
         .ok_or(UkiError::NoName)
 }
 
+/// Buffered version of [`get_boot_label`].
+///
+/// See [`get_boot_label`] for details. This version works with any [`Read`] + [`Seek`] source.
+pub fn get_boot_label_buffered<R: Read + Seek>(image: &mut R) -> Result<String, UkiError> {
+    let osrel = get_text_section_buffered(image, ".osrel")?;
+    OsReleaseInfo::parse(&osrel)
+        .get_boot_label()
+        .ok_or(UkiError::NoName)
+}
+
 /// Gets the contents of the .cmdline section of a UKI.
 pub fn get_cmdline(image: &[u8]) -> Result<&str, UkiError> {
     get_text_section(image, ".cmdline")
 }
 
+/// Buffered version of [`get_cmdline`]. See [`get_cmdline`] for details.
+pub fn get_cmdline_buffered<R: Read + Seek>(image: &mut R) -> Result<String, UkiError> {
+    get_text_section_buffered(image, ".cmdline")
+}
+
 #[cfg(test)]
 mod test {
     use core::mem::size_of;
@@ -264,19 +352,40 @@ ID=pretty-os
 "#,
         );
 
+        // Test slice-based functions
         assert_eq!(
             get_boot_label(uki.as_ref()).unwrap(),
             "prettyOS Rocky Racoon"
         );
+
+        // Test buffered functions produce same results
+        let mut cursor = std::io::Cursor::new(&uki);
+        assert_eq!(
+            get_boot_label_buffered(&mut cursor).unwrap(),
+            "prettyOS Rocky Racoon"
+        );
     }
 
     #[test]
     fn test_bad_pe() {
         fn pe_err(img: &[u8]) {
-            assert_eq!(get_boot_label(img), Err(UkiError::PortableExecutableError));
+            assert!(matches!(
+                get_boot_label(img),
+                Err(UkiError::PortableExecutableError)
+            ));
         }
         fn no_sec(img: &[u8]) {
-            assert_eq!(get_boot_label(img), Err(UkiError::MissingSection(".osrel")));
+            assert!(matches!(
+                get_boot_label(img),
+                Err(UkiError::MissingSection(s)) if s == ".osrel"
+            ));
+
+            // Test buffered version
+            let mut cursor = std::io::Cursor::new(img);
+            assert!(matches!(
+                get_boot_label_buffered(&mut cursor),
+                Err(UkiError::MissingSection(s)) if s == ".osrel"
+            ));
         }
 
         pe_err(b"");
@@ -319,4 +428,93 @@ ID=pretty-os
             &[],
         ));
     }
+
+    #[test]
+    fn test_section_functions() {
+        let osrel_data = b"PRETTY_NAME='TestOS'\nVERSION_ID=1.0\n";
+        let cmdline_data = b"root=/dev/sda1 quiet";
+
+        let osrel_offset = data_offset(2);
+        let cmdline_offset = osrel_offset + osrel_data.len();
+
+        let uki = peify(
+            b"",
+            &[
+                SectionHeader {
+                    name: *b".osrel\0\0",
+                    virtual_size: U32::new(osrel_data.len() as u32),
+                    pointer_to_raw_data: U32::new(osrel_offset as u32),
+                    ..Default::default()
+                },
+                SectionHeader {
+                    name: *b".cmdline",
+                    virtual_size: U32::new(cmdline_data.len() as u32),
+                    pointer_to_raw_data: U32::new(cmdline_offset as u32),
+                    ..Default::default()
+                },
+            ],
+            &[osrel_data, cmdline_data],
+        );
+
+        // Test slice-based functions
+        let osrel_section = get_section(&uki, ".osrel").unwrap().unwrap();
+        assert_eq!(osrel_section, osrel_data);
+
+        let cmdline_section = get_section(&uki, ".cmdline").unwrap().unwrap();
+        assert_eq!(cmdline_section, cmdline_data);
+
+        let osrel_text = get_text_section(&uki, ".osrel").unwrap();
+        assert_eq!(osrel_text, "PRETTY_NAME='TestOS'\nVERSION_ID=1.0\n");
+
+        let cmdline_text = get_cmdline(&uki).unwrap();
+        assert_eq!(cmdline_text, "root=/dev/sda1 quiet");
+
+        // Test buffered functions produce same results
+        let mut cursor = std::io::Cursor::new(&uki);
+        let osrel_section_buf = get_section_buffered(&mut cursor, ".osrel").unwrap();
+        assert_eq!(osrel_section_buf, osrel_data);
+
+        cursor.set_position(0);
+        let cmdline_section_buf = get_section_buffered(&mut cursor, ".cmdline").unwrap();
+        assert_eq!(cmdline_section_buf, cmdline_data);
+
+        cursor.set_position(0);
+        let osrel_text_buf = get_text_section_buffered(&mut cursor, ".osrel").unwrap();
+        assert_eq!(osrel_text_buf, "PRETTY_NAME='TestOS'\nVERSION_ID=1.0\n");
+
+        cursor.set_position(0);
+        let cmdline_text_buf = get_cmdline_buffered(&mut cursor).unwrap();
+        assert_eq!(cmdline_text_buf, "root=/dev/sda1 quiet");
+
+        // Test missing section
+        cursor.set_position(0);
+        let missing_result = get_section_buffered(&mut cursor, ".missing");
+        assert!(matches!(missing_result, Err(UkiError::MissingSection(s)) if s == ".missing"));
+    }
+
+    #[test]
+    fn test_invalid_utf8() {
+        let invalid_utf8 = b"\xff\xfe\xfd";
+        let osrel_offset = data_offset(1);
+
+        let uki = peify(
+            b"",
+            &[SectionHeader {
+                name: *b".osrel\0\0",
+                virtual_size: U32::new(invalid_utf8.len() as u32),
+                pointer_to_raw_data: U32::new(osrel_offset as u32),
+                ..Default::default()
+            }],
+            &[invalid_utf8],
+        );
+
+        // Test slice-based function
+        let result = get_text_section(&uki, ".osrel");
+        assert!(matches!(result, Err(UkiError::UnicodeError(s)) if s == ".osrel"));
+
+        // Test buffered function gives same error
+        let mut cursor = std::io::Cursor::new(&uki);
+        let result_buf = get_text_section_buffered(&mut cursor, ".osrel");
+        assert!(matches!(result_buf, Err(UkiError::UnicodeError(s)) if s == ".osrel"));
+    }
 }