From 339e2d72dc2f883c34b573d92a17422579622422 Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Sun, 30 Dec 2018 17:43:11 +0100
Subject: [PATCH] Fix more memory leaks

---
 misc/main.cpp           |  2 ++
 src/App.h               | 10 ++++++++++
 src/Loop.h              |  8 ++++++++
 src/LoopData.h          | 10 ++++++++++
 src/PerMessageDeflate.h |  1 +
 src/WebSocketContext.h  |  9 +++++++++
 6 files changed, 40 insertions(+)

diff --git a/misc/main.cpp b/misc/main.cpp
index fc45e94..be96d45 100644
--- a/misc/main.cpp
+++ b/misc/main.cpp
@@ -67,6 +67,8 @@ int main(int argc, char **argv) {
         }
     }).run();
 
+    uWS::Loop::defaultLoop()->free();
+
 //    return 0;
 
 //    AsyncFileStreamer *asyncFileStreamer = new AsyncFileStreamer("/home/alexhultman/v0.15/public");
diff --git a/src/App.h b/src/App.h
index d427e89..94bf5fc 100644
--- a/src/App.h
+++ b/src/App.h
@@ -46,6 +46,8 @@ private:
     /* The app always owns at least one http context, but creates websocket contexts on demand */
     HttpContext<SSL> *httpContext;
 
+    std::vector<WebSocketContext<SSL, true> *> webSocketContexts;
+
     using SOCKET_TYPE = typename StaticDispatch<SSL>::SOCKET_TYPE;
     using StaticDispatch<SSL>::static_dispatch;
 public:
@@ -57,6 +59,10 @@ public:
     ~TemplatedApp() {
         /* Let's just put everything here */
         httpContext->free();
+
+        for (auto *webSocketContext : webSocketContexts) {
+            webSocketContext->free();
+        }
     }
 
     TemplatedApp(const TemplatedApp &other) {
@@ -83,6 +89,9 @@ public:
         /* Every route has its own websocket context with its own behavior and user data type */
         auto *webSocketContext = WebSocketContext<SSL, true>::create(Loop::defaultLoop(), (typename StaticDispatch<SSL>::SOCKET_CONTEXT_TYPE *) httpContext);
 
+        /* We need to clear this later on */
+        webSocketContexts.push_back(webSocketContext);
+
         /* Quick fix to disable any compression if set */
 #ifdef UWS_NO_ZLIB
         behavior.compression = uWS::DISABLED;
@@ -184,6 +193,7 @@ public:
             }
         });
 
+        // never called
         return *this;
     }
 
diff --git a/src/Loop.h b/src/Loop.h
index 2af4387..a184c5d 100644
--- a/src/Loop.h
+++ b/src/Loop.h
@@ -62,6 +62,10 @@ private:
 
     Loop() = delete;
 
+    ~Loop() {
+        std::cout << "Loop destructor called" << std::endl;
+    }
+
     Loop *init() {
         new (us_loop_ext((us_loop *) this)) LoopData;
         return this;
@@ -96,7 +100,11 @@ public:
 
     /* Freeing the default loop should be done once */
     void free() {
+        LoopData *loopData = (LoopData *) us_loop_ext((us_loop *) this);
+        loopData->~LoopData();
         us_loop_free((us_loop *) this);
+
+        std::cout << "Loop::free" << std::endl;
     }
 
     /* Set postCb callback */
diff --git a/src/LoopData.h b/src/LoopData.h
index f87b150..04bbc86 100644
--- a/src/LoopData.h
+++ b/src/LoopData.h
@@ -38,6 +38,16 @@ private:
     std::function<void(Loop *)> postHandler;
 
 public:
+    ~LoopData() {
+        /* If we have had App.ws called with compression we need to clear this */
+        if (zlibContext) {
+            delete zlibContext;
+            delete inflationStream;
+            delete deflationStream;
+        }
+        delete [] corkBuffer;
+    }
+
     /* Good 16k for SSL perf. */
     static const int CORK_BUFFER_SIZE = 16 * 1024;
 
diff --git a/src/PerMessageDeflate.h b/src/PerMessageDeflate.h
index 47e10cb..003f178 100644
--- a/src/PerMessageDeflate.h
+++ b/src/PerMessageDeflate.h
@@ -124,6 +124,7 @@ struct InflationStream {
 
     ~InflationStream() {
         std::cout << "Destructing inflationstream" << std::endl;
+        inflateEnd(&inflationStream);
     }
 
     std::string_view inflate(ZlibContext *zlibContext, std::string_view compressed, size_t maxPayloadLength) {
diff --git a/src/WebSocketContext.h b/src/WebSocketContext.h
index 03e4647..7b61ebb 100644
--- a/src/WebSocketContext.h
+++ b/src/WebSocketContext.h
@@ -335,6 +335,15 @@ private:
         return this;
     }
 
+    void free() {
+        std::cout << "websocket context free" << std::endl;
+
+        WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext((SOCKET_CONTEXT_TYPE *) this);
+        webSocketContextData->~WebSocketContextData();
+
+        us_socket_context_free((SOCKET_CONTEXT_TYPE *) this);
+    }
+
 public:
     /* WebSocket contexts are always child contexts to a HTTP context so no SSL options are needed as they are inherited */
     static WebSocketContext *create(Loop *loop, SOCKET_CONTEXT_TYPE *parentSocketContext) {