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

Author: Christopher Dombroski

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