Core/Network: Socket refactors

* Devirtualize calls to Read and Update by marking concrete implementations as final
* Removed derived class template argument
* Specialize boost::asio::basic_stream_socket for boost::asio::io_context instead of type-erased any_io_executor
* Make socket initialization easier composable (before entering Read loop)
* Remove use of deprecated boost::asio::null_buffers and boost::beast::ssl_stream

4 files changed, 151 insertions, 156 deletions

diff --git a/src/server/game/Server/WorldSocket.cpp b/src/server/game/Server/WorldSocket.cpp
index 4b2ad6c3c5c..2b868ff3884 100644
--- a/src/server/game/Server/WorldSocket.cpp
+++ b/src/server/game/Server/WorldSocket.cpp
@@ -26,6 +26,7 @@
 #include "GameTime.h"
 #include "HMAC.h"
 #include "IPLocation.h"
+#include "IpBanCheckConnectionInitializer.h"
 #include "PacketLog.h"
 #include "ProtobufJSON.h"
 #include "RealmList.h"
@@ -33,6 +34,7 @@
 #include "RealmList.pb.h"
 #include "ScriptMgr.h"
 #include "SessionKeyGenerator.h"
+#include "SslStream.h"
 #include "World.h"
 #include "WorldPacket.h"
 #include "WorldSession.h"
@@ -49,8 +51,6 @@ struct CompressedWorldPacket
 
 #pragma pack(pop)
 
-std::string const WorldSocket::ServerConnectionInitialize("WORLD OF WARCRAFT CONNECTION - SERVER TO CLIENT - V2");
-std::string const WorldSocket::ClientConnectionInitialize("WORLD OF WARCRAFT CONNECTION - CLIENT TO SERVER - V2");
 uint32 const WorldSocket::MinSizeForCompression = 0x400;
 
 std::array<uint8, 32> const WorldSocket::AuthCheckSeed = { 0xDE, 0x3A, 0x2A, 0x8E, 0x6B, 0x89, 0x52, 0x66, 0x88, 0x9D, 0x7E, 0x7A, 0x77, 0x1D, 0x5D, 0x1F,
@@ -62,7 +62,7 @@ std::array<uint8, 32> const WorldSocket::ContinuedSessionSeed = { 0x56, 0x5C, 0x
 std::array<uint8, 32> const WorldSocket::EncryptionKeySeed = { 0x71, 0xC9, 0xED, 0x5A, 0xA7, 0x0E, 0x4D, 0xFF, 0x4C, 0x36, 0xA6, 0x5A, 0x3E, 0x46, 0x8A, 0x4A,
     0x5D, 0xA1, 0x48, 0xC8, 0x30, 0x47, 0x4A, 0xDE, 0xF6, 0x0D, 0x6C, 0xBE, 0x6F, 0xE4, 0x55, 0x73 };
 
-WorldSocket::WorldSocket(boost::asio::ip::tcp::socket&& socket) : Socket(std::move(socket)),
+WorldSocket::WorldSocket(Trinity::Net::IoContextTcpSocket&& socket) : BaseSocket(std::move(socket)),
     _type(CONNECTION_TYPE_REALM), _key(0), _serverChallenge(), _sessionKey(), _encryptKey(), _OverSpeedPings(0),
     _worldSession(nullptr), _authed(false), _canRequestHotfixes(true), _headerBuffer(sizeof(IncomingPacketHeader)), _sendBufferSize(4096), _compressionStream(nullptr)
 {
@@ -77,127 +77,127 @@ WorldSocket::~WorldSocket()
     }
 }
 
-void WorldSocket::Start()
+struct WorldSocketProtocolInitializer final : Trinity::Net::SocketConnectionInitializer
 {
-    std::string ip_address = GetRemoteIpAddress().to_string();
-    LoginDatabasePreparedStatement* stmt = LoginDatabase.GetPreparedStatement(LOGIN_SEL_IP_INFO);
-    stmt->setString(0, ip_address);
+    static constexpr std::string_view ServerConnectionInitialize = "WORLD OF WARCRAFT CONNECTION - SERVER TO CLIENT - V2\n";
+    static constexpr std::string_view ClientConnectionInitialize = "WORLD OF WARCRAFT CONNECTION - CLIENT TO SERVER - V2\n";
 
-    _queryProcessor.AddCallback(LoginDatabase.AsyncQuery(stmt).WithPreparedCallback([self = shared_from_this()](PreparedQueryResult result)
-    {
-        self->CheckIpCallback(std::move(result));
-    }));
-}
+    explicit WorldSocketProtocolInitializer(WorldSocket* socket) : _socket(socket) { }
 
-void WorldSocket::CheckIpCallback(PreparedQueryResult result)
-{
-    if (result)
+    void Start() override
     {
-        bool banned = false;
-        do
-        {
-            Field* fields = result->Fetch();
-            if (fields[0].GetUInt64() != 0)
-                banned = true;
+        _packetBuffer.Resize(ClientConnectionInitialize.length());
 
-        } while (result->NextRow());
+        AsyncRead();
 
-        if (banned)
-        {
-            TC_LOG_ERROR("network", "WorldSocket::CheckIpCallback: Sent Auth Response (IP {} banned).", GetRemoteIpAddress().to_string());
-            DelayedCloseSocket();
-            return;
-        }
+        MessageBuffer initializer;
+        initializer.Write(ServerConnectionInitialize.data(), ServerConnectionInitialize.length());
+
+        // - IoContext.run thread, safe.
+        _socket->QueuePacket(std::move(initializer));
     }
 
-    _packetBuffer.Resize(ClientConnectionInitialize.length() + 1);
+    void AsyncRead()
+    {
+        _socket->AsyncRead(
+            [socketRef = _socket->weak_from_this(), self = static_pointer_cast<WorldSocketProtocolInitializer>(this->shared_from_this())]
+            {
+                if (!socketRef.expired())
+                    return self->ReadHandler();
+
+                return Trinity::Net::SocketReadCallbackResult::Stop;
+            });
+    }
 
-    AsyncReadWithCallback(&WorldSocket::InitializeHandler);
+    Trinity::Net::SocketReadCallbackResult ReadHandler();
 
-    MessageBuffer initializer;
-    initializer.Write(ServerConnectionInitialize.c_str(), ServerConnectionInitialize.length());
-    initializer.Write("\n", 1);
+    void HandleDataReady();
 
-    // - IoContext.run thread, safe.
-    QueuePacket(std::move(initializer));
+private:
+    WorldSocket* _socket;
+    MessageBuffer _packetBuffer;
+};
+
+void WorldSocket::Start()
+{
+    // build initializer chain
+    std::array<std::shared_ptr<Trinity::Net::SocketConnectionInitializer>, 3> initializers =
+    { {
+        std::make_shared<Trinity::Net::IpBanCheckConnectionInitializer<WorldSocket>>(this),
+        std::make_shared<WorldSocketProtocolInitializer>(this),
+        std::make_shared<Trinity::Net::ReadConnectionInitializer<WorldSocket>>(this),
+    } };
+
+    Trinity::Net::SocketConnectionInitializer::SetupChain(initializers)->Start();
 }
 
-void WorldSocket::InitializeHandler(boost::system::error_code const& error, std::size_t transferedBytes)
+Trinity::Net::SocketReadCallbackResult WorldSocketProtocolInitializer::ReadHandler()
 {
-    if (error)
+    MessageBuffer& packet = _socket->GetReadBuffer();
+    if (packet.GetActiveSize() > 0 && _packetBuffer.GetRemainingSpace() > 0)
     {
-        CloseSocket();
-        return;
+        // need to receive the header
+        std::size_t readHeaderSize = std::min(packet.GetActiveSize(), _packetBuffer.GetRemainingSpace());
+        _packetBuffer.Write(packet.GetReadPointer(), readHeaderSize);
+        packet.ReadCompleted(readHeaderSize);
+
+        if (_packetBuffer.GetRemainingSpace() == 0)
+        {
+            HandleDataReady();
+            return Trinity::Net::SocketReadCallbackResult::Stop;
+        }
+
+        // Couldn't receive the whole header this time.
+        ASSERT(packet.GetActiveSize() == 0);
     }
 
-    GetReadBuffer().WriteCompleted(transferedBytes);
+    return Trinity::Net::SocketReadCallbackResult::KeepReading;
+}
 
-    MessageBuffer& packet = GetReadBuffer();
-    if (packet.GetActiveSize() > 0)
+void WorldSocketProtocolInitializer::HandleDataReady()
+{
+    try
     {
-        if (_packetBuffer.GetRemainingSpace() > 0)
+        ByteBuffer buffer(std::move(_packetBuffer));
+        if (buffer.ReadString(ClientConnectionInitialize.length()) != ClientConnectionInitialize)
         {
-            // need to receive the header
-            std::size_t readHeaderSize = std::min(packet.GetActiveSize(), _packetBuffer.GetRemainingSpace());
-            _packetBuffer.Write(packet.GetReadPointer(), readHeaderSize);
-            packet.ReadCompleted(readHeaderSize);
-
-            if (_packetBuffer.GetRemainingSpace() > 0)
-            {
-                // Couldn't receive the whole header this time.
-                ASSERT(packet.GetActiveSize() == 0);
-                AsyncReadWithCallback(&WorldSocket::InitializeHandler);
-                return;
-            }
-
-            try
-            {
-                ByteBuffer buffer(std::move(_packetBuffer));
-                std::string initializer(buffer.ReadString(ClientConnectionInitialize.length()));
-                if (initializer != ClientConnectionInitialize)
-                {
-                    CloseSocket();
-                    return;
-                }
+            _socket->CloseSocket();
+            return;
+        }
+    }
+    catch (ByteBufferException const& ex)
+    {
+        TC_LOG_ERROR("network", "WorldSocket::InitializeHandler ByteBufferException {} occured while parsing initial packet from {}",
+            ex.what(), _socket->GetRemoteIpAddress().to_string());
+        _socket->CloseSocket();
+        return;
+    }
 
-                uint8 terminator;
-                buffer >> terminator;
-                if (terminator != '\n')
-                {
-                    CloseSocket();
-                    return;
-                }
-            }
-            catch (ByteBufferException const& ex)
-            {
-                TC_LOG_ERROR("network", "WorldSocket::InitializeHandler ByteBufferException {} occured while parsing initial packet from {}",
-                    ex.what(), GetRemoteIpAddress().to_string());
-                CloseSocket();
-                return;
-            }
+    if (!_socket->InitializeCompression())
+        return;
 
-            _compressionStream = new z_stream();
-            _compressionStream->zalloc = (alloc_func)nullptr;
-            _compressionStream->zfree = (free_func)nullptr;
-            _compressionStream->opaque = (voidpf)nullptr;
-            _compressionStream->avail_in = 0;
-            _compressionStream->next_in = nullptr;
-            int32 z_res = deflateInit2(_compressionStream, sWorld->getIntConfig(CONFIG_COMPRESSION), Z_DEFLATED, -15, 8, Z_DEFAULT_STRATEGY);
-            if (z_res != Z_OK)
-            {
-                CloseSocket();
-                TC_LOG_ERROR("network", "Can't initialize packet compression (zlib: deflateInit) Error code: {} ({})", z_res, zError(z_res));
-                return;
-            }
+    _socket->SendAuthSession();
+    if (next)
+        next->Start();
+}
 
-            _packetBuffer.Reset();
-            HandleSendAuthSession();
-            AsyncRead();
-            return;
-        }
+bool WorldSocket::InitializeCompression()
+{
+    _compressionStream = new z_stream();
+    _compressionStream->zalloc = (alloc_func)nullptr;
+    _compressionStream->zfree = (free_func)nullptr;
+    _compressionStream->opaque = (voidpf)nullptr;
+    _compressionStream->avail_in = 0;
+    _compressionStream->next_in = nullptr;
+    int32 z_res = deflateInit2(_compressionStream, sWorld->getIntConfig(CONFIG_COMPRESSION), Z_DEFLATED, -15, 8, Z_DEFAULT_STRATEGY);
+    if (z_res != Z_OK)
+    {
+        CloseSocket();
+        TC_LOG_ERROR("network", "Can't initialize packet compression (zlib: deflateInit) Error code: {} ({})", z_res, zError(z_res));
+        return false;
     }
 
-    AsyncReadWithCallback(&WorldSocket::InitializeHandler);
+    return true;
 }
 
 bool WorldSocket::Update()
@@ -206,7 +206,7 @@ bool WorldSocket::Update()
     MessageBuffer buffer(_sendBufferSize);
     while (_bufferQueue.Dequeue(queued))
     {
-        uint32 packetSize = queued->size() + 2 /*opcode*/;
+        uint32 packetSize = queued->size() + 4 /*opcode*/;
         if (packetSize > MinSizeForCompression && queued->NeedsEncryption())
             packetSize = deflateBound(_compressionStream, packetSize) + sizeof(CompressedWorldPacket);
 
@@ -240,7 +240,7 @@ bool WorldSocket::Update()
     return true;
 }
 
-void WorldSocket::HandleSendAuthSession()
+void WorldSocket::SendAuthSession()
 {
     Trinity::Crypto::GetRandomBytes(_serverChallenge);
 
@@ -260,11 +260,8 @@ void WorldSocket::OnClose()
     }
 }
 
-void WorldSocket::ReadHandler()
+Trinity::Net::SocketReadCallbackResult WorldSocket::ReadHandler()
 {
-    if (!IsOpen())
-        return;
-
     MessageBuffer& packet = GetReadBuffer();
     while (packet.GetActiveSize() > 0)
     {
@@ -286,7 +283,7 @@ void WorldSocket::ReadHandler()
             if (!ReadHeaderHandler())
             {
                 CloseSocket();
-                return;
+                return Trinity::Net::SocketReadCallbackResult::Stop;
             }
         }
 
@@ -314,11 +311,16 @@ void WorldSocket::ReadHandler()
             if (result != ReadDataHandlerResult::WaitingForQuery)
                 CloseSocket();
 
-            return;
+            return Trinity::Net::SocketReadCallbackResult::Stop;
         }
     }
 
-    AsyncRead();
+    return Trinity::Net::SocketReadCallbackResult::KeepReading;
+}
+
+void WorldSocket::QueueQuery(QueryCallback&& queryCallback)
+{
+    _queryProcessor.AddCallback(std::move(queryCallback));
 }
 
 void WorldSocket::SetWorldSession(WorldSession* session)
@@ -510,14 +512,13 @@ WorldSocket::ReadDataHandlerResult WorldSocket::ReadDataHandler()
 
 void WorldSocket::LogOpcodeText(OpcodeClient opcode, std::unique_lock<std::mutex> const& guard) const
 {
-    if (!guard)
+    if (!guard || !_worldSession)
     {
         TC_LOG_TRACE("network.opcode", "C->S: {} {}", GetRemoteIpAddress().to_string(), GetOpcodeNameForLogging(opcode));
     }
     else
     {
-        TC_LOG_TRACE("network.opcode", "C->S: {} {}", (_worldSession ? _worldSession->GetPlayerInfo() : GetRemoteIpAddress().to_string()),
-            GetOpcodeNameForLogging(opcode));
+        TC_LOG_TRACE("network.opcode", "C->S: {} {}", _worldSession->GetPlayerInfo(), GetOpcodeNameForLogging(opcode));
     }
 }
 
@@ -688,7 +689,7 @@ void WorldSocket::HandleAuthSession(std::shared_ptr<WorldPackets::Auth::AuthSess
     stmt->setInt32(0, int32(sRealmList->GetCurrentRealmId().Realm));
     stmt->setString(1, joinTicket->gameaccount());
 
-    _queryProcessor.AddCallback(LoginDatabase.AsyncQuery(stmt).WithPreparedCallback([this, authSession = std::move(authSession), joinTicket = std::move(joinTicket)](PreparedQueryResult result) mutable
+    QueueQuery(LoginDatabase.AsyncQuery(stmt).WithPreparedCallback([this, authSession = std::move(authSession), joinTicket = std::move(joinTicket)](PreparedQueryResult result) mutable
     {
         HandleAuthSessionCallback(std::move(authSession), std::move(joinTicket), std::move(result));
     }));
@@ -898,19 +899,20 @@ void WorldSocket::HandleAuthSessionCallback(std::shared_ptr<WorldPackets::Auth::
     sScriptMgr->OnAccountLogin(account.Game.Id);
 
     _authed = true;
-    _worldSession = new WorldSession(account.Game.Id, std::move(*joinTicket->mutable_gameaccount()), account.BattleNet.Id, shared_from_this(), account.Game.Security,
-        account.Game.Expansion, mutetime, account.Game.OS, account.Game.TimezoneOffset, account.Game.Build, buildVariant, account.Game.Locale,
+    _worldSession = new WorldSession(account.Game.Id, std::move(*joinTicket->mutable_gameaccount()), account.BattleNet.Id,
+        static_pointer_cast<WorldSocket>(shared_from_this()), account.Game.Security, account.Game.Expansion, mutetime,
+        account.Game.OS, account.Game.TimezoneOffset, account.Game.Build, buildVariant, account.Game.Locale,
         account.Game.Recruiter, account.Game.IsRectuiter);
 
     // Initialize Warden system only if it is enabled by config
     if (wardenActive)
         _worldSession->InitWarden(_sessionKey);
 
-    _queryProcessor.AddCallback(_worldSession->LoadPermissionsAsync().WithPreparedCallback([this](PreparedQueryResult result)
+    QueueQuery(_worldSession->LoadPermissionsAsync().WithPreparedCallback([this](PreparedQueryResult result)
     {
         LoadSessionPermissionsCallback(std::move(result));
     }));
-    AsyncRead();
+    AsyncRead(Trinity::Net::InvokeReadHandlerCallback<WorldSocket>{ .Socket = this });
 }
 
 void WorldSocket::LoadSessionPermissionsCallback(PreparedQueryResult result)
@@ -938,7 +940,7 @@ void WorldSocket::HandleAuthContinuedSession(std::shared_ptr<WorldPackets::Auth:
     LoginDatabasePreparedStatement* stmt = LoginDatabase.GetPreparedStatement(LOGIN_SEL_ACCOUNT_INFO_CONTINUED_SESSION);
     stmt->setUInt32(0, accountId);
 
-    _queryProcessor.AddCallback(LoginDatabase.AsyncQuery(stmt).WithPreparedCallback([this, authSession = std::move(authSession)](PreparedQueryResult result) mutable
+    QueueQuery(LoginDatabase.AsyncQuery(stmt).WithPreparedCallback([this, authSession = std::move(authSession)](PreparedQueryResult result) mutable
     {
         HandleAuthContinuedSessionCallback(std::move(authSession), std::move(result));
     }));
@@ -985,7 +987,7 @@ void WorldSocket::HandleAuthContinuedSessionCallback(std::shared_ptr<WorldPacket
     memcpy(_encryptKey.data(), encryptKeyGen.GetDigest().data(), 32);
 
     SendPacketAndLogOpcode(*WorldPackets::Auth::EnterEncryptedMode(_encryptKey, true).Write());
-    AsyncRead();
+    AsyncRead(Trinity::Net::InvokeReadHandlerCallback<WorldSocket>{ .Socket = this });
 }
 
 void WorldSocket::HandleConnectToFailed(WorldPackets::Auth::ConnectToFailed& connectToFailed)
@@ -1032,7 +1034,7 @@ void WorldSocket::HandleEnterEncryptedModeAck()
     if (_type == CONNECTION_TYPE_REALM)
         sWorld->AddSession(_worldSession);
     else
-        sWorld->AddInstanceSocket(shared_from_this(), _key);
+        sWorld->AddInstanceSocket(static_pointer_cast<WorldSocket>(shared_from_this()), _key);
 }
 
 void WorldSocket::SendAuthResponseError(uint32 code)
diff --git a/src/server/game/Server/WorldSocket.h b/src/server/game/Server/WorldSocket.h
index 04451dc26e8..b5388532b4d 100644
--- a/src/server/game/Server/WorldSocket.h
+++ b/src/server/game/Server/WorldSocket.h
@@ -15,8 +15,8 @@
  * with this program. If not, see <http://www.gnu.org/licenses/>.
  */
 
-#ifndef __WORLDSOCKET_H__
-#define __WORLDSOCKET_H__
+#ifndef TRINITYCORE_WORLD_SOCKET_H
+#define TRINITYCORE_WORLD_SOCKET_H
 
 #include "AsyncCallbackProcessor.h"
 #include "AuthDefines.h"
@@ -77,7 +77,7 @@ struct PacketHeader
     uint32 Size;
     uint8 Tag[12];
 
-    bool IsValidSize() { return Size < 0x10000; }
+    bool IsValidSize() const { return Size < 0x10000; }
 };
 
 struct IncomingPacketHeader : PacketHeader
@@ -87,10 +87,8 @@ struct IncomingPacketHeader : PacketHeader
 
 #pragma pack(pop)
 
-class TC_GAME_API WorldSocket : public Socket<WorldSocket>
+class TC_GAME_API WorldSocket final : public Trinity::Net::Socket<>
 {
-    static std::string const ServerConnectionInitialize;
-    static std::string const ClientConnectionInitialize;
     static uint32 const MinSizeForCompression;
 
     static std::array<uint8, 32> const AuthCheckSeed;
@@ -98,14 +96,16 @@ class TC_GAME_API WorldSocket : public Socket<WorldSocket>
     static std::array<uint8, 32> const ContinuedSessionSeed;
     static std::array<uint8, 32> const EncryptionKeySeed;
 
-    typedef Socket<WorldSocket> BaseSocket;
+    using BaseSocket = Socket;
 
 public:
-    WorldSocket(boost::asio::ip::tcp::socket&& socket);
+    WorldSocket(Trinity::Net::IoContextTcpSocket&& socket);
     ~WorldSocket();
 
     WorldSocket(WorldSocket const& right) = delete;
+    WorldSocket(WorldSocket&& right) = delete;
     WorldSocket& operator=(WorldSocket const& right) = delete;
+    WorldSocket& operator=(WorldSocket&& right) = delete;
 
     void Start() override;
     bool Update() override;
@@ -118,9 +118,15 @@ public:
     void SetWorldSession(WorldSession* session);
     void SetSendBufferSize(std::size_t sendBufferSize) { _sendBufferSize = sendBufferSize; }
 
-protected:
     void OnClose() override;
-    void ReadHandler() override;
+    Trinity::Net::SocketReadCallbackResult ReadHandler() override;
+
+    void QueueQuery(QueryCallback&& queryCallback);
+
+    void SendAuthSession();
+    bool InitializeCompression();
+
+protected:
     bool ReadHeaderHandler();
 
     enum class ReadDataHandlerResult
@@ -132,9 +138,6 @@ protected:
 
     ReadDataHandlerResult ReadDataHandler();
 private:
-    void CheckIpCallback(PreparedQueryResult result);
-    void InitializeHandler(boost::system::error_code const& error, std::size_t transferedBytes);
-
     /// writes network.opcode log
     /// accessing WorldSession is not threadsafe, only do it when holding _worldSessionLock
     void LogOpcodeText(OpcodeClient opcode, std::unique_lock<std::mutex> const& guard) const;
@@ -143,7 +146,6 @@ private:
     void WritePacketToBuffer(EncryptablePacket const& packet, MessageBuffer& buffer);
     uint32 CompressPacket(uint8* buffer, WorldPacket const& packet);
 
-    void HandleSendAuthSession();
     void HandleAuthSession(std::shared_ptr<WorldPackets::Auth::AuthSession> authSession);
     void HandleAuthSessionCallback(std::shared_ptr<WorldPackets::Auth::AuthSession> authSession,
         std::shared_ptr<JSON::RealmList::RealmJoinTicket> joinTicket, PreparedQueryResult result);
diff --git a/src/server/game/Server/WorldSocketMgr.cpp b/src/server/game/Server/WorldSocketMgr.cpp
index 58f242b52f8..cc44e5bad4b 100644
--- a/src/server/game/Server/WorldSocketMgr.cpp
+++ b/src/server/game/Server/WorldSocketMgr.cpp
@@ -22,21 +22,16 @@
 #include "WorldSocket.h"
 #include <boost/system/error_code.hpp>
 
-static void OnSocketAccept(boost::asio::ip::tcp::socket&& sock, uint32 threadIndex)
-{
-    sWorldSocketMgr.OnSocketOpen(std::move(sock), threadIndex);
-}
-
-class WorldSocketThread : public NetworkThread<WorldSocket>
+class WorldSocketThread : public Trinity::Net::NetworkThread<WorldSocket>
 {
 public:
-    void SocketAdded(std::shared_ptr<WorldSocket> sock) override
+    void SocketAdded(std::shared_ptr<WorldSocket> const& sock) override
     {
         sock->SetSendBufferSize(sWorldSocketMgr.GetApplicationSendBufferSize());
         sScriptMgr->OnSocketOpen(sock);
     }
 
-    void SocketRemoved(std::shared_ptr<WorldSocket> sock) override
+    void SocketRemoved(std::shared_ptr<WorldSocket>const& sock) override
     {
         sScriptMgr->OnSocketClose(sock);
     }
@@ -75,7 +70,10 @@ bool WorldSocketMgr::StartNetwork(Trinity::Asio::IoContext& ioContext, std::stri
     if (!BaseSocketMgr::StartNetwork(ioContext, bindIp, port, threadCount))
         return false;
 
-    _acceptor->AsyncAcceptWithCallback<&OnSocketAccept>();
+    _acceptor->AsyncAccept([this](Trinity::Net::IoContextTcpSocket&& sock, uint32 threadIndex)
+    {
+        OnSocketOpen(std::move(sock), threadIndex);
+    });
 
     sScriptMgr->OnNetworkStart();
     return true;
@@ -88,7 +86,7 @@ void WorldSocketMgr::StopNetwork()
     sScriptMgr->OnNetworkStop();
 }
 
-void WorldSocketMgr::OnSocketOpen(boost::asio::ip::tcp::socket&& sock, uint32 threadIndex)
+void WorldSocketMgr::OnSocketOpen(Trinity::Net::IoContextTcpSocket&& sock, uint32 threadIndex)
 {
     // set some options here
     if (_socketSystemSendBufferSize >= 0)
@@ -117,7 +115,7 @@ void WorldSocketMgr::OnSocketOpen(boost::asio::ip::tcp::socket&& sock, uint32 th
     BaseSocketMgr::OnSocketOpen(std::move(sock), threadIndex);
 }
 
-NetworkThread<WorldSocket>* WorldSocketMgr::CreateThreads() const
+Trinity::Net::NetworkThread<WorldSocket>* WorldSocketMgr::CreateThreads() const
 {
     return new WorldSocketThread[GetNetworkThreadCount()];
 }
diff --git a/src/server/game/Server/WorldSocketMgr.h b/src/server/game/Server/WorldSocketMgr.h
index 84b190575a4..8859da81074 100644
--- a/src/server/game/Server/WorldSocketMgr.h
+++ b/src/server/game/Server/WorldSocketMgr.h
@@ -15,21 +15,15 @@
  * with this program. If not, see <http://www.gnu.org/licenses/>.
  */
 
-/** \addtogroup u2w User to World Communication
- *  @{
- *  \file WorldSocketMgr.h
- *  \author Derex <derex101@gmail.com>
- */
-
-#ifndef __WORLDSOCKETMGR_H
-#define __WORLDSOCKETMGR_H
+#ifndef TRINITYCORE_WORLD_SOCKET_MGR_H
+#define TRINITYCORE_WORLD_SOCKET_MGR_H
 
 #include "SocketMgr.h"
 
 class WorldSocket;
 
 /// Manages all sockets connected to peers and network threads
-class TC_GAME_API WorldSocketMgr : public SocketMgr<WorldSocket>
+class TC_GAME_API WorldSocketMgr : public Trinity::Net::SocketMgr<WorldSocket>
 {
     typedef SocketMgr<WorldSocket> BaseSocketMgr;
 
@@ -44,14 +38,14 @@ public:
     /// Stops all network threads, It will wait for all running threads .
     void StopNetwork() override;
 
-    void OnSocketOpen(boost::asio::ip::tcp::socket&& sock, uint32 threadIndex) override;
+    void OnSocketOpen(Trinity::Net::IoContextTcpSocket&& sock, uint32 threadIndex) override;
 
     std::size_t GetApplicationSendBufferSize() const { return _socketApplicationSendBufferSize; }
 
 protected:
     WorldSocketMgr();
 
-    NetworkThread<WorldSocket>* CreateThreads() const override;
+    Trinity::Net::NetworkThread<WorldSocket>* CreateThreads() const override;
 
 private:
     int32 _socketSystemSendBufferSize;
@@ -62,4 +56,3 @@ private:
 #define sWorldSocketMgr WorldSocketMgr::Instance()
 
 #endif
-/// @}