Core/NetworkIO: Applied 4.3.4 changes to "new" network code

7 files changed, 61 insertions, 31 deletions

diff --git a/src/server/authserver/Main.cpp b/src/server/authserver/Main.cpp
index 5c06a1e1c34..e614b2b79fa 100644
--- a/src/server/authserver/Main.cpp
+++ b/src/server/authserver/Main.cpp
@@ -126,9 +126,9 @@ int main(int argc, char** argv)
     }
 
     std::string bindIp = sConfigMgr->GetStringDefault("BindIP", "0.0.0.0");
-    AsyncAcceptor<Battlenet::Session> bnetServer(_ioService, bindIp, bnport);
 
     sAuthSocketMgr.StartNetwork(_ioService, bindIp, port);
+    sBattlenetSessionMgr.StartNetwork(_ioService, bindIp, bnport);
 
     // Set signal handlers
     boost::asio::signal_set signals(_ioService, SIGINT, SIGTERM);
diff --git a/src/server/authserver/Server/BattlenetBitStream.h b/src/server/authserver/Server/BattlenetBitStream.h
index 82c2a0a6d5d..6a5e366ede5 100644
--- a/src/server/authserver/Server/BattlenetBitStream.h
+++ b/src/server/authserver/Server/BattlenetBitStream.h
@@ -64,7 +64,7 @@ namespace Battlenet
             _buffer.resize(length, 0);
         }
 
-        BitStream(MessageBuffer&& buffer) : _writePos(buffer.GetReadyDataSize() * 8), _readPos(0), _buffer(buffer.Move())
+        BitStream(MessageBuffer&& buffer) : _writePos(buffer.GetActiveSize() * 8), _readPos(0), _buffer(buffer.Move())
         {
         }
 
diff --git a/src/server/authserver/Server/BattlenetSession.cpp b/src/server/authserver/Server/BattlenetSession.cpp
index 2c2026253dc..7abb9df9db9 100644
--- a/src/server/authserver/Server/BattlenetSession.cpp
+++ b/src/server/authserver/Server/BattlenetSession.cpp
@@ -56,7 +56,7 @@ Battlenet::Session::ModuleHandler const Battlenet::Session::ModuleHandlers[MODUL
     &Battlenet::Session::HandleResumeModule,
 };
 
-Battlenet::Session::Session(tcp::socket&& socket) : Socket(std::move(socket), std::size_t(BufferSizes::Read)), _accountId(0), _accountName(), _locale(),
+Battlenet::Session::Session(tcp::socket&& socket) : Socket(std::move(socket)), _accountId(0), _accountName(), _locale(),
     _os(), _build(0), _gameAccountId(0), _gameAccountName(), _accountSecurityLevel(SEC_PLAYER), I(), s(), v(), b(), B(), K(),
     _reconnectProof(), _crypt(), _authed(false)
 {
@@ -537,9 +537,9 @@ bool Battlenet::Session::HandleRealmJoinRequest(PacketHeader& header, BitStream&
     return true;
 }
 
-void Battlenet::Session::ReadDataHandler()
+void Battlenet::Session::ReadHandler()
 {
-    BitStream packet(MoveData());
+    BitStream packet(std::move(GetReadBuffer()));
     _crypt.DecryptRecv(packet.GetBuffer(), packet.GetSize());
 
     while (!packet.IsRead())
@@ -581,13 +581,14 @@ void Battlenet::Session::ReadDataHandler()
         }
     }
 
-    AsyncReadData(size_t(BufferSizes::Read));
+    GetReadBuffer().Resize(size_t(BufferSizes::Read));
+    AsyncRead();
 }
 
 void Battlenet::Session::Start()
 {
     TC_LOG_TRACE("server.battlenet", "Battlenet::Session::Start");
-    AsyncReadData(size_t(BufferSizes::Read));
+    AsyncRead();
 }
 
 void Battlenet::Session::AsyncWrite(ServerPacket* packet)
@@ -602,16 +603,15 @@ void Battlenet::Session::AsyncWrite(ServerPacket* packet)
 
     packet->Write();
 
-    std::lock_guard<std::mutex> guard(_writeLock);
+    MessageBuffer buffer;
+    buffer.Write(packet->GetData(), packet->GetSize());
+    delete packet;
 
-    _crypt.EncryptSend(packet->GetData(), packet->GetSize());
+    std::unique_lock<std::mutex> guard(_writeLock);
 
-    bool needsWriteStart = _writeQueue.empty();
+    _crypt.EncryptSend(buffer.GetReadPointer(), buffer.GetActiveSize());
 
-    _writeQueue.push(packet);
-
-    if (needsWriteStart)
-        BattlenetSocket::AsyncWrite(_writeQueue.front());
+    QueuePacket(std::move(buffer), guard);
 }
 
 inline void ReplaceResponse(Battlenet::ServerPacket** oldResponse, Battlenet::ServerPacket* newResponse)
diff --git a/src/server/authserver/Server/BattlenetSession.h b/src/server/authserver/Server/BattlenetSession.h
index f32f97d15e7..2e6c32cbabf 100644
--- a/src/server/authserver/Server/BattlenetSession.h
+++ b/src/server/authserver/Server/BattlenetSession.h
@@ -51,9 +51,9 @@ namespace Battlenet
         Read = 0x4000
     };
 
-    class Session : public Socket<Session, ServerPacket*>
+    class Session : public Socket<Session>
     {
-        typedef Socket<Session, ServerPacket*> BattlenetSocket;
+        typedef Socket<Session> BattlenetSocket;
 
     public:
         explicit Session(tcp::socket&& socket);
@@ -79,11 +79,8 @@ namespace Battlenet
 
         void AsyncWrite(ServerPacket* packet);
 
-        bool IsDataReady() const override { return GetDataSize() > 0; }
-
     protected:
-        void ReadHeaderHandler() override { }
-        void ReadDataHandler() override;
+        void ReadHandler() override;
 
     private:
         void _SetVSFields(std::string const& rI);
diff --git a/src/server/authserver/Server/BattlenetSessionManager.cpp b/src/server/authserver/Server/BattlenetSessionManager.cpp
index e16057ecbf6..91ba2b65094 100644
--- a/src/server/authserver/Server/BattlenetSessionManager.cpp
+++ b/src/server/authserver/Server/BattlenetSessionManager.cpp
@@ -16,3 +16,22 @@
  */
 
 #include "BattlenetSessionManager.h"
+
+bool Battlenet::SessionManager::StartNetwork(boost::asio::io_service& service, std::string const& bindIp, uint16 port)
+{
+    if (!BaseSocketMgr::StartNetwork(service, bindIp, port))
+        return false;
+
+    _acceptor->AsyncAcceptManaged(&OnSocketAccept);
+    return true;
+}
+
+NetworkThread<Battlenet::Session>* Battlenet::SessionManager::CreateThreads() const
+{
+    return new NetworkThread<Session>[GetNetworkThreadCount()];
+}
+
+void Battlenet::SessionManager::OnSocketAccept(tcp::socket&& sock)
+{
+    sBattlenetSessionMgr.OnSocketOpen(std::forward<tcp::socket>(sock));
+}
diff --git a/src/server/authserver/Server/BattlenetSessionManager.h b/src/server/authserver/Server/BattlenetSessionManager.h
index 454bac96ddb..b5a54438ef1 100644
--- a/src/server/authserver/Server/BattlenetSessionManager.h
+++ b/src/server/authserver/Server/BattlenetSessionManager.h
@@ -19,6 +19,7 @@
 #define BattlenetSessionManager_h__
 
 #include "BattlenetSession.h"
+#include "SocketMgr.h"
 
 namespace Battlenet
 {
@@ -37,8 +38,10 @@ namespace Battlenet
 
 #pragma pack(pop)
 
-    class SessionManager
+    class SessionManager : SocketMgr<Session>
     {
+        typedef SocketMgr<Session> BaseSocketMgr;
+
     public:
         static SessionManager& Instance()
         {
@@ -46,12 +49,19 @@ namespace Battlenet
             return instance;
         }
 
+        bool StartNetwork(boost::asio::io_service& service, std::string const& bindIp, uint16 port) override;
+
         // noop for now, will be needed later to broadcast realmlist updates for example
         void AddSession(Session* /*session*/) { }
 
         void RemoveSession(Session* /*session*/) { }
 
+    protected:
+        NetworkThread<Session>* CreateThreads() const override;
+
     private:
+        static void OnSocketAccept(tcp::socket&& sock);
+
         std::map<SessionInfo, Session> _sessions;
     };
 }
diff --git a/src/server/game/Server/WorldSocket.cpp b/src/server/game/Server/WorldSocket.cpp
index cb42732b464..83bf2c2fff8 100644
--- a/src/server/game/Server/WorldSocket.cpp
+++ b/src/server/game/Server/WorldSocket.cpp
@@ -36,7 +36,7 @@ std::string const WorldSocket::ClientConnectionInitialize("WORLD OF WARCRAFT CON
 WorldSocket::WorldSocket(tcp::socket&& socket)
     : Socket(std::move(socket)), _authSeed(rand32()), _OverSpeedPings(0), _worldSession(nullptr), _initialized(false)
 {
-    _headerBuffer.Resize(sizeof(ClientPktHeader));
+    _headerBuffer.Resize(2);
 }
 
 void WorldSocket::Start()
@@ -46,7 +46,7 @@ void WorldSocket::Start()
     MessageBuffer initializer;
     ServerPktHeader header(ServerConnectionInitialize.size(), 0);
     initializer.Write(header.header, header.getHeaderLength() - 2);
-    initializer.Write((void*)ServerConnectionInitialize.c_str(), ServerConnectionInitialize.length());
+    initializer.Write(ServerConnectionInitialize.c_str(), ServerConnectionInitialize.length());
 
     std::unique_lock<std::mutex> dummy(_writeLock, std::defer_lock);
     QueuePacket(std::move(initializer), dummy);
@@ -123,15 +123,17 @@ void WorldSocket::ReadHandler()
 
 bool WorldSocket::ReadHeaderHandler()
 {
-    ASSERT(_headerBuffer.GetActiveSize() == sizeof(ClientPktHeader));
+    ASSERT(_headerBuffer.GetActiveSize() == (_initialized ? sizeof(ClientPktHeader) : 2));
 
-    _authCrypt.DecryptRecv(_headerBuffer.GetReadPointer(), sizeof(ClientPktHeader));
+    _authCrypt.DecryptRecv(_headerBuffer.GetReadPointer(), _headerBuffer.GetActiveSize());
 
     ClientPktHeader* header = reinterpret_cast<ClientPktHeader*>(_headerBuffer.GetReadPointer());
     EndianConvertReverse(header->size);
-    EndianConvert(header->cmd);
 
-    if (!header->IsValidSize() || !header->IsValidOpcode())
+    if (_initialized)
+        EndianConvert(header->cmd);
+
+    if (!header->IsValidSize() || (_initialized && !header->IsValidOpcode()))
     {
         if (_worldSession)
         {
@@ -147,7 +149,9 @@ bool WorldSocket::ReadHeaderHandler()
         return false;
     }
 
-    header->size -= sizeof(header->cmd);
+    if (_initialized)
+        header->size -= sizeof(header->cmd);
+
     _packetBuffer.Resize(header->size);
     return true;
 }
@@ -221,9 +225,7 @@ bool WorldSocket::ReadDataHandler()
     }
     else
     {
-        ClientPktHeader* header = reinterpret_cast<ClientPktHeader*>(_headerBuffer.GetReadPointer());
-
-        std::string initializer(reinterpret_cast<char const*>(_packetBuffer.GetReadPointer()), header->size);
+        std::string initializer(reinterpret_cast<char const*>(_packetBuffer.GetReadPointer()), std::min(_packetBuffer.GetActiveSize(), ClientConnectionInitialize.length()));
         if (initializer != ClientConnectionInitialize)
         {
             CloseSocket();
@@ -231,6 +233,8 @@ bool WorldSocket::ReadDataHandler()
         }
 
         _initialized = true;
+        _headerBuffer.Resize(sizeof(ClientPktHeader));
+        _packetBuffer.Reset();
         HandleSendAuthSession();
     }