From cd80da58d8458e3d04572d571a5fb4e5765b0cee Mon Sep 17 00:00:00 2001
From: Mark Thompson <129641948+NotherNgineer@users.noreply.github.com>
Date: Sun, 18 Feb 2024 12:02:42 -0600
Subject: [PATCH] Improved flash image validation (#1916)

---
 .../application/apps/ui_flash_utility.cpp     | 30 ++++++++++++++++---
 .../application/apps/ui_flash_utility.hpp     |  1 +
 2 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/firmware/application/apps/ui_flash_utility.cpp b/firmware/application/apps/ui_flash_utility.cpp
index 639f707d..a91a3951 100644
--- a/firmware/application/apps/ui_flash_utility.cpp
+++ b/firmware/application/apps/ui_flash_utility.cpp
@@ -28,27 +28,49 @@ namespace ui {
 
 static const char16_t* firmware_folder = u"/FIRMWARE";
 
+// Firmware image validation
+static const char* hackrf_magic = "HACKRFFW";
+#define FIRMWARE_INFO_AREA_OFFSET 0x400
+#define FIRST_CHECKSUM_NIGHTLY 240125
+
 Thread* FlashUtilityView::thread{nullptr};
 static constexpr size_t max_filename_length = 26;
 
 bool valid_firmware_file(std::filesystem::path::string_type path) {
     File firmware_file;
+    bool require_checksum{false};
     uint32_t read_buffer[128];
-    uint32_t checksum{(uint32_t)~FLASH_EXPECTED_CHECKSUM};  // initializing to invalid checksum in case file can't be read
+    uint32_t checksum{FLASH_CHECKSUM_ERROR};  // initializing to invalid checksum in case file can't be read
+
+    static_assert((FIRMWARE_INFO_AREA_OFFSET % sizeof(read_buffer)) == 0, "Read buffer size must divide evenly into FIRMWARE_INFO_AREA_OFFSET");
 
     // test read of the whole file just to validate checksum (baseband flash code will re-read when flashing)
     auto result = firmware_file.open(path.c_str());
     if (!result.is_valid()) {
         checksum = 0;
-        for (uint32_t i = 0; i < FLASH_ROM_SIZE / sizeof(read_buffer); i++) {
+        for (uint32_t offset = 0; offset < FLASH_ROM_SIZE; offset += sizeof(read_buffer)) {
             auto readResult = firmware_file.read(&read_buffer, sizeof(read_buffer));
 
-            // if file is smaller than 1MB, assume it's a downgrade to an old FW version and ignore the checksum
             if ((!readResult) || (readResult.value() != sizeof(read_buffer))) {
-                checksum = FLASH_EXPECTED_CHECKSUM;
+                // File was smaller than 1MB:
+                // If version is such that the file SHOULD have been 1MB, call it a checksum error (otherwise say it's OK).
+                checksum = (require_checksum) ? FLASH_CHECKSUM_ERROR : FLASH_EXPECTED_CHECKSUM;
                 break;
             }
 
+            // Did we just read the firmware info area?
+            if (offset == FIRMWARE_INFO_AREA_OFFSET) {
+                // If there's no info area (missing HACKRFFW signature), it could be an ancient FW version (so skipping check)
+                if (memcmp(read_buffer, hackrf_magic, 8) == 0) {
+                    char* version_string = (char*)&read_buffer[4];
+
+                    // Require a 1MB firmware image with a valid checksum if release version >=v2.x or nightly >n_240125
+                    if (((version_string[0] == 'v') && (std::atoi(&version_string[1]) >= 2)) ||
+                        ((version_string[0] == 'n') && (version_string[1] == '_') && (std::atoi(&version_string[2]) >= FIRST_CHECKSUM_NIGHTLY)))
+                        require_checksum = true;
+                }
+            }
+
             checksum += simple_checksum((uint32_t)read_buffer, sizeof(read_buffer));
         }
     }
diff --git a/firmware/application/apps/ui_flash_utility.hpp b/firmware/application/apps/ui_flash_utility.hpp
index de03ebcc..50a97c84 100644
--- a/firmware/application/apps/ui_flash_utility.hpp
+++ b/firmware/application/apps/ui_flash_utility.hpp
@@ -35,6 +35,7 @@
 #define FLASH_ROM_SIZE 1048576
 #define FLASH_STARTING_ADDRESS 0x00000000
 #define FLASH_EXPECTED_CHECKSUM 0x00000000
+#define FLASH_CHECKSUM_ERROR 0xFFFFFFFF
 
 namespace ui {