From 674a6e98c93907d1ac697ae9e0d697d4c6a0cb77 Mon Sep 17 00:00:00 2001
From: maurus_fritsche <maurus.fritsche@hotmail.com>
Date: Mon, 10 Feb 2020 19:45:55 +0100
Subject: [PATCH] Implemented HeartBeat in WebSocketsServer

---
 src/WebSocketsServer.cpp | 63 ++++++++++++++++++++++++++++++++++++++++
 src/WebSocketsServer.h   |  9 ++++++
 2 files changed, 72 insertions(+)

diff --git a/src/WebSocketsServer.cpp b/src/WebSocketsServer.cpp
index d65eec5..da382b6 100644
--- a/src/WebSocketsServer.cpp
+++ b/src/WebSocketsServer.cpp
@@ -30,6 +30,9 @@ WebSocketsServer::WebSocketsServer(uint16_t port, String origin, String protocol
     _origin   = origin;
     _protocol = protocol;
     _runnning = false;
+    _pingInterval = 0;
+    _pongTimeout  = 0;
+    _disconnectTimeoutCount = 0;
 
     _server = new WEBSOCKETS_NETWORK_SERVER_CLASS(port);
 
@@ -91,6 +94,10 @@ void WebSocketsServer::begin(void) {
 #if(WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266_ASYNC)
         client->cHttpLine = "";
 #endif
+
+        client->pingInterval = _pingInterval;
+        client->pongTimeout  = _pongTimeout;
+        client->disconnectTimeoutCount = _disconnectTimeoutCount;
     }
 
 #ifdef ESP8266
@@ -486,6 +493,12 @@ bool WebSocketsServer::newClient(WEBSOCKETS_NETWORK_CLASS * TCPclient) {
             client->tcp->readStringUntil('\n', &(client->cHttpLine), std::bind(&WebSocketsServer::handleHeader, this, client, &(client->cHttpLine)));
 #endif
 
+            client->pingInterval = _pingInterval;
+            client->pongTimeout  = _pongTimeout;
+            client->disconnectTimeoutCount = _disconnectTimeoutCount;
+            client->lastPing     = millis();
+            client->pongReceived = false;
+
             return true;
             break;
         }
@@ -677,6 +690,9 @@ void WebSocketsServer::handleClientData(void) {
                         break;
                 }
             }
+            
+            handleHBPing(client);
+            handleHBTimeout(client);
         }
 #if(WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266)
         delay(0);
@@ -854,3 +870,50 @@ void WebSocketsServer::handleHeader(WSclient_t * client, String * headerLine) {
         }
     }
 }
+
+/**
+ * send heartbeat ping to server in set intervals
+ */
+void WebSocketsServer::handleHBPing(WSclient_t * client) {
+    if(client->pingInterval == 0)
+        return;
+    uint32_t pi = millis() - client->lastPing;
+    if(pi > client->pingInterval) {
+    DEBUG_WEBSOCKETS("[WS-Server][%d] sending HB ping\n", client->num);
+        if(sendPing(client->num)) {
+            client->lastPing     = millis();
+            client->pongReceived = false;
+        }
+    }
+}
+
+/**
+ * enable ping/pong heartbeat process
+ * @param pingInterval uint32_t how often ping will be sent
+ * @param pongTimeout uint32_t millis after which pong should timout if not received
+ * @param disconnectTimeoutCount uint8_t how many timeouts before disconnect, 0=> do not disconnect
+ */
+void WebSocketsServer::enableHeartbeat(uint32_t pingInterval, uint32_t pongTimeout, uint8_t disconnectTimeoutCount) {
+    _pingInterval = pingInterval;
+    _pongTimeout = pongTimeout;
+    _disconnectTimeoutCount = disconnectTimeoutCount;
+    
+    WSclient_t * client;
+    for(uint8_t i = 0; i < WEBSOCKETS_SERVER_CLIENT_MAX; i++) {
+        client = &_clients[i];
+        WebSockets::enableHeartbeat(client, pingInterval, pongTimeout, disconnectTimeoutCount);
+    }
+}
+
+/**
+ * disable ping/pong heartbeat process
+ */
+void WebSocketsServer::disableHeartbeat() {
+    _pingInterval = 0;
+    
+    WSclient_t * client;
+    for(uint8_t i = 0; i < WEBSOCKETS_SERVER_CLIENT_MAX; i++) {
+        client = &_clients[i];
+        client->pingInterval = 0;
+    }
+}
\ No newline at end of file
diff --git a/src/WebSocketsServer.h b/src/WebSocketsServer.h
index 4874a70..74ea85f 100644
--- a/src/WebSocketsServer.h
+++ b/src/WebSocketsServer.h
@@ -91,6 +91,9 @@ class WebSocketsServer : protected WebSockets {
     void setAuthorization(const char * auth);
 
     int connectedClients(bool ping = false);
+    
+    void enableHeartbeat(uint32_t pingInterval, uint32_t pongTimeout, uint8_t disconnectTimeoutCount);
+    void disableHeartbeat();
 
 #if(WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266) || (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP8266_ASYNC) || (WEBSOCKETS_NETWORK_TYPE == NETWORK_ESP32)
     IPAddress remoteIP(uint8_t num);
@@ -112,6 +115,10 @@ class WebSocketsServer : protected WebSockets {
     WebSocketServerHttpHeaderValFunc _httpHeaderValidationFunc;
 
     bool _runnning;
+    
+    uint32_t _pingInterval;
+    uint32_t _pongTimeout;
+    uint8_t _disconnectTimeoutCount;
 
     bool newClient(WEBSOCKETS_NETWORK_CLASS * TCPclient);
 
@@ -126,6 +133,8 @@ class WebSocketsServer : protected WebSockets {
 #endif
 
     void handleHeader(WSclient_t * client, String * headerLine);
+    
+    void handleHBPing(WSclient_t * client);    // send ping in specified intervals
 
     /**
          * called if a non Websocket connection is coming in.