From 2fda054bea45aa6d1be3ab4265406b0846247357 Mon Sep 17 00:00:00 2001
From: Sathishkumar T <tsatiz@gmail.com>
Date: Tue, 24 Jul 2018 23:10:18 +0530
Subject: [PATCH] [OTA Timeout] Added ability set OTA timeout in the OTA client
 (#1669)

---
 libraries/ArduinoOTA/src/ArduinoOTA.cpp | 12 +++++++++---
 libraries/ArduinoOTA/src/ArduinoOTA.h   | 12 +++++++-----
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/libraries/ArduinoOTA/src/ArduinoOTA.cpp b/libraries/ArduinoOTA/src/ArduinoOTA.cpp
index 1e100991..5df367d5 100644
--- a/libraries/ArduinoOTA/src/ArduinoOTA.cpp
+++ b/libraries/ArduinoOTA/src/ArduinoOTA.cpp
@@ -9,7 +9,7 @@
 #include "Update.h"
 
 
-//#define OTA_DEBUG Serial
+// #define OTA_DEBUG Serial
 
 ArduinoOTAClass::ArduinoOTAClass()
 : _port(0)
@@ -20,6 +20,7 @@ ArduinoOTAClass::ArduinoOTAClass()
 , _size(0)
 , _cmd(0)
 , _ota_port(0)
+, _ota_timeout(1000)
 , _start_callback(NULL)
 , _end_callback(NULL)
 , _error_callback(NULL)
@@ -260,8 +261,9 @@ void ArduinoOTAClass::_runUpdate() {
     }
 
     uint32_t written = 0, total = 0, tried = 0;
+
     while (!Update.isFinished() && client.connected()) {
-        size_t waited = 1000;
+        size_t waited = _ota_timeout;
         size_t available = client.available();
         while (!available && waited){
             delay(1);
@@ -387,6 +389,10 @@ int ArduinoOTAClass::getCommand() {
     return _cmd;
 }
 
+void ArduinoOTAClass::setTimeout(int timeoutInMillis) {
+    _ota_timeout = timeoutInMillis;
+}
+
 #if !defined(NO_GLOBAL_INSTANCES) && !defined(NO_GLOBAL_ARDUINOOTA)
 ArduinoOTAClass ArduinoOTA;
-#endif
+#endif
\ No newline at end of file
diff --git a/libraries/ArduinoOTA/src/ArduinoOTA.h b/libraries/ArduinoOTA/src/ArduinoOTA.h
index ee8d5897..db0ead63 100644
--- a/libraries/ArduinoOTA/src/ArduinoOTA.h
+++ b/libraries/ArduinoOTA/src/ArduinoOTA.h
@@ -7,7 +7,6 @@
 
 #define INT_BUFFER_SIZE 16
 
-
 typedef enum {
   OTA_IDLE,
   OTA_WAITAUTH,
@@ -25,9 +24,9 @@ typedef enum {
 class ArduinoOTAClass
 {
   public:
-	typedef std::function<void(void)> THandlerFunction;
-	typedef std::function<void(ota_error_t)> THandlerFunction_Error;
-	typedef std::function<void(unsigned int, unsigned int)> THandlerFunction_Progress;
+    typedef std::function<void(void)> THandlerFunction;
+    typedef std::function<void(ota_error_t)> THandlerFunction_Error;
+    typedef std::function<void(unsigned int, unsigned int)> THandlerFunction_Progress;
 
     ArduinoOTAClass();
     ~ArduinoOTAClass();
@@ -75,6 +74,8 @@ class ArduinoOTAClass
     //Gets update command type after OTA has started. Either U_FLASH or U_SPIFFS
     int getCommand();
 
+    void setTimeout(int timeoutInMillis);
+
   private:
     int _port;
     String _password;
@@ -88,6 +89,7 @@ class ArduinoOTAClass
     int _size;
     int _cmd;
     int _ota_port;
+    int _ota_timeout;
     IPAddress _ota_ip;
     String _md5;
 
@@ -106,4 +108,4 @@ class ArduinoOTAClass
 extern ArduinoOTAClass ArduinoOTA;
 #endif
 
-#endif /* __ARDUINO_OTA_H */
+#endif /* __ARDUINO_OTA_H */
\ No newline at end of file