--capsule: Find any number of PD binaries in capsule

The PDs may not be all the same version, so we should show them all.

Signed-off-by: Daniel Schaefer <dhs@frame.work>

Author: JohnAZoidberg

framework_lib/src/capsule_content.rs

@@ -6,6 +6,7 @@
 
 use crate::alloc::string::ToString;
 use alloc::string::String;
+use alloc::vec::Vec;
 use core::convert::TryInto;
 
 use crate::ccgx::binary::{CCG5_PD_LEN, CCG6_PD_LEN, CCG8_PD_LEN};
@@ -51,19 +52,38 @@ pub fn find_ec_in_bios_cap(data: &[u8]) -> Option<&[u8]> {
 }
 
 pub fn find_pd_in_bios_cap(data: &[u8]) -> Option<&[u8]> {
-    // Just search for the first couple of bytes in PD binaries
-    // TODO: There's a second one but unless the capsule is bad, we can assume
-    // they're the same version
-    let ccg5_needle = &[0x00, 0x20, 0x00, 0x20, 0x11, 0x00];
-    let ccg6_needle = &[0x00, 0x40, 0x00, 0x20, 0x11, 0x00];
-    let ccg8_needle = &[0x00, 0x80, 0x00, 0x20, 0xAD, 0x0C];
-    if let Some(found_pd1) = util::find_sequence(data, ccg5_needle) {
-        Some(&data[found_pd1..found_pd1 + CCG5_PD_LEN])
-    } else if let Some(found_pd1) = util::find_sequence(data, ccg6_needle) {
-        Some(&data[found_pd1..found_pd1 + CCG6_PD_LEN])
-    } else if let Some(found_pd1) = util::find_sequence(data, ccg8_needle) {
-        Some(&data[found_pd1..found_pd1 + CCG8_PD_LEN])
-    } else {
-        None
+    find_all_pds_in_bios_cap(data).into_iter().next()
+}
+
+/// PD binary signatures and their corresponding lengths
+const CCG5_NEEDLE: &[u8] = &[0x00, 0x20, 0x00, 0x20, 0x11, 0x00];
+const CCG6_NEEDLE: &[u8] = &[0x00, 0x40, 0x00, 0x20, 0x11, 0x00];
+const CCG8_NEEDLE: &[u8] = &[0x00, 0x80, 0x00, 0x20, 0xAD, 0x0C];
+
+/// Find all PD firmware binaries embedded in a BIOS capsule
+pub fn find_all_pds_in_bios_cap(data: &[u8]) -> Vec<&[u8]> {
+    let mut results = Vec::new();
+
+    // Search for CCG5 PDs
+    for offset in util::find_all_sequences(data, CCG5_NEEDLE) {
+        if offset + CCG5_PD_LEN <= data.len() {
+            results.push(&data[offset..offset + CCG5_PD_LEN]);
+        }
     }
+
+    // Search for CCG6 PDs
+    for offset in util::find_all_sequences(data, CCG6_NEEDLE) {
+        if offset + CCG6_PD_LEN <= data.len() {
+            results.push(&data[offset..offset + CCG6_PD_LEN]);
+        }
+    }
+
+    // Search for CCG8 PDs
+    for offset in util::find_all_sequences(data, CCG8_NEEDLE) {
+        if offset + CCG8_PD_LEN <= data.len() {
+            results.push(&data[offset..offset + CCG8_PD_LEN]);
+        }
+    }
+
+    results
 }

framework_lib/src/commandline/mod.rs

@@ -28,7 +28,7 @@ use crate::built_info;
 use crate::camera::check_camera_version;
 use crate::capsule;
 use crate::capsule_content::{
-    find_bios_version, find_ec_in_bios_cap, find_pd_in_bios_cap, find_retimer_version,
+    find_all_pds_in_bios_cap, find_bios_version, find_ec_in_bios_cap, find_retimer_version,
 };
 use crate::ccgx::device::{FwMode, PdController, PdPort};
 #[cfg(feature = "hidapi")]
@@ -1691,9 +1691,10 @@ pub fn run_with_args(args: &Cli, _allupdate: bool) -> i32 {
                     println!("Embedded EC");
                     analyze_ec_fw(ec_bin);
                 }
-                if let Some(pd_bin) = find_pd_in_bios_cap(&data) {
+                let pd_bins = find_all_pds_in_bios_cap(&data);
+                for (i, pd_bin) in pd_bins.iter().enumerate() {
                     found_any = true;
-                    println!("Embedded PD");
+                    println!("Embedded PD {}", i + 1);
                     analyze_ccgx_pd_fw(pd_bin);
                 }
                 if !found_any {
@@ -2436,8 +2437,9 @@ pub fn analyze_capsule(data: &[u8]) -> Option<capsule::EfiCapsuleHeader> {
                     println!("  RW Version:   {:>18}", ver.version);
                 }
             }
-            if let Some(pd_bin) = find_pd_in_bios_cap(data) {
-                println!("Embedded PD");
+            let pd_bins = find_all_pds_in_bios_cap(data);
+            for (i, pd_bin) in pd_bins.iter().enumerate() {
+                println!("Embedded PD {}", i + 1);
                 analyze_ccgx_pd_fw(pd_bin);
             }
         }

framework_lib/src/util.rs

@@ -258,6 +258,22 @@ pub fn find_sequence(haystack: &[u8], needle: &[u8]) -> Option<usize> {
         .position(|window| window == needle)
 }
 
+/// Find all occurrences of a sequence of bytes in a long slice of bytes
+pub fn find_all_sequences(haystack: &[u8], needle: &[u8]) -> Vec<usize> {
+    let mut results = Vec::new();
+    let mut pos = 0;
+    while pos + needle.len() <= haystack.len() {
+        if let Some(offset) = find_sequence(&haystack[pos..], needle) {
+            let absolute_pos = pos + offset;
+            results.push(absolute_pos);
+            pos = absolute_pos + needle.len();
+        } else {
+            break;
+        }
+    }
+    results
+}
+
 /// Assert length of an EC response from the windows driver
 /// It's always 20 more than expected. TODO: Figure out why
 pub fn assert_win_len<N: Num + std::fmt::Debug + Ord + NumCast + Copy>(left: N, right: N) {