Support WiFi OTA (#6352)

* Support WiFi OTA

* Fix trunk warnings

* Make getVersion() check for project name too

---------

Co-authored-by: Ben Meadors <benmmeadors@gmail.com>

src/mesh/NodeDB.cpp

@@ -51,6 +51,10 @@
 #include <utility/bonding.h>
 #endif
 
+#if defined(ARCH_ESP32) && !MESHTASTIC_EXCLUDE_WIFI
+#include <WiFiOTA.h>
+#endif
+
 NodeDB *nodeDB = nullptr;
 
 // we have plenty of ram so statically alloc this tempbuf (for now)
@@ -635,6 +639,12 @@ void NodeDB::installDefaultConfig(bool preserveKey = false)
     config.display.wake_on_tap_or_motion = true;
 #endif
 
+#if defined(ARCH_ESP32) && !MESHTASTIC_EXCLUDE_WIFI
+    if (WiFiOTA::isUpdated()) {
+        WiFiOTA::recoverConfig(&config.network);
+    }
+#endif
+
     initConfigIntervals();
 }
 

src/modules/AdminModule.cpp

@@ -10,6 +10,9 @@
 #if defined(ARCH_ESP32) && !MESHTASTIC_EXCLUDE_BLUETOOTH
 #include "BleOta.h"
 #endif
+#if defined(ARCH_ESP32) && !MESHTASTIC_EXCLUDE_WIFI
+#include "WiFiOTA.h"
+#endif
 #include "Router.h"
 #include "configuration.h"
 #include "main.h"
@@ -194,19 +197,23 @@ bool AdminModule::handleReceivedProtobuf(const meshtastic_MeshPacket &mp, meshta
     }
     case meshtastic_AdminMessage_reboot_ota_seconds_tag: {
         int32_t s = r->reboot_ota_seconds;
-#if defined(ARCH_ESP32) && !MESHTASTIC_EXCLUDE_BLUETOOTH
-        if (BleOta::getOtaAppVersion().isEmpty()) {
-            LOG_INFO("No OTA firmware available, scheduling regular reboot in %d seconds", s);
-            screen->startAlert("Rebooting...");
-        } else {
+#if defined(ARCH_ESP32)
+#if !MESHTASTIC_EXCLUDE_BLUETOOTH
+        if (!BleOta::getOtaAppVersion().isEmpty()) {
             screen->startFirmwareUpdateScreen();
             BleOta::switchToOtaApp();
-            LOG_INFO("Reboot to OTA in %d seconds", s);
+            LOG_INFO("Rebooting to BLE OTA");
         }
-#else
-        LOG_INFO("Not on ESP32, scheduling regular reboot in %d seconds", s);
-        screen->startAlert("Rebooting...");
 #endif
+#if !MESHTASTIC_EXCLUDE_WIFI
+        if (WiFiOTA::trySwitchToOTA()) {
+            screen->startFirmwareUpdateScreen();
+            WiFiOTA::saveConfig(&config.network);
+            LOG_INFO("Rebooting to WiFi OTA");
+        }
+#endif
+#endif
+        LOG_INFO("Reboot in %d seconds", s);
         rebootAtMsec = (s < 0) ? 0 : (millis() + s * 1000);
         break;
     }

src/platform/esp32/WiFiOTA.cpp

@@ -0,0 +1,92 @@
+#include "WiFiOTA.h"
+#include "configuration.h"
+#include <Preferences.h>
+#include <esp_ota_ops.h>
+
+namespace WiFiOTA
+{
+
+static const char *nvsNamespace = "ota-wifi";
+static const char *appProjectName = "OTA-WiFi";
+
+static bool updated = false;
+
+bool isUpdated()
+{
+    return updated;
+}
+
+void initialize()
+{
+    Preferences prefs;
+    prefs.begin(nvsNamespace);
+    if (prefs.getBool("updated")) {
+        LOG_INFO("First boot after OTA update");
+        updated = true;
+        prefs.putBool("updated", false);
+    }
+    prefs.end();
+}
+
+void recoverConfig(meshtastic_Config_NetworkConfig *network)
+{
+    LOG_INFO("Recovering WiFi settings after OTA update");
+
+    Preferences prefs;
+    prefs.begin(nvsNamespace, true);
+    String ssid = prefs.getString("ssid");
+    String psk = prefs.getString("psk");
+    prefs.end();
+
+    network->wifi_enabled = true;
+    strncpy(network->wifi_ssid, ssid.c_str(), sizeof(network->wifi_ssid));
+    strncpy(network->wifi_psk, psk.c_str(), sizeof(network->wifi_psk));
+}
+
+void saveConfig(meshtastic_Config_NetworkConfig *network)
+{
+    LOG_INFO("Saving WiFi settings for upcoming OTA update");
+
+    Preferences prefs;
+    prefs.begin(nvsNamespace);
+    prefs.putString("ssid", network->wifi_ssid);
+    prefs.putString("psk", network->wifi_psk);
+    prefs.putBool("updated", false);
+    prefs.end();
+}
+
+const esp_partition_t *getAppPartition()
+{
+    return esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_OTA_1, NULL);
+}
+
+bool getAppDesc(const esp_partition_t *part, esp_app_desc_t *app_desc)
+{
+    if (esp_ota_get_partition_description(part, app_desc) != ESP_OK)
+        return false;
+    if (strcmp(app_desc->project_name, appProjectName) != 0)
+        return false;
+    return true;
+}
+
+bool trySwitchToOTA()
+{
+    const esp_partition_t *part = getAppPartition();
+    esp_app_desc_t app_desc;
+    if (!getAppDesc(part, &app_desc))
+        return false;
+    if (esp_ota_set_boot_partition(part) != ESP_OK)
+        return false;
+    return true;
+}
+
+String getVersion()
+{
+    const esp_partition_t *part = getAppPartition();
+    esp_app_desc_t app_desc;
+    if (!getAppDesc(part, &app_desc))
+        return String();
+    return String(app_desc.version);
+}
+
+} // namespace WiFiOTA

src/platform/esp32/WiFiOTA.h

@@ -0,0 +1,18 @@
+#ifndef WIFIOTA_H
+#define WIFIOTA_H
+
+#include "mesh-pb-constants.h"
+#include <Arduino.h>
+
+namespace WiFiOTA
+{
+void initialize();
+bool isUpdated();
+
+void recoverConfig(meshtastic_Config_NetworkConfig *network);
+void saveConfig(meshtastic_Config_NetworkConfig *network);
+bool trySwitchToOTA();
+String getVersion();
+} // namespace WiFiOTA
+
+#endif // WIFIOTA_H

src/platform/esp32/main-esp32.cpp

@@ -9,6 +9,8 @@
 #include "nimble/NimbleBluetooth.h"
 #endif
 
+#include <WiFiOTA.h>
+
 #if HAS_WIFI
 #include "mesh/wifi/WiFiAPClient.h"
 #endif
@@ -139,12 +141,19 @@ void esp32Setup()
 #if !MESHTASTIC_EXCLUDE_BLUETOOTH
     String BLEOTA = BleOta::getOtaAppVersion();
     if (BLEOTA.isEmpty()) {
-        LOG_INFO("No OTA firmware available");
+        LOG_INFO("No BLE OTA firmware available");
     } else {
-        LOG_INFO("OTA firmware version %s", BLEOTA.c_str());
+        LOG_INFO("BLE OTA firmware version %s", BLEOTA.c_str());
     }
-#else
-    LOG_INFO("No OTA firmware available");
+#endif
+#if !MESHTASTIC_EXCLUDE_WIFI
+    String version = WiFiOTA::getVersion();
+    if (version.isEmpty()) {
+        LOG_INFO("No WiFi OTA firmware available");
+    } else {
+        LOG_INFO("WiFi OTA firmware version %s", version.c_str());
+    }
+    WiFiOTA::initialize();
 #endif
 
     // enableModemSleep();