From f6f3fca5d6b97ea5836304a1c4b4abf96c0a44ac Mon Sep 17 00:00:00 2001
From: Intel <chemicstry@gmail.com>
Date: Mon, 5 May 2014 13:10:59 -0400
Subject: [PATCH] Fix mem leak and crash

---
 src/server/poolserver/Stratum/Client.cpp    |  6 ++++--
 src/server/poolserver/Stratum/Client.h      | 20 +++++++++++++++-----
 src/server/poolserver/Stratum/Server.h      |  9 +++++----
 src/server/shared/MySQL/DatabaseOperation.h |  1 +
 4 files changed, 25 insertions(+), 11 deletions(-)

diff --git a/src/server/poolserver/Stratum/Client.cpp b/src/server/poolserver/Stratum/Client.cpp
index 392af9a..bddcedb 100644
--- a/src/server/poolserver/Stratum/Client.cpp
+++ b/src/server/poolserver/Stratum/Client.cpp
@@ -9,7 +9,7 @@
 
 namespace Stratum
 {
-    void Client::Start()
+    bool Client::Start()
     {
         // Get IP
         tcp::endpoint remote_ep = _socket.remote_endpoint();
@@ -19,11 +19,13 @@ namespace Stratum
         if (_server->IsBanned(_ip)) {
             sLog.Warn(LOG_STRATUM, "Blocked banned client from: %s", remote_ad.to_v4().to_string().c_str());
             Disconnect();
-            return;
+            return false;
         }
         
         // Start reading socket
         StartRead();
+        
+        return true;
     }
     
     void Client::SendJob(bool clean)
diff --git a/src/server/poolserver/Stratum/Client.h b/src/server/poolserver/Stratum/Client.h
index ff873aa..ebf3032 100644
--- a/src/server/poolserver/Stratum/Client.h
+++ b/src/server/poolserver/Stratum/Client.h
@@ -25,19 +25,24 @@ namespace Stratum
     class Client : public boost::enable_shared_from_this<Client>
     {
     public:
-        Client(Server* server, asio::io_service& io_service, uint64 id) : _server(server), _socket(io_service), _ioStrand(io_service), _id(id), _subscribed(false), _jobid(0), _shareLimiter(this)
+        Client(Server* server, asio::io_service& io_service, uint64 id) : _io_service(io_service), _server(server), _socket(io_service), _ioStrand(io_service), _id(id), _subscribed(false), _jobid(0), _shareLimiter(this)
         {
             _diff = sConfig.Get<uint32>("StratumMinDifficulty");
             _minDiff = _diff;
         }
         
+        ~Client()
+        {
+            sLog.Info(LOG_STRATUM, "%u: I'm going out! Cya!", _ip);
+        }
+        
         tcp::socket& GetSocket()
         {
             return _socket;
         }
         
         // Start client up!
-        void Start();
+        bool Start();
         
         void StartRead()
         {
@@ -46,7 +51,7 @@ namespace Stratum
                 _socket,
                 _recvBuffer,
                 asio::transfer_at_least(1),
-                _ioStrand.wrap(boost::bind(&Client::_OnReceive, this, asio::placeholders::error, asio::placeholders::bytes_transferred)));
+                _ioStrand.wrap(boost::bind(&Client::_OnReceive, shared_from_this(), asio::placeholders::error, asio::placeholders::bytes_transferred)));
         }
         
         void SendJob(bool clean);
@@ -61,7 +66,7 @@ namespace Stratum
             boost::asio::async_write(
                 _socket,
                 boost::asio::buffer(data.c_str(), data.length()),
-                _ioStrand.wrap(boost::bind(&Client::_OnSend, this, boost::asio::placeholders::error)));
+                _ioStrand.wrap(boost::bind(&Client::_OnSend, shared_from_this(), boost::asio::placeholders::error)));
         }
         
         void OnMiningSubmit(JSON msg);
@@ -126,13 +131,18 @@ namespace Stratum
         
         void CloseSocket()
         {
-            _socket.close();
+            boost::system::error_code ec;
+            _socket.shutdown(asio::ip::tcp::socket::shutdown_both, ec);
+            _socket.close(ec);
         }
     public:
         void _OnReceive(const boost::system::error_code& error, size_t bytes_transferred);
         void _OnSend(const boost::system::error_code& error);
         
     private:
+        // ASIO
+        asio::io_service& _io_service;
+        
         // Networking
         asio::streambuf _recvBuffer;
         std::string _recvMessage;
diff --git a/src/server/poolserver/Stratum/Server.h b/src/server/poolserver/Stratum/Server.h
index a8a8917..4b2cd12 100644
--- a/src/server/poolserver/Stratum/Server.h
+++ b/src/server/poolserver/Stratum/Server.h
@@ -98,8 +98,8 @@ namespace Stratum
         // Disconnects client
         void Disconnect(ClientPtr client)
         {
-            client->CloseSocket();
             _clients.erase(client);
+            client->CloseSocket();
             sLog.Debug(LOG_STRATUM, "Stratum client disconnected from %s. Total clients: %u", asio::ip::address_v4(client->GetIP()).to_string().c_str(), _clients.size());
         }
         
@@ -134,9 +134,10 @@ namespace Stratum
         void _OnAccept(ClientPtr client, const boost::system::error_code& error)
         {
             if (!error) {
-                client->Start();
-                _clients.insert(client);
-                sLog.Debug(LOG_STRATUM, "New stratum client accepted from %s. Total clients: %u", asio::ip::address_v4(client->GetIP()).to_string().c_str(), _clients.size());
+                if (client->Start()) {
+                    _clients.insert(client);
+                    sLog.Debug(LOG_STRATUM, "New stratum client accepted from %s. Total clients: %u", asio::ip::address_v4(client->GetIP()).to_string().c_str(), _clients.size());
+                }
             } else {
                 sLog.Debug(LOG_STRATUM, "Failed to accept stratum client");
             }
diff --git a/src/server/shared/MySQL/DatabaseOperation.h b/src/server/shared/MySQL/DatabaseOperation.h
index 3e8f05b..47639d0 100644
--- a/src/server/shared/MySQL/DatabaseOperation.h
+++ b/src/server/shared/MySQL/DatabaseOperation.h
@@ -14,6 +14,7 @@ namespace MySQL
     {
     public:
         DatabaseOperation(): _conn(NULL) {}
+        virtual ~DatabaseOperation() {}
         virtual void Execute() = 0;
         void SetConnection(DatabaseConnection* conn) { _conn = conn; }
     protected: