Keep admitted WiFi OTA uploads alive through receive gaps

arch/esp32/AsyncElegantOTA/src/AsyncElegantOTA.cpp

@@ -59,7 +59,9 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
                 return request->requestAuthentication();
             }
         }
-        if (request->isSent()) return; // Preserve an earlier upload error.
+        // 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();
         AsyncWebServerResponse *response = request->beginResponse(complete?200:500, "text/plain", complete?"OK":"FAIL");
@@ -68,6 +70,7 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
         request->send(response);
     }, [&](AsyncWebServerRequest *request, String filename, size_t index, uint8_t *data, size_t len, bool final) {
         //Upload handler chunks in data
+        if (request->isSent() || request->getResponse()) return;
         if(_authRequired){
             if(!request->authenticate(_username.c_str(), _password.c_str())){
                 return request->requestAuthentication();
@@ -102,6 +105,11 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
                 return request->send(400, "text/plain", "MD5 parameter invalid");
             }
 
+            // The HTTP listener defaults to a three-second receive timeout.
+            // A valid OTA body can pause longer on a weak WiFi link; retain a
+            // bounded timeout on this admitted upload without changing peers.
+            request->client()->setRxTimeout(30);
+
             #if defined(ESP8266)
                 int cmd = (filename == "filesystem") ? U_FS : U_FLASH;
                 Update.runAsync(true);
@@ -112,7 +120,6 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
                 int cmd = (filename == "filesystem") ? U_SPIFFS : U_FLASH;
                 if (!Update.begin(UPDATE_SIZE_UNKNOWN, cmd)) { // Start with max available size
             #endif
-                Update.printError(Serial);
                 ota_upload_busy.store(false);
                 return request->send(400, "text/plain", "OTA could not begin");
             }
@@ -144,13 +151,12 @@ void AsyncElegantOtaClass::begin(AsyncWebServer *server, const char* username, c
         // Write chunked data to the free sketch space
         if(len){
             if (Update.write(data, len) != len) {
-                return request->send(400, "text/plain", "OTA could not begin");
+                return request->send(400, "text/plain", "OTA flash write failed");
             }
         }
             
         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
-                Update.printError(Serial);
                 return request->send(400, "text/plain", "Could not end OTA");
             }
             *static_cast<bool*>(request->_tempObject) = true;

test/fixtures/async_ota_completion.cpp

@@ -4,10 +4,13 @@
 #include <cstring>
 #include <functional>
 #include <string>
+#include <vector>
 using String = std::string;
 constexpr int HTTP_GET=0, HTTP_POST=1, U_FLASH=0, U_SPIFFS=100;
 constexpr unsigned UPDATE_SIZE_UNKNOWN=~0U;
 constexpr int pdPASS=1;
+constexpr int ERR_OK=0;
+#define async_tcp_log_d(...) ((void)0)
 #define pdMS_TO_TICKS(value) (value)
 const char ELEGANT_HTML[] = "page";
 const unsigned ELEGANT_HTML_SIZE=sizeof(ELEGANT_HTML);
@@ -32,26 +35,52 @@ int xTaskCreate(Task task, const char*, unsigned, void*, unsigned, void*) {
 struct Parameter { String text="valid"; String value() { return text; } };
 struct Response {
   int code;
+  String body;
   void addHeader(const char*, const char*) {}
 };
+unsigned fixture_now=0;
+struct Client {
+  unsigned _rx_timeout=0, _rx_last_packet=fixture_now;
+  bool closed=false;
+  std::vector<unsigned> timeout_changes;
+  std::function<void()> disconnected;
+  void setRxTimeout(unsigned seconds) { _rx_timeout=seconds; timeout_changes.push_back(seconds); }
+  void received() { _rx_last_packet=fixture_now; }
+  void _close() { closed=true; if (disconnected) disconnected(); }
+  int poll() {
+    const unsigned now=fixture_now;
+@RX_TIMEOUT@
+    return ERR_OK;
+  }
+};
+void accept(Client* c) {
+@ACCEPT_TIMEOUT@
+}
 struct Request {
   void* _tempObject=nullptr;
   bool md5_present=true, authorized=true, sent=false;
   int code=0;
+  String body;
+  Response* queued_response=nullptr;
   std::function<void()> disconnected;
   Parameter md5;
-  ~Request() { free(_tempObject); }
+  Client connection;
+  Request() { accept(&connection); }
+  ~Request() { free(_tempObject); delete queued_response; }
+  Client* client() { return &connection; }
   bool authenticate(const char*, const char*) { return authorized; }
-  void requestAuthentication() { code=401; sent=true; }
+  void requestAuthentication() { send(401,"text/plain","unauthorized"); }
   bool hasParam(const char*, bool) { return md5_present; }
   Parameter* getParam(const char*, bool) { return &md5; }
-  Response* beginResponse(int code, const char*, const char*) { return new Response{code}; }
-  Response* beginResponse_P(int code, const char*, const char*, unsigned) { return new Response{code}; }
-  void send(Response* response) { code=response->code; sent=true; delete response; }
-  void send(int status, const char*, const char*) { code=status; sent=true; }
-  void send(int status, const char*, const String&) { code=status; sent=true; }
+  Response* beginResponse(int code, const char*, const char* text) { return new Response{code,text}; }
+  Response* beginResponse_P(int code, const char*, const char* text, unsigned) { return new Response{code,text}; }
+  void send(Response* response) { delete queued_response; queued_response=response; code=response->code; body=response->body; }
+  void send(int status, const char*, const char* text) { send(new Response{status,text}); }
+  void send(int status, const char*, const String& text) { send(new Response{status,text}); }
+  Response* getResponse() const { return queued_response; }
+  void flush_response() { assert(queued_response); sent=true; connection.setRxTimeout(0); }
   bool isSent() const { return sent; }
-  void onDisconnect(std::function<void()> fn) { disconnected=fn; }
+  void onDisconnect(std::function<void()> fn) { disconnected=fn; connection.disconnected=fn; }
 };
 using AsyncWebServerRequest=Request;
 static AsyncWebServerRequest* ota_upload_owner=nullptr;
@@ -65,15 +94,17 @@ struct AsyncWebServer {
   }
 };
 struct UpdateFake {
-  bool error=false, fail_end=false;
-  int aborts=0;
+  bool error=false, fail_begin=false, fail_end=false, fail_write=false;
+  int aborts=0, writes=0;
   bool hasError() const { return error; }
   bool setMD5(const char* md5) { return String(md5)=="valid"; }
-  bool begin(unsigned, int) { error=false; return true; }
-  size_t write(uint8_t*, size_t size) { return size; }
+  bool begin(unsigned, int) { error=fail_begin; return !error; }
+  size_t write(uint8_t*, size_t size) { ++writes; error=fail_write; return fail_write ? size-1 : size; }
   bool end(bool) { error=fail_end; return !error; }
   void abort() { error=true; ++aborts; }
-  void printError(int) {}
+  // A connected USB host which does not drain TX can block raw Serial. Every
+  // upload error path must complete without attempting that raw print.
+  void printError(int) { assert(false && "Raw USB printing can block OTA callbacks"); }
 } Update;
 struct AsyncElegantOtaClass {
   AsyncWebServer* _server=nullptr;
@@ -91,7 +122,9 @@ int main() {
   AsyncElegantOtaClass ota;
   ota.begin(&server,"","");
   auto complete=[&](Request& request) {
-    in_network_callback=true; server.completion(&request); in_network_callback=false;
+    in_network_callback=true;
+    server.completion(&request); request.flush_response();
+    in_network_callback=false;
   };
   auto upload=[&](Request& request) {
     uint8_t data[8]={};
@@ -99,6 +132,80 @@ int main() {
     server.upload(&request,"firmware",0,data,sizeof(data),true);
     in_network_callback=false;
   };
+#ifdef TEST_OTA_BODY_TIMEOUT
+  {
+  uint8_t chunk[4]={};
+  auto receive=[&](Request& request,size_t index,bool final) {
+    assert(!request.connection.closed);
+    request.connection.received();
+    in_network_callback=true;
+    server.upload(&request,"firmware",index,chunk,sizeof(chunk),final);
+    in_network_callback=false;
+  };
+  auto idle=[&](Request& request,unsigned ms) {
+    fixture_now+=ms;
+    in_network_callback=true; request.connection.poll(); in_network_callback=false;
+  };
+  Request ordinary;
+  assert(ordinary.connection._rx_timeout==3);
+  idle(ordinary,4000);
+  assert(ordinary.connection.closed && !ota_upload_owner);
+  Request admitted;
+  receive(admitted,0,false);
+  assert(ota_upload_owner==&admitted && !admitted.sent
+      && admitted.connection.timeout_changes==std::vector<unsigned>({3,30}));
+  Request concurrent;
+  upload(concurrent);
+  assert(!concurrent.isSent() && concurrent.getResponse()
+      && concurrent.connection.timeout_changes==std::vector<unsigned>({3}));
+  complete(concurrent);
+  assert(concurrent.code==409 && ota_upload_owner==&admitted
+      && concurrent.connection.timeout_changes==std::vector<unsigned>({3,0}));
+  Request untouched;
+  assert(untouched.connection._rx_timeout==3);
+  idle(admitted,4000);
+  assert(!admitted.connection.closed && ota_upload_owner==&admitted);
+  idle(admitted,25999);
+  assert(!admitted.connection.closed);
+  idle(admitted,1);
+  assert(admitted.connection.closed && Update.aborts==1 && !ota_upload_owner
+      && !ota_upload_busy.load() && !pending_task && ESP.restarts==0);
+  Request missing_md5;
+  missing_md5.md5_present=false;
+  upload(missing_md5);
+  assert(!missing_md5.isSent() && missing_md5.getResponse()
+      && missing_md5.connection.timeout_changes==std::vector<unsigned>({3}));
+  complete(missing_md5);
+  assert(missing_md5.code==400
+      && missing_md5.connection.timeout_changes==std::vector<unsigned>({3,0})
+      && !ota_upload_owner && !ota_upload_busy.load());
+  Request bad_md5;
+  bad_md5.md5.text="invalid";
+  upload(bad_md5); complete(bad_md5);
+  assert(bad_md5.code==400 && !ota_upload_owner && !ota_upload_busy.load()
+      && bad_md5.connection.timeout_changes==std::vector<unsigned>({3,0}));
+  ota.begin(&server,"admin","password");
+  Request unauthorized;
+  unauthorized.authorized=false;
+  upload(unauthorized); complete(unauthorized);
+  assert(unauthorized.code==401 && !ota_upload_owner && !ota_upload_busy.load()
+      && unauthorized.connection.timeout_changes==std::vector<unsigned>({3,0}));
+  ota.begin(&server,"","");
+  Request successful;
+  receive(successful,0,false);
+  idle(successful,4000);
+  assert(!successful.connection.closed && ota_upload_owner==&successful);
+  receive(successful,sizeof(chunk),true);
+  complete(successful);
+  assert(successful.code==200 && successful.body=="OK"
+      && successful.connection._rx_timeout==0 && !pending_task);
+  in_network_callback=true; successful.disconnected(); in_network_callback=false;
+  assert(pending_task && !ota_upload_owner);
+  auto task=pending_task; pending_task=nullptr; task(nullptr);
+  assert(ESP.restarts==1);
+  return 0;
+  }
+#endif
   // A POST without an actual completed upload cannot report OK or reboot.
   Request empty;
   complete(empty);
@@ -107,6 +214,13 @@ int main() {
   invalid.md5_present=false;
   upload(invalid); complete(invalid);
   assert(invalid.code==400 && !pending_task && ESP.restarts==0);
+  Request begin_failed;
+  Update.fail_begin=true;
+  upload(begin_failed); complete(begin_failed);
+  assert(begin_failed.code==400 && begin_failed.body=="OTA could not begin"
+      && !ota_upload_owner && !ota_upload_busy.load() && Update.aborts==0
+      && !pending_task && ESP.restarts==0);
+  Update.fail_begin=false;
   Request corrupt;
   Update.fail_end=true;
   upload(corrupt); complete(corrupt);
@@ -128,6 +242,23 @@ int main() {
   upload(stopped); complete(stopped);
   assert(stopped.code==503 && !ota_upload_owner && !pending_task);
   assert(ota.setEnabled(true));
+  Request write_failed;
+  Update.fail_write=true;
+  server.upload(&write_failed,"firmware",0,chunk,sizeof(chunk),false);
+  assert(!write_failed.isSent() && write_failed.getResponse()
+      && write_failed.code==400 && write_failed.body=="OTA flash write failed");
+  const int writes_after_error=Update.writes;
+  server.upload(&write_failed,"firmware",sizeof(chunk),chunk,sizeof(chunk),true);
+  assert(Update.writes==writes_after_error);
+  server.upload(&write_failed,"firmware",0,chunk,sizeof(chunk),false);
+  assert(Update.writes==writes_after_error && write_failed.code==400
+      && write_failed.body=="OTA flash write failed");
+  complete(write_failed);
+  assert(write_failed.code==400 && write_failed.body=="OTA flash write failed"
+      && !pending_task);
+  in_network_callback=true; write_failed.disconnected(); in_network_callback=false;
+  assert(Update.aborts==3 && !ota_upload_owner && !ota_upload_busy.load());
+  Update.fail_write=false;
   Request valid;
   upload(valid); complete(valid);
   assert(valid.code==200 && valid.disconnected && !pending_task && ESP.restarts==0);

test/test_async_web_lifecycle.py

@@ -1,6 +1,7 @@
 """Execute the vendored uploader and pinned TCP listener restart path."""
 import importlib.util
 from pathlib import Path
+import re
 import urllib.request
 import unittest
 
@@ -16,7 +17,7 @@ spec.loader.exec_module(fix)
 class AsyncWebLifecycleTest(unittest.TestCase):
     compile_and_run=wifi_start.WiFiOtaStartTest.compile_and_run
 
-    def test_upload_response_drains_before_deferred_reboot(self):
+    def uploader_fixture(self):
         source=(ROOT/"arch/esp32/AsyncElegantOTA/src/AsyncElegantOTA.cpp").read_text()
         fixture=(ROOT/"test/fixtures/async_ota_completion.cpp").read_text()
         state="""static std::atomic<bool> ota_reboot_pending{false};
@@ -25,9 +26,46 @@ static std::atomic<bool> ota_upload_busy{false};"""
         # The two preprocessor branches for Update.begin each open the same
         # logical block, so extract this method by its next declaration.
         begin=source[source.index("void AsyncElegantOtaClass::begin("):source.index("// deprecated")]
-        self.compile_and_run(fixture.replace("@REBOOT_STATE@",state)
+        cache=ROOT/".pio/libdeps/Station_G2_repeater_observer_mqtt"
+        web_file=cache/"ESPAsyncWebServer/src/WebServer.cpp"
+        tcp_file=cache/"AsyncTCP/src/AsyncTCP.cpp"
+        web=web_file.read_text() if web_file.exists() else urllib.request.urlopen(
+            "https://raw.githubusercontent.com/ESP32Async/ESPAsyncWebServer/v3.12.0/src/WebServer.cpp",
+            timeout=30).read().decode("ascii")
+        tcp=tcp_file.read_text() if tcp_file.exists() else urllib.request.urlopen(
+            "https://raw.githubusercontent.com/ESP32Async/AsyncTCP/v3.5.0/src/AsyncTCP.cpp",
+            timeout=30).read().decode("ascii")
+        # Exercise the pinned listener's actual default and AsyncTCP's close
+        # condition while running the production multipart upload callbacks.
+        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("@BEGIN@",begin)
-            .replace("@RESTART@",extract_braced(source,"void AsyncElegantOtaClass::restart()")),"-DESP32=1")
+            .replace("@RESTART@",extract_braced(source,"void AsyncElegantOtaClass::restart()"))
+            .replace("@ACCEPT_TIMEOUT@",accept).replace("@RX_TIMEOUT@",rx))
+
+    def test_upload_response_drains_before_deferred_reboot(self):
+        fixture=self.uploader_fixture()
+        self.compile_and_run(fixture,"-DESP32=1")
+        old=fixture.replace(" || request->getResponse()","")
+        self.assertNotEqual(fixture,old,"The negative control must remove queued-response guards")
+        with self.assertRaises(AssertionError):
+            self.compile_and_run(old,"-DESP32=1")
+        for body in ("OTA could not begin","Could not end OTA","OTA flash write failed"):
+            with self.subTest(raw_serial_error=body):
+                error_return='return request->send(400, "text/plain", "'+body+'");'
+                blocked=fixture.replace(error_return,'Update.printError(Serial);\n'+error_return)
+                self.assertNotEqual(fixture,blocked,"The negative control must restore a raw Serial print")
+                with self.assertRaises(AssertionError):
+                    self.compile_and_run(blocked,"-DESP32=1")
+
+    def test_admitted_upload_survives_body_gap_but_abandoned_owner_expires(self):
+        fixture=self.uploader_fixture()
+        self.compile_and_run(fixture,"-DESP32=1","-DTEST_OTA_BODY_TIMEOUT=1")
+        old=fixture.replace("request->client()->setRxTimeout(30);","")
+        self.assertNotEqual(fixture,old,"The negative control must remove the production timeout fix")
+        with self.assertRaises(AssertionError):
+            self.compile_and_run(old,"-DESP32=1","-DTEST_OTA_BODY_TIMEOUT=1")
 
     def test_tcp_listener_can_rebind_with_completed_connections(self):
         cached=ROOT/".pio/libdeps/heltec_v4_repeater_observer_mqtt/AsyncTCP/src/AsyncTCP.cpp"