diff --git a/src/client.cpp b/src/client.cpp index 3c0dbf2..ac84754 100644 --- a/src/client.cpp +++ b/src/client.cpp @@ -6,6 +6,8 @@ #include #include #include +#include +#include #include #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);