Introduce EifSector to represent coupled header and buffer

Signed-off-by: Maayan Keshet <[email protected]>

Author: maayank

eif_extract/src/main.rs

@@ -19,11 +19,11 @@ fn extract_ramdisks(eif_path: &str, output_dir: &str, prefix: &str) -> io::Resul
 
     println!("Reading section data...");
     for section_result in sections {
-        let (section, _, data) = section_result.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
-        if section.section_type == EifSectionType::EifSectionRamdisk {
+        let section = section_result.map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
+        if section.section_header.section_type == EifSectionType::EifSectionRamdisk {
             let output_file_path = format!("{}/{}{}.dat", output_dir, prefix, ramdisk_count);
             let mut output_file = File::create(&output_file_path)?;
-            output_file.write_all(&data)?;
+            output_file.write_all(&section.section_data)?;
             println!("Saved ramdisk to {}", output_file_path);
             ramdisk_count += 1;
         }

src/defs/mod.rs

@@ -222,6 +222,14 @@ impl EifSectionHeader {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct EifSection {
+    /// The header of the EIF section
+    pub section_header: EifSectionHeader,
+    /// Buffer containing the section data
+    pub section_data: Vec<u8>,
+}
+
 /// Array containing the signatures of at least one PCR.
 /// For now, it only contains the signature of PRC0.
 pub type EifSignature = Vec<PcrSignature>;

src/utils/eif_reader.rs

@@ -3,7 +3,8 @@
 
 use crate::defs::eif_hasher::EifHasher;
 use crate::defs::{
-    EifHeader, EifIdentityInfo, EifSectionHeader, EifSectionType, PcrInfo, PcrSignature,
+    EifHeader, EifIdentityInfo, EifSectionType, EifSectionHeader, EifSection, PcrInfo,
+    PcrSignature,
 };
 use aws_nitro_enclaves_cose::{crypto::Openssl, CoseSign1};
 use crc::{Crc, CRC_32_ISO_HDLC};
@@ -75,33 +76,36 @@ impl Sections {
 }
 
 impl Iterator for Sections {
-    type Item = Result<(EifSectionHeader, Vec<u8>, Vec<u8>), String>; // Header, HeaderBuf, Buf
+    type Item = Result<EifSection, String>; // Changed to EifSection
 
     fn next(&mut self) -> Option<Self::Item> {
         let mut section_buf = vec![0u8; EifSectionHeader::size()];
         if self.eif_file.read_exact(&mut section_buf).is_err() {
             return None;
         }
 
-        let section = match EifSectionHeader::from_be_bytes(&section_buf) {
+        let section_header = match EifSectionHeader::from_be_bytes(&section_buf) {
             Ok(sec) => sec,
             Err(e) => return Some(Err(format!("Error extracting EIF section header: {:?}", e))),
         };
 
-        let mut buf = vec![0u8; section.section_size as usize];
+        let mut section_data = vec![0u8; section_header.section_size as usize];
         self.curr_seek += EifSectionHeader::size() as u64;
         if self.eif_file.seek(SeekFrom::Start(self.curr_seek)).is_err() {
             return Some(Err("Failed to seek after EIF header".to_string()));
         }
-        if self.eif_file.read_exact(&mut buf).is_err() {
+        if self.eif_file.read_exact(&mut section_data).is_err() {
             return Some(Err("Error while reading section from EIF".to_string()));
         }
-        self.curr_seek += section.section_size as u64;
+        self.curr_seek += section_header.section_size as u64;
         if self.eif_file.seek(SeekFrom::Start(self.curr_seek)).is_err() {
             return Some(Err("Failed to seek after EIF section".to_string()));
         }
 
-        Some(Ok((section, section_buf, buf)))
+        Some(Ok(EifSection {
+            section_header,
+            section_data,
+        }))
     }
 }
 
@@ -165,41 +169,40 @@ impl EifReader {
 
         // Read all sections and treat by type
         for section_result in sections {
-            let (section, section_buf, buf) = section_result
-                .map_err(|e| e.to_string())?;
-            eif_crc.update(&section_buf);
+            let section = section_result.map_err(|e| e.to_string())?;
+            eif_crc.update(&section.section_header.to_be_bytes());
 
-            match section.section_type {
+            match section.section_header.section_type {
                 EifSectionType::EifSectionKernel | EifSectionType::EifSectionCmdline => {
-                    image_hasher.write_all(&buf).map_err(|e| {
+                    image_hasher.write_all(&section.section_data).map_err(|e| {
                         format!("Failed to write EIF section to image_hasher: {:?}", e)
                     })?;
-                    bootstrap_hasher.write_all(&buf).map_err(|e| {
+                    bootstrap_hasher.write_all(&section.section_data).map_err(|e| {
                         format!("Failed to write EIF section to bootstrap_hasher: {:?}", e)
                     })?;
                 }
                 EifSectionType::EifSectionRamdisk => {
-                    image_hasher.write_all(&buf).map_err(|e| {
+                    image_hasher.write_all(&section.section_data).map_err(|e| {
                         format!("Failed to write ramdisk section to image_hasher: {:?}", e)
                     })?;
                     if ramdisk_idx == 0 {
-                        bootstrap_hasher.write_all(&buf).map_err(|e| {
+                        bootstrap_hasher.write_all(&section.section_data).map_err(|e| {
                             format!(
                                 "Failed to write ramdisk section to bootstrap_hasher: {:?}",
                                 e
                             )
                         })?;
                     } else {
-                        app_hasher.write_all(&buf).map_err(|e| {
+                        app_hasher.write_all(&section.section_data).map_err(|e| {
                             format!("Failed to write ramdisk section to app_hasher: {:?}", e)
                         })?;
                     }
                     ramdisk_idx += 1;
                 }
                 EifSectionType::EifSectionSignature => {
-                    signature_section = Some(buf.clone());
+                    signature_section = Some(section.section_data.clone());
                     // Deserialize PCR0 signature structure and write it to the hasher
-                    let des_sign: Vec<PcrSignature> = from_slice(&buf[..])
+                    let des_sign: Vec<PcrSignature> = from_slice(&section.section_data[..])
                         .map_err(|e| format!("Error deserializing certificate: {:?}", e))?;
 
                     let cert = openssl::x509::X509::from_pem(&des_sign[0].signing_certificate)
@@ -212,7 +215,7 @@ impl EifReader {
                     })?;
                 }
                 EifSectionType::EifSectionMetadata => {
-                    metadata = serde_json::from_slice(&buf[..])
+                    metadata = serde_json::from_slice(&section.section_data[..])
                         .map_err(|e| format!("Error deserializing metadata: {:?}", e))?;
                 }
                 EifSectionType::EifSectionInvalid => {