From 88a72ae0ca78da55bac3044390ac8d5e3f684cdd Mon Sep 17 00:00:00 2001
From: Christopher Dombroski <chris@cdombroski.com>
Date: Sun, 29 Dec 2024 14:37:28 -0800
Subject: [PATCH] Calculate end of bootloader to update SHA256

The current implementation uses the last 32 bytes of the bootloader
file. When Secure Boot V2 is enabled, the bootloader is padded. The
new implementation walks through the segments to find the end and adds
the 16-byte aligned 1-byte checksum to update the SHA256 instead of
incorrectly updating the padding.

Closes #715
---
 espflash/src/image_format.rs | 32 ++++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/espflash/src/image_format.rs b/espflash/src/image_format.rs
index 2af83ece..eea82fc0 100644
--- a/espflash/src/image_format.rs
+++ b/espflash/src/image_format.rs
@@ -135,11 +135,23 @@ impl<'a> IdfBootloaderFormat<'a> {
         };
 
         // fetch the generated header from the bootloader
-        let mut header: ImageHeader = *from_bytes(&bootloader[0..size_of::<ImageHeader>()]);
+        let mut calc_bootloader_size = 0;
+        let bootloader_header_size = size_of::<ImageHeader>();
+        calc_bootloader_size += bootloader_header_size;
+        let mut header: ImageHeader = *from_bytes(&bootloader[0..bootloader_header_size]);
         if header.magic != ESP_MAGIC {
             return Err(Error::InvalidBootloader);
         }
 
+        for _ in 0..header.segment_count {
+            let segment: SegmentHeader = *from_bytes(
+                &bootloader
+                    [calc_bootloader_size..calc_bootloader_size + size_of::<SegmentHeader>()],
+            );
+            println!("{:X?}", segment);
+            calc_bootloader_size += segment.length as usize + size_of::<SegmentHeader>();
+        }
+
         // update the header if a user has specified any custom arguments
         if let Some(mode) = flash_settings.mode {
             header.flash_mode = mode as u8;
@@ -157,11 +169,23 @@ impl<'a> IdfBootloaderFormat<'a> {
         );
 
         // re-calculate hash of the bootloader - needed since we modified the header
-        let bootloader_len = bootloader.len();
+        // the hash is at the end of the bootloader, but the bootloader bytes are padded.
+        // the real end of the bootloader is the end of the segments plus the 16-byte aligned 1-byte checksum.
+        let bootloader_sha_start =
+            calc_bootloader_size + (16 - ((calc_bootloader_size + 1) % 16)) + 1;
+        println!("{}", bootloader_sha_start);
         let mut hasher = Sha256::new();
-        hasher.update(&bootloader[..bootloader_len - 32]);
+        hasher.update(&bootloader[..bootloader_sha_start]);
         let hash = hasher.finalize();
-        bootloader.to_mut()[bootloader_len - 32..].copy_from_slice(&hash);
+        println!(
+            "before: {:X?}",
+            &bootloader[bootloader_sha_start..bootloader_sha_start + 32]
+        );
+        bootloader.to_mut()[bootloader_sha_start..bootloader_sha_start + 32].copy_from_slice(&hash);
+        println!(
+            "after:  {:X?}",
+            &bootloader[bootloader_sha_start..bootloader_sha_start + 32]
+        );
 
         // write the header of the app
         // use the same settings as the bootloader