From 48604bf622c53e851597d55282ae0ec120277aa5 Mon Sep 17 00:00:00 2001
From: Asbelos <asbelos@btinternet.com>
Date: Wed, 6 Dec 2023 17:12:49 +0000
Subject: [PATCH] ESP32 websockets

---
 CommandDistributor.cpp |  4 ++--
 Websockets.cpp         | 35 +++++++++++++++++++++--------------
 Websockets.h           |  1 +
 WifiESP32.cpp          | 14 +++++++++++---
 4 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/CommandDistributor.cpp b/CommandDistributor.cpp
index 866dbaf..cd78896 100644
--- a/CommandDistributor.cpp
+++ b/CommandDistributor.cpp
@@ -63,8 +63,6 @@ template<typename... Targs> void CommandDistributor::broadcastReply(clientType t
 // Parse is called by Withrottle or Ethernet interface to determine which
 // protocol the client is using and call the appropriate part of dcc++Ex
 void  CommandDistributor::parse(byte clientId,byte * buffer, RingStream * stream) {
-  if (Diag::WIFI && Diag::CMD)
-    DIAG(F("Parse C=%d T=%d B=%s"),clientId, clients[clientId], buffer);
   ring=stream;
 
   // First check if the client is not known
@@ -206,7 +204,9 @@ void CommandDistributor::setClockTime(int16_t clocktime, int8_t clockrate, byte
         // CAH. DIAG removed because LCD does it anyway. 
         LCD(6,F("Clk Time:%d Sp %d"), clocktime, clockrate);
         // look for an event for this time
+#ifdef EXRAIL_ACTIVE        
         RMFT2::clockEvent(clocktime,1);
+#endif        
         // Now tell everyone else what the time is.
         CommandDistributor::broadcastClockTime(clocktime, clockrate);
         lastclocktime = clocktime;
diff --git a/Websockets.cpp b/Websockets.cpp
index a051586..c9d964a 100644
--- a/Websockets.cpp
+++ b/Websockets.cpp
@@ -51,6 +51,10 @@
 #include "libsha1.h"
 #include "Websockets.h"
 #include "DIAG.h"
+#ifdef ARDUINO_ARCH_ESP32
+  // ESP32 runtime or definitions has strlcat_P missing 
+  #define strlcat_P strlcat
+#endif  
 static const char b64_table[] = {
   'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
   'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P',
@@ -179,21 +183,24 @@ byte * Websockets::unmask(byte clientId,RingStream *ring, byte * buffer) {
    return (dataLength>=126)? 4:2;
  }
 
- void Websockets::writeOutboundHeader(Print * stream,uint16_t dataLength) {
-    // write the outbound header 
-    // length patched if necessary.
+int Websockets::fillOutboundHeader(uint16_t dataLength, byte * buffer) {
     // text opcode, flag(126= use 2 length bytes, no mask bit) , length
-    if (dataLength>=126) {
-      const byte prefix[]={0x81,126,
-           (byte)(dataLength & 0xFF), (byte)(dataLength>>8)}; 
-      stream->write(prefix,sizeof(prefix));
+    buffer[0]=0x81;
+    if (dataLength<126) {
+         buffer[1]=(byte)dataLength;
+         return 2; 
     }
-    else {
-      const byte prefix[]={0x81,(byte)dataLength}; 
-      stream->write(prefix,sizeof(prefix));
-    }
-     
- }
-
+    buffer[1]=126;
+    buffer[2]=(byte)(dataLength & 0xFF);
+    buffer[3]= (byte)(dataLength>>8);
+    return 4;  
+}
+    
+ void Websockets::writeOutboundHeader(Print * stream,uint16_t dataLength) {
+    byte prefix[4];
+    int headerlen=fillOutboundHeader(dataLength,prefix);
+    stream->write(prefix,sizeof(headerlen));
+  }
+ 
 
 
diff --git a/Websockets.h b/Websockets.h
index a49d815..0792b47 100644
--- a/Websockets.h
+++ b/Websockets.h
@@ -26,6 +26,7 @@ class Websockets {
     static bool checkConnectionString(byte clientId,byte * cmd, RingStream * outbound );
     static byte * unmask(byte clientId,RingStream *ring, byte * buffer);
     static int16_t getOutboundHeaderSize(uint16_t dataLength);
+    static int fillOutboundHeader(uint16_t dataLength, byte * buffer);
     static void writeOutboundHeader(Print * stream,uint16_t dataLength);
     static const byte WEBSOCK_CLIENT_MARKER=0x80;
 };
diff --git a/WifiESP32.cpp b/WifiESP32.cpp
index f0a857f..d6a099c 100644
--- a/WifiESP32.cpp
+++ b/WifiESP32.cpp
@@ -2,6 +2,8 @@
     © 2023 Paul M. Antoine
     © 2021 Harald Barth
     © 2023 Nathan Kellenicki
+    © 2023 Chris Harlow
+    
 
     This file is part of CommandStation-EX
 
@@ -30,6 +32,7 @@
 #include "RingStream.h"
 #include "CommandDistributor.h"
 #include "WiThrottle.h"
+#include "Websockets.h"
 /*
 #include "soc/rtc_wdt.h"
 #include "esp_task_wdt.h"
@@ -330,6 +333,8 @@ void WifiESP::loop() {
 
     // something to write out?
     clientId=outboundRing->read();
+    bool useWebsocket=clientId & Websockets::WEBSOCK_CLIENT_MARKER;
+    clientId &= ~ Websockets::WEBSOCK_CLIENT_MARKER;
     if (clientId >= 0) {
       // We have data to send in outboundRing
       // and we have a valid clientId.
@@ -337,19 +342,22 @@ void WifiESP::loop() {
       // and then look if it can be sent because
       // we can not leave it in the ring for ever
       int count=outboundRing->count();
+      auto wsHeaderLen=useWebsocket? Websockets::getOutboundHeaderSize(count) : 0;
       {
-	char buffer[count+1]; // one extra for '\0'
+        
+	byte buffer[wsHeaderLen+count+1]; // one extra for '\0'
+  if (useWebsocket) Websockets::fillOutboundHeader(count, buffer);
 	for(int i=0;i<count;i++) {
 	  int c = outboundRing->read();
 	  if (c >= 0) // Panic check, should never be false
-	    buffer[i] = (char)c;
+	    buffer[i+wsHeaderLen] = (char)c;
 	  else {
 	    DIAG(F("Ringread fail at %d"),i);
 	    break;
 	  }
 	}
 	// buffer filled, end with '\0' so we can use it as C string
-	buffer[count]='\0';
+	buffer[wsHeaderLen+count]='\0';
 	if((unsigned int)clientId <= clients.size() && clients[clientId].ok()) {
 	  if (Diag::CMD || Diag::WITHROTTLE)
 	    DIAG(F("SEND %d:%s"), clientId, buffer);