@@ -57,6 +57,14 @@ bool UpdaterClass::begin(size_t size, int command) {
5757 return false ;
5858 }
5959
60+ if (!ESP.checkFlashConfig (false )) {
61+ _error = UPDATE_ERROR_FLASH_CONFIG;
62+ #ifdef DEBUG_UPDATER
63+ printError (DEBUG_UPDATER);
64+ #endif
65+ return false ;
66+ }
67+
6068 _reset ();
6169 _error = 0 ;
6270
@@ -116,9 +124,13 @@ bool UpdaterClass::begin(size_t size, int command) {
116124 return true ;
117125}
118126
119- void UpdaterClass::setMD5 (const char * expected_md5){
120- if (strlen (expected_md5) != 32 ) return ;
127+ bool UpdaterClass::setMD5 (const char * expected_md5){
128+ if (strlen (expected_md5) != 32 )
129+ {
130+ return false ;
131+ }
121132 _target_md5 = expected_md5;
133+ return true ;
122134}
123135
124136bool UpdaterClass::end (bool evenIfRemaining){
@@ -152,13 +164,22 @@ bool UpdaterClass::end(bool evenIfRemaining){
152164#ifdef DEBUG_UPDATER
153165 DEBUG_UPDATER.printf (" MD5 Failed: expected:%s, calculated:%s\n " , _target_md5.c_str (), _md5.toString ().c_str ());
154166#endif
167+ _reset ();
155168 return false ;
156169 }
157170#ifdef DEBUG_UPDATER
158171 else DEBUG_UPDATER.printf (" MD5 Success: %s\n " , _target_md5.c_str ());
159172#endif
160173 }
161174
175+ if (!_verifyEnd ()) {
176+ #ifdef DEBUG_UPDATER
177+ printError (DEBUG_UPDATER);
178+ #endif
179+ _reset ();
180+ return false ;
181+ }
182+
162183 if (_command == U_FLASH) {
163184 eboot_command ebcmd;
164185 ebcmd.action = ACTION_COPY_RAW;
@@ -233,12 +254,70 @@ size_t UpdaterClass::write(uint8_t *data, size_t len) {
233254 return len;
234255}
235256
257+ bool UpdaterClass::_verifyHeader (uint8_t data) {
258+ if (_command == U_FLASH) {
259+ // check for valid first magic byte (is always 0xE9)
260+ if (data != 0xE9 ) {
261+ _error = UPDATE_ERROR_MAGIC_BYTE;
262+ _currentAddress = (_startAddress + _size);
263+ return false ;
264+ }
265+ return true ;
266+ } else if (_command == U_SPIFFS) {
267+ // no check of SPIFFS possible with first byte.
268+ return true ;
269+ }
270+ return false ;
271+ }
272+
273+ bool UpdaterClass::_verifyEnd () {
274+ if (_command == U_FLASH) {
275+
276+ uint8_t buf[4 ];
277+ if (!ESP.flashRead (_startAddress, (uint32_t *) &buf[0 ], 4 )) {
278+ _error = UPDATE_ERROR_READ;
279+ _currentAddress = (_startAddress);
280+ return false ;
281+ }
282+
283+ // check for valid first magic byte
284+ if (buf[0 ] != 0xE9 ) {
285+ _error = UPDATE_ERROR_MAGIC_BYTE;
286+ _currentAddress = (_startAddress);
287+ return false ;
288+ }
289+
290+ uint32_t bin_flash_size = ESP.magicFlashChipSize ((buf[3 ] & 0xf0 ) >> 4 );
291+
292+ // check if new bin fits to SPI flash
293+ if (bin_flash_size > ESP.getFlashChipRealSize ()) {
294+ _error = UPDATE_ERROR_NEW_FLASH_CONFIG;
295+ _currentAddress = (_startAddress);
296+ return false ;
297+ }
298+
299+ return true ;
300+ } else if (_command == U_SPIFFS) {
301+ // SPIFFS is already over written checks make no sense any more.
302+ return true ;
303+ }
304+ return false ;
305+ }
306+
236307size_t UpdaterClass::writeStream (Stream &data) {
237308 size_t written = 0 ;
238309 size_t toRead = 0 ;
239310 if (hasError () || !isRunning ())
240311 return 0 ;
241312
313+ if (!_verifyHeader (data.peek ())) {
314+ #ifdef DEBUG_UPDATER
315+ printError (DEBUG_UPDATER);
316+ #endif
317+ _reset ();
318+ return 0 ;
319+ }
320+
242321 while (remaining ()) {
243322 toRead = data.readBytes (_buffer + _bufferLen, (FLASH_SECTOR_SIZE - _bufferLen));
244323 if (toRead == 0 ) { // Timeout
@@ -250,8 +329,9 @@ size_t UpdaterClass::writeStream(Stream &data) {
250329#ifdef DEBUG_UPDATER
251330 printError (DEBUG_UPDATER);
252331#endif
332+ _reset ();
333+ return written;
253334 }
254- return written;
255335 }
256336 _bufferLen += toRead;
257337 if ((_bufferLen == remaining () || _bufferLen == FLASH_SECTOR_SIZE) && !_writeBuffer ())
@@ -270,6 +350,8 @@ void UpdaterClass::printError(Stream &out){
270350 out.println (" Flash Write Failed" );
271351 } else if (_error == UPDATE_ERROR_ERASE){
272352 out.println (" Flash Erase Failed" );
353+ } else if (_error == UPDATE_ERROR_READ){
354+ out.println (" Flash Read Failed" );
273355 } else if (_error == UPDATE_ERROR_SPACE){
274356 out.println (" Not Enough Space" );
275357 } else if (_error == UPDATE_ERROR_SIZE){
@@ -278,6 +360,12 @@ void UpdaterClass::printError(Stream &out){
278360 out.println (" Stream Read Timeout" );
279361 } else if (_error == UPDATE_ERROR_MD5){
280362 out.println (" MD5 Check Failed" );
363+ } else if (_error == UPDATE_ERROR_FLASH_CONFIG){
364+ out.printf (" Flash config wrong real: %d IDE: %d\n " , ESP.getFlashChipRealSize (), ESP.getFlashChipSize ());
365+ } else if (_error == UPDATE_ERROR_NEW_FLASH_CONFIG){
366+ out.printf (" new Flash config wrong real: %d\n " , ESP.getFlashChipRealSize ());
367+ } else if (_error == UPDATE_ERROR_MAGIC_BYTE){
368+ out.println (" Magic byte is wrong, not 0xE9" );
281369 } else {
282370 out.println (" UNKNOWN" );
283371 }
0 commit comments