feat(client): add OTA server mode and update notification handling

src/client.cpp

@@ -6,6 +6,8 @@
 #include <mbedtls/gcm.h>
 #include <mbedtls/sha256.h>
 #include <Preferences.h>
+#include <esp_ota_ops.h>
+#include <esp_partition.h>
 
 #include <esp_log.h>
 #include "version.h"
@@ -39,6 +41,8 @@ CRGB leds[NUM_LEDS];
 #define ECDH_UUID           "beb5483e-36e1-4688-b7f5-ea07361f26dc"
 #define CLIENT_REGISTRY_UUID "beb5483e-36e1-4688-b7f5-ea07361f26dd"
 #define CLIENT_CONTROL_UUID  "beb5483e-36e1-4688-b7f5-ea07361f26de"
+#define VERSION_UUID        "beb5483e-36e1-4688-b7f5-ea07361f26db"
+#define OTA_UUID            "beb5483e-36e1-4688-b7f5-ea07361f26da"
 #define ESP_UUID_LEN  16
 #define BROKER_PLAIN_SIZE  20
 #define BROKER_PKT_SIZE    (4 + 8 + BROKER_PLAIN_SIZE + 16)
@@ -81,6 +85,30 @@ uint32_t connectionStableMs = 0;
 uint32_t relayFlashUntil = 0;
 bool otaInProgress = false;
 bool otaEnded = false;
+bool otaServerMode = false;
+
+static NimBLEServer *pOtaServer = nullptr;
+static NimBLECharacteristic *pOtaChar = nullptr;
+static NimBLECharacteristic *pOtaUuidChar = nullptr;
+
+esp_ota_handle_t otaHandle = 0;
+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;
+bool otaReady = false;
+bool otaConfirmed = false;
+bool otaEarlyConfirm = false;
+uint32_t otaEndedAt = 0;
+#define OTA_CONFIRM_TIMEOUT_MS 10000
+
+#define OTA_CMD_START  0x10
+#define OTA_STATUS_ACK 0x00
+#define OTA_STATUS_NAK 0x01
+#define OTA_STATUS_COMPLETE 0x10
+#define OTA_STATUS_ERROR    0x11
 
 static bool bleReady = false;
 static bool serverFound = false;
@@ -192,6 +220,182 @@ static void disconnectClient() {
     memset(aesKey, 0, AES_KEY_LEN);
 }
 
+static void sendOtaNotify(uint8_t status, uint32_t seq) {
+    if (pOtaChar == nullptr) return;
+    uint8_t buf[5];
+    buf[0] = status;
+    buf[1] = (seq >> 24) & 0xFF;
+    buf[2] = (seq >> 16) & 0xFF;
+    buf[3] = (seq >> 8) & 0xFF;
+    buf[4] = seq & 0xFF;
+    pOtaChar->setValue(buf, 5);
+    bool ok = pOtaChar->notify();
+    if (!ok) {
+        ESP_LOGW(TAG, "OTA: notify() failed (status=0x%02X seq=%lu)", status, seq);
+    }
+}
+
+static void otaReset() {
+    otaInProgress = false;
+    otaHandle = 0;
+    otaPartition = nullptr;
+    otaExpectedSize = 0;
+    otaWritten = 0;
+    otaReady = false;
+    otaConfirmed = false;
+    otaEarlyConfirm = false;
+    if (otaShaStarted) {
+        mbedtls_sha256_free(&otaShaCtx);
+        otaShaStarted = false;
+    }
+}
+
+static void abortOta(const char *reason) {
+    ESP_LOGE(TAG, "OTA: abort — %s", reason);
+    if (otaHandle != 0) {
+        esp_ota_end(otaHandle);
+    }
+    otaReset();
+    sendOtaNotify(OTA_STATUS_ERROR, 0);
+}
+
+class OtaCharacteristicCallbacks : public NimBLECharacteristicCallbacks {
+    void onWrite(NimBLECharacteristic *pChar, NimBLEConnInfo &connInfo) override {
+        std::string value = pChar->getValue();
+        if (value.empty()) return;
+        const uint8_t *data = (const uint8_t *)value.data();
+        size_t len = value.length();
+
+        if (len == 5 && data[0] == OTA_CMD_START) {
+            if (otaInProgress) {
+                esp_ota_end(otaHandle);
+                otaReset();
+            }
+            otaExpectedSize = ((uint32_t)data[1] << 24) | ((uint32_t)data[2] << 16) |
+                              ((uint32_t)data[3] << 8)  | (uint32_t)data[4];
+            otaPartition = esp_ota_get_next_update_partition(nullptr);
+            if (otaPartition == nullptr) {
+                ESP_LOGE(TAG, "OTA: no update partition found");
+                sendOtaNotify(OTA_STATUS_ERROR, 0);
+                return;
+            }
+            esp_err_t err = esp_ota_begin(otaPartition, otaExpectedSize, &otaHandle);
+            if (err != ESP_OK) {
+                ESP_LOGE(TAG, "OTA: begin failed (%d)", err);
+                sendOtaNotify(OTA_STATUS_ERROR, 0);
+                return;
+            }
+            otaInProgress = true;
+            otaWritten = 0;
+            mbedtls_sha256_init(&otaShaCtx);
+            mbedtls_sha256_starts_ret(&otaShaCtx, 0);
+            otaShaStarted = true;
+            ESP_LOGI(TAG, "OTA: started, expecting %lu bytes", otaExpectedSize);
+            return;
+        }
+
+        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_LOGD(TAG, "OTA: SHA256 verify (written=%lu expected=%lu)", otaWritten, otaExpectedSize);
+            if (memcmp(computedSha, otaExpectedSha, 32) != 0) {
+                ESP_LOGE(TAG, "OTA: SHA256 mismatch — aborting");
+                esp_ota_abort(otaHandle);
+                otaReset();
+                sendOtaNotify(OTA_STATUS_ERROR, 0);
+                return;
+            }
+
+            ESP_LOGI(TAG, "OTA: SHA256 OK, waiting for 0xFF confirm");
+            otaReady = true;
+            sendOtaNotify(OTA_STATUS_COMPLETE, 0);
+            return;
+        }
+
+        if (!otaInProgress) {
+            ESP_LOGD(TAG, "OTA: write ignored (not in progress, len=%d data0=0x%02X)", len, data[0]);
+            return;
+        }
+
+        if (len == 1 && data[0] == 0xFF) {
+            ESP_LOGD(TAG, "OTA: 0xFF received (ready=%d confirmed=%d inProgress=%d)",
+                     otaReady, otaConfirmed, otaInProgress);
+            if (otaReady) {
+                ESP_LOGI(TAG, "OTA: 0xFF confirm accepted");
+                otaConfirmed = true;
+            } else {
+                ESP_LOGD(TAG, "OTA: 0xFF arrived early, deferring");
+                otaEarlyConfirm = true;
+            }
+            return;
+        }
+
+        if (data[0] < 0x10) {
+            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) {
+                ESP_LOGE(TAG, "OTA: write error at seq %lu (%d)", seq, err);
+                sendOtaNotify(OTA_STATUS_NAK, seq);
+                abortOta("write error");
+                return;
+            }
+            otaWritten += dataLen;
+            sendOtaNotify(OTA_STATUS_ACK, seq);
+        }
+    }
+};
+
+static bool startOtaServer() {
+    if (otaServerMode) return true;
+
+    char otaName[32];
+    snprintf(otaName, sizeof(otaName), "SylDa_Client_OTA_%02X%02X", deviceUuid[0], deviceUuid[1]);
+
+    NimBLEDevice::init(otaName);
+    NimBLEDevice::setMTU(248);
+
+    pOtaServer = NimBLEDevice::createServer();
+
+    NimBLEService *pOtaService = pOtaServer->createService(SERVICE_UUID);
+
+    pOtaChar = pOtaService->createCharacteristic(
+        OTA_UUID,
+        NIMBLE_PROPERTY::WRITE_NR | NIMBLE_PROPERTY::NOTIFY
+    );
+    pOtaChar->setCallbacks(new OtaCharacteristicCallbacks());
+
+    pOtaUuidChar = pOtaService->createCharacteristic(
+        UUID_CHAR_UUID,
+        NIMBLE_PROPERTY::READ
+    );
+    pOtaUuidChar->setValue(deviceUuid, ESP_UUID_LEN);
+
+    NimBLECharacteristic *pVerChar = pOtaService->createCharacteristic(
+        VERSION_UUID,
+        NIMBLE_PROPERTY::READ
+    );
+    pVerChar->setValue(FIRMWARE_VERSION_CLIENT);
+
+    pOtaService->start();
+
+    NimBLEAdvertising *pAdvertising = NimBLEDevice::getAdvertising();
+    pAdvertising->setName(otaName);
+    pAdvertising->setMinInterval(32);
+    pAdvertising->setMaxInterval(64);
+    pAdvertising->start();
+
+    otaServerMode = true;
+    ESP_LOGI(TAG, "OTA: server started as %s", otaName);
+    return true;
+}
+
 static bool performEcdhHandshake(NimBLERemoteService *pSvc) {
     if (!pskProvisioned) {
         ESP_LOGE(TAG, "ECDH: no PSK");
@@ -374,6 +578,90 @@ static bool connectToServer() {
         }
     }
 
+    NimBLERemoteCharacteristic *pControlChar = pSvc->getCharacteristic(CLIENT_CONTROL_UUID);
+    if (pControlChar == nullptr) {
+        ESP_LOGW(TAG, "BLE: CLIENT_CONTROL characteristic not found");
+    } else {
+        if (!pControlChar->subscribe(true, [pControlChar](NimBLERemoteCharacteristic *pChar, const uint8_t *data, size_t len, bool isNotify) {
+            if (len != BROKER_PKT_SIZE) {
+                ESP_LOGW(TAG, "BROKER: bad notify length %d", (int)len);
+                return;
+            }
+
+            uint8_t nonce[12];
+            memcpy(nonce, data, 4);
+            memcpy(nonce + 4, data + 4, 8);
+
+            uint8_t plainPayload[BROKER_PLAIN_SIZE];
+            uint8_t ciphertext[BROKER_PLAIN_SIZE];
+            memcpy(ciphertext, data + 12, BROKER_PLAIN_SIZE);
+
+            mbedtls_gcm_context gcm;
+            mbedtls_gcm_init(&gcm);
+            int ret = mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES, aesKey, AES_KEY_LEN * 8);
+            if (ret != 0) {
+                ESP_LOGE(TAG, "BROKER: setkey failed (ret=%d)", ret);
+                mbedtls_gcm_free(&gcm);
+                return;
+            }
+
+            ret = mbedtls_gcm_auth_decrypt(&gcm,
+                BROKER_PLAIN_SIZE, nonce, 12, NULL, 0,
+                data + 12 + BROKER_PLAIN_SIZE, 16, ciphertext, plainPayload);
+            mbedtls_gcm_free(&gcm);
+
+            if (ret != 0) {
+                ESP_LOGE(TAG, "BROKER: AES-GCM decrypt failed (ret=%d)", ret);
+                return;
+            }
+
+            uint8_t cmd = plainPayload[0];
+            ESP_LOGI(TAG, "BROKER: received cmd 0x%02X", cmd);
+
+            if (cmd == BROKER_CMD_UPDATE_AVAILABLE) {
+                ESP_LOGI(TAG, "BROKER: update available, starting OTA server");
+                startOtaServer();
+
+                uint8_t clientUuid[ESP_UUID_LEN];
+                memcpy(clientUuid, plainPayload + 1, ESP_UUID_LEN);
+
+                uint8_t plainResp[BROKER_PLAIN_SIZE];
+                memset(plainResp, 0, BROKER_PLAIN_SIZE);
+                plainResp[0] = BROKER_CMD_CLIENT_READY;
+                memcpy(plainResp + 1, clientUuid, ESP_UUID_LEN);
+
+                uint8_t pkt[BROKER_PKT_SIZE];
+                memset(pkt, 0, BROKER_PKT_SIZE);
+                uint32_t s = ++seq;
+                pkt[0] = (s >> 24) & 0xFF;
+                pkt[1] = (s >> 16) & 0xFF;
+                pkt[2] = (s >> 8)  & 0xFF;
+                pkt[3] = s & 0xFF;
+
+                for (int i = 0; i < 8; i++) pkt[4 + i] = (uint8_t)(esp_random() & 0xFF);
+
+                uint8_t nonceResp[12];
+                memcpy(nonceResp, pkt, 4);
+                memcpy(nonceResp + 4, pkt + 4, 8);
+
+                mbedtls_gcm_context gcmResp;
+                mbedtls_gcm_init(&gcmResp);
+                ret = mbedtls_gcm_setkey(&gcmResp, MBEDTLS_CIPHER_ID_AES, aesKey, AES_KEY_LEN * 8);
+                if (ret == 0) {
+                    mbedtls_gcm_crypt_and_tag(&gcmResp, MBEDTLS_GCM_ENCRYPT,
+                        BROKER_PLAIN_SIZE, nonceResp, 12, NULL, 0,
+                        plainResp, pkt + 12, 16, pkt + 12 + BROKER_PLAIN_SIZE);
+                    mbedtls_gcm_free(&gcmResp);
+
+                    pControlChar->writeValue(pkt, BROKER_PKT_SIZE, false);
+                    ESP_LOGI(TAG, "BROKER: sent CLIENT_READY to server");
+                }
+            }
+        }, true)) {
+            ESP_LOGW(TAG, "BROKER: subscribe to CLIENT_CONTROL failed");
+        }
+    }
+
     return true;
 }
 
@@ -583,9 +871,46 @@ void loop() {
     bool relay1State = !t1Active || btnActive;
     bool relay2State = !t2Active;
 
+    if (otaServerMode && otaReady) {
+        esp_err_t err = esp_ota_end(otaHandle);
+        if (err != ESP_OK) {
+            ESP_LOGE(TAG, "OTA: esp_ota_end failed (%d) — aborting", err);
+            otaReset();
+        } else {
+            ESP_LOGI(TAG, "OTA: esp_ota_end OK, waiting for 0xFF confirm");
+            otaEnded = true;
+            otaEndedAt = millis();
+        }
+        otaReady = false;
+    }
+
+    if (otaServerMode && otaEnded && otaEarlyConfirm) {
+        ESP_LOGD(TAG, "OTA: processing deferred 0xFF confirm");
+        otaConfirmed = true;
+        otaEarlyConfirm = false;
+    }
+
+    if (otaServerMode && otaEnded && !otaConfirmed && (millis() - otaEndedAt > OTA_CONFIRM_TIMEOUT_MS)) {
+        ESP_LOGW(TAG, "OTA: 0xFF confirm timeout (%lu ms) — discarding update", millis() - otaEndedAt);
+        otaReset();
+    }
+
+    if (otaServerMode && otaConfirmed) {
+        ESP_LOGI(TAG, "OTA: confirmed, setting boot partition");
+        esp_err_t err = esp_ota_set_boot_partition(otaPartition);
+        if (err != ESP_OK) {
+            ESP_LOGE(TAG, "OTA: set boot partition failed (%d)", err);
+        } else {
+            ESP_LOGI(TAG, "OTA: done, rebooting...");
+            delay(500);
+            esp_restart();
+        }
+        otaReset();
+    }
+
     CRGB ledColor;
 
-    if (otaInProgress) {
+    if (otaServerMode || otaInProgress) {
         ledColor = otaEnded
             ? CRGB::Yellow
             : (((millis() / 100) % 2) ? CRGB::Green : CRGB::Black);