New Update library, example, upload and more

Proper error handling in the uploading python script
Much faster OTA example sketch with better results
New Update class that simplifies updating the firmware from any source
Updated Esp.updateSketch() to use the new class

Author: John Doe

cores/esp8266/Arduino.h

@@ -233,6 +233,7 @@ void loop(void);
 
 #include "HardwareSerial.h"
 #include "Esp.h"
+#include "Updater.h"
 #include "debug.h"
 
 #define min(a,b) ((a)<(b)?(a):(b))

cores/esp8266/Esp.cpp

@@ -30,7 +30,7 @@ extern struct rst_info resetInfo;
 }
 
 
-// #define DEBUG_SERIAL Serial
+//#define DEBUG_SERIAL Serial
 
 
 /**
@@ -358,96 +358,38 @@ uint32_t EspClass::getFreeSketchSpace() {
     return freeSpaceEnd - freeSpaceStart;
 }
 
-bool EspClass::updateSketch(Stream& in, uint32_t size, bool restartOnFail) {
-
-    if (size > getFreeSketchSpace()){
-        if(restartOnFail) ESP.restart();
-        return false;
-    }
-
-    uint32_t usedSize = getSketchSize();
-    uint32_t freeSpaceStart = (usedSize + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
-    uint32_t roundedSize = (size + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
-
-#ifdef DEBUG_SERIAL
-    DEBUG_SERIAL.printf("erase @0x%x size=0x%x\r\n", freeSpaceStart, roundedSize);
-#endif
-
-    noInterrupts();
-    int rc = SPIEraseAreaEx(freeSpaceStart, roundedSize);
-    interrupts();
-    if (rc){
-        if(restartOnFail) ESP.restart();
-        return false;
-    }
-
+bool EspClass::updateSketch(Stream& in, uint32_t size, bool restartOnFail, bool restartOnSuccess) {
+  if(!Update.begin(size)){
 #ifdef DEBUG_SERIAL
-    DEBUG_SERIAL.println("erase done");
+    DEBUG_SERIAL.print("Update ");
+    Update.printError(DEBUG_SERIAL);
 #endif
+    if(restartOnFail) ESP.restart();
+    return false;
+  }
 
-    uint32_t addr = freeSpaceStart;
-    uint32_t left = size;
-
-    const uint32_t bufferSize = FLASH_SECTOR_SIZE;
-    std::unique_ptr<uint8_t> buffer(new uint8_t[bufferSize]);
-
+  if(Update.writeStream(in) != size){
 #ifdef DEBUG_SERIAL
-    DEBUG_SERIAL.println("writing");
+    DEBUG_SERIAL.print("Update ");
+    Update.printError(DEBUG_SERIAL);
 #endif
-    while (left > 0) {
-        size_t willRead = (left < bufferSize) ? left : bufferSize;
-        size_t rd = in.readBytes(buffer.get(), willRead);
-        if (rd != willRead) {
-#ifdef DEBUG_SERIAL
-            DEBUG_SERIAL.printf("stream read less: %u/%u\n", rd, willRead);
-#endif
-            if(rd == 0){ //we got nothing from the client
-              //we should actually give it a bit of a chance to send us something
-              //connection could be slow ;)
-              if(restartOnFail) ESP.restart();
-              return false;
-            }
-            //we at least got some data, lets write it to the flash
-            willRead = rd;
-        }
-        
-        if(addr == freeSpaceStart) {
-            // check for valid first magic byte
-            if(*((uint8 *) buffer.get()) != 0xE9) {
-                if(restartOnFail) ESP.restart();
-                return false;
-            }
-        }
-
-        noInterrupts();
-        rc = SPIWrite(addr, buffer.get(), willRead);
-        interrupts();
-        if (rc) {
-#ifdef DEBUG_SERIAL
-            DEBUG_SERIAL.println("write failed");
-#endif            
-            if(restartOnFail) ESP.restart();
-            return false;
-        }
+    if(restartOnFail) ESP.restart();
+    return false;
+  }
 
-        addr += willRead;
-        left -= willRead;
+  if(!Update.end()){
 #ifdef DEBUG_SERIAL
-        DEBUG_SERIAL.print(".");
+    DEBUG_SERIAL.print("Update ");
+    Update.printError(DEBUG_SERIAL);
 #endif
-    }
+    if(restartOnFail) ESP.restart();
+    return false;
+  }
 
 #ifdef DEBUG_SERIAL
-    DEBUG_SERIAL.println("\r\nrestarting");
-#endif
-    eboot_command ebcmd;
-    ebcmd.action = ACTION_COPY_RAW;
-    ebcmd.args[0] = freeSpaceStart;
-    ebcmd.args[1] = 0x00000;
-    ebcmd.args[2] = size;
-    eboot_command_write(&ebcmd);
-
-    ESP.restart();
-    return true; // never happens
+    DEBUG_SERIAL.println("Update SUCCESS");
+#endif  
+    if(restartOnSuccess) ESP.restart();
+    return true;
 }
 

cores/esp8266/Esp.h

@@ -116,7 +116,7 @@ class EspClass {
 
         uint32_t getSketchSize();
         uint32_t getFreeSketchSpace();
-        bool updateSketch(Stream& in, uint32_t size, bool restartOnFail = false);
+        bool updateSketch(Stream& in, uint32_t size, bool restartOnFail = false, bool restartOnSuccess = true);
 
         String getResetInfo();
         struct rst_info * getResetInfoPtr();

cores/esp8266/Updater.cpp

@@ -0,0 +1,190 @@
+#include "Updater.h"
+#include "Arduino.h"
+#include "eboot_command.h"
+extern "C"{
+  #include "mem.h"
+}
+#define DEBUG_UPDATER Serial
+
+extern "C" uint32_t _SPIFFS_start;
+
+UpdaterClass::UpdaterClass() : _error(0), _buffer(0), _bufferLen(0), _size(0), _startAddress(0), _currentAddress(0) {}
+
+bool UpdaterClass::begin(size_t size){
+  if(_size > 0){
+#ifdef DEBUG_UPDATER
+    DEBUG_UPDATER.println("already running");
+#endif
+    return false;
+  }
+  
+  if(size == 0){
+    _error = UPDATE_ERROR_SIZE;
+#ifdef DEBUG_UPDATER
+    printError(DEBUG_UPDATER);
+#endif
+    return false;
+  }
+  
+  if(_buffer) os_free(_buffer);
+  
+  _bufferLen = 0;
+  _startAddress = 0;
+  _currentAddress = 0;
+  _size = 0;
+  _error = 0;
+  
+  uint32_t usedSize = ESP.getSketchSize();
+  uint32_t freeSpaceStart = (usedSize + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
+  uint32_t freeSpaceEnd = (uint32_t)&_SPIFFS_start - 0x40200000;
+  uint32_t roundedSize = (size + FLASH_SECTOR_SIZE - 1) & (~(FLASH_SECTOR_SIZE - 1));
+  
+  if(roundedSize > (freeSpaceEnd - freeSpaceStart)){
+    _error = UPDATE_ERROR_SPACE;
+#ifdef DEBUG_UPDATER
+    printError(DEBUG_UPDATER);
+#endif
+    return false;
+  }
+  noInterrupts();
+  int rc = SPIEraseAreaEx(freeSpaceStart, roundedSize);
+  interrupts();
+  if (rc){
+    _error = UPDATE_ERROR_ERASE;
+#ifdef DEBUG_UPDATER
+    printError(DEBUG_UPDATER);
+#endif
+    return false;
+  }
+  _startAddress = freeSpaceStart;
+  _currentAddress = _startAddress;
+  _size = size;
+  _buffer = (uint8_t*)os_malloc(FLASH_SECTOR_SIZE);
+  
+  return true;
+}
+
+bool UpdaterClass::end(){
+  if(_size == 0){
+#ifdef DEBUG_UPDATER
+    DEBUG_UPDATER.println("no update");
+#endif
+    return false;
+  }
+  
+  if(_buffer) os_free(_buffer);
+  _bufferLen = 0;
+  
+  if(hasError() || !isFinished()){
+#ifdef DEBUG_UPDATER
+    DEBUG_UPDATER.printf("premature end: res:%u, pos:%u/%u\n", getError(), progress(), _size);
+#endif
+    _currentAddress = 0;
+    _startAddress = 0;
+    _size = 0;
+    return false;
+  }
+  
+  eboot_command ebcmd;
+  ebcmd.action = ACTION_COPY_RAW;
+  ebcmd.args[0] = _startAddress;
+  ebcmd.args[1] = 0x00000;
+  ebcmd.args[2] = _size;
+  eboot_command_write(&ebcmd);
+
+  _currentAddress = 0;
+  _startAddress = 0;
+  _size = 0;
+  _error = UPDATE_ERROR_OK;
+  return true;
+}
+
+bool UpdaterClass::_writeBuffer(){
+  WDT_FEED();
+  noInterrupts();
+  int rc = SPIWrite(_currentAddress, _buffer, _bufferLen);
+  interrupts();
+  if (rc) {
+    _error = UPDATE_ERROR_WRITE;
+#ifdef DEBUG_UPDATER
+    printError(DEBUG_UPDATER);
+#endif
+    return false;
+  }
+  _currentAddress += _bufferLen;
+  _bufferLen = 0;
+  return true;
+}
+
+size_t UpdaterClass::write(uint8_t *data, size_t len){
+  size_t left = len;
+  if(hasError())
+    return 0;
+  
+  while((_bufferLen + left) > FLASH_SECTOR_SIZE){
+    size_t toBuff = FLASH_SECTOR_SIZE - _bufferLen;
+    memcpy(_buffer + _bufferLen, data + (len - left), toBuff);
+    _bufferLen += toBuff;
+    if(!_writeBuffer()){
+      return len - left;
+    }
+    left -= toBuff;
+    yield();
+  }
+  //lets see whats left
+  memcpy(_buffer + _bufferLen, data + (len - left), left);
+  _bufferLen += left;
+  if(_bufferLen == remaining()){
+    //we are at the end of the update, so should write what's left to flash
+    if(!_writeBuffer()){
+      return len - left;
+    }
+  }
+  return len;
+}
+
+size_t UpdaterClass::writeStream(Stream &data){
+  size_t written = 0;
+  size_t toRead = 0;
+  if(hasError())
+    return 0;
+  
+  while(remaining()){
+    toRead = FLASH_SECTOR_SIZE - _bufferLen;
+    toRead = data.readBytes(_buffer + _bufferLen, toRead);
+    if(toRead == 0){ //Timeout
+      _error = UPDATE_ERROR_STREAM;
+#ifdef DEBUG_UPDATER
+      printError(DEBUG_UPDATER);
+#endif
+      return written;
+    }
+    _bufferLen += toRead;
+    if((_bufferLen == remaining() || _bufferLen == FLASH_SECTOR_SIZE) && !_writeBuffer())
+      return written;
+    written += toRead;
+    yield();
+  }
+  return written;
+}
+
+void UpdaterClass::printError(Stream &out){
+  out.printf("ERROR[%u]: ", _error);
+  if(_error == UPDATE_ERROR_OK){
+    out.println("No Error");
+  } else if(_error == UPDATE_ERROR_WRITE){
+    out.println("Flash Write Failed");
+  } else if(_error == UPDATE_ERROR_ERASE){
+    out.println("Flash Erase Failed");
+  } else if(_error == UPDATE_ERROR_SPACE){
+    out.println("Not Enough Space");
+  } else if(_error == UPDATE_ERROR_SIZE){
+    out.println("Bad Size Given");
+  } else if(_error == UPDATE_ERROR_STREAM){
+    out.println("Stream Read Timeout");
+  } else {
+    out.println("UNKNOWN");
+  }
+}
+
+UpdaterClass Update;