ESP8266HTTPUpdateServer: fix responses after uploading

- fix response not being delivered to the browser after upload is done (esp8266#2221) - if Update.begin fails, don’t attempt to write data - if update is not successful, send error message from Update to the client - move strings into PROGMEM
oligoform · May 10, 2017 · ace0622 · ace0622
1 parent b623613
commit ace0622
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 23 deletions.
diff --git a/libraries/ESP8266HTTPUpdateServer/src/ESP8266HTTPUpdateServer.cpp b/libraries/ESP8266HTTPUpdateServer/src/ESP8266HTTPUpdateServer.cpp
@@ -3,17 +3,18 @@
 #include <WiFiServer.h>
 #include <ESP8266WebServer.h>
 #include <WiFiUdp.h>
+#include "StreamString.h"
 #include "ESP8266HTTPUpdateServer.h"
 
 
-const char* ESP8266HTTPUpdateServer::_serverIndex =
-R"(<html><body><form method='POST' action='' enctype='multipart/form-data'>
+static const char serverIndex[] PROGMEM =
+  R"(<html><body><form method='POST' action='' enctype='multipart/form-data'>
                   <input type='file' name='update'>
                   <input type='submit' value='Update'>
                </form>
-         </body></html>)";
-const char* ESP8266HTTPUpdateServer::_failedResponse = R"(Update Failed!)";
-const char* ESP8266HTTPUpdateServer::_successResponse = "<META http-equiv=\"refresh\" content=\"15;URL=\">Update Success! Rebooting...";
+         </body></html>\n)";
+static const char successResponse[] PROGMEM = 
+  "<META http-equiv=\"refresh\" content=\"15;URL=\">Update Success! Rebooting...\n";
 
 ESP8266HTTPUpdateServer::ESP8266HTTPUpdateServer(bool serial_debug)
 {
@@ -34,20 +35,29 @@ void ESP8266HTTPUpdateServer::setup(ESP8266WebServer *server, const char * path,
     _server->on(path, HTTP_GET, [&](){
       if(_username != NULL && _password != NULL && !_server->authenticate(_username, _password))
         return _server->requestAuthentication();
-      _server->send(200, "text/html", _serverIndex);
+      _server->send_P(200, PSTR("text/html"), serverIndex);
     });
 
     // handler for the /update form POST (once file upload finishes)
     _server->on(path, HTTP_POST, [&](){
       if(!_authenticated)
         return _server->requestAuthentication();
-      _server->send(200, "text/html", Update.hasError() ? _failedResponse : _successResponse);
-      ESP.restart();
+      if (Update.hasError()) {
+        _server->send(200, F("text/html"), String(F("Update error: ")) + _updaterError);
+      } else {
+        _server->client().setNoDelay(true);
+        _server->send_P(200, PSTR("text/html"), successResponse);
+        delay(100);
+        _server->client().stop();
+        ESP.restart();
+      }
     },[&](){
       // handler for the file upload, get's the sketch bytes, and writes
       // them through the Update object
       HTTPUpload& upload = _server->upload();
+
       if(upload.status == UPLOAD_FILE_START){
+        _updaterError = String();
         if (_serial_output)
           Serial.setDebugOutput(true);
 
@@ -63,19 +73,18 @@ void ESP8266HTTPUpdateServer::setup(ESP8266WebServer *server, const char * path,
           Serial.printf("Update: %s\n", upload.filename.c_str());
         uint32_t maxSketchSpace = (ESP.getFreeSketchSpace() - 0x1000) & 0xFFFFF000;
         if(!Update.begin(maxSketchSpace)){//start with max available size
-          if (_serial_output) Update.printError(Serial);
+          _setUpdaterError();
         }
-      } else if(_authenticated && upload.status == UPLOAD_FILE_WRITE){
+      } else if(_authenticated && upload.status == UPLOAD_FILE_WRITE && !_updaterError.length()){
         if (_serial_output) Serial.printf(".");
         if(Update.write(upload.buf, upload.currentSize) != upload.currentSize){
-          if (_serial_output) Update.printError(Serial);
-
+          _setUpdaterError();
         }
-      } else if(_authenticated && upload.status == UPLOAD_FILE_END){
+      } else if(_authenticated && upload.status == UPLOAD_FILE_END && !_updaterError.length()){
         if(Update.end(true)){ //true to set the size to the current progress
           if (_serial_output) Serial.printf("Update Success: %u\nRebooting...\n", upload.totalSize);
         } else {
-          if (_serial_output) Update.printError(Serial);
+          _setUpdaterError();
         }
         if (_serial_output) Serial.setDebugOutput(false);
       } else if(_authenticated && upload.status == UPLOAD_FILE_ABORTED){
@@ -85,3 +94,11 @@ void ESP8266HTTPUpdateServer::setup(ESP8266WebServer *server, const char * path,
       delay(0);
     });
 }
+
+void ESP8266HTTPUpdateServer::_setUpdaterError()
+{
+  if (_serial_output) Update.printError(Serial);
+  StreamString str;
+  Update.printError(str);
+  _updaterError = str.c_str();
+}
diff --git a/libraries/ESP8266HTTPUpdateServer/src/ESP8266HTTPUpdateServer.h b/libraries/ESP8266HTTPUpdateServer/src/ESP8266HTTPUpdateServer.h
@@ -5,15 +5,6 @@ class ESP8266WebServer;
 
 class ESP8266HTTPUpdateServer
 {
-  private:
-    bool _serial_output;
-    ESP8266WebServer *_server;
-    static const char *_serverIndex;
-    static const char *_failedResponse;
-    static const char *_successResponse;
-    char * _username;
-    char * _password;
-    bool _authenticated;
   public:
     ESP8266HTTPUpdateServer(bool serial_debug=false);
 
@@ -33,6 +24,17 @@ class ESP8266HTTPUpdateServer
     }
 
     void setup(ESP8266WebServer *server, const char * path, const char * username, const char * password);
+
+  protected:
+    void _setUpdaterError();
+
+  private:
+    bool _serial_output;
+    ESP8266WebServer *_server;
+    char * _username;
+    char * _password;
+    bool _authenticated;
+    String _updaterError;
 };