From 4c4b484d5573511fa934d0db8cccb7e89c88afc7 Mon Sep 17 00:00:00 2001
From: William Vinnicombe <william.vinnicombe@raspberrypi.com>
Date: Tue, 14 Jan 2025 17:17:42 +0000
Subject: [PATCH] Improve checking for malicious flash data

Add data_max_size to prevent overwriting the bootloader with data from flash
---
 bootloaders/encrypted/enc_bootloader.c | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/bootloaders/encrypted/enc_bootloader.c b/bootloaders/encrypted/enc_bootloader.c
index dc828d62b..02f81ff09 100644
--- a/bootloaders/encrypted/enc_bootloader.c
+++ b/bootloaders/encrypted/enc_bootloader.c
@@ -92,31 +92,35 @@ int main() {
 
     rc = rom_get_partition_table_info((uint32_t*)workarea, 0x8, PT_INFO_PARTITION_LOCATION_AND_FLAGS | PT_INFO_SINGLE_PARTITION | (boot_partition << 24));
 
-    uint32_t data_start_addr;
-    uint32_t data_end_addr;
+    uint32_t data_start_addr = 0;
+    uint32_t data_end_addr = 0;
+    uint32_t data_max_size = 0;
     if (rc != 3) {
         printf("No boot partition - assuming bin at start of flash\n");
         data_start_addr = 0;
         data_end_addr = 0x70000; // must fit into 0x20000000 -> 0x20070000
+        data_max_size = data_end_addr - data_start_addr;
     } else {
         uint16_t first_sector_number = (((uint32_t*)workarea)[1] & PICOBIN_PARTITION_LOCATION_FIRST_SECTOR_BITS) >> PICOBIN_PARTITION_LOCATION_FIRST_SECTOR_LSB;
         uint16_t last_sector_number = (((uint32_t*)workarea)[1] & PICOBIN_PARTITION_LOCATION_LAST_SECTOR_BITS) >> PICOBIN_PARTITION_LOCATION_LAST_SECTOR_LSB;
         data_start_addr = first_sector_number * 0x1000;
         data_end_addr = (last_sector_number + 1) * 0x1000;
+        data_max_size = data_end_addr - data_start_addr;
 
-        printf("Partition Start %x, End %x\n", data_start_addr, data_end_addr);
+        printf("Partition Start %x, End %x, Max Size %x\n", data_start_addr, data_end_addr, data_max_size);
     }
 
     printf("Decrypting the chosen image\n");
     uint32_t first_mb_start = 0;
+    bool first_mb_start_found = false;
     uint32_t first_mb_end = 0;
     uint32_t last_mb_start = 0;
-    for (uint16_t i=0; i <= 0x1000; i += 4) {
+    for (uint16_t i=0; i < 0x1000; i += 4) {
         if (*(uint32_t*)(XIP_BASE + data_start_addr + i) == 0xffffded3) {
             printf("Found first block start\n");
             first_mb_start = i;
-        }
-        if (*(uint32_t*)(XIP_BASE + data_start_addr + i) == 0xab123579) {
+            first_mb_start_found = true;
+        } else if (first_mb_start_found && (*(uint32_t*)(XIP_BASE + data_start_addr + i) == 0xab123579)) {
             printf("Found first block end\n");
             first_mb_end = i + 4;
             last_mb_start = *(uint32_t*)(XIP_BASE + data_start_addr + i-4) + first_mb_start;
@@ -124,6 +128,12 @@ int main() {
         }
     }
 
+    if (last_mb_start > data_max_size) {
+        // todo - harden this check
+        printf("ERROR: Encrypted binary is too big for it's partition - resetting\n");
+        reset_usb_boot(0, 0);
+    }
+
     if (*(uint32_t*)(XIP_BASE + data_start_addr + last_mb_start) == 0xffffded3) {
         printf("Found last block start where expected\n");
     } else {