memory leak fix

ITpC · Jun 28, 2019 · ca0b688 · ca0b688
1 parent e64095f
commit ca0b688
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 9 deletions.
diff --git a/include/ClientWebSocket.h b/include/ClientWebSocket.h
@@ -113,8 +113,6 @@ namespace LAppS
         if(tls)
         {
           wolfSSL_shutdown(TLSSocket);
-          wolfSSL_free(TLSSocket);
-          TLSSocket=nullptr;
         }      
         mState=WSClient::State::CLOSED;
       }
@@ -534,11 +532,26 @@ namespace LAppS
         // handshake
 
         int ret=force_send(httpUpgradeRequest);
-        if(ret == 0) throw std::system_error(ECONNRESET,std::system_category(),"ClientWebSocket::ClientWebSocket() - error on send, - peer has shutdown connection");
+
+        if(ret == 0)
+        {
+          if(tls&&TLSSocket)
+          {
+            wolfSSL_free(TLSSocket);
+            TLSSocket=nullptr;
+          }
+          throw std::system_error(ECONNRESET,std::system_category(),"ClientWebSocket::ClientWebSocket() - error on send, - peer has shutdown connection");
+        }
+
         if(ret>0)
         {
           mState=WSClient::State::HANDSHAKE;
         }else{
+          if(tls&&TLSSocket)
+          {
+            wolfSSL_free(TLSSocket);
+            TLSSocket=nullptr;
+          }
           throw std::system_error(EBADE,std::system_category(),"ClientWebSocket handshake failed on write");
         }
       }
@@ -607,7 +620,7 @@ namespace LAppS
       send(event.message->data(),event.message->size(),WebSocketProtocol::OpCode::PONG);
     }
 
-    ~ClientWebSocket()
+    ~ClientWebSocket() override
     {
       // prevent SIGPIPE on thread calling destructor.
       sigset_t sigsetmask;
@@ -631,6 +644,12 @@ namespace LAppS
           this->close();
         }
       }
+
+      if(tls&&TLSSocket)
+      {
+        wolfSSL_free(TLSSocket);
+        TLSSocket=nullptr;
+      }
     }
 
     const int send(const uint8_t* src, const size_t len, const WebSocketProtocol::OpCode& opcode)

diff --git a/include/modules/cws.h b/include/modules/cws.h
@@ -255,20 +255,25 @@ extern "C" {
       return 2;
     }
 
-    auto udptr=static_cast<int32_t*>(lua_newuserdata(L,sizeof(int32_t)));
-    std::string uri{lua_tostring(L,argc-1)};
     try
     {
-      (*udptr)=WSCPool.create(std::move(uri));
+      std::string uri{lua_tostring(L,argc-1)};
+
+      auto wscs=WSCPool.create(std::move(uri));
+
+      auto udptr=static_cast<int32_t*>(lua_newuserdata(L,sizeof(int32_t)));
+
+      (*udptr)=wscs;
 
       luaL_getmetatable(L,"cws");
       lua_pushvalue(L,3);
       lua_setfield(L,-2,std::to_string(*udptr).c_str());
       lua_setmetatable(L, -2);
       return 1;
-    }
-    catch(const std::exception& e)
+      
+    }catch(const std::exception& e)
     {
+      lua_pushnil(L);
       lua_pushnil(L);
       std::string message("Can't create socket: ");
       message.append(e.what());