Core/Network: Refactored socket code, moved common operations to base Socket class

4 files changed, 76 insertions, 130 deletions

diff --git a/src/server/game/CMakeLists.txt b/src/server/game/CMakeLists.txt
index 84d6f87d75c..532900c0438 100644
--- a/src/server/game/CMakeLists.txt
+++ b/src/server/game/CMakeLists.txt
@@ -123,6 +123,7 @@ include_directories(
   ${CMAKE_SOURCE_DIR}/src/server/shared/Dynamic/LinkedReference
   ${CMAKE_SOURCE_DIR}/src/server/shared/Dynamic
   ${CMAKE_SOURCE_DIR}/src/server/shared/Logging
+  ${CMAKE_SOURCE_DIR}/src/server/shared/Networking
   ${CMAKE_SOURCE_DIR}/src/server/shared/Packets
   ${CMAKE_SOURCE_DIR}/src/server/shared/Threading
   ${CMAKE_SOURCE_DIR}/src/server/shared/Utilities
diff --git a/src/server/game/Server/WorldSession.cpp b/src/server/game/Server/WorldSession.cpp
index f30991b385e..eb83a5f55be 100644
--- a/src/server/game/Server/WorldSession.cpp
+++ b/src/server/game/Server/WorldSession.cpp
@@ -130,7 +130,7 @@ WorldSession::WorldSession(uint32 id, std::shared_ptr<WorldSocket> sock, Account
 
     if (sock)
     {
-        m_Address = sock->GetRemoteIpAddress();
+        m_Address = sock->GetRemoteIpAddress().to_string();
         ResetTimeOutTime();
         LoginDatabase.PExecute("UPDATE account SET online = 1 WHERE id = %u;", GetAccountId());     // One-time query
     }
diff --git a/src/server/game/Server/WorldSocket.cpp b/src/server/game/Server/WorldSocket.cpp
index 575c145687c..87340593c79 100644
--- a/src/server/game/Server/WorldSocket.cpp
+++ b/src/server/game/Server/WorldSocket.cpp
@@ -31,7 +31,7 @@ using boost::asio::ip::tcp;
 using boost::asio::streambuf;
 
 WorldSocket::WorldSocket(tcp::socket&& socket)
-    : _socket(std::move(socket)), _authSeed(rand32()), _OverSpeedPings(0), _worldSession(nullptr)
+    : Socket(std::move(socket), sizeof(ClientPktHeader)), _authSeed(rand32()), _OverSpeedPings(0), _worldSession(nullptr)
 {
 }
 
@@ -58,105 +58,89 @@ void WorldSocket::HandleSendAuthSession()
     AsyncWrite(packet);
 }
 
-void WorldSocket::AsyncReadHeader()
+void WorldSocket::ReadHeaderHandler(boost::system::error_code error, size_t transferedBytes)
 {
-    auto self(shared_from_this());
-    _socket.async_read_some(boost::asio::buffer(_readBuffer, sizeof(ClientPktHeader)), [this, self](boost::system::error_code error, size_t transferedBytes)
+    if (!error && transferedBytes == sizeof(ClientPktHeader))
     {
-        if (!error && transferedBytes == sizeof(ClientPktHeader))
-        {
-            ClientPktHeader* header = (ClientPktHeader*)&_readBuffer;
-
-            if (_worldSession)
-                _authCrypt.DecryptRecv((uint8*)header, sizeof(ClientPktHeader));
+        _authCrypt.DecryptRecv(GetReadBuffer(), sizeof(ClientPktHeader));
 
-            EndianConvertReverse(header->size);
-            EndianConvert(header->cmd);
+        ClientPktHeader* header = reinterpret_cast<ClientPktHeader*>(GetReadBuffer());
+        EndianConvertReverse(header->size);
+        EndianConvert(header->cmd);
 
-            AsyncReadData(header->size - sizeof(header->cmd));
-        }
-        else
-        {
-            // _socket.is_open() till returns true even after calling close()
-            CloseSocket();
-        }
-    });
+        AsyncReadData(header->size - sizeof(header->cmd), sizeof(ClientPktHeader));
+    }
+    else
+        CloseSocket();
 }
 
-void WorldSocket::AsyncReadData(size_t dataSize)
+void WorldSocket::ReadDataHandler(boost::system::error_code error, size_t transferedBytes)
 {
-    auto self(shared_from_this());
-    _socket.async_read_some(boost::asio::buffer(&_readBuffer[sizeof(ClientPktHeader)], dataSize), [this, dataSize, self](boost::system::error_code error, size_t transferedBytes)
+    ClientPktHeader* header = reinterpret_cast<ClientPktHeader*>(GetReadBuffer());
+
+    if (!error && transferedBytes == (header->size - sizeof(header->cmd)))
     {
-        if (!error && transferedBytes == dataSize)
-        {
-            ClientPktHeader* header = (ClientPktHeader*)&_readBuffer;
+        header->size -= sizeof(header->cmd);
 
-            header->size -= sizeof(header->cmd);
+        uint16 opcode = uint16(header->cmd);
 
-            uint16 opcode = (uint16)header->cmd;
+        std::string opcodeName = GetOpcodeNameForLogging(opcode);
 
-            std::string opcodeName = GetOpcodeNameForLogging(opcode);
+        WorldPacket packet(opcode, header->size);
 
-            WorldPacket packet(opcode, header->size);
+        if (header->size > 0)
+        {
+            packet.resize(header->size);
 
-            if (header->size > 0)
-            {
-                packet.resize(header->size);
+            std::memcpy(packet.contents(), &(GetReadBuffer()[sizeof(ClientPktHeader)]), header->size);
+        }
 
-                std::memcpy(packet.contents(), &_readBuffer[sizeof(ClientPktHeader)], header->size);
-            }
+        if (sPacketLog->CanLogPacket())
+            sPacketLog->LogPacket(packet, CLIENT_TO_SERVER);
 
-            if (sPacketLog->CanLogPacket())
-                sPacketLog->LogPacket(packet, CLIENT_TO_SERVER);
+        TC_LOG_TRACE("network.opcode", "C->S: %s %s", (_worldSession ? _worldSession->GetPlayerInfo() : GetRemoteIpAddress().to_string()).c_str(), GetOpcodeNameForLogging(opcode).c_str());
 
-            TC_LOG_TRACE("network.opcode", "C->S: %s %s", (_worldSession ? _worldSession->GetPlayerInfo() : GetRemoteIpAddress()).c_str(), GetOpcodeNameForLogging(opcode).c_str());
+        switch (opcode)
+        {
+            case CMSG_PING:
+                HandlePing(packet);
+                break;
+            case CMSG_AUTH_SESSION:
+                if (_worldSession)
+                {
+                    TC_LOG_ERROR("network", "WorldSocket::ProcessIncoming: received duplicate CMSG_AUTH_SESSION from %s", _worldSession->GetPlayerInfo().c_str());
+                    break;
+                }
 
-            switch (opcode)
+                sScriptMgr->OnPacketReceive(shared_from_this(), packet);
+                HandleAuthSession(packet);
+                break;
+            case CMSG_KEEP_ALIVE:
+                TC_LOG_DEBUG("network", "%s", opcodeName.c_str());
+                sScriptMgr->OnPacketReceive(shared_from_this(), packet);
+                break;
+            default:
             {
-                case CMSG_PING:
-                    HandlePing(packet);
-                    break;
-                case CMSG_AUTH_SESSION:
-                    if (_worldSession)
-                    {
-                        TC_LOG_ERROR("network", "WorldSocket::ProcessIncoming: received duplicate CMSG_AUTH_SESSION from %s", _worldSession->GetPlayerInfo().c_str());
-                        break;
-                    }
-
-                    sScriptMgr->OnPacketReceive(shared_from_this(), packet);
-                    HandleAuthSession(packet);
-                    break;
-                case CMSG_KEEP_ALIVE:
-                    TC_LOG_DEBUG("network", "%s", opcodeName.c_str());
-                    sScriptMgr->OnPacketReceive(shared_from_this(), packet);
-                    break;
-                default:
+                if (!_worldSession)
                 {
-                    if (!_worldSession)
-                    {
-                        TC_LOG_ERROR("network.opcode", "ProcessIncoming: Client not authed opcode = %u", uint32(opcode));
-                        break;
-                    }
-
-                    // Our Idle timer will reset on any non PING opcodes.
-                    // Catches people idling on the login screen and any lingering ingame connections.
-                    _worldSession->ResetTimeOutTime();
-
-                    // Copy the packet to the heap before enqueuing
-                    _worldSession->QueuePacket(new WorldPacket(packet));
+                    TC_LOG_ERROR("network.opcode", "ProcessIncoming: Client not authed opcode = %u", uint32(opcode));
                     break;
                 }
-            }
 
-            AsyncReadHeader();
-        }
-        else
-        {
-            // _socket.is_open() till returns true even after calling close()
-            CloseSocket();
+                // Our Idle timer will reset on any non PING opcodes.
+                // Catches people idling on the login screen and any lingering ingame connections.
+                _worldSession->ResetTimeOutTime();
+
+                // Copy the packet to the heap before enqueuing
+                _worldSession->QueuePacket(new WorldPacket(packet));
+                break;
+            }
         }
-    });
+
+        AsyncReadHeader();
+    }
+    else
+        CloseSocket();
 }
 
 void WorldSocket::AsyncWrite(WorldPacket const& packet)
@@ -164,7 +148,7 @@ void WorldSocket::AsyncWrite(WorldPacket const& packet)
     if (sPacketLog->CanLogPacket())
         sPacketLog->LogPacket(packet, SERVER_TO_CLIENT);
 
-    TC_LOG_TRACE("network.opcode", "S->C: %s %s", (_worldSession ? _worldSession->GetPlayerInfo() : GetRemoteIpAddress()).c_str(), GetOpcodeNameForLogging(packet.GetOpcode()).c_str());
+    TC_LOG_TRACE("network.opcode", "S->C: %s %s", (_worldSession ? _worldSession->GetPlayerInfo() : GetRemoteIpAddress().to_string()).c_str(), GetOpcodeNameForLogging(packet.GetOpcode()).c_str());
 
     ServerPktHeader header(packet.size() + 2, packet.GetOpcode());
 
@@ -185,25 +169,6 @@ void WorldSocket::AsyncWrite(WorldPacket const& packet)
         AsyncWrite(_writeQueue.front());
 }
 
-void WorldSocket::AsyncWrite(std::vector<uint8> const& data)
-{
-    auto self(shared_from_this());
-    boost::asio::async_write(_socket, boost::asio::buffer(data), [this, self](boost::system::error_code error, std::size_t /*length*/)
-    {
-        if (!error)
-        {
-            std::lock_guard<std::mutex> deleteGuard(_writeLock);
-
-            _writeQueue.pop();
-
-            if (!_writeQueue.empty())
-                AsyncWrite(_writeQueue.front());
-        }
-        else
-            CloseSocket();
-    });
-}
-
 void WorldSocket::HandleAuthSession(WorldPacket& recvPacket)
 {
     uint8 digest[20];
@@ -223,7 +188,7 @@ void WorldSocket::HandleAuthSession(WorldPacket& recvPacket)
     if (sWorld->IsClosed())
     {
         SendAuthResponseError(AUTH_REJECT);
-        TC_LOG_ERROR("network", "WorldSocket::HandleAuthSession: World closed, denying client (%s).", GetRemoteIpAddress().c_str());
+        TC_LOG_ERROR("network", "WorldSocket::HandleAuthSession: World closed, denying client (%s).", GetRemoteIpAddress().to_string().c_str());
         return;
     }
 
@@ -270,7 +235,7 @@ void WorldSocket::HandleAuthSession(WorldPacket& recvPacket)
         expansion = world_expansion;
 
     // For hook purposes, we get Remoteaddress at this point.
-    std::string address = GetRemoteIpAddress();
+    std::string address = GetRemoteIpAddress().to_string();
 
     // As we don't know if attempted login process by ip works, we update last_attempt_ip right away
     stmt = LoginDatabase.GetPreparedStatement(LOGIN_UPD_LAST_ATTEMPT_IP);
@@ -471,7 +436,7 @@ void WorldSocket::HandlePing(WorldPacket& recvPacket)
                 if (_worldSession && !_worldSession->HasPermission(rbac::RBAC_PERM_SKIP_CHECK_OVERSPEED_PING))
                 {
                     TC_LOG_ERROR("network", "WorldSocket::HandlePing: %s kicked for over-speed pings (address: %s)",
-                        _worldSession->GetPlayerInfo().c_str(), GetRemoteIpAddress().c_str());
+                        _worldSession->GetPlayerInfo().c_str(), GetRemoteIpAddress().to_string().c_str());
 
                     CloseSocket();
                     return;
@@ -489,8 +454,7 @@ void WorldSocket::HandlePing(WorldPacket& recvPacket)
     }
     else
     {
-        TC_LOG_ERROR("network", "WorldSocket::HandlePing: peer sent CMSG_PING, but is not authenticated or got recently kicked, address = %s",
-            GetRemoteIpAddress().c_str());
+        TC_LOG_ERROR("network", "WorldSocket::HandlePing: peer sent CMSG_PING, but is not authenticated or got recently kicked, address = %s", GetRemoteIpAddress().to_string().c_str());
 
         CloseSocket();
         return;
@@ -500,13 +464,3 @@ void WorldSocket::HandlePing(WorldPacket& recvPacket)
     packet << ping;
     return AsyncWrite(packet);
 }
-
-void WorldSocket::CloseSocket()
-{
-    boost::system::error_code socketError;
-    _socket.close(socketError);
-    if (socketError)
-        TC_LOG_DEBUG("network", "WorldSocket::CloseSocket: Player '%s' (%s) errored when closing socket: %i (%s)",
-            _worldSession ? _worldSession->GetPlayerInfo().c_str() : "unknown", GetRemoteIpAddress().c_str(),
-            socketError.value(), socketError.message().c_str());
-}
diff --git a/src/server/game/Server/WorldSocket.h b/src/server/game/Server/WorldSocket.h
index 0f3fc553872..5839c2194c4 100644
--- a/src/server/game/Server/WorldSocket.h
+++ b/src/server/game/Server/WorldSocket.h
@@ -21,6 +21,7 @@
 
 #include "Common.h"
 #include "AuthCrypt.h"
+#include "Socket.h"
 #include "Util.h"
 #include "WorldPacket.h"
 #include "WorldSession.h"
@@ -42,7 +43,7 @@ struct ClientPktHeader
 
 #pragma pack(pop)
 
-class WorldSocket : public std::enable_shared_from_this<WorldSocket>
+class WorldSocket : public Socket<WorldSocket>
 {
 public:
     WorldSocket(tcp::socket&& socket);
@@ -50,15 +51,15 @@ public:
     WorldSocket(WorldSocket const& right) = delete;
     WorldSocket& operator=(WorldSocket const& right) = delete;
 
-    void Start();
+    void Start() override;
 
-    std::string GetRemoteIpAddress() const { return _socket.remote_endpoint().address().to_string(); };
-    uint16 GetRemotePort() const { return _socket.remote_endpoint().port(); }
+    void AsyncWrite(WorldPacket const& packet);
 
-    void CloseSocket();
-    bool IsOpen() const { return _socket.is_open(); }
+    using Socket<WorldSocket>::AsyncWrite;
 
-    void AsyncWrite(WorldPacket const& packet);
+protected:
+    void ReadHeaderHandler(boost::system::error_code error, size_t transferedBytes) override;
+    void ReadDataHandler(boost::system::error_code error, size_t transferedBytes) override;
 
 private:
     void HandleSendAuthSession();
@@ -67,16 +68,6 @@ private:
 
     void HandlePing(WorldPacket& recvPacket);
 
-    void AsyncReadHeader();
-    void AsyncReadData(size_t dataSize);
-    void AsyncWrite(std::vector<uint8> const& data);
-
-    tcp::socket _socket;
-
-    char _readBuffer[4096];
-    std::mutex _writeLock;
-    std::queue<std::vector<uint8> > _writeQueue;
-
     uint32 _authSeed;
     AuthCrypt _authCrypt;