Files
HaloKeymind/examples/esp32_partition_migrator/main.cpp
T

582 lines
24 KiB
C++

#include <Arduino.h>
#include <WiFi.h>
#include <AsyncTCP.h>
#include <ESPAsyncWebServer.h>
#include <AsyncElegantOTA.h>
#include <Preferences.h>
#include <SPIFFS.h>
#include <esp_ota_ops.h>
#include <esp_attr.h>
#include <esp_flash.h>
#include <esp_partition.h>
#include <helpers/ESP32PartitionMigrationPolicy.h>
#include <helpers/esp32/WiFiRadioPolicy.h>
namespace migration = mesh::esp32_partition_migration;
namespace {
constexpr char kApSsid[] = "MeshCore-Migrate";
constexpr char kApPassword[] = "meshcore-migrate";
#ifndef MESHCORE_MIGRATION_DELAY_MS
#define MESHCORE_MIGRATION_DELAY_MS 4500
#endif
#ifndef MESHCORE_MIGRATION_RESTART_DELAY_MS
#define MESHCORE_MIGRATION_RESTART_DELAY_MS 250
#endif
constexpr uint32_t kMigrationDelayMs = MESHCORE_MIGRATION_DELAY_MS;
constexpr uint32_t kMigrationRestartDelayMs = MESHCORE_MIGRATION_RESTART_DELAY_MS;
constexpr size_t kCopyBufferBytes = 4096;
constexpr size_t kIdentityFileBytes = 96; // public key (32) followed by private key (64)
constexpr char kMigrationNvsNamespace[] = "mesh-pt-migrate";
constexpr char kMigrationIdentityKey[] = "identity";
// NVS key names are limited to 15 characters.
constexpr char kMigrationIdentityPendingKey[] = "id-pending";
static_assert(sizeof(kMigrationNvsNamespace) - 1 <= 15,
"ESP32 NVS namespace names are limited to 15 characters");
AsyncWebServer server(80);
bool migration_started = false;
bool migration_complete = false;
bool reboot_requested = false;
uint32_t migration_at = 0;
uint32_t reboot_at = 0;
char status_text[160] = "Starting";
// The Arduino loop task can use external RAM on an ESP32-S3. Flash
// erase/write disables the external-memory cache, so every buffer passed to a
// flash operation must be explicitly placed in internal DRAM.
DRAM_ATTR uint8_t copy_buffer[kCopyBufferBytes];
DRAM_ATTR uint8_t partition_table_bytes[
migration::kExpandedPartitionTablePrefixBytes];
DRAM_ATTR uint8_t partition_table_verified[
migration::kExpandedPartitionTablePrefixBytes];
// Keep the temporary flash hooks in DRAM. The normal ESP-IDF hooks keep the
// two CPU cores and their flash caches safe while an erase/write is running;
// only the address-validation callback is narrowed for this one deliberately
// dangerous sector. esp_flash_os_functions_t is documented for advanced
// callers which need to replace individual hooks.
DRAM_ATTR esp_flash_os_functions_t partition_table_flash_hooks;
using RegionProtectedFn = esp_err_t (*)(void*, size_t, size_t);
DRAM_ATTR RegionProtectedFn original_region_protected = nullptr;
DRAM_ATTR void* original_flash_hook_data = nullptr;
// esp_partition_write() reads the descriptor while the flash cache is off.
// Keep the synthetic copy destination out of a task stack that may be
// allocated in PSRAM on ESP32-S3 builds.
DRAM_ATTR esp_partition_t copy_destination;
struct PartitionRefs {
const esp_partition_t* nvs = nullptr;
const esp_partition_t* otadata = nullptr;
const esp_partition_t* app0 = nullptr;
const esp_partition_t* app1 = nullptr;
const esp_partition_t* spiffs = nullptr;
};
uint32_t crc32(const uint8_t* data, size_t size, uint32_t crc = 0xFFFFFFFFU) {
for (size_t i = 0; i < size; ++i) {
crc ^= data[i];
for (unsigned bit = 0; bit < 8; ++bit) {
crc = (crc >> 1) ^ ((crc & 1U) ? 0xEDB88320U : 0U);
}
}
return crc;
}
// This callback is invoked before ESP-IDF enters its cache-off flash critical
// section. It permits only the partition-table sector, and delegates every
// other address to ESP-IDF's original protection policy.
esp_err_t IRAM_ATTR partitionTableRegionProtected(void*, size_t address, size_t size) {
const size_t table_start = migration::kPartitionTableAddress;
const size_t table_end = table_start + migration::kPartitionTableBytes;
if (address >= table_start && size <= table_end - address) {
return ESP_OK;
}
if (!original_region_protected) return ESP_ERR_INVALID_STATE;
return original_region_protected(original_flash_hook_data, address, size);
}
const char* errName(esp_err_t err) {
const char* name = esp_err_to_name(err);
return name ? name : "unknown ESP error";
}
bool findPartitions(PartitionRefs& refs, migration::PartitionGeometry& geometry) {
refs.nvs = esp_partition_find_first(ESP_PARTITION_TYPE_DATA,
ESP_PARTITION_SUBTYPE_DATA_NVS, nullptr);
refs.otadata = esp_partition_find_first(ESP_PARTITION_TYPE_DATA,
ESP_PARTITION_SUBTYPE_DATA_OTA, nullptr);
refs.app0 = esp_partition_find_first(ESP_PARTITION_TYPE_APP,
ESP_PARTITION_SUBTYPE_APP_OTA_0, nullptr);
refs.app1 = esp_partition_find_first(ESP_PARTITION_TYPE_APP,
ESP_PARTITION_SUBTYPE_APP_OTA_1, nullptr);
refs.spiffs = esp_partition_find_first(ESP_PARTITION_TYPE_DATA,
ESP_PARTITION_SUBTYPE_DATA_SPIFFS, nullptr);
if (!refs.nvs || !refs.otadata || !refs.app0 || !refs.app1 || !refs.spiffs) {
return false;
}
geometry = {
refs.nvs->address, refs.nvs->size,
refs.otadata->address, refs.otadata->size,
refs.app0->address, refs.app0->size,
refs.app1->address, refs.app1->size,
refs.spiffs->address, refs.spiffs->size,
};
return true;
}
// esp_partition_* validates only the geometry in the supplied descriptor. A
// descriptor for the future data region is therefore enough to write it while
// the old partition table is still active. This avoids touching old SPIFFS
// until the target table is published.
esp_partition_t rawPartition(uint32_t address, uint32_t size, const char* label) {
esp_partition_t part = {};
part.type = ESP_PARTITION_TYPE_DATA;
part.subtype = ESP_PARTITION_SUBTYPE_DATA_UNDEFINED;
part.address = address;
part.size = size;
strncpy(part.label, label, sizeof(part.label) - 1);
part.label[sizeof(part.label) - 1] = 0;
part.encrypted = false;
return part;
}
esp_partition_t rawApp0Partition(uint32_t address, uint32_t size, const char* label) {
esp_partition_t part = rawPartition(address, size, label);
part.type = ESP_PARTITION_TYPE_APP;
part.subtype = ESP_PARTITION_SUBTYPE_APP_OTA_0;
return part;
}
bool rangesOverlap(uint32_t first_address, uint32_t first_size,
uint32_t second_address, uint32_t second_size) {
const uint64_t first_end = static_cast<uint64_t>(first_address) + first_size;
const uint64_t second_end = static_cast<uint64_t>(second_address) + second_size;
return first_address < second_end && second_address < first_end;
}
bool eraseRaw(const esp_partition_t& destination) {
for (uint32_t offset = 0; offset < destination.size;) {
uint32_t span = destination.size - offset;
// A large flash erase can run long enough to trip the task watchdog on an
// S3. Erase one sector at a time and explicitly yield between operations.
// This is slower but keeps the board alive while its future app slot is
// being prepared.
if (span > migration::kSectorBytes) span = migration::kSectorBytes;
const esp_err_t result = esp_partition_erase_range(&destination, offset, span);
if (result != ESP_OK) {
snprintf(status_text, sizeof(status_text), "Erase failed at 0x%lx: %s",
(unsigned long)(destination.address + offset), errName(result));
return false;
}
offset += span;
delay(1);
}
return true;
}
bool copyAndVerify(const esp_partition_t& source, uint32_t destination_address,
uint32_t bytes, const char* destination_label) {
if (bytes == 0 || bytes > source.size || bytes % migration::kSectorBytes != 0) {
strcpy(status_text, "Invalid migration copy geometry");
return false;
}
copy_destination = rawPartition(destination_address, bytes, destination_label);
if (!eraseRaw(copy_destination)) return false;
uint32_t source_crc = 0xFFFFFFFFU;
for (uint32_t offset = 0; offset < bytes; offset += sizeof(copy_buffer)) {
esp_err_t result = esp_partition_read(&source, offset, copy_buffer,
sizeof(copy_buffer));
if (result != ESP_OK) {
snprintf(status_text, sizeof(status_text), "Read failed at 0x%lx: %s",
(unsigned long)(source.address + offset), errName(result));
return false;
}
source_crc = crc32(copy_buffer, sizeof(copy_buffer), source_crc);
result = esp_partition_write(&copy_destination, offset, copy_buffer,
sizeof(copy_buffer));
if (result != ESP_OK) {
snprintf(status_text, sizeof(status_text), "Write failed at 0x%lx: %s",
(unsigned long)(copy_destination.address + offset), errName(result));
return false;
}
delay(1);
}
uint32_t destination_crc = 0xFFFFFFFFU;
for (uint32_t offset = 0; offset < bytes; offset += sizeof(copy_buffer)) {
const esp_err_t result = esp_partition_read(&copy_destination, offset, copy_buffer,
sizeof(copy_buffer));
if (result != ESP_OK) {
snprintf(status_text, sizeof(status_text), "Verify read failed at 0x%lx: %s",
(unsigned long)(copy_destination.address + offset), errName(result));
return false;
}
destination_crc = crc32(copy_buffer, sizeof(copy_buffer), destination_crc);
delay(1);
}
source_crc = ~source_crc;
destination_crc = ~destination_crc;
if (source_crc != destination_crc) {
snprintf(status_text, sizeof(status_text),
"CRC mismatch while copying %s (%08lx != %08lx)", destination_label,
(unsigned long)source_crc, (unsigned long)destination_crc);
return false;
}
return true;
}
// The larger destination SPIFFS partition is allowed to reformat itself on
// its first mount. A raw SPIFFS image is not reliably expandable, so stage
// the one irreplaceable file in NVS first. NVS stays at the same address in
// both layouts. The source file format is exactly the historical
// IdentityStore layout: 32 public-key bytes followed by 64 private-key bytes.
bool stageLegacyIdentity() {
if (!SPIFFS.begin(false)) {
strcpy(status_text, "Could not mount legacy SPIFFS to save identity");
return false;
}
File identity = SPIFFS.open("/identity/_main.id", "r");
const bool read_ok = identity && identity.size() >= kIdentityFileBytes
&& identity.read(copy_buffer, kIdentityFileBytes) == kIdentityFileBytes;
if (identity) identity.close();
SPIFFS.end();
if (!read_ok) {
strcpy(status_text, "Refused: legacy private-key file is unavailable");
return false;
}
Preferences migration_nvs;
if (!migration_nvs.begin(kMigrationNvsNamespace, false)) {
strcpy(status_text, "Could not open NVS identity staging");
return false;
}
const bool saved = migration_nvs.putBytes(kMigrationIdentityKey, copy_buffer,
kIdentityFileBytes) == kIdentityFileBytes
&& migration_nvs.putBool(kMigrationIdentityPendingKey, true);
migration_nvs.end();
if (!saved) strcpy(status_text, "Could not save private key for migration");
return saved;
}
bool restoreStagedIdentity() {
Preferences migration_nvs;
if (!migration_nvs.begin(kMigrationNvsNamespace, false)) {
strcpy(status_text, "Could not open NVS identity recovery");
return false;
}
const bool pending = migration_nvs.getBool(kMigrationIdentityPendingKey, false);
const size_t stored_bytes = pending
? migration_nvs.getBytes(kMigrationIdentityKey, copy_buffer, sizeof(copy_buffer))
: 0;
migration_nvs.end();
if (!pending) return true;
if (stored_bytes != kIdentityFileBytes) {
strcpy(status_text, "Refused: staged private key is incomplete");
return false;
}
// `true` intentionally formats only if the raw legacy SPIFFS image cannot
// mount at the expanded size. All settings except the identity are allowed
// to be recreated; the staged identity is immediately written back below.
if (!SPIFFS.begin(true)) {
strcpy(status_text, "Could not initialize expanded SPIFFS");
return false;
}
if (!SPIFFS.exists("/identity") && !SPIFFS.mkdir("/identity")) {
SPIFFS.end();
strcpy(status_text, "Could not create identity folder");
return false;
}
File identity = SPIFFS.open("/identity/_main.id", "w");
const bool wrote = identity && identity.write(copy_buffer, kIdentityFileBytes)
== kIdentityFileBytes;
if (identity) {
identity.flush();
identity.close();
}
File verify = SPIFFS.open("/identity/_main.id", "r");
const bool verified = wrote && verify && verify.size() >= kIdentityFileBytes
&& verify.read(partition_table_verified, kIdentityFileBytes) == kIdentityFileBytes
&& memcmp(copy_buffer, partition_table_verified, kIdentityFileBytes) == 0;
if (verify) verify.close();
SPIFFS.end();
if (!verified) {
strcpy(status_text, "Private-key restore verification failed");
return false;
}
if (!migration_nvs.begin(kMigrationNvsNamespace, false)) {
strcpy(status_text, "Could not finalize NVS identity recovery");
return false;
}
migration_nvs.remove(kMigrationIdentityKey);
const bool cleared = migration_nvs.remove(kMigrationIdentityPendingKey);
migration_nvs.end();
if (!cleared) {
strcpy(status_text, "Private key restored; NVS cleanup needs retry");
return false;
}
return true;
}
bool publishExpandedPartitionTable(const migration::TargetPlan& plan) {
// Preserve ESP-IDF's normal OS flash hooks. In particular, their start/end
// hooks suspend the other core and safely disable/re-enable caches. Do not
// call esp_flash_app_disable_protect(): that internal coredump-only helper
// removes those hooks entirely and is unsafe from a running application.
esp_flash_t* const chip = esp_flash_default_chip;
const esp_flash_os_functions_t* const original_hooks = chip ? chip->os_func : nullptr;
if (!chip || !original_hooks || !original_hooks->region_protected) {
strcpy(status_text, "ESP-IDF flash protection hooks are unavailable");
return false;
}
original_region_protected = original_hooks->region_protected;
original_flash_hook_data = chip->os_func_data;
partition_table_flash_hooks = *original_hooks;
partition_table_flash_hooks.region_protected = partitionTableRegionProtected;
chip->os_func = &partition_table_flash_hooks;
bool ok = false;
// esp_partition_write disables the flash cache. The generated prefix is
// otherwise stored in DROM, so copy it to explicitly internal DRAM first.
memcpy(partition_table_bytes, plan.partition_table_prefix,
plan.partition_table_prefix_bytes);
// Do not use esp_partition_write for the table itself. Once the old sector
// is erased, a partition-manager path must not be allowed to consult that
// erased metadata. esp_flash_* operates on the main flash chip directly.
esp_err_t result = esp_flash_erase_region(chip,
migration::kPartitionTableAddress,
migration::kPartitionTableBytes);
if (result == ESP_OK) {
result = esp_flash_write(chip, partition_table_bytes,
migration::kPartitionTableAddress,
sizeof(partition_table_bytes));
if (result == ESP_OK) {
result = esp_flash_read(chip, partition_table_verified,
migration::kPartitionTableAddress,
sizeof(partition_table_verified));
ok = result == ESP_OK && memcmp(partition_table_verified,
plan.partition_table_prefix, plan.partition_table_prefix_bytes) == 0;
if (!ok) strcpy(status_text, "Partition table verification failed");
} else {
snprintf(status_text, sizeof(status_text), "Partition table write failed: %s",
errName(result));
}
} else {
snprintf(status_text, sizeof(status_text), "Partition table erase failed: %s",
errName(result));
}
chip->os_func = original_hooks;
original_region_protected = nullptr;
original_flash_hook_data = nullptr;
return ok;
}
void runMigration() {
migration_started = true;
PartitionRefs refs;
migration::PartitionGeometry geometry = {};
const uint32_t flash_bytes = ESP.getFlashChipSize();
const migration::TargetPlan* const plan = migration::targetForFlash(flash_bytes);
if (!plan || !findPartitions(refs, geometry)
|| !migration::canMigrateGeneric(flash_bytes, geometry)) {
snprintf(status_text, sizeof(status_text),
"Refused: unsupported layout or flash size (%lu bytes)",
(unsigned long)flash_bytes);
Serial.println(status_text);
return;
}
const esp_partition_t* running = esp_ota_get_running_partition();
// Partition handles are opaque; their addresses, rather than their pointer
// identities, determine the slot. This keeps the A/B decision correct if
// ESP-IDF returns a distinct descriptor for the currently-running image.
if (!running || (running->address != refs.app0->address &&
running->address != refs.app1->address)) {
strcpy(status_text, "Refused: migration image is not in a legacy OTA slot");
Serial.println(status_text);
return;
}
Serial.println("Migration: staging private key; do not interrupt power");
if (!stageLegacyIdentity()) {
Serial.println(status_text);
return;
}
Serial.println("Migration: private key safely staged in NVS");
// Always make target app0 contain this migration image. If it was uploaded
// into a different legacy slot, this preserves a Wi-Fi endpoint after the
// table change. Never erase a running source range while copying it.
if (running->address != plan->layout.app0_address) {
if (running->size > plan->layout.app0_size) {
strcpy(status_text, "Refused: migration slot is larger than future app0");
Serial.println(status_text);
return;
}
if (rangesOverlap(running->address, running->size,
plan->layout.app0_address, running->size)) {
strcpy(status_text, "Refused: target app0 overlaps the running image");
Serial.println(status_text);
return;
}
Serial.println("Migration: placing Wi-Fi bridge image in app0");
if (!copyAndVerify(*running, plan->layout.app0_address,
running->size, "future-app0")) {
Serial.println(status_text);
return;
}
Serial.println("Migration: bridge image copied to app0");
}
// The OTA-select data records a slot identity. Select a descriptor for the
// target app0 before replacing the table so the next boot runs this bridge
// at the target table's app0 address, never the blank new app1.
const esp_partition_t target_app0 = rawApp0Partition(plan->layout.app0_address,
plan->layout.app0_size,
"app0");
const esp_err_t select_result = esp_ota_set_boot_partition(&target_app0);
if (select_result != ESP_OK) {
snprintf(status_text, sizeof(status_text), "Could not select app0: %s",
errName(select_result));
Serial.println(status_text);
return;
}
Serial.println("Migration: app0 selected for restart");
Serial.println("Migration: publishing expanded partition table");
// This bridge is built without native USB CDC, so a connected USB host cannot
// post a flash-backed event while the partition-sector operation disables the
// flash cache. Stop UART0 as well: the S3's serial event path is otherwise
// still able to interrupt the raw flash operation. The bridge itself is
// deliberately Wi-Fi-only; the final normal repeater build restores its
// standard USB behavior.
Serial.flush();
Serial.end();
const bool table_published = publishExpandedPartitionTable(*plan);
Serial.begin(115200);
delay(50);
if (!table_published) {
Serial.println(status_text);
return;
}
migration_complete = true;
if (kMigrationRestartDelayMs == 0) {
strcpy(status_text, "Partition table verified; waiting for test reboot");
} else {
strcpy(status_text, "Migration complete; restarting Wi-Fi uploader");
reboot_at = millis() + kMigrationRestartDelayMs;
}
Serial.println(status_text);
}
void sendHome(AsyncWebServerRequest* request) {
const char* mode = status_text;
String page;
page.reserve(1000);
page += "<!doctype html><meta name=viewport content='width=device-width,initial-scale=1'>";
page += "<h2>MeshCore Wi-Fi partition migration</h2><p>";
page += mode;
page += "</p>";
if (strstr(status_text, "ready") != nullptr) {
page += "<p>The expanded partition layout is active and the private device "
"identity was staged and restored before this uploader was exposed.</p>"
"<p><a href='/update'>Upload the full application image for this board.</a>.</p>";
} else if (migration_complete) {
page += "<p>Partition-table bytes were read back successfully. The test harness "
"is waiting for an explicit reboot.</p><p><a href='/reboot'>Restart now</a></p>";
} else if (!migration_started) {
page += "<p>Copying private data and replacing the partition table starts shortly. "
"Keep USB power connected; this page will disappear while the board restarts.</p>";
} else {
page += "<p>Keep power connected. Refresh after two minutes if the board did not restart.</p>";
}
request->send(200, "text/html", page);
}
void startServer() {
// A legacy ESP-NOW repeater can leave the AP protocol mask in proprietary
// LR mode. SoftAP then reports success but ordinary phones and laptops
// cannot discover it. Reuse the normal WebConfig recipe: AP+STA mode,
// an explicit interoperable protocol mask, and the project AP channel.
WiFi.mode(WIFI_AP_STA);
WiFi.setAutoReconnect(false);
WiFi.disconnect(false, true);
delay(100);
WiFi.setSleep(false);
const IPAddress address(192, 168, 4, 1);
const IPAddress netmask(255, 255, 255, 0);
if (!WiFi.softAPConfig(address, address, netmask)
|| !WiFi.softAP(kApSsid, kApPassword, mesh::wifi::accessPointChannel())
|| mesh::wifi::applyAccessPointProtocolMask() != ESP_OK
|| esp_wifi_set_protocol(WIFI_IF_STA, mesh::wifi::kProtocolMask) != ESP_OK
|| esp_wifi_set_max_tx_power(78) != ESP_OK) {
strcpy(status_text, "Wi-Fi AP failed; restart the board and retry");
return;
}
wifi_config_t ap_config = {};
int8_t max_tx_power = 0;
const esp_err_t config_result = esp_wifi_get_config(WIFI_IF_AP, &ap_config);
const esp_err_t power_result = esp_wifi_get_max_tx_power(&max_tx_power);
Serial.printf("Wi-Fi AP active: %s at %s (channel %u, hidden %u, power %.2f dBm, config %d, power %d)\n",
WiFi.softAPSSID().c_str(), WiFi.softAPIP().toString().c_str(),
(unsigned)ap_config.ap.channel, (unsigned)ap_config.ap.ssid_hidden,
max_tx_power / 4.0, (int)config_result, (int)power_result);
server.on("/", HTTP_GET, sendHome);
server.on("/reboot", HTTP_GET, [](AsyncWebServerRequest* request) {
if (!migration_complete) {
request->send(409, "text/plain", "Migration has not completed");
return;
}
request->send(200, "text/plain", "Restarting migration bridge");
reboot_requested = true;
});
AsyncElegantOTA.begin(&server);
server.begin();
}
} // namespace
void setup() {
Serial.begin(115200);
delay(250);
PartitionRefs refs;
migration::PartitionGeometry geometry = {};
const uint32_t flash_bytes = ESP.getFlashChipSize();
if (findPartitions(refs, geometry)
&& migration::isTargetLayout(flash_bytes, geometry)) {
if (restoreStagedIdentity()) {
strcpy(status_text, "Expanded layout ready");
}
} else if (findPartitions(refs, geometry)
&& migration::canMigrateGeneric(flash_bytes, geometry)) {
strcpy(status_text, "Legacy layout verified; migration begins shortly");
migration_at = millis() + kMigrationDelayMs;
} else {
strcpy(status_text, "Refused: unsupported legacy source layout");
}
startServer();
Serial.printf("%s. Join %s (password: %s), then open http://192.168.4.1/\n",
status_text, kApSsid, kApPassword);
}
void loop() {
if (!migration_started && migration_at != 0
&& static_cast<int32_t>(millis() - migration_at) >= 0) {
runMigration();
}
if ((reboot_at != 0 && static_cast<int32_t>(millis() - reboot_at) >= 0)
|| reboot_requested) {
delay(100);
ESP.restart();
}
delay(10);
}