From 6833d8278935b54fd9a076421a6bec462d314c17 Mon Sep 17 00:00:00 2001
From: Asbelos <asbelos@btinternet.com>
Date: Mon, 26 Oct 2020 12:59:40 +0000
Subject: [PATCH] Buffer overflow resilience

And diag output of replies
---
 RingStream.cpp         | 25 +++++++++++++++++++++++--
 RingStream.h           |  4 ++++
 WifiInboundHandler.cpp | 25 +++++++++++++++++++++----
 WifiInboundHandler.h   |  5 +++--
 4 files changed, 51 insertions(+), 8 deletions(-)

diff --git a/RingStream.cpp b/RingStream.cpp
index eee0753..cea17fb 100644
--- a/RingStream.cpp
+++ b/RingStream.cpp
@@ -28,11 +28,12 @@ RingStream::RingStream( const uint16_t len)
   _pos_read=0;
   _buffer[0]=0;
   _overflow=false;
+  _mark=0;
 }
 
-size_t RingStream::write(uint8_t byte) {
+size_t RingStream::write(uint8_t b) {
   if (_overflow) return 0;
-  _buffer[_pos_write] = byte;
+  _buffer[_pos_write] = b;
   ++_pos_write;
   if (_pos_write>=_len) _pos_write=0;
   if (_pos_write==_pos_read) {
@@ -63,3 +64,23 @@ int RingStream::count() {
   }
   return counter;
   }
+
+int RingStream::freeSpace() {
+  if (_pos_read>_pos_write) return _pos_read-_pos_write-2;
+  else return _len - _pos_write + _pos_read-2;  
+}
+
+
+void RingStream::mark() {
+    _mark=_pos_write;
+}
+
+bool RingStream::commit() {
+  if (_overflow) {
+    _pos_write=_mark;
+    _overflow=false;
+    return false; // commit failed
+  }
+  write((uint8_t)0);
+  return true; // commit worked
+}
diff --git a/RingStream.h b/RingStream.h
index 7a129e0..52caf9e 100644
--- a/RingStream.h
+++ b/RingStream.h
@@ -30,12 +30,16 @@ class RingStream : public Print {
     using Print::write;
     int read();
     int count();
+    int freeSpace();
+    void mark();
+    bool commit();
 
  private:
    int _len;
    int _pos_write;
    int _pos_read;
    bool _overflow;
+   int _mark;
    byte * _buffer;
 };
 
diff --git a/WifiInboundHandler.cpp b/WifiInboundHandler.cpp
index 8a779ee..6aa8738 100644
--- a/WifiInboundHandler.cpp
+++ b/WifiInboundHandler.cpp
@@ -65,9 +65,12 @@ void WifiInboundHandler::loop1() {
          cmd[count]=0;
          if (Diag::WIFI) DIAG(F("%e\n"),cmd); 
          
+         outboundRing->mark();  // remember start of outbound data 
          outboundRing->print(clientId);
          CommandDistributor::parse(clientId,cmd,outboundRing);
-         outboundRing->write((byte)0);
+         // The commit call will either write the null byte at the end of the output,
+         // OR rollback to the mark because the commend generated more than fits rthe buffer 
+         outboundRing->commit();
          return;
       }
    }
@@ -95,8 +98,11 @@ WifiInboundHandler::INBOUND_STATE WifiInboundHandler::loop2() {
         }
         
         if (ch=='>') { 
-           if (Diag::WIFI) DIAG(F("[[XMIT %d]]"),currentReplySize);
-           for (int i=0;i<currentReplySize;i++) wifiStream->write(outboundRing->read());
+           for (int i=0;i<currentReplySize;i++) {
+             int cout=outboundRing->read();
+             wifiStream->write(cout);
+             if (Diag::WIFI) StringFormatter::printEscape(cout); // DIAG in disguise
+           }
            outboundRing->read(); // drop the end marker
            clientPendingCIPSEND=-1;
            pendingCipsend=false;
@@ -159,6 +165,12 @@ WifiInboundHandler::INBOUND_STATE WifiInboundHandler::loop2() {
             break;
           }
           if (Diag::WIFI) DIAG(F("\nWifi inbound data(%d:%d):"),runningClientId,dataLength); 
+          if (inboundRing->freeSpace()<=(dataLength+1)) {
+            // This input would overflow the inbound ring, ignore it  
+            loopState=IPD_IGNORE_DATA;
+            break;
+          }
+          inboundRing->mark();
           inboundRing->print(runningClientId); // prefix inbound with client id
           loopState=IPD_DATA;
           break; 
@@ -170,11 +182,16 @@ WifiInboundHandler::INBOUND_STATE WifiInboundHandler::loop2() {
          inboundRing->write(ch);    
         dataLength--;
         if (dataLength == 0) {
-           inboundRing->write((byte)0);    
+           inboundRing->commit();    
           loopState = ANYTHING;
         }
         break;
 
+      case IPD_IGNORE_DATA: // ignoring data that would not fit in inbound ring
+        dataLength--;
+        if (dataLength == 0) loopState = ANYTHING;
+        break;
+
       case GOT_CLIENT_ID:  // got x before CLOSE or CONNECTED
         loopState=(ch==',') ? GOT_CLIENT_ID2: SKIPTOEND;
         break;
diff --git a/WifiInboundHandler.h b/WifiInboundHandler.h
index fa41aa4..e947bdb 100644
--- a/WifiInboundHandler.h
+++ b/WifiInboundHandler.h
@@ -32,6 +32,7 @@ class WifiInboundHandler {
           IPD5,        // got +IPD,c 
           IPD6_LENGTH, // got +IPD,c, reading length 
           IPD_DATA,    // got +IPD,c,ll,: collecting data
+          IPD_IGNORE_DATA, // got +IPD,c,ll,: ignoring the data that won't fit inblound Ring
 
           GOT_CLIENT_ID,  // clientid prefix to CONNECTED / CLOSED
           GOT_CLIENT_ID2,  // clientid prefix to CONNECTED / CLOSED
@@ -44,8 +45,8 @@ class WifiInboundHandler {
    INBOUND_STATE loop2();
    Stream * wifiStream;
    
-   static const int INBOUND_RING = 200;
-   static const int OUTBOUND_RING = 1024;
+   static const int INBOUND_RING = 512;
+   static const int OUTBOUND_RING = 2048;
  
    RingStream * inboundRing;
    RingStream * outboundRing;