Fix Updater potential overflow, add host tests

cores/esp8266/Updater.cpp

@@ -104,7 +104,9 @@ bool UpdaterClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
   _reset();
   clearError(); //  _error = 0
 
+#ifndef HOST_MOCK
   wifi_set_sleep_type(NONE_SLEEP_T);
+#endif
 
   //address where we will start writing the update
   uintptr_t updateStartAddress = 0;
@@ -115,7 +117,11 @@ bool UpdaterClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
 
   if (command == U_FLASH) {
     //address of the end of the space available for sketch and update
+#ifdef HOST_MOCK
+    uintptr_t updateEndAddress = (uintptr_t) -1;
+#else
     uintptr_t updateEndAddress = (uintptr_t)&_FS_start - 0x40200000;
+#endif
 
     updateStartAddress = (updateEndAddress > roundedSize)? (updateEndAddress - roundedSize) : 0;
 
@@ -132,10 +138,12 @@ bool UpdaterClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
     }
   }
   else if (command == U_FS) {
+#ifndef HOST_MOCK
     if((uintptr_t)&_FS_start + roundedSize > (uintptr_t)&_FS_end) {
       _setError(UPDATE_ERROR_SPACE);
       return false;
     }
+#endif
 
 #ifdef ATOMIC_FS_UPDATE
     //address of the end of the space available for update
@@ -148,7 +156,11 @@ bool UpdaterClass::begin(size_t size, int command, int ledPin, uint8_t ledOn) {
       return false;
     }
 #else
+  #ifdef HOST_MOCK
+    updateStartAddress = 0;
+  #else
     updateStartAddress = (uintptr_t)&_FS_start - 0x40200000;
+  #endif
 #endif
   }
   else {
@@ -378,9 +390,7 @@ size_t UpdaterClass::write(uint8_t *data, size_t len) {
   if(hasError() || !isRunning())
     return 0;
 
-  if(len > remaining()){
-    //len = remaining();
-    //fail instead
+  if(progress() + _bufferLen + len > _size) {
     _setError(UPDATE_ERROR_SPACE);
     return 0;
   }

tests/host/Makefile

@@ -70,6 +70,7 @@ CORE_CPP_FILES := $(addprefix $(CORE_PATH)/,\
 	Schedule.cpp \
 	HardwareSerial.cpp \
 	crc32.cpp \
+	Updater.cpp \
 	) \
 	$(addprefix $(LIBRARIES_PATH)/ESP8266SdFat/src/, \
 		FatLib/FatFile.cpp \
@@ -98,6 +99,7 @@ MOCK_CPP_FILES_COMMON := $(addprefix common/,\
 	MockUART.cpp \
 	MockTools.cpp \
 	MocklwIP.cpp \
+	MockDigital.cpp \
 )
 
 MOCK_CPP_FILES := $(MOCK_CPP_FILES_COMMON) $(addprefix common/,\
@@ -136,7 +138,8 @@ TEST_CPP_FILES := \
 	core/test_md5builder.cpp \
 	core/test_string.cpp \
 	core/test_PolledTimeout.cpp \
-	core/test_Print.cpp
+	core/test_Print.cpp \
+	core/test_Updater.cpp
 
 PREINCLUDES := \
 	-include common/mock.h \

tests/host/common/MockDigital.cpp

@@ -0,0 +1,56 @@
+/*
+  digital.c - wiring digital implementation for esp8266
+
+  Copyright (c) 2015 Hristo Gochkov. All rights reserved.
+  This file is part of the esp8266 core for Arduino environment.
+
+  This library is free software; you can redistribute it and/or
+  modify it under the terms of the GNU Lesser General Public
+  License as published by the Free Software Foundation; either
+  version 2.1 of the License, or (at your option) any later version.
+
+  This library is distributed in the hope that it will be useful,
+  but WITHOUT ANY WARRANTY; without even the implied warranty of
+  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+  Lesser General Public License for more details.
+
+  You should have received a copy of the GNU Lesser General Public
+  License along with this library; if not, write to the Free Software
+  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
+*/
+#define ARDUINO_MAIN
+#include "wiring_private.h"
+#include "pins_arduino.h"
+#include "c_types.h"
+#include "eagle_soc.h"
+#include "ets_sys.h"
+#include "user_interface.h"
+#include "core_esp8266_waveform.h"
+#include "interrupts.h"
+
+extern "C" {
+
+static uint8_t _mode[17];
+static uint8_t _gpio[17];
+
+extern void pinMode(uint8_t pin, uint8_t mode) {
+  if (pin < 17) {
+    _mode[pin] = mode;
+  }
+}
+
+extern void digitalWrite(uint8_t pin, uint8_t val) {
+  if (pin < 17) {
+    _gpio[pin] = val;
+  }
+}
+
+extern int digitalRead(uint8_t pin) {
+  if (pin < 17) {
+    return _gpio[pin] != 0;
+  } else {
+    return 0;
+  }
+}
+
+};

tests/host/core/test_Updater.cpp

@@ -0,0 +1,55 @@
+/*
+ test_Updater.cpp - Updater tests
+ Copyright © 2019 Earle F. Philhower, III
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+ */
+
+#include <catch.hpp>
+#include <Updater.h>
+
+
+// Use a SPIFFS file because we can't instantiate a virtual class like Print
+TEST_CASE("Updater fails when writes overflow requested size", "[core][Updater]")
+{
+    UpdaterClass *u;
+    uint8_t buff[6000];
+    u = new UpdaterClass();
+    REQUIRE(u->begin(6000));
+    REQUIRE(u->write(buff, 1000));
+    REQUIRE(u->write(buff, 1000));
+    REQUIRE(u->write(buff, 1000));
+    REQUIRE(u->write(buff, 1000));
+    REQUIRE(u->write(buff, 1000));
+    REQUIRE(u->write(buff, 1000));
+    REQUIRE(u->remaining() == 0);
+    // All bytes written, should fail next
+    REQUIRE(!u->write(buff, 1000));
+    delete u;
+
+    // Updater to a 4K aligned size, check nomissing over/underflow
+    u = new UpdaterClass();
+    REQUIRE(u->begin(8192));
+    REQUIRE(u->remaining() == 8192);
+    REQUIRE(u->write(buff, 4096));
+    REQUIRE(u->write(buff, 4096));
+    REQUIRE(!u->write(buff, 1));
+    delete u;
+
+    // Issue #4674
+    u = new UpdaterClass();
+    REQUIRE(u->begin(5000));
+    REQUIRE(u->write(buff, 2048));
+    REQUIRE(u->write(buff, 2048));
+    // Should fail, would write 6KB
+    REQUIRE(!u->write(buff, 2048));
+    delete u;
+}