From 5c2c257526ef32214f8dcab7144216fcdd115a7b Mon Sep 17 00:00:00 2001
From: Markus Sattler <help.markus+git@gmail.com>
Date: Sat, 23 Jan 2016 18:34:30 +0100
Subject: [PATCH] fix server disconnect problems

---
 src/WebSockets.cpp       | 12 ++++++++++++
 src/WebSockets.h         |  2 +-
 src/WebSocketsClient.cpp | 10 ++++++++--
 src/WebSocketsServer.cpp | 18 ++++++++++++++++++
 4 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/src/WebSockets.cpp b/src/WebSockets.cpp
index d45748e..205b7e7 100644
--- a/src/WebSockets.cpp
+++ b/src/WebSockets.cpp
@@ -282,6 +282,10 @@ void WebSockets::handleWebsocket(WSclient_t * client) {
  * @param size
  */
 bool WebSockets::handleWebsocketWaitFor(WSclient_t * client, size_t size) {
+    if(!client->tcp || !client->tcp->connected()) {
+        return false;
+    }
+
     if(size > WEBSOCKETS_MAX_HEADER_SIZE) {
         DEBUG_WEBSOCKETS("[WS][%d][handleWebsocketWaitFor] size: %d to big!\n", client->num, size);
         return false;
@@ -309,6 +313,10 @@ bool WebSockets::handleWebsocketWaitFor(WSclient_t * client, size_t size) {
 
 void WebSockets::handleWebsocketCb(WSclient_t * client) {
 
+    if(!client->tcp || !client->tcp->connected()) {
+        return;
+    }
+
     uint8_t * buffer = client->cWsHeader;
 
     WSMessageHeader_t * header = &client->cWsHeaderDecode;
@@ -390,6 +398,7 @@ void WebSockets::handleWebsocketCb(WSclient_t * client) {
 }
 
 void WebSockets::handleWebsocketPayloadCb(WSclient_t * client, bool ok, uint8_t * payload) {
+
     WSMessageHeader_t * header = &client->cWsHeaderDecode;
     if(ok) {
         if(header->payloadLen > 0) {
@@ -513,6 +522,9 @@ String WebSockets::base64_encode(uint8_t * data, size_t length) {
  */
 bool WebSockets::readCb(WSclient_t * client, uint8_t * out, size_t n, WSreadWaitCb cb) {
 #if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266_ASYNC)
+    if(!client->tcp || !client->tcp->connected()) {
+        return false;
+    }
 
     client->tcp->readBytes(out, n, std::bind([](WSclient_t * client, bool ok, WSreadWaitCb cb) {
         if(cb) {
diff --git a/src/WebSockets.h b/src/WebSockets.h
index 96cf09c..c203da3 100644
--- a/src/WebSockets.h
+++ b/src/WebSockets.h
@@ -27,7 +27,7 @@
 
 #include <Arduino.h>
 
-//#define DEBUG_WEBSOCKETS(...) os_printf( __VA_ARGS__ )
+#define DEBUG_WEBSOCKETS(...) os_printf( __VA_ARGS__ )
 
 #ifndef DEBUG_WEBSOCKETS
 #define DEBUG_WEBSOCKETS(...)
diff --git a/src/WebSocketsClient.cpp b/src/WebSocketsClient.cpp
index 0351012..987ff85 100644
--- a/src/WebSocketsClient.cpp
+++ b/src/WebSocketsClient.cpp
@@ -89,6 +89,7 @@ void WebSocketsClient::beginSSL(String host, uint16_t port, String url, String f
  * called in arduino loop
  */
 void WebSocketsClient::loop(void) {
+#if (WEBSOCKETS_NETWORK_TYPE != NETWORK_ESP8266_ASYNC)
     if(!clientIsConnected(&_client)) {
 
 #if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266)
@@ -146,10 +147,9 @@ void WebSocketsClient::loop(void) {
             delay(10); //some little delay to not flood the server
         }
     } else {
-#if (WEBSOCKETS_NETWORK_TYPE != NETWORK_ESP8266_ASYNC)
         handleClientData();
-#endif
     }
+#endif
 }
 
 /**
@@ -269,11 +269,17 @@ void WebSocketsClient::clientDisconnect(WSclient_t * client) {
 
     if(client->tcp) {
         if(client->tcp->connected()) {
+#if (WEBSOCKETS_NETWORK_TYPE != NETWORK_ESP8266_ASYNC)
             client->tcp->flush();
+#endif
             client->tcp->stop();
         }
         event = true;
+#if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266_ASYNC)
+        client->status = WSC_NOT_CONNECTED;
+#else
         delete client->tcp;
+#endif
         client->tcp = NULL;
     }
 
diff --git a/src/WebSocketsServer.cpp b/src/WebSocketsServer.cpp
index 1d82605..cddb606 100644
--- a/src/WebSocketsServer.cpp
+++ b/src/WebSocketsServer.cpp
@@ -296,6 +296,7 @@ bool WebSocketsServer::newClient(WEBSOCKETS_NETWORK_CLASS * TCPclient) {
 
         // state is not connected or tcp connection is lost
         if(!clientIsConnected(client)) {
+
             client->tcp = TCPclient;
 
 #if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266)
@@ -316,6 +317,18 @@ bool WebSocketsServer::newClient(WEBSOCKETS_NETWORK_CLASS * TCPclient) {
 
 
 #if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266_ASYNC)
+            client->tcp->onDisconnect(std::bind([](WebSocketsServer * server, AsyncTCPbuffer * obj, WSclient_t * client) -> bool {
+                DEBUG_WEBSOCKETS("[WS-Server][%d] Disconnect client\n", client->num);
+
+                AsyncTCPbuffer ** sl = &server->_clients[client->num].tcp;
+                if(*sl == obj) {
+                    client->status = WSC_NOT_CONNECTED;
+                    *sl = NULL;
+                }
+                return true;
+            },  this, std::placeholders::_1, client));
+
+
             client->tcp->readStringUntil('\n', &(client->cHttpLine), std::bind(&WebSocketsServer::handleHeader, this, client, &(client->cHttpLine)));
 #endif
 
@@ -375,7 +388,11 @@ void WebSocketsServer::clientDisconnect(WSclient_t * client) {
 #endif
             client->tcp->stop();
         }
+#if (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266_ASYNC)
+        client->status = WSC_NOT_CONNECTED;
+#else
         delete client->tcp;
+#endif
         client->tcp = NULL;
     }
 
@@ -425,6 +442,7 @@ bool WebSocketsServer::clientIsConnected(WSclient_t * client) {
 
     if(client->tcp) {
         // do cleanup
+        DEBUG_WEBSOCKETS("[WS-Server][%d] client list cleanup.\n", client->num);
         clientDisconnect(client);
     }