diff --git a/libraries/Update/src/Update.h b/libraries/Update/src/Update.h index 6f733d41..6cdcff24 100644 --- a/libraries/Update/src/Update.h +++ b/libraries/Update/src/Update.h @@ -137,6 +137,15 @@ class UpdateClass { return written; } + /* + check if there is a firmware on the other OTA partition that you can bootinto + */ + bool canRollBack(); + /* + set the other OTA partition as bootable (reboot to enable) + */ + bool rollBack(); + private: void _reset(); void _abort(uint8_t err); diff --git a/libraries/Update/src/Updater.cpp b/libraries/Update/src/Updater.cpp index 11676d27..7f9b4c57 100644 --- a/libraries/Update/src/Updater.cpp +++ b/libraries/Update/src/Updater.cpp @@ -35,6 +35,34 @@ static const char * _err2str(uint8_t _error){ return ("UNKNOWN"); } +static bool _partitionIsBootable(const esp_partition_t* partition){ + uint8_t buf[4]; + if(!partition){ + return false; + } + if(!ESP.flashRead(partition->address, (uint32_t*)buf, 4)) { + return false; + } + + if(buf[0] != ESP_IMAGE_HEADER_MAGIC) { + return false; + } + return true; +} + +static bool _enablePartition(const esp_partition_t* partition){ + uint8_t buf[4]; + if(!partition){ + return false; + } + if(!ESP.flashRead(partition->address, (uint32_t*)buf, 4)) { + return false; + } + buf[0] = ESP_IMAGE_HEADER_MAGIC; + + return ESP.flashWrite(partition->address, (uint32_t*)buf, 4); +} + UpdateClass::UpdateClass() : _error(0) , _buffer(0) @@ -56,6 +84,22 @@ void UpdateClass::_reset() { _command = U_FLASH; } +bool UpdateClass::canRollBack(){ + if(_buffer){ //Update is running + return false; + } + const esp_partition_t* partition = esp_ota_get_next_update_partition(NULL); + return _partitionIsBootable(partition); +} + +bool UpdateClass::rollBack(){ + if(_buffer){ //Update is running + return false; + } + const esp_partition_t* partition = esp_ota_get_next_update_partition(NULL); + return _partitionIsBootable(partition) && !esp_ota_set_boot_partition(partition); +} + bool UpdateClass::begin(size_t size, int command) { if(_size > 0){ log_w("already running"); @@ -121,6 +165,17 @@ void UpdateClass::abort(){ } bool UpdateClass::_writeBuffer(){ + //first bytes of new firmware + if(!_progress && _command == U_FLASH){ + //check magic + if(_buffer[0] != ESP_IMAGE_HEADER_MAGIC){ + _abort(UPDATE_ERROR_MAGIC_BYTE); + return false; + } + //remove magic byte from the firmware now and write it upon success + //this ensures that partially written firmware will not be bootable + _buffer[0] = 0xFF; + } if(!ESP.flashEraseSector((_partition->address + _progress)/SPI_FLASH_SEC_SIZE)){ _abort(UPDATE_ERROR_ERASE); return false; @@ -129,6 +184,10 @@ bool UpdateClass::_writeBuffer(){ _abort(UPDATE_ERROR_WRITE); return false; } + //restore magic or md5 will fail + if(!_progress && _command == U_FLASH){ + _buffer[0] = ESP_IMAGE_HEADER_MAGIC; + } _md5.add(_buffer, _bufferLen); _progress += _bufferLen; _bufferLen = 0; @@ -150,17 +209,11 @@ bool UpdateClass::_verifyHeader(uint8_t data) { bool UpdateClass::_verifyEnd() { if(_command == U_FLASH) { - uint8_t buf[4]; - if(!ESP.flashRead(_partition->address, (uint32_t*)buf, 4)) { + if(!_partitionIsBootable(_partition)) { _abort(UPDATE_ERROR_READ); return false; } - if(buf[0] != ESP_IMAGE_HEADER_MAGIC) { - _abort(UPDATE_ERROR_MAGIC_BYTE); - return false; - } - if(esp_ota_set_boot_partition(_partition)){ _abort(UPDATE_ERROR_ACTIVATE); return false; @@ -168,6 +221,7 @@ bool UpdateClass::_verifyEnd() { _reset(); return true; } else if(_command == U_SPIFFS) { + _reset(); return true; } return false;