From e613e42249830f6221f11c4951c5b1840a728dbd Mon Sep 17 00:00:00 2001
From: Me No Dev <me-no-dev@vdnsbg.com>
Date: Sun, 8 Nov 2015 02:01:06 +0200
Subject: [PATCH] Fix example sketch, espota output and failing updater

I get 100% success with OTA now
---
 cores/esp8266/Updater.h                       |  2 +-
 .../DNS_SD_Arduino_OTA/DNS_SD_Arduino_OTA.ino | 66 ++++++++++---------
 tools/espota.py                               |  8 ++-
 3 files changed, 42 insertions(+), 34 deletions(-)

diff --git a/cores/esp8266/Updater.h b/cores/esp8266/Updater.h
index 670b0e7e94..587831ef2e 100644
--- a/cores/esp8266/Updater.h
+++ b/cores/esp8266/Updater.h
@@ -103,7 +103,7 @@ class UpdaterClass {
         }
         if(remaining() == 0)
           return written;
-        yield();
+        delay(1);
         available = data.available();
       }
       return written;
diff --git a/libraries/ESP8266mDNS/examples/DNS_SD_Arduino_OTA/DNS_SD_Arduino_OTA.ino b/libraries/ESP8266mDNS/examples/DNS_SD_Arduino_OTA/DNS_SD_Arduino_OTA.ino
index 0689a8531a..583a1b4062 100644
--- a/libraries/ESP8266mDNS/examples/DNS_SD_Arduino_OTA/DNS_SD_Arduino_OTA.ino
+++ b/libraries/ESP8266mDNS/examples/DNS_SD_Arduino_OTA/DNS_SD_Arduino_OTA.ino
@@ -1,20 +1,20 @@
 #include <ESP8266WiFi.h>
 #include <ESP8266mDNS.h>
-#include <WiFiUdp.h>
 
 const char* host = "esp8266-ota";
 const char* ssid = "**********";
 const char* pass = "**********";
-const uint16_t aport = 8266;
+const uint16_t ota_port = 8266;
+const char* ota_pass = "1234";
 
-WiFiServer TelnetServer(aport);
-WiFiClient Telnet;
 WiFiUDP OTA;
+WiFiServer MonitorServer(ota_port);
+WiFiClient Monitor;
 
 void setup() {
   Serial.begin(115200);
   Serial.println("");
-  Serial.println("Arduino OTA Test");
+  Serial.println("Bare Minimum Arduino OTA");
 
   Serial.printf("Sketch size: %u\n", ESP.getSketchSize());
   Serial.printf("Free size: %u\n", ESP.getFreeSketchSpace());
@@ -22,45 +22,52 @@ void setup() {
   WiFi.begin(ssid, pass);
   if(WiFi.waitForConnectResult() == WL_CONNECTED){
     MDNS.begin(host);
-    MDNS.addService("arduino", "tcp", aport);
-    OTA.begin(aport);
-    TelnetServer.begin();
-    TelnetServer.setNoDelay(true);
+    MDNS.enableArduino(ota_port, true);
+    OTA.begin(ota_port);
+    MonitorServer.begin();
+    MonitorServer.setNoDelay(true);
     Serial.print("IP address: ");
     Serial.println(WiFi.localIP());
+  } else {
+    Serial.println("WiFi Connect Failed");
+    delay(10000);
+    ESP.reset();
   }
 }
 
 void loop() {
-  //OTA Sketch
   if (OTA.parsePacket()) {
     IPAddress remote = OTA.remoteIP();
+    String pass = OTA.readStringUntil(' ');
     int cmd  = OTA.parseInt();
     int port = OTA.parseInt();
     int size   = OTA.parseInt();
-
+    
+    if(!pass.equals(String(ota_pass))){
+      Serial.println("ERROR: Wrong Password");
+      return;
+    }
+    
     Serial.print("Update Start: ip:");
     Serial.print(remote);
     Serial.printf(", port:%d, size:%d\n", port, size);
     uint32_t startTime = millis();
 
-    WiFiUDP::stopAll();
-
     if(!Update.begin(size)){
-      Serial.println("Update Begin Error");
+      Update.printError(Serial);
       return;
     }
 
+    WiFiUDP::stopAll();
+    WiFiClient::stopAll();
+
     WiFiClient client;
     if (client.connect(remote, port)) {
-
       uint32_t written;
       while(!Update.isFinished()){
         written = Update.write(client);
         if(written > 0) client.print(written, DEC);
       }
-      Serial.setDebugOutput(false);
-
       if(Update.end()){
         client.println("OK");
         Serial.printf("Update Success: %u\nRebooting...\n", millis() - startTime);
@@ -74,28 +81,27 @@ void loop() {
     }
   }
   //IDE Monitor (connected to Serial)
-  if (TelnetServer.hasClient()){
-    if (!Telnet || !Telnet.connected()){
-      if(Telnet) Telnet.stop();
-      Telnet = TelnetServer.available();
+  if (MonitorServer.hasClient()){
+    if (!Monitor || !Monitor.connected()){
+      if(Monitor) Monitor.stop();
+      Monitor = MonitorServer.available();
     } else {
-      WiFiClient toKill = TelnetServer.available();
-      toKill.stop();
+      MonitorServer.available().stop();
     }
   }
-  if (Telnet && Telnet.connected() && Telnet.available()){
-    while(Telnet.available())
-      Serial.write(Telnet.read());
+  if (Monitor && Monitor.connected() && Monitor.available()){
+    while(Monitor.available())
+      Serial.write(Monitor.read());
   }
   if(Serial.available()){
     size_t len = Serial.available();
     uint8_t * sbuf = (uint8_t *)malloc(len);
     Serial.readBytes(sbuf, len);
-    if (Telnet && Telnet.connected()){
-      Telnet.write((uint8_t *)sbuf, len);
-      yield();
+    if (Monitor && Monitor.connected()){
+      Monitor.write((uint8_t *)sbuf, len);
+      delay(0);
     }
     free(sbuf);
   }
-  delay(1);
+  delay(0);
 }
diff --git a/tools/espota.py b/tools/espota.py
index bea72910da..cacbc12d54 100755
--- a/tools/espota.py
+++ b/tools/espota.py
@@ -53,7 +53,7 @@ def serve(remoteAddr, remotePort, password, filename, command = FLASH):
   sent = sock2.sendto(message, remote_address)
   sock2.close()
 
-  logging.info('Waiting for device...\n')
+  logging.info('Waiting for device...')
   try:
     sock.settimeout(10)
     connection, client_address = sock.accept()
@@ -78,13 +78,15 @@ def serve(remoteAddr, remotePort, password, filename, command = FLASH):
         connection.sendall(chunk)
         res = connection.recv(4)
       except:
-        logging.error('\nError Uploading')
+        sys.stderr.write('\n')
+        logging.error('Error Uploading')
         connection.close()
         f.close()
         sock.close()
         return 1
 
-    logging.info('\nWaiting for result...\n')
+    sys.stderr.write('\n')
+    logging.info('Waiting for result...')
     try:
       connection.settimeout(60)
       data = connection.recv(32)