From ea8bab45714db939335b91b63bc03714571681a9 Mon Sep 17 00:00:00 2001
From: eddy <eddy@macmini.local>
Date: Thu, 30 Jul 2026 19:13:27 +0200
Subject: [PATCH] feat: incremental SHA256 verification before
 esp_ota_set_boot_partition

---
 src/main.cpp | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/main.cpp b/src/main.cpp
index 4acb47a..73aa58e 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -3,6 +3,7 @@
 #include <NimBLEAdvertising.h>
 
 #include <mbedtls/md.h>
+#include <mbedtls/sha256.h>
 #include <Preferences.h>
 #include <esp_ota_ops.h>
 #include <esp_partition.h>
@@ -50,6 +51,8 @@ const esp_partition_t *otaPartition = nullptr;
 uint32_t otaExpectedSize = 0;
 uint32_t otaWritten = 0;
 uint8_t otaExpectedSha[32];
+mbedtls_sha256_context otaShaCtx;
+bool otaShaStarted = false;
 
 uint8_t hmacKey[HMAC_KEY_LEN];
 uint8_t deviceUuid[ESP_UUID_LEN];
@@ -129,6 +132,10 @@ static void otaReset() {
     otaPartition = nullptr;
     otaExpectedSize = 0;
     otaWritten = 0;
+    if (otaShaStarted) {
+        mbedtls_sha256_free(&otaShaCtx);
+        otaShaStarted = false;
+    }
 }
 
 static void sendOtaNotify(NimBLECharacteristic *otaChar, uint8_t status, uint32_t seq) {
@@ -179,6 +186,9 @@ class OtaCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
             }
             otaInProgress = true;
             otaWritten = 0;
+            mbedtls_sha256_init(&otaShaCtx);
+            mbedtls_sha256_starts_ret(&otaShaCtx, 0);
+            otaShaStarted = true;
             digitalWrite(RELAY1_PIN, LOW);
             digitalWrite(RELAY2_PIN, LOW);
             relay1State = false;
@@ -191,6 +201,10 @@ class OtaCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
 
         if (len == 32 && otaInProgress) {
             memcpy(otaExpectedSha, data, 32);
+            uint8_t computedSha[32];
+            mbedtls_sha256_finish_ret(&otaShaCtx, computedSha);
+            mbedtls_sha256_free(&otaShaCtx);
+            otaShaStarted = false;
             esp_err_t err = esp_ota_end(otaHandle);
             if (err != ESP_OK) {
                 Serial.printf("OTA: end failed (%d)\n", err);
@@ -198,6 +212,12 @@ class OtaCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
                 sendOtaNotify(pChar, OTA_STATUS_ERROR, 0);
                 return;
             }
+            if (memcmp(computedSha, otaExpectedSha, 32) != 0) {
+                Serial.println("OTA: SHA256 mismatch — aborting");
+                otaReset();
+                sendOtaNotify(pChar, OTA_STATUS_ERROR, 0);
+                return;
+            }
             err = esp_ota_set_boot_partition(otaPartition);
             if (err != ESP_OK) {
                 Serial.printf("OTA: set boot partition failed (%d)\n", err);
@@ -207,7 +227,7 @@ class OtaCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
             }
             digitalWrite(LED_PIN, HIGH);
             sendOtaNotify(pChar, OTA_STATUS_COMPLETE, 0);
-            Serial.println("OTA: complete, rebooting...");
+            Serial.println("OTA: SHA256 OK, rebooting...");
             delay(500);
             esp_restart();
             return;
@@ -219,6 +239,7 @@ class OtaCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
             uint32_t seq = ((uint32_t)data[0] << 24) | ((uint32_t)data[1] << 16) |
                            ((uint32_t)data[2] << 8)  | (uint32_t)data[3];
             size_t dataLen = len - 4;
+            mbedtls_sha256_update_ret(&otaShaCtx, data + 4, dataLen);
             esp_err_t err = esp_ota_write(otaHandle, data + 4, dataLen);
             if (err != ESP_OK) {
                 Serial.printf("OTA: write err at seq %lu (%d)\n", seq, err);