Add single-click auto-update functionality

dsbaars · dsbaars · commit 023ff2913143 · 2024-09-11T20:27:40.000+02:00
diff --git a/data b/data
@@ -1 +1 @@
-Subproject commit 1c2d8dcdd0efd846f39b0f24977c982a96f16392
+Subproject commit 6c40b54273b7f7c7d6c2624d3c2a066435f27756
diff --git a/src/lib/ota.cpp b/src/lib/ota.cpp
@@ -4,6 +4,8 @@ TaskHandle_t taskOtaHandle = NULL;
 bool isOtaUpdating = false;
 QueueHandle_t otaQueue;
 
+
+
 void setupOTA()
 {
   if (preferences.getBool("otaEnabled", DEFAULT_OTA_ENABLED))
@@ -65,7 +67,7 @@ void onOTAStart()
   vTaskSuspend(workerTaskHandle);
   vTaskSuspend(taskScreenRotateTaskHandle);
 
-  vTaskSuspend(ledTaskHandle);
+//  vTaskSuspend(ledTaskHandle);
   vTaskSuspend(buttonTaskHandle);
 
   // stopWebServer();
@@ -81,15 +83,26 @@ void handleOTATask(void *parameter)
   {
     if (xQueueReceive(otaQueue, &msg, 0) == pdTRUE)
     {
-      int result = downloadUpdateHandler(msg.updateType);
+      if (msg.updateType == UPDATE_ALL) {
+        int resultWebUi = downloadUpdateHandler(UPDATE_WEBUI);
+        int resultFw = downloadUpdateHandler(UPDATE_FIRMWARE);
+
+        if (resultWebUi == 0 && resultFw == 0) {
+          ESP.restart();
+        } else {
+          queueLedEffect(LED_FLASH_ERROR);
+          vTaskDelay(pdMS_TO_TICKS(3000));
+          ESP.restart();
+        }
+      }
     }
 
     ArduinoOTA.handle(); // Allow OTA updates to occur
     vTaskDelay(pdMS_TO_TICKS(2000));
   }
 }
 
-String getLatestRelease(const String &fileToDownload)
+ReleaseInfo getLatestRelease(const String &fileToDownload)
 {
   String releaseUrl = "https://api.github.com/repos/btclock/btclock_v3/releases/latest";
   WiFiClientSecure client;
@@ -100,7 +113,7 @@ String getLatestRelease(const String &fileToDownload)
 
   int httpCode = http.GET();
 
-  String downloadUrl = "";
+  ReleaseInfo info = {"", ""};
 
   if (httpCode > 0)
   {
@@ -113,15 +126,26 @@ String getLatestRelease(const String &fileToDownload)
 
     for (JsonObject asset : assets)
     {
-      if (asset["name"] == fileToDownload)
+      String assetName = asset["name"].as<String>();
+      if (assetName == fileToDownload)
+      {
+        info.fileUrl = asset["browser_download_url"].as<String>();
+      }
+      else if (assetName == fileToDownload + ".sha256")
+      {
+        info.checksumUrl = asset["browser_download_url"].as<String>();
+      }
+
+      if (!info.fileUrl.isEmpty() && !info.checksumUrl.isEmpty())
       {
-        downloadUrl = asset["browser_download_url"].as<String>();
         break;
       }
     }
-    Serial.printf("Latest release URL: %s\r\n", downloadUrl.c_str());
+    Serial.printf("Latest release URL: %s\r\n", info.fileUrl.c_str());
+    Serial.printf("Checksum URL: %s\r\n", info.checksumUrl.c_str());
   }
-  return downloadUrl;
+  http.end();
+  return info;
 }
 
 int downloadUpdateHandler(char updateType)
@@ -131,7 +155,7 @@ int downloadUpdateHandler(char updateType)
   HTTPClient http;
   http.setFollowRedirects(HTTPC_STRICT_FOLLOW_REDIRECTS);
 
-  String latestRelease = "";
+  ReleaseInfo latestRelease;
 
   switch (updateType)
   {
@@ -143,25 +167,22 @@ int downloadUpdateHandler(char updateType)
   case UPDATE_WEBUI:
   {
     latestRelease = getLatestRelease("littlefs.bin");
-    updateWebUi(latestRelease, U_SPIFFS);
-    return 0;
+    // updateWebUi(latestRelease.fileUrl, U_SPIFFS);
+    // return 0;
   }
   break;
   }
 
-  if (latestRelease.isEmpty())
-  {
-    return 503;
-  }
+
   // First, download the expected SHA256
-  String expectedSHA256 = downloadSHA256(getFirmwareFilename());
+  String expectedSHA256 = downloadSHA256(latestRelease.checksumUrl);
   if (expectedSHA256.isEmpty())
   {
     Serial.println("Failed to get SHA256 checksum. Aborting update.");
     return false;
   }
 
-  http.begin(client, latestRelease);
+  http.begin(client, latestRelease.fileUrl);
   http.setUserAgent(USER_AGENT);
 
   int httpCode = http.GET();
@@ -215,19 +236,21 @@ int downloadUpdateHandler(char updateType)
       
       Update.onProgress(onOTAProgress);
 
-      int updateType = (updateType == UPDATE_WEBUI) ? U_SPIFFS : U_FLASH;
-
       if (Update.begin(contentLength, updateType))
       {
-        size_t written = Update.writeStream(*stream);
+        onOTAStart();
+        size_t written = Update.write(firmware, contentLength);
 
         if (written == contentLength)
         {
           Serial.println("Written : " + String(written) + " successfully");
+          free(firmware);
         }
         else
         {
           Serial.println("Written only : " + String(written) + "/" + String(contentLength) + ". Retry?");
+          free(firmware);
+          return 503;
         }
 
         if (Update.end())
@@ -236,26 +259,33 @@ int downloadUpdateHandler(char updateType)
           if (Update.isFinished())
           {
             Serial.println("Update successfully completed. Rebooting.");
-            ESP.restart();
+//            ESP.restart();
           }
           else
           {
             Serial.println("Update not finished? Something went wrong!");
+            free(firmware);
+            return 503;
           }
         }
         else
         {
           Serial.println("Error Occurred. Error #: " + String(Update.getError()));
+          free(firmware);
+          return 503;
         }
       }
       else
       {
         Serial.println("Not enough space to begin OTA");
+        free(firmware);
+        return 503;
       }
     }
     else
     {
       Serial.println("Invalid content length");
+      return 503;
     }
   }
   else
@@ -265,7 +295,7 @@ int downloadUpdateHandler(char updateType)
   }
   http.end();
 
-  return 200;
+  return 0;
 }
 
 void updateWebUi(String latestRelease, int command)
@@ -380,9 +410,8 @@ bool getIsOTAUpdating()
   return isOtaUpdating;
 }
 
-String downloadSHA256(const String &filename)
+String downloadSHA256(const String &sha256Url)
 {
-  String sha256Url = getLatestRelease(filename + ".sha256");
   if (sha256Url.isEmpty())
   {
     Serial.println("Failed to get SHA256 file URL");
diff --git a/src/lib/ota.hpp b/src/lib/ota.hpp
@@ -15,6 +15,11 @@ typedef struct {
 
 extern QueueHandle_t otaQueue;
 
+struct ReleaseInfo {
+  String fileUrl;
+  String checksumUrl;
+};
+
 void setupOTA();
 void onOTAStart();
 void handleOTATask(void *parameter);
@@ -23,9 +28,10 @@ void onOTAProgress(unsigned int progress, unsigned int total);
 void onOTAError(ota_error_t error);
 void onOTAComplete();
 int downloadUpdateHandler(char updateType);
-String getLatestRelease(const String& fileToDownload);
+ReleaseInfo getLatestRelease(const String& fileToDownload);
 
 bool getIsOTAUpdating();
 
 void updateWebUi(String latestRelease, int command);
-String downloadSHA256(const String& filename);
+String downloadSHA256(const String& filename);
+
diff --git a/src/lib/shared.hpp b/src/lib/shared.hpp
@@ -9,6 +9,7 @@
 #include <GxEPD2.h>
 #include <GxEPD2_BW.h>
 #include <mbedtls/md.h>
+#include <Update.h>
 
 #include <mutex>
 #include <utils.hpp>
@@ -69,9 +70,9 @@ const int usPerMinute = 60 * usPerSecond;
 
 extern const char *github_root_ca;
 
-const PROGMEM char UPDATE_FIRMWARE = 0;
-const PROGMEM char UPDATE_WEBUI = 1;
-
+const PROGMEM char UPDATE_FIRMWARE = U_FLASH;
+const PROGMEM char UPDATE_WEBUI = U_SPIFFS;
+const PROGMEM char UPDATE_ALL = 99;
 
 struct ScreenMapping {
     int value;
diff --git a/src/lib/webserver.cpp b/src/lib/webserver.cpp
@@ -86,8 +86,7 @@ void setupWebserver()
   {
     server.on("/upload/firmware", HTTP_POST, onFirmwareUpdate, asyncFirmwareUpdateHandler);
     server.on("/upload/webui", HTTP_POST, onFirmwareUpdate, asyncWebuiUpdateHandler);
-    // server.on("/update/webui", HTTP_GET, onUpdateWebUi);
-    // server.on("/update/firmware", HTTP_GET, onUpdateFirmware);
+    server.on("/api/firmware/auto_update", HTTP_GET, onAutoUpdateFirmware);
   }
 
   server.on("/api/restart", HTTP_GET, onApiRestart);
@@ -142,29 +141,16 @@ void onFirmwareUpdate(AsyncWebServerRequest *request)
   request->send(response);
 }
 
-void onUpdateWebUi(AsyncWebServerRequest *request)
+void onAutoUpdateFirmware(AsyncWebServerRequest *request)
 {
-  UpdateMessage msg = {UPDATE_WEBUI};
+  UpdateMessage msg = {UPDATE_ALL};
   if (xQueueSend(otaQueue, &msg, 0) == pdTRUE)
   {
-    request->send(200, "text/plain", "WebUI update triggered");
+    request->send(200, "application/json", "{\"msg\":\"Firmware update triggered\"}");
   }
   else
   {
-    request->send(503, "text/plain", "Update already in progress");
-  }
-}
-
-void onUpdateFirmware(AsyncWebServerRequest *request)
-{
-  UpdateMessage msg = {UPDATE_FIRMWARE};
-  if (xQueueSend(otaQueue, &msg, 0) == pdTRUE)
-  {
-    request->send(200, "text/plain", "Firmware update triggered");
-  }
-  else
-  {
-    request->send(503, "text/plain", "Update already in progress");
+    request->send(503,"application/json", "{\"msg\":\"Update already in progress\"}"); 
   }
 }
 
diff --git a/src/lib/webserver.hpp b/src/lib/webserver.hpp
@@ -26,10 +26,6 @@ bool processEpdColorSettings(AsyncWebServerRequest *request);
 void onApiStatus(AsyncWebServerRequest *request);
 void onApiSystemStatus(AsyncWebServerRequest *request);
 void onApiSetWifiTxPower(AsyncWebServerRequest *request);
-void onUpdateWebUi(AsyncWebServerRequest *request);
-void onUpdateFirmware(AsyncWebServerRequest *request);
-
-
 
 void onApiScreenNext(AsyncWebServerRequest *request);
 void onApiScreenPrevious(AsyncWebServerRequest *request);
@@ -58,6 +54,7 @@ void onFirmwareUpdate(AsyncWebServerRequest *request);
 void asyncFirmwareUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final);
 void asyncFileUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final, int command);
 void asyncWebuiUpdateHandler(AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final);
+void onAutoUpdateFirmware(AsyncWebServerRequest *request);
 
 void onIndex(AsyncWebServerRequest *request);
 void onNotFound(AsyncWebServerRequest *request);

Original file line number	Diff line number	Diff line change
`@@ -4,6 +4,8 @@ TaskHandle_t taskOtaHandle = NULL;`
`4`	`4`	`bool isOtaUpdating = false;`
`5`	`5`	`QueueHandle_t otaQueue;`
`6`	`6`
	`7`	`+`
	`8`	`+`
`7`	`9`	`void setupOTA()`
`8`	`10`	`{`
`9`	`11`	`if (preferences.getBool("otaEnabled", DEFAULT_OTA_ENABLED))`
`@@ -65,7 +67,7 @@ void onOTAStart()`
`65`	`67`	`vTaskSuspend(workerTaskHandle);`
`66`	`68`	`vTaskSuspend(taskScreenRotateTaskHandle);`
`67`	`69`
`68`		`- vTaskSuspend(ledTaskHandle);`
	`70`	`+// vTaskSuspend(ledTaskHandle);`
`69`	`71`	`vTaskSuspend(buttonTaskHandle);`
`70`	`72`
`71`	`73`	`// stopWebServer();`
`@@ -81,15 +83,26 @@ void handleOTATask(void *parameter)`
`81`	`83`	`{`
`82`	`84`	`if (xQueueReceive(otaQueue, &msg, 0) == pdTRUE)`
`83`	`85`	`{`
`84`		`- int result = downloadUpdateHandler(msg.updateType);`
	`86`	`+ if (msg.updateType == UPDATE_ALL) {`
	`87`	`+ int resultWebUi = downloadUpdateHandler(UPDATE_WEBUI);`
	`88`	`+ int resultFw = downloadUpdateHandler(UPDATE_FIRMWARE);`
	`89`	`+`
	`90`	`+ if (resultWebUi == 0 && resultFw == 0) {`
	`91`	`+ ESP.restart();`
	`92`	`+ } else {`
	`93`	`+ queueLedEffect(LED_FLASH_ERROR);`
	`94`	`+ vTaskDelay(pdMS_TO_TICKS(3000));`
	`95`	`+ ESP.restart();`
	`96`	`+ }`
	`97`	`+ }`
`85`	`98`	`}`
`86`	`99`
`87`	`100`	`ArduinoOTA.handle(); // Allow OTA updates to occur`
`88`	`101`	`vTaskDelay(pdMS_TO_TICKS(2000));`
`89`	`102`	`}`
`90`	`103`	`}`
`91`	`104`
`92`		`-String getLatestRelease(const String &fileToDownload)`
	`105`	`+ReleaseInfo getLatestRelease(const String &fileToDownload)`
`93`	`106`	`{`
`94`	`107`	`String releaseUrl = "https://api.github.com/repos/btclock/btclock_v3/releases/latest";`
`95`	`108`	`WiFiClientSecure client;`
`@@ -100,7 +113,7 @@ String getLatestRelease(const String &fileToDownload)`
`100`	`113`
`101`	`114`	`int httpCode = http.GET();`
`102`	`115`
`103`		`- String downloadUrl = "";`
	`116`	`+ ReleaseInfo info = {"", ""};`
`104`	`117`
`105`	`118`	`if (httpCode > 0)`
`106`	`119`	`{`
`@@ -113,15 +126,26 @@ String getLatestRelease(const String &fileToDownload)`
`113`	`126`
`114`	`127`	`for (JsonObject asset : assets)`
`115`	`128`	`{`
`116`		`- if (asset["name"] == fileToDownload)`
	`129`	`+ String assetName = asset["name"].as<String>();`
	`130`	`+ if (assetName == fileToDownload)`
	`131`	`+ {`
	`132`	`+ info.fileUrl = asset["browser_download_url"].as<String>();`
	`133`	`+ }`
	`134`	`+ else if (assetName == fileToDownload + ".sha256")`
	`135`	`+ {`
	`136`	`+ info.checksumUrl = asset["browser_download_url"].as<String>();`
	`137`	`+ }`
	`138`	`+`
	`139`	`+ if (!info.fileUrl.isEmpty() && !info.checksumUrl.isEmpty())`
`117`	`140`	`{`
`118`		`- downloadUrl = asset["browser_download_url"].as<String>();`
`119`	`141`	`break;`
`120`	`142`	`}`
`121`	`143`	`}`
`122`		`- Serial.printf("Latest release URL: %s\r\n", downloadUrl.c_str());`
	`144`	`+ Serial.printf("Latest release URL: %s\r\n", info.fileUrl.c_str());`
	`145`	`+ Serial.printf("Checksum URL: %s\r\n", info.checksumUrl.c_str());`
`123`	`146`	`}`
`124`		`- return downloadUrl;`
	`147`	`+ http.end();`
	`148`	`+ return info;`
`125`	`149`	`}`
`126`	`150`
`127`	`151`	`int downloadUpdateHandler(char updateType)`
`@@ -131,7 +155,7 @@ int downloadUpdateHandler(char updateType)`
`131`	`155`	`HTTPClient http;`
`132`	`156`	`http.setFollowRedirects(HTTPC_STRICT_FOLLOW_REDIRECTS);`
`133`	`157`
`134`		`- String latestRelease = "";`
	`158`	`+ ReleaseInfo latestRelease;`
`135`	`159`
`136`	`160`	`switch (updateType)`
`137`	`161`	`{`
`@@ -143,25 +167,22 @@ int downloadUpdateHandler(char updateType)`
`143`	`167`	`case UPDATE_WEBUI:`
`144`	`168`	`{`
`145`	`169`	`latestRelease = getLatestRelease("littlefs.bin");`
`146`		`- updateWebUi(latestRelease, U_SPIFFS);`
`147`		`- return 0;`
	`170`	`+ // updateWebUi(latestRelease.fileUrl, U_SPIFFS);`
	`171`	`+ // return 0;`
`148`	`172`	`}`
`149`	`173`	`break;`
`150`	`174`	`}`
`151`	`175`
`152`		`- if (latestRelease.isEmpty())`
`153`		`- {`
`154`		`- return 503;`
`155`		`- }`
	`176`	`+`
`156`	`177`	`// First, download the expected SHA256`
`157`		`- String expectedSHA256 = downloadSHA256(getFirmwareFilename());`
	`178`	`+ String expectedSHA256 = downloadSHA256(latestRelease.checksumUrl);`
`158`	`179`	`if (expectedSHA256.isEmpty())`
`159`	`180`	`{`
`160`	`181`	`Serial.println("Failed to get SHA256 checksum. Aborting update.");`
`161`	`182`	`return false;`
`162`	`183`	`}`
`163`	`184`
`164`		`- http.begin(client, latestRelease);`
	`185`	`+ http.begin(client, latestRelease.fileUrl);`
`165`	`186`	`http.setUserAgent(USER_AGENT);`
`166`	`187`
`167`	`188`	`int httpCode = http.GET();`
`@@ -215,19 +236,21 @@ int downloadUpdateHandler(char updateType)`
`215`	`236`
`216`	`237`	`Update.onProgress(onOTAProgress);`
`217`	`238`
`218`		`- int updateType = (updateType == UPDATE_WEBUI) ? U_SPIFFS : U_FLASH;`
`219`		`-`
`220`	`239`	`if (Update.begin(contentLength, updateType))`
`221`	`240`	`{`
`222`		`- size_t written = Update.writeStream(*stream);`
	`241`	`+ onOTAStart();`
	`242`	`+ size_t written = Update.write(firmware, contentLength);`
`223`	`243`
`224`	`244`	`if (written == contentLength)`
`225`	`245`	`{`
`226`	`246`	`Serial.println("Written : " + String(written) + " successfully");`
	`247`	`+ free(firmware);`
`227`	`248`	`}`
`228`	`249`	`else`
`229`	`250`	`{`
`230`	`251`	`Serial.println("Written only : " + String(written) + "/" + String(contentLength) + ". Retry?");`
	`252`	`+ free(firmware);`
	`253`	`+ return 503;`
`231`	`254`	`}`
`232`	`255`
`233`	`256`	`if (Update.end())`
`@@ -236,26 +259,33 @@ int downloadUpdateHandler(char updateType)`
`236`	`259`	`if (Update.isFinished())`
`237`	`260`	`{`
`238`	`261`	`Serial.println("Update successfully completed. Rebooting.");`
`239`		`- ESP.restart();`
	`262`	`+// ESP.restart();`
`240`	`263`	`}`
`241`	`264`	`else`
`242`	`265`	`{`
`243`	`266`	`Serial.println("Update not finished? Something went wrong!");`
	`267`	`+ free(firmware);`
	`268`	`+ return 503;`
`244`	`269`	`}`
`245`	`270`	`}`
`246`	`271`	`else`
`247`	`272`	`{`
`248`	`273`	`Serial.println("Error Occurred. Error #: " + String(Update.getError()));`
	`274`	`+ free(firmware);`
	`275`	`+ return 503;`
`249`	`276`	`}`
`250`	`277`	`}`
`251`	`278`	`else`
`252`	`279`	`{`
`253`	`280`	`Serial.println("Not enough space to begin OTA");`
	`281`	`+ free(firmware);`
	`282`	`+ return 503;`
`254`	`283`	`}`
`255`	`284`	`}`
`256`	`285`	`else`
`257`	`286`	`{`
`258`	`287`	`Serial.println("Invalid content length");`
	`288`	`+ return 503;`
`259`	`289`	`}`
`260`	`290`	`}`
`261`	`291`	`else`
`@@ -265,7 +295,7 @@ int downloadUpdateHandler(char updateType)`
`265`	`295`	`}`
`266`	`296`	`http.end();`
`267`	`297`
`268`		`- return 200;`
	`298`	`+ return 0;`
`269`	`299`	`}`
`270`	`300`
`271`	`301`	`void updateWebUi(String latestRelease, int command)`
`@@ -380,9 +410,8 @@ bool getIsOTAUpdating()`
`380`	`410`	`return isOtaUpdating;`
`381`	`411`	`}`
`382`	`412`
`383`		`-String downloadSHA256(const String &filename)`
	`413`	`+String downloadSHA256(const String &sha256Url)`
`384`	`414`	`{`
`385`		`- String sha256Url = getLatestRelease(filename + ".sha256");`
`386`	`415`	`if (sha256Url.isEmpty())`
`387`	`416`	`{`
`388`	`417`	`Serial.println("Failed to get SHA256 file URL");`
Original file line number	Diff line number	Diff line change
`@@ -86,8 +86,7 @@ void setupWebserver()`
`86`	`86`	`{`
`87`	`87`	`server.on("/upload/firmware", HTTP_POST, onFirmwareUpdate, asyncFirmwareUpdateHandler);`
`88`	`88`	`server.on("/upload/webui", HTTP_POST, onFirmwareUpdate, asyncWebuiUpdateHandler);`
`89`		`- // server.on("/update/webui", HTTP_GET, onUpdateWebUi);`
`90`		`- // server.on("/update/firmware", HTTP_GET, onUpdateFirmware);`
	`89`	`+ server.on("/api/firmware/auto_update", HTTP_GET, onAutoUpdateFirmware);`
`91`	`90`	`}`
`92`	`91`
`93`	`92`	`server.on("/api/restart", HTTP_GET, onApiRestart);`
`@@ -142,29 +141,16 @@ void onFirmwareUpdate(AsyncWebServerRequest *request)`
`142`	`141`	`request->send(response);`
`143`	`142`	`}`
`144`	`143`
`145`		`-void onUpdateWebUi(AsyncWebServerRequest *request)`
	`144`	`+void onAutoUpdateFirmware(AsyncWebServerRequest *request)`
`146`	`145`	`{`
`147`		`- UpdateMessage msg = {UPDATE_WEBUI};`
	`146`	`+ UpdateMessage msg = {UPDATE_ALL};`
`148`	`147`	`if (xQueueSend(otaQueue, &msg, 0) == pdTRUE)`
`149`	`148`	`{`
`150`		`- request->send(200, "text/plain", "WebUI update triggered");`
	`149`	`+ request->send(200, "application/json", "{\"msg\":\"Firmware update triggered\"}");`
`151`	`150`	`}`
`152`	`151`	`else`
`153`	`152`	`{`
`154`		`- request->send(503, "text/plain", "Update already in progress");`
`155`		`- }`
`156`		`-}`
`157`		`-`
`158`		`-void onUpdateFirmware(AsyncWebServerRequest *request)`
`159`		`-{`
`160`		`- UpdateMessage msg = {UPDATE_FIRMWARE};`
`161`		`- if (xQueueSend(otaQueue, &msg, 0) == pdTRUE)`
`162`		`- {`
`163`		`- request->send(200, "text/plain", "Firmware update triggered");`
`164`		`- }`
`165`		`- else`
`166`		`- {`
`167`		`- request->send(503, "text/plain", "Update already in progress");`
	`153`	`+ request->send(503,"application/json", "{\"msg\":\"Update already in progress\"}");`
`168`	`154`	`}`
`169`	`155`	`}`
`170`	`156`