Fix buffered write and fix freezes when losing client connection

src/ESPTelnet.cpp

@@ -32,7 +32,11 @@ void ESPTelnet::handleInput() {
 
 void ESPTelnet::println() {
   if (client && isConnected()) {
-    client.println();
+    if (!client.println()) {
+      onFailedWrite();
+    } else {
+      onSuccessfullyWrite();
+    }
   }
 }
 
@@ -57,10 +61,10 @@ size_t ESPTelnet::printf(const char* format, ...) {
     va_start(arg, format);
     vsnprintf(temp, len + 1, format, arg);
     va_end(arg);
-    len = client.write((uint8_t*)temp, len);
+    len = write((uint8_t*)temp, len);
     free(temp);
   } else {
-    len = client.write((uint8_t*)loc_buf, len);
+    len = write((uint8_t*)loc_buf, len);
   }
 
   return len;

src/ESPTelnet.h

@@ -18,28 +18,44 @@ class ESPTelnet : public ESPTelnetBase {
   template<typename T>
   void print(const T& data) {
     if (client && isConnected()) {
-      client.print(data);
+      if (!client.print(data)) {
+        onFailedWrite();
+      } else {
+        onSuccessfullyWrite();
+      }
     }
   }
 
   template<typename T>
   void println(const T& data) {
     if (client && isConnected()) {
-      client.println(data);
+      if (!client.println(data)) {
+        onFailedWrite();
+      } else {
+        onSuccessfullyWrite();
+      }
     }
   }
 
   template<typename T>
   void print(const T& data, int base) {
     if (client && isConnected()) {
-      client.print(data, base);
+      if (!client.print(data, base)) {
+        onFailedWrite();
+      } else {
+        onSuccessfullyWrite();
+      }
     }
   }
 
   template<typename T>
   void println(const T& data, int base) {
     if (client && isConnected()) {
-      client.println(data, base);
+      if (!client.println(data, base)) {
+        onFailedWrite();
+      } else {
+        onSuccessfullyWrite();
+      }
     }
   }
 

src/ESPTelnetBase.cpp

@@ -119,10 +119,66 @@ int ESPTelnetBase::getKeepAliveInterval() {
 
 /////////////////////////////////////////////////////////////////
 
+void ESPTelnetBase::flush() {
+  if (!client || !isConnected()) {
+    return;
+  }
+
+// only the ESP8266 has a "bool flush()" method
+// https://github.com/esp8266/Arduino/blob/master/libraries/ESP8266WiFi/src/WiFiClient.cpp#L306
+
+#ifdef ARDUINO_ARCH_ESP8266
+  if (!client.flush(this->getKeepAliveInterval())) {
+    onFailedWrite();
+  } else {
+    onSuccessfullyWrite();
+  }
+#else
+  client.flush();
+#endif
+}
+
+/////////////////////////////////////////////////////////////////
+
+size_t ESPTelnetBase::write(uint8_t data) {
+  if (!client || !isConnected()) {
+    return 0;
+  }
+
+  size_t written = client.write(data);
+  if (!written) {
+    onFailedWrite();
+  } else {
+    onSuccessfullyWrite();
+  }
+
+  return written;
+}
+
+/////////////////////////////////////////////////////////////////
+
+size_t ESPTelnetBase::write(const uint8_t* data, size_t size) {
+  if (!client || !isConnected()) {
+    return 0;
+  }
+
+  size_t written = client.write(data, size);
+  if (written != size) {
+    onFailedWrite();
+  } else {
+    onSuccessfullyWrite();
+  }
+
+  return written;
+}
+
+/////////////////////////////////////////////////////////////////
+
 void ESPTelnetBase::connectClient(WiFiClient c, bool triggerEvent) {
   client = c;
   ip = client.remoteIP().toString();
   client.setNoDelay(true);
+  client.setTimeout(this->getKeepAliveInterval());
   if (triggerEvent && on_connect != NULL) on_connect(ip);
   emptyClientStream();
   connected = true;
@@ -182,8 +238,27 @@ String ESPTelnetBase::getLastAttemptIP() const {
 
 /////////////////////////////////////////////////////////////////
 
+void ESPTelnetBase::onFailedWrite() {
+  failedWrites++;
+
+  if (failedWrites >= MAX_ERRORS_ON_WRITE) {
+    failedWrites = 0;
+    disconnectClient();
+  }
+}
+
+/////////////////////////////////////////////////////////////////
+
+void ESPTelnetBase::onSuccessfullyWrite() {
+  if (failedWrites > 0) {
+    failedWrites = 0;
+  }
+}
+
+/////////////////////////////////////////////////////////////////
+
 void ESPTelnetBase::emptyClientStream() {
-  client.flush();
+  flush();
   delay(50);
   while (client.available()) {
     client.read();

src/ESPTelnetBase.h

@@ -11,6 +11,7 @@
 #define ASCII_LF 10
 #define ASCII_CR 13
 #define KEEP_ALIVE_INTERVAL_MS 1000
+#define MAX_ERRORS_ON_WRITE 3
 
 /////////////////////////////////////////////////////////////////
 
@@ -48,6 +49,10 @@ class ESPTelnetBase {
   void setKeepAliveInterval(int ms);
   int getKeepAliveInterval();
 
+  virtual void flush();
+  virtual size_t write(uint8_t);
+  virtual size_t write(const uint8_t* data, size_t size);
+
   String getIP() const;
   String getLastAttemptIP() const;
 
@@ -57,6 +62,8 @@ class ESPTelnetBase {
   void onDisconnect(CallbackFunction f);
   void onInputReceived(CallbackFunction f);
 
+
+
  protected:
   TCPServer server = TCPServer(23);  // must be initalized here
   TCPClient client;
@@ -68,13 +75,17 @@ class ESPTelnetBase {
   uint16_t server_port = 23;
   int keep_alive_interval = KEEP_ALIVE_INTERVAL_MS;
   long last_status_check;
+  unsigned int failedWrites = 0;
 
   CallbackFunction on_connect = NULL;
   CallbackFunction on_reconnect = NULL;
   CallbackFunction on_disconnect = NULL;
   CallbackFunction on_connection_attempt = NULL;
   CallbackFunction on_input = NULL;
 
+  virtual void onFailedWrite();
+  virtual void onSuccessfullyWrite();
+
   void emptyClientStream();
   bool _isIPSet(IPAddress ip);
   virtual void handleInput() = 0;

src/ESPTelnetStream.cpp

@@ -42,19 +42,19 @@ int ESPTelnetStream::peek() {
 /////////////////////////////////////////////////////////////////
 
 void ESPTelnetStream::flush() {
-  if (client && isConnected()) {
-    client.flush();
-  }
+  return ESPTelnetBase::flush();
 }
 
 /////////////////////////////////////////////////////////////////
 
 size_t ESPTelnetStream::write(uint8_t data) {
-  if (client && isConnected()) {
-    return client.write(data);
-  } else {
-    return 0;
-  }
+  return ESPTelnetBase::write(data);
+}
+
+/////////////////////////////////////////////////////////////////
+
+size_t ESPTelnetStream::write(const uint8_t* data, size_t size) {
+  return ESPTelnetBase::write(data, size);
 }
 
 /////////////////////////////////////////////////////////////////

src/ESPTelnetStream.h

@@ -17,9 +17,10 @@ class ESPTelnetStream : public ESPTelnetBase, public Stream {
   int available();
   int read();
   int peek();
-  void flush();
+  void flush() override;
 
-  size_t write(uint8_t);
+  size_t write(uint8_t) override;
+  size_t write(const uint8_t* data, size_t size) override;
 
  protected:
   void handleInput();