From 7a83fc758e857103ec2dbb47a399b4caebfe84dc Mon Sep 17 00:00:00 2001
From: "chris@subtlebytes.com" <chris@subtlebytes.com>
Date: Mon, 20 Jan 2025 13:39:17 -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
---
 CHANGELOG.md                 |  2 ++
 espflash/Cargo.toml          |  3 +--
 espflash/src/image_format.rs | 33 +++++++++++++++++++++++++++++----
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 445ee4ca..fef4ab93 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,6 +13,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Update the app image SHA in the correct location for padded images (#715)
+
 ### Removed
 
 ## [3.3.0] - 2025-01-13
diff --git a/espflash/Cargo.toml b/espflash/Cargo.toml
index 9dea8f38..39a3d614 100644
--- a/espflash/Cargo.toml
+++ b/espflash/Cargo.toml
@@ -43,7 +43,7 @@ directories = { version = "5.0.1", optional = true }
 env_logger = { version = "0.11.6", optional = true }
 esp-idf-part = "0.5.0"
 flate2 = "1.0.35"
-hex = { version = "0.4.3", features = ["serde"], optional = true }
+hex = { version = "0.4.3", features = ["serde"]}
 indicatif = { version = "0.17.9", optional = true }
 lazy_static = { version = "1.5.0", optional = true }
 log = "0.4.22"
@@ -79,7 +79,6 @@ cli = [
     "dep:dialoguer",
     "dep:directories",
     "dep:env_logger",
-    "dep:hex",
     "dep:indicatif",
     "dep:lazy_static",
     "dep:parse_int",
diff --git a/espflash/src/image_format.rs b/espflash/src/image_format.rs
index 2af83ece..38baac1c 100644
--- a/espflash/src/image_format.rs
+++ b/espflash/src/image_format.rs
@@ -4,6 +4,7 @@ use std::{borrow::Cow, io::Write, iter::once, mem::size_of};
 
 use bytemuck::{bytes_of, from_bytes, Pod, Zeroable};
 use esp_idf_part::{Partition, PartitionTable, Type};
+use log;
 use sha2::{Digest, Sha256};
 
 use crate::{
@@ -135,11 +136,22 @@ 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>()],
+            );
+            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,24 @@ 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.
+        // the checksum is stored in the last byte so that the file is a multiple of 16 bytes.
+        calc_bootloader_size += 1; // add checksum size
+        calc_bootloader_size = calc_bootloader_size + ((16 - (calc_bootloader_size % 16)) % 16);
+        let bootloader_sha_start = calc_bootloader_size;
+        calc_bootloader_size += 32; // add sha256 size
+        let bootloader_sha_end = calc_bootloader_size;
+
         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);
+        log::info!(
+            "Updating bootloader SHA256 from {} to {}",
+            hex::encode(&bootloader[bootloader_sha_start..bootloader_sha_end]),
+            hex::encode(hash)
+        );
+        bootloader.to_mut()[bootloader_sha_start..bootloader_sha_end].copy_from_slice(&hash);
 
         // write the header of the app
         // use the same settings as the bootloader