From c972197643340b8c8234bd1d36fe367a45b6eec1 Mon Sep 17 00:00:00 2001
From: Kevin Hester <kevinh@geeksville.com>
Date: Sun, 27 Dec 2020 16:58:32 +0800
Subject: [PATCH] fix #598 don't corrupt the heap when a TCP connection drops

---
 src/Observer.h              |  9 +++++++++
 src/esp32/WiFiServerAPI.cpp | 11 ++++++++---
 src/mesh/PhoneAPI.cpp       | 19 ++++++++++++-------
 src/mesh/PhoneAPI.h         | 13 +++++++++----
 4 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/src/Observer.h b/src/Observer.h
index 701495b38..85c40b1ea 100644
--- a/src/Observer.h
+++ b/src/Observer.h
@@ -16,6 +16,10 @@ template <class T> class Observer
   public:
     virtual ~Observer();
 
+    /// Stop watching our current obserable
+    void unobserve();
+
+    /// Start watching a specified observable
     void observe(Observable<T> *o);
 
   private:
@@ -82,6 +86,11 @@ template <class T> class Observable
 };
 
 template <class T> Observer<T>::~Observer()
+{
+    unobserve();
+}
+
+template <class T> void Observer<T>::unobserve()
 {
     if (observed)
         observed->removeObserver(this);
diff --git a/src/esp32/WiFiServerAPI.cpp b/src/esp32/WiFiServerAPI.cpp
index 632116bf3..3b11c2ae4 100644
--- a/src/esp32/WiFiServerAPI.cpp
+++ b/src/esp32/WiFiServerAPI.cpp
@@ -32,9 +32,14 @@ void WiFiServerAPI::loop()
 {
     if (client.connected()) {
         StreamAPI::loop();
-    } else {
-        DEBUG_MSG("Client dropped connection, closing TCP server\n");
-        delete this;
+    } else if(isConnected) {
+        // If our API link was up, shut it down
+
+        DEBUG_MSG("Client dropped connection, closing API client\n");
+        // Note: we can't call delete here because this object includes other state
+        // besides the stream API.  Instead kill it later when we start a new instance
+        // delete this;
+        close();
     }
 }
 
diff --git a/src/mesh/PhoneAPI.cpp b/src/mesh/PhoneAPI.cpp
index 754f0d380..a0cad4273 100644
--- a/src/mesh/PhoneAPI.cpp
+++ b/src/mesh/PhoneAPI.cpp
@@ -7,22 +7,27 @@
 #include <assert.h>
 
 #if FromRadio_size > MAX_TO_FROM_RADIO_SIZE
-    #error FromRadio is too big
+#error FromRadio is too big
 #endif
 
 #if ToRadio_size > MAX_TO_FROM_RADIO_SIZE
-    #error ToRadio is too big
+#error ToRadio is too big
 #endif
 
-PhoneAPI::PhoneAPI()
-{
-}
+PhoneAPI::PhoneAPI() {}
 
 void PhoneAPI::init()
 {
     observe(&service.fromNumChanged);
 }
 
+void PhoneAPI::close() {
+    unobserve();
+    state = STATE_SEND_NOTHING;
+    isConnected = false;
+    onConnectionChanged(isConnected);
+}
+
 void PhoneAPI::checkConnectionTimeout()
 {
     if (isConnected) {
@@ -102,7 +107,7 @@ size_t PhoneAPI::getFromRadio(uint8_t *buf)
     if (!available()) {
         // DEBUG_MSG("getFromRadio, !available\n");
         return 0;
-    } 
+    }
 
     DEBUG_MSG("getFromRadio, state=%d\n", state);
 
@@ -170,7 +175,7 @@ size_t PhoneAPI::getFromRadio(uint8_t *buf)
         if (packetForPhone) {
 
             printPacket("phone downloaded packet", packetForPhone);
-            
+
             // Encapsulate as a FromRadio packet
             fromRadioScratch.which_variant = FromRadio_packet_tag;
             fromRadioScratch.variant.packet = *packetForPhone;
diff --git a/src/mesh/PhoneAPI.h b/src/mesh/PhoneAPI.h
index 129ecb5db..b895adc7e 100644
--- a/src/mesh/PhoneAPI.h
+++ b/src/mesh/PhoneAPI.h
@@ -21,7 +21,7 @@ class PhoneAPI
     : public Observer<uint32_t> // FIXME, we shouldn't be inheriting from Observer, instead use CallbackObserver as a member
 {
     enum State {
-        STATE_LEGACY,       // Temporary default state - until Android apps are all updated, uses the old BLE API
+        STATE_LEGACY,       // (no longer used) old default state - until Android apps are all updated, uses the old BLE API
         STATE_SEND_NOTHING, // (Eventual) Initial state, don't send anything until the client starts asking for config
         STATE_SEND_MY_INFO, // send our my info record
         STATE_SEND_RADIO,
@@ -31,7 +31,7 @@ class PhoneAPI
         STATE_SEND_PACKETS // send packets or debug strings
     };
 
-    State state = STATE_LEGACY;
+    State state = STATE_SEND_NOTHING;
 
     /**
      * Each packet sent to the phone has an incrementing count
@@ -53,14 +53,16 @@ class PhoneAPI
     /** the last msec we heard from the client on the other side of this link */
     uint32_t lastContactMsec = 0;
 
-    bool isConnected = false;
-
   public:
     PhoneAPI();
 
     /// Do late init that can't happen at constructor time
     virtual void init();
 
+    // Call this when the client drops the connection, resets the state to STATE_SEND_NOTHING
+    // Unregisters our observer
+    void close();
+
     /**
      * Handle a ToRadio protobuf
      */
@@ -87,6 +89,9 @@ class PhoneAPI
     void handleSetRadio(const RadioConfig &r);
 
   protected:
+    /// Are we currently connected to a client?
+    bool isConnected = false;
+
     /// Our fromradio packet while it is being assembled
     FromRadio fromRadioScratch;