Commit WiFi OTA only after the complete upload request

arch/esp32/AsyncElegantOTA/src/AsyncElegantOTA.cpp

@@ -8,6 +8,10 @@ static std::atomic<bool> ota_reboot_pending{false};
 // All upload callbacks run on AsyncTCP's task. Update is a global singleton;
 // retain one request as its owner until that request disconnects.
 static AsyncWebServerRequest* ota_upload_owner = nullptr;
+struct OtaUploadState {
+    bool file_final;
+    bool committed;
+};
 
 AsyncElegantOtaClass AsyncElegantOTA;
 
@@ -62,8 +66,18 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
         // send() queues a response until the request body has finished. Keep
         // an earlier upload error even before it has reached the wire.
         if (request->isSent() || request->getResponse()) return;
-        const bool complete = request->_tempObject
-            && *static_cast<bool*>(request->_tempObject) && !Update.hasError();
+        auto* state = static_cast<OtaUploadState*>(request->_tempObject);
+        // A file boundary can precede more multipart fields or another file.
+        // Update.end() selects the next boot partition, so commit only after
+        // the whole POST has completed without an earlier upload error.
+        if (state && state->file_final && !state->committed
+            && ota_upload_owner == request && !Update.hasError()) {
+            if (!Update.end(true)) {
+                return request->send(400, "text/plain", "Could not end OTA");
+            }
+            state->committed = true;
+        }
+        const bool complete = state && state->committed && !Update.hasError();
         AsyncWebServerResponse *response = request->beginResponse(complete?200:500, "text/plain", complete?"OK":"FAIL");
         response->addHeader("Connection", "close");
         response->addHeader("Access-Control-Allow-Origin", "*");
@@ -128,7 +142,7 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
 
             // AsyncWebServer frees this per-request allocation on disconnect.
             // hasError() alone also reports success for a POST with no upload.
-            request->_tempObject = calloc(1, sizeof(bool));
+            request->_tempObject = calloc(1, sizeof(OtaUploadState));
             if (!request->_tempObject) {
                 Update.abort();
                 ota_upload_busy.store(false);
@@ -138,7 +152,7 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
             request->onDisconnect([this, request]() {
                 if (ota_upload_owner != request) return;
                 ota_upload_owner = nullptr;
-                if (*static_cast<bool*>(request->_tempObject)) {
+                if (static_cast<OtaUploadState*>(request->_tempObject)->committed) {
                     // Let the response drain and the network callback return
                     // before rebooting from the separate restart task.
                     restart();
@@ -158,13 +172,8 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
             }
         }
             
-        if (final) { // if the final flag is set then this is the last frame of data
-            if (!Update.end(true)) { //true to set the size to the current progress
-                return request->send(400, "text/plain", "Could not end OTA");
-            }
-            *static_cast<bool*>(request->_tempObject) = true;
-        }else{
-            return;
+        if (final) {
+            static_cast<OtaUploadState*>(request->_tempObject)->file_final = true;
         }
     });
 }

test/fixtures/async_ota_completion.cpp

@@ -32,6 +32,7 @@ int xTaskCreate(Task task, const char*, unsigned, void*, unsigned, void*) {
   assert(!pending_task); pending_task=task; return pdPASS;
 }
 @REBOOT_STATE@
+@UPLOAD_STATE@
 // Independent MD5 vectors: eight zero bytes, and that payload with byte zero
 // changed to one. The fake models Update's lifecycle rather than its digest
 // algorithm; both vectors are fixed independently of the uploader callbacks.
@@ -100,7 +101,7 @@ struct AsyncWebServer {
 };
 struct UpdateFake {
   bool error=false, running=false, fail_begin=false, fail_end=false, fail_write=false;
-  int aborts=0, writes=0, begins=0, md5_checks=0;
+  int aborts=0, writes=0, begins=0, md5_checks=0, ends=0, boot_commits=0;
   String target_md5;
   std::vector<uint8_t> received;
   bool hasError() const { return error; }
@@ -126,6 +127,7 @@ struct UpdateFake {
     return size;
   }
   bool end(bool) {
+    ++ends;
     if (error || !running) return false;
     if (fail_end) { error=true; running=false; return false; }
     assert(received.size()==8);
@@ -140,6 +142,9 @@ struct UpdateFake {
       if (target_md5!=actual_md5) { error=true; running=false; return false; }
     }
     running=false;
+    // Pinned UpdateClass::_verifyEnd() activates the boot partition when end()
+    // succeeds. Track that irreversible effect separately from file receipt.
+    ++boot_commits;
     return true;
   }
   void abort() { error=true; running=false; ++aborts; }
@@ -173,6 +178,47 @@ int main() {
     server.upload(&request,"firmware",0,data,sizeof(data),true);
     in_network_callback=false;
   };
+#if defined(TEST_OTA_SECOND_FILE) || defined(TEST_OTA_INCOMPLETE_POST)
+  {
+  Request rejected;
+  upload(rejected);
+  assert(ota_upload_owner==&rejected && ota_upload_busy.load()
+      && !rejected.isSent() && !rejected.getResponse());
+#ifdef TEST_OTA_SECOND_FILE
+  // AsyncWebServer starts each multipart file at index zero. A second file
+  // must keep its queued 409 instead of committing the already received file.
+  upload(rejected);
+  assert(!rejected.isSent() && rejected.getResponse() && rejected.code==409);
+  complete(rejected);
+  assert(rejected.code==409 && rejected.body=="OTA upload already active");
+#else
+  // The file-final boundary can arrive without the closing POST boundary.
+  // Do not invoke the completion callback for this interrupted request.
+  assert(rejected.code==0);
+#endif
+  in_network_callback=true; rejected.disconnected(); in_network_callback=false;
+  // Execute any incorrectly scheduled restart so the negative control checks
+  // both scheduling and the boot effect, rather than leaving work pending.
+  if (pending_task) { auto task=pending_task; pending_task=nullptr; task(nullptr); }
+  assert(Update.ends==0 && Update.boot_commits==0 && Update.md5_checks==0
+      && Update.aborts==1 && !Update.running && !ota_upload_owner
+      && !ota_upload_busy.load() && !ota_reboot_pending.load()
+      && !pending_task && ESP.restarts==0);
+  // Cleanup permits a valid upload on the same listener without restarting it.
+  Request valid;
+  upload(valid);
+  assert(Update.ends==0 && Update.boot_commits==0 && Update.running);
+  complete(valid);
+  assert(valid.code==200 && valid.body=="OK" && Update.ends==1
+      && Update.boot_commits==1 && Update.md5_checks==1 && !Update.running
+      && ota_upload_owner==&valid && !pending_task && ESP.restarts==0);
+  in_network_callback=true; valid.disconnected(); in_network_callback=false;
+  assert(!ota_upload_owner && !ota_upload_busy.load() && pending_task && ESP.restarts==0);
+  auto task=pending_task; pending_task=nullptr; task(nullptr);
+  assert(ESP.restarts==1);
+  return 0;
+  }
+#endif
 #ifdef TEST_OTA_BODY_TIMEOUT
   {
   uint8_t chunk[4]={};
@@ -285,7 +331,7 @@ int main() {
   upload(mismatched_md5); complete(mismatched_md5);
   assert(mismatched_md5.code==400 && mismatched_md5.body=="Could not end OTA"
       && ota_upload_owner==&mismatched_md5 && ota_upload_busy.load()
-      && !*static_cast<bool*>(mismatched_md5._tempObject)
+      && !static_cast<OtaUploadState*>(mismatched_md5._tempObject)->committed
       && Update.md5_checks==1 && !pending_task && ESP.restarts==0);
   in_network_callback=true; mismatched_md5.disconnected(); in_network_callback=false;
   assert(Update.aborts==2 && !ota_upload_owner && !ota_upload_busy.load()
@@ -300,7 +346,7 @@ int main() {
   complete(corrupt_bytes);
   assert(corrupt_bytes.code==400 && corrupt_bytes.body=="Could not end OTA"
       && ota_upload_owner==&corrupt_bytes && ota_upload_busy.load()
-      && !*static_cast<bool*>(corrupt_bytes._tempObject)
+      && !static_cast<OtaUploadState*>(corrupt_bytes._tempObject)->committed
       && Update.md5_checks==2 && !pending_task && ESP.restarts==0);
   in_network_callback=true; corrupt_bytes.disconnected(); in_network_callback=false;
   assert(Update.aborts==3 && !ota_upload_owner && !ota_upload_busy.load()
@@ -344,9 +390,14 @@ int main() {
   assert(Update.aborts==6 && !ota_upload_owner && !ota_upload_busy.load());
   Update.fail_write=false;
   Request valid;
-  upload(valid); complete(valid);
+  const int ends_before_valid=Update.ends, commits_before_valid=Update.boot_commits;
+  upload(valid);
+  assert(Update.ends==ends_before_valid && Update.boot_commits==commits_before_valid
+      && Update.running && !static_cast<OtaUploadState*>(valid._tempObject)->committed);
+  complete(valid);
   assert(valid.code==200 && valid.disconnected && !pending_task && ESP.restarts==0
-      && Update.md5_checks==3 && !Update.running);
+      && Update.md5_checks==3 && !Update.running
+      && Update.ends==ends_before_valid+1 && Update.boot_commits==commits_before_valid+1);
   // The response drains before disconnect; its callback only schedules work.
   in_network_callback=true; valid.disconnected(); in_network_callback=false;
   assert(pending_task && ESP.restarts==0);

test/test_async_web_lifecycle.py

@@ -40,6 +40,7 @@ static std::atomic<bool> ota_upload_busy{false};"""
         accept=re.search(r"c->setRxTimeout\(\d+\);",web).group(0)
         rx=tcp[tcp.index("  // RX Timeout"):tcp.index("  // Everything is fine")]
         return (fixture.replace("@REBOOT_STATE@",state)
+            .replace("@UPLOAD_STATE@",extract_braced(source,"struct OtaUploadState")+";")
             .replace("@BEGIN@",begin)
             .replace("@RESTART@",extract_braced(source,"void AsyncElegantOtaClass::restart()"))
             .replace("@ACCEPT_TIMEOUT@",accept).replace("@RX_TIMEOUT@",rx))
@@ -67,6 +68,34 @@ static std::atomic<bool> ota_upload_busy{false};"""
         with self.assertRaises(AssertionError):
             self.compile_and_run(old,"-DESP32=1","-DTEST_OTA_BODY_TIMEOUT=1")
 
+    def premature_commit_fixture(self, fixture):
+        commit=extract_braced(fixture,"if (state && state->file_final")
+        end=extract_braced(commit,"if (!Update.end(true))")
+        final="static_cast<OtaUploadState*>(request->_tempObject)->file_final = true;"
+        # Restore the old policy using the actual end/error block: commit at
+        # file-final, before the complete multipart POST has been accepted.
+        old=fixture.replace(commit,"",1).replace(final,end+"\n"
+            +"static_cast<OtaUploadState*>(request->_tempObject)->committed = true;\n"
+            +final,1)
+        self.assertNotEqual(fixture,old,"The negative control must restore premature file commit")
+        return old
+
+    def test_second_file_rejection_does_not_commit_first_file_or_reboot(self):
+        fixture=self.uploader_fixture()
+        flags=("-DESP32=1","-DTEST_OTA_SECOND_FILE=1")
+        self.compile_and_run(fixture,*flags)
+        with self.assertRaisesRegex(AssertionError,r"Assertion (?:.* failed|failed:)") as failure:
+            self.compile_and_run(self.premature_commit_fixture(fixture),*flags)
+        self.assertRegex(str(failure.exception),r"Update\.ends==0 && Update\.boot_commits==0")
+
+    def test_file_final_disconnect_before_post_completion_does_not_commit_or_reboot(self):
+        fixture=self.uploader_fixture()
+        flags=("-DESP32=1","-DTEST_OTA_INCOMPLETE_POST=1")
+        self.compile_and_run(fixture,*flags)
+        with self.assertRaisesRegex(AssertionError,r"Assertion (?:.* failed|failed:)") as failure:
+            self.compile_and_run(self.premature_commit_fixture(fixture),*flags)
+        self.assertRegex(str(failure.exception),r"Update\.ends==0 && Update\.boot_commits==0")
+
     def test_md5_survives_begin_and_rejected_digest_releases_update(self):
         fixture=self.uploader_fixture()
         self.compile_and_run(fixture,"-DESP32=1")