websockets part 3

CommandDistributor.cpp

@@ -75,7 +75,7 @@ void  CommandDistributor::parse(byte clientId,byte * buffer, RingStream * stream
   if (clients[clientId] == NONE_TYPE) {
     auto websock=Websockets::checkConnectionString(clientId,buffer,stream);
     if (websock) {
-      clients[clientId]=COMMAND_TYPE;
+      clients[clientId]=WEBSOCK_CONNECTING_TYPE;
       return;
     }
     if (buffer[0] == '<')
@@ -84,16 +84,28 @@ void  CommandDistributor::parse(byte clientId,byte * buffer, RingStream * stream
       clients[clientId]=WITHROTTLE_TYPE;
   }
 
-  // mark buffer that is sent to parser
-  ring->mark(clientId);
+  // after first inbound transmission the websocket is connected
+  if (clients[clientId]==WEBSOCK_CONNECTING_TYPE)
+      clients[clientId]=WEBSOCKET_TYPE;
       
+      
+  // mark buffer that is sent to parser
   // When type is known, send the string
   // to the right parser
   if (clients[clientId] == COMMAND_TYPE) {
+    ring->mark(clientId);
     DCCEXParser::parse(stream, buffer, ring);
   } else if (clients[clientId] == WITHROTTLE_TYPE) {
+    ring->mark(clientId);
     WiThrottle::getThrottle(clientId)->parse(ring, buffer);
   }
+  else if (clients[clientId] == WEBSOCKET_TYPE) {
+    buffer=Websockets::unmask(clientId,ring, buffer);
+    if (!buffer) return; // unmask may have handled it alrerday (ping/pong)
+    // mark ring with client flagged as websocket for transmission later
+    ring->mark(clientId | Websockets::WEBSOCK_CLIENT_MARKER);
+    DCCEXParser::parse(stream, buffer, ring);
+    }
 
   if (ring->peekTargetMark()!=RingStream::NO_CLIENT) {
     // The commit call will either write the length bytes

CommandDistributor.h

@@ -36,7 +36,7 @@
 
 class CommandDistributor {
 public:
-  enum clientType: byte {NONE_TYPE,COMMAND_TYPE,WITHROTTLE_TYPE};
+  enum clientType: byte {NONE_TYPE,COMMAND_TYPE,WITHROTTLE_TYPE,WEBSOCK_CONNECTING_TYPE,WEBSOCKET_TYPE};
 private:
   static void broadcastToClients(clientType type);
   static StringBuffer * broadcastBufferWriter;

Websockets.cpp

@@ -1,3 +1,50 @@
+/*
+ *  © 2023 Chris Harlow
+ *  All rights reserved.
+ *
+ *  This file is part of CommandStation-EX
+ *
+ *  This is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  It is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with CommandStation.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+
+/**************************************************
+  HOW IT WORKS
+
+  1) Refer to https://developer.mozilla.org/en-US/docs/Web/API/WebSockets_API/Writing_WebSocket_servers
+  
+  2) When a new client sends in a socket stream, the
+     CommandDistributor pass it to this code 
+     checkConnectionString() to check for an HTTP
+     protocol GET requesting a change to websocket protocol.
+     If that is found, the relevant answer is generated and queued and
+     the CommandDistributor marks this client as a websocket client awaiting connection.
+     Once the outbound handshake has completed, the CommandDistributor promotes the client
+     from awaiting connection to connected websocket so that all
+     future traffic for this client is handled with websocket protocol.
+
+  3) When an input is received from a client marked as websocket,
+     CommandDistributor calls  unmask() to strip off the websocket header and
+     un-mask the input bytes. The command distributor will flag the
+     clientid in the ringstream so that anyone transmitting this
+     output will know to handle it differently.   
+
+   4) when the  Wifi/Ethernet handler needs to transmit the result from the
+      output ring, it recognises the websockets flag and adds the websocket 
+      header to the output dynamically. 
+
+ *************************************************************/
 #include <Arduino.h>
 #include "FSH.h"
 #include "RingStream.h"
@@ -59,7 +106,7 @@ bool Websockets::checkConnectionString(byte clientId,byte * cmd, RingStream * ou
     SHA1Update(&ctx, (unsigned char *)replyKey, strlen(replyKey));
     SHA1Final(sha1HashBin, &ctx);
       
-    // ghenerate the response and embed the base64 encode 
+    // generate the response and embed the base64 encode 
     // of the key
     outbound->mark(clientId);
     outbound->print("HTTP/1.1 101 Switching Protocols\r\n"
@@ -84,3 +131,70 @@ bool Websockets::checkConnectionString(byte clientId,byte * cmd, RingStream * ou
     return true;          
 }
 
+byte * Websockets::unmask(byte clientId,RingStream *ring, byte * buffer) {
+ // buffer should have a websocket header
+ //byte opcode=buffer[0] & 0x0f;
+ DIAG(F("Websock in: %x %x %x %x %x %x %x %x"),
+      buffer[0],buffer[1],buffer[2],buffer[3],
+       buffer[4],buffer[5],buffer[6]);
+
+ byte opcode=buffer[0];
+ bool maskbit=buffer[1]&0x80;
+ int16_t payloadLength=buffer[1]&0x7f;
+ 
+ byte * mask;
+ if (payloadLength<126) {
+     mask=buffer+2;
+     }
+ else {
+     payloadLength=(buffer[3]<<8)|(buffer[2]);
+     mask=buffer+4;
+ }
+ DIAG(F("Websock op=%x mb=%b pl=%d m=%x %x %x %x"), opcode, maskbit, payloadLength, 
+      mask[0],mask[1],mask[2], mask[3]);
+
+ if (opcode==0x89) { // ping
+     DIAG(F("Websocket ping"));
+     buffer[0]=0x8a;  // pong.. and send it back
+     ring->mark(clientId &0x7f); // dont readjust
+     ring->print((char *)buffer);
+     ring->commit();
+     return nullptr; 
+     }
+     
+ if (opcode!=0x81) {
+  DIAG(F("Websocket unknown opcode 0x%x"),opcode);
+  return nullptr;
+ }
+ 
+ byte * payload=mask+4;
+ for (int i=0;i<payloadLength;i++) {
+     payload[i]^=mask[i%4];
+ }
+ 
+ return payload; // payload will be parsed as normal
+     
+ }
+
+ int16_t Websockets::getOutboundHeaderSize(uint16_t dataLength) {
+   return (dataLength>=126)? 4:2;
+ }
+
+ void Websockets::writeOutboundHeader(Print * stream,uint16_t dataLength) {
+    // write the outbound header 
+    // length patched if necessary.
+    // 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));
+    }
+    else {
+      const byte prefix[]={0x81,(byte)dataLength}; 
+      stream->write(prefix,sizeof(prefix));
+    }
+     
+ }
+
+
+

Websockets.h

@@ -1,3 +1,22 @@
+/*
+ *  © 2023 Chris Harlow
+ *  All rights reserved.
+ *
+ *  This file is part of CommandStation-EX
+ *
+ *  This is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  It is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with CommandStation.  If not, see <https://www.gnu.org/licenses/>.
+ */
 #ifndef Websockets_h
 #define Websockets_h
 #include <Arduino.h>
@@ -5,6 +24,10 @@
 class Websockets {
     public:
     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 void writeOutboundHeader(Print * stream,uint16_t dataLength);
+    static const byte WEBSOCK_CLIENT_MARKER=0x80;
 };
 
 #endif

WifiInboundHandler.cpp

@@ -1,8 +1,7 @@
 /*
  *  © 2021 Fred Decker
  *  © 2021 Fred Decker
- *  © 2020-2021 Chris Harlow
- *  © 2020, Chris Harlow. All rights reserved.
+ *  © 2020-2023, Chris Harlow. All rights reserved.
  *  © 2020, Harald Barth.
  *  
  *  This file is part of Asbelos DCC API
@@ -26,6 +25,7 @@
 #include "RingStream.h"
 #include "CommandDistributor.h"
 #include "DIAG.h"
+#include "Websockets.h"
 
 WifiInboundHandler * WifiInboundHandler::singleton;
 
@@ -67,8 +67,13 @@ void WifiInboundHandler::loop1() {
     
 
     if (pendingCipsend && millis()-lastCIPSEND > CIPSENDgap) {
-         if (Diag::WIFI) DIAG( F("WiFi: [[CIPSEND=%d,%d]]"), clientPendingCIPSEND, currentReplySize);
-         StringFormatter::send(wifiStream, F("AT+CIPSEND=%d,%d\r\n"),  clientPendingCIPSEND, currentReplySize);
+         // add allowances for websockets
+         bool websocket=clientPendingCIPSEND & Websockets::WEBSOCK_CLIENT_MARKER;
+         byte realClient=clientPendingCIPSEND & ~Websockets::WEBSOCK_CLIENT_MARKER;
+         int16_t realSize=currentReplySize;
+         if (websocket) realSize+=Websockets::getOutboundHeaderSize(currentReplySize);
+         if (Diag::WIFI) DIAG( F("WiFi: [[CIPSEND=%d,%d]]"), realClient, realSize);
+         StringFormatter::send(wifiStream, F("AT+CIPSEND=%d,%d\r\n"),  realClient,realSize);
          pendingCipsend=false;
          return;
       }
@@ -112,9 +117,12 @@ WifiInboundHandler::INBOUND_STATE WifiInboundHandler::loop2() {
         }
         
         if (ch=='>') { 
-           if (Diag::WIFI) DIAG(F("[XMIT %d]"),currentReplySize); 
+           bool websocket=clientPendingCIPSEND & Websockets::WEBSOCK_CLIENT_MARKER; 
+           if (Diag::WIFI) DIAG(F("[XMIT %d ws=%b]"),currentReplySize,websocket);
+           if (websocket) Websockets::writeOutboundHeader(wifiStream,currentReplySize); 
            for (int i=0;i<currentReplySize;i++) {
              int cout=outboundRing->read();
+             if (websocket && (cout=='\n')) cout='\r';
              wifiStream->write(cout);
              if (Diag::WIFI) StringFormatter::printEscape(cout); // DIAG in disguise
            }

libsha1.cpp

@@ -1,3 +1,7 @@
+// For DCC-EX: This file downloaded from:
+//  https://github.com/Links2004/arduinoWebSockets
+// All due credit to Steve Reid
+
 /* from valgrind tests */
 
 /* ================ sha1.c ================ */

libsha1.h

@@ -1,3 +1,7 @@
+// For DCC-EX: This file downloaded from:
+//  https://github.com/Links2004/arduinoWebSockets
+// All due credit to Steve Reid
+
 /* ================ sha1.h ================ */
 /*
 SHA-1 in C