From d908b4c2fcbb2f4c6f977c41cacf381631bff4ad Mon Sep 17 00:00:00 2001 From: blinkysc <37940565+blinkysc@users.noreply.github.com> Date: Thu, 15 Jan 2026 07:47:58 -0600 Subject: [PATCH] refactor(Core/Network): Port TrinityCore socket optimizations (#24384) Co-authored-by: blinkysc Co-authored-by: Shauren --- .../apps/authserver/Server/AuthSession.cpp | 12 +- .../apps/authserver/Server/AuthSession.h | 8 +- .../apps/authserver/Server/AuthSocketMgr.h | 4 +- .../apps/worldserver/RemoteAccess/RASession.h | 8 +- .../Scripting/ScriptDefines/ServerScript.cpp | 4 +- .../Scripting/ScriptDefines/ServerScript.h | 4 +- src/server/game/Scripting/ScriptMgr.h | 4 +- src/server/game/Server/WorldSocket.cpp | 12 +- src/server/game/Server/WorldSocket.h | 8 +- src/server/game/Server/WorldSocketMgr.cpp | 8 +- src/server/game/Server/WorldSocketMgr.h | 6 +- src/server/shared/Network/AsyncAcceptor.h | 17 ++- src/server/shared/Network/NetworkThread.h | 8 +- src/server/shared/Network/Socket.h | 69 ++++++--- src/server/shared/Network/SocketMgr.h | 4 +- tools/socket_stress_heavy.py | 141 ++++++++++++++++++ 16 files changed, 242 insertions(+), 75 deletions(-) create mode 100644 tools/socket_stress_heavy.py diff --git a/src/server/apps/authserver/Server/AuthSession.cpp b/src/server/apps/authserver/Server/AuthSession.cpp index 29c0c643c..e5315c9ec 100644 --- a/src/server/apps/authserver/Server/AuthSession.cpp +++ b/src/server/apps/authserver/Server/AuthSession.cpp @@ -162,7 +162,7 @@ void AccountInfo::LoadResult(Field* fields) Utf8ToUpperOnlyLatin(Login); } -AuthSession::AuthSession(tcp::socket&& socket) : +AuthSession::AuthSession(IoContextTcpSocket&& socket) : Socket(std::move(socket)), _status(STATUS_CHALLENGE), _build(0), _expversion(0) { } void AuthSession::Start() @@ -216,7 +216,7 @@ void AuthSession::CheckIpCallback(PreparedQueryResult result) AsyncRead(); } -void AuthSession::ReadHandler() +SocketReadCallbackResult AuthSession::ReadHandler() { MessageBuffer& packet = GetReadBuffer(); @@ -234,7 +234,7 @@ void AuthSession::ReadHandler() if (_status != itr->second.status) { CloseSocket(); - return; + return SocketReadCallbackResult::Stop; } uint16 size = uint16(itr->second.packetSize); @@ -248,7 +248,7 @@ void AuthSession::ReadHandler() if (size > MAX_ACCEPTED_CHALLENGE_SIZE) { CloseSocket(); - return; + return SocketReadCallbackResult::Stop; } } @@ -258,13 +258,13 @@ void AuthSession::ReadHandler() if (!(*this.*itr->second.handler)()) { CloseSocket(); - return; + return SocketReadCallbackResult::Stop; } packet.ReadCompleted(size); } - AsyncRead(); + return SocketReadCallbackResult::KeepReading; } void AuthSession::SendPacket(ByteBuffer& packet) diff --git a/src/server/apps/authserver/Server/AuthSession.h b/src/server/apps/authserver/Server/AuthSession.h index 1c5e28eee..0c61e5e2f 100644 --- a/src/server/apps/authserver/Server/AuthSession.h +++ b/src/server/apps/authserver/Server/AuthSession.h @@ -60,22 +60,22 @@ struct AccountInfo AccountTypes SecurityLevel = SEC_PLAYER; }; -class AuthSession : public Socket +class AuthSession final : public Socket { typedef Socket AuthSocket; public: static std::unordered_map InitHandlers(); - AuthSession(tcp::socket&& socket); + AuthSession(IoContextTcpSocket&& socket); void Start() override; - bool Update() override; + bool Update() final; void SendPacket(ByteBuffer& packet); protected: - void ReadHandler() override; + SocketReadCallbackResult ReadHandler() final; private: bool HandleLogonChallenge(); diff --git a/src/server/apps/authserver/Server/AuthSocketMgr.h b/src/server/apps/authserver/Server/AuthSocketMgr.h index 162a6363c..99cf080e2 100644 --- a/src/server/apps/authserver/Server/AuthSocketMgr.h +++ b/src/server/apps/authserver/Server/AuthSocketMgr.h @@ -54,9 +54,9 @@ protected: return threads; } - static void OnSocketAccept(tcp::socket&& sock, uint32 threadIndex) + static void OnSocketAccept(IoContextTcpSocket&& sock, uint32 threadIndex) { - Instance().OnSocketOpen(std::forward(sock), threadIndex); + Instance().OnSocketOpen(std::move(sock), threadIndex); } }; diff --git a/src/server/apps/worldserver/RemoteAccess/RASession.h b/src/server/apps/worldserver/RemoteAccess/RASession.h index 27e2ab4a6..1f06b32ae 100644 --- a/src/server/apps/worldserver/RemoteAccess/RASession.h +++ b/src/server/apps/worldserver/RemoteAccess/RASession.h @@ -18,18 +18,16 @@ #ifndef __RASESSION_H__ #define __RASESSION_H__ -#include +#include "Socket.h" #include #include -using boost::asio::ip::tcp; - const std::size_t bufferSize = 4096; class RASession : public std::enable_shared_from_this { public: - RASession(tcp::socket&& socket) : + RASession(IoContextTcpSocket&& socket) : _socket(std::move(socket)), _commandExecuting(nullptr) { } void Start(); @@ -47,7 +45,7 @@ private: static void CommandPrint(void* callbackArg, std::string_view text); static void CommandFinished(void* callbackArg, bool); - tcp::socket _socket; + IoContextTcpSocket _socket; boost::asio::streambuf _readBuffer; boost::asio::streambuf _writeBuffer; std::promise* _commandExecuting; diff --git a/src/server/game/Scripting/ScriptDefines/ServerScript.cpp b/src/server/game/Scripting/ScriptDefines/ServerScript.cpp index 9794127f0..74d49c555 100644 --- a/src/server/game/Scripting/ScriptDefines/ServerScript.cpp +++ b/src/server/game/Scripting/ScriptDefines/ServerScript.cpp @@ -29,14 +29,14 @@ void ScriptMgr::OnNetworkStop() CALL_ENABLED_HOOKS(ServerScript, SERVERHOOK_ON_NETWORK_STOP, script->OnNetworkStop()); } -void ScriptMgr::OnSocketOpen(std::shared_ptr socket) +void ScriptMgr::OnSocketOpen(std::shared_ptr const& socket) { ASSERT(socket); CALL_ENABLED_HOOKS(ServerScript, SERVERHOOK_ON_SOCKET_OPEN, script->OnSocketOpen(socket)); } -void ScriptMgr::OnSocketClose(std::shared_ptr socket) +void ScriptMgr::OnSocketClose(std::shared_ptr const& socket) { ASSERT(socket); diff --git a/src/server/game/Scripting/ScriptDefines/ServerScript.h b/src/server/game/Scripting/ScriptDefines/ServerScript.h index 5f3a7787b..5125f2317 100644 --- a/src/server/game/Scripting/ScriptDefines/ServerScript.h +++ b/src/server/game/Scripting/ScriptDefines/ServerScript.h @@ -46,11 +46,11 @@ public: virtual void OnNetworkStop() { } // Called when a remote socket establishes a connection to the server. Do not store the socket object. - virtual void OnSocketOpen(std::shared_ptr /*socket*/) { } + virtual void OnSocketOpen(std::shared_ptr const& /*socket*/) { } // Called when a socket is closed. Do not store the socket object, and do not rely on the connection // being open; it is not. - virtual void OnSocketClose(std::shared_ptr /*socket*/) { } + virtual void OnSocketClose(std::shared_ptr const& /*socket*/) { } /** * @brief This hook called when a packet is sent to a client. The packet object is a copy of the original packet, so reading and modifying it is safe. diff --git a/src/server/game/Scripting/ScriptMgr.h b/src/server/game/Scripting/ScriptMgr.h index 0dadd18f6..056a44a57 100644 --- a/src/server/game/Scripting/ScriptMgr.h +++ b/src/server/game/Scripting/ScriptMgr.h @@ -155,8 +155,8 @@ public: /* SpellScriptLoader */ public: /* ServerScript */ void OnNetworkStart(); void OnNetworkStop(); - void OnSocketOpen(std::shared_ptr socket); - void OnSocketClose(std::shared_ptr socket); + void OnSocketOpen(std::shared_ptr const& socket); + void OnSocketClose(std::shared_ptr const& socket); bool CanPacketReceive(WorldSession* session, WorldPacket const& packet); bool CanPacketSend(WorldSession* session, WorldPacket const& packet); diff --git a/src/server/game/Server/WorldSocket.cpp b/src/server/game/Server/WorldSocket.cpp index 452a16ad1..e8da0fb6e 100644 --- a/src/server/game/Server/WorldSocket.cpp +++ b/src/server/game/Server/WorldSocket.cpp @@ -116,7 +116,7 @@ void EncryptableAndCompressiblePacket::CompressIfNeeded() SetOpcode(SMSG_COMPRESSED_UPDATE_OBJECT); } -WorldSocket::WorldSocket(tcp::socket&& socket) +WorldSocket::WorldSocket(IoContextTcpSocket&& socket) : Socket(std::move(socket)), _OverSpeedPings(0), _worldSession(nullptr), _authed(false), _sendBufferSize(4096), _loggingPackets(false) { Acore::Crypto::GetRandomBytes(_authSeed); @@ -238,10 +238,10 @@ void WorldSocket::OnClose() } } -void WorldSocket::ReadHandler() +SocketReadCallbackResult WorldSocket::ReadHandler() { if (!IsOpen()) - return; + return SocketReadCallbackResult::Stop; MessageBuffer& packet = GetReadBuffer(); while (packet.GetActiveSize() > 0) @@ -264,7 +264,7 @@ void WorldSocket::ReadHandler() if (!ReadHeaderHandler()) { CloseSocket(); - return; + return SocketReadCallbackResult::Stop; } } @@ -295,11 +295,11 @@ void WorldSocket::ReadHandler() CloseSocket(); } - return; + return SocketReadCallbackResult::Stop; } } - AsyncRead(); + return SocketReadCallbackResult::KeepReading; } bool WorldSocket::ReadHeaderHandler() diff --git a/src/server/game/Server/WorldSocket.h b/src/server/game/Server/WorldSocket.h index 347e9b506..cc47d5e7c 100644 --- a/src/server/game/Server/WorldSocket.h +++ b/src/server/game/Server/WorldSocket.h @@ -67,19 +67,19 @@ struct ClientPktHeader struct AuthSession; -class AC_GAME_API WorldSocket : public Socket +class AC_GAME_API WorldSocket final : public Socket { typedef Socket BaseSocket; public: - WorldSocket(tcp::socket&& socket); + WorldSocket(IoContextTcpSocket&& socket); ~WorldSocket(); WorldSocket(WorldSocket const& right) = delete; WorldSocket& operator=(WorldSocket const& right) = delete; void Start() override; - bool Update() override; + bool Update() final; void SendPacket(WorldPacket const& packet); @@ -90,7 +90,7 @@ public: protected: void OnClose() override; - void ReadHandler() override; + SocketReadCallbackResult ReadHandler() final; bool ReadHeaderHandler(); enum class ReadDataHandlerResult diff --git a/src/server/game/Server/WorldSocketMgr.cpp b/src/server/game/Server/WorldSocketMgr.cpp index 07e9d4dce..3b5053e1d 100644 --- a/src/server/game/Server/WorldSocketMgr.cpp +++ b/src/server/game/Server/WorldSocketMgr.cpp @@ -25,13 +25,13 @@ class WorldSocketThread : public NetworkThread { public: - void SocketAdded(std::shared_ptr sock) override + void SocketAdded(std::shared_ptr const& sock) override { sock->SetSendBufferSize(sWorldSocketMgr.GetApplicationSendBufferSize()); sScriptMgr->OnSocketOpen(sock); } - void SocketRemoved(std::shared_ptr sock) override + void SocketRemoved(std::shared_ptr const& sock) override { sScriptMgr->OnSocketClose(sock); } @@ -81,7 +81,7 @@ void WorldSocketMgr::StopNetwork() sScriptMgr->OnNetworkStop(); } -void WorldSocketMgr::OnSocketOpen(tcp::socket&& sock, uint32 threadIndex) +void WorldSocketMgr::OnSocketOpen(IoContextTcpSocket&& sock, uint32 threadIndex) { // set some options here if (_socketSystemSendBufferSize >= 0) @@ -109,7 +109,7 @@ void WorldSocketMgr::OnSocketOpen(tcp::socket&& sock, uint32 threadIndex) } } - BaseSocketMgr::OnSocketOpen(std::forward(sock), threadIndex); + BaseSocketMgr::OnSocketOpen(std::move(sock), threadIndex); } NetworkThread* WorldSocketMgr::CreateThreads() const diff --git a/src/server/game/Server/WorldSocketMgr.h b/src/server/game/Server/WorldSocketMgr.h index a93a12fa6..efd2f91d7 100644 --- a/src/server/game/Server/WorldSocketMgr.h +++ b/src/server/game/Server/WorldSocketMgr.h @@ -41,7 +41,7 @@ public: /// Stops all network threads, It will wait for all running threads . void StopNetwork() override; - void OnSocketOpen(tcp::socket&& sock, uint32 threadIndex) override; + void OnSocketOpen(IoContextTcpSocket&& sock, uint32 threadIndex) override; std::size_t GetApplicationSendBufferSize() const { return _socketApplicationSendBufferSize; } @@ -50,9 +50,9 @@ protected: NetworkThread* CreateThreads() const override; - static void OnSocketAccept(tcp::socket&& sock, uint32 threadIndex) + static void OnSocketAccept(IoContextTcpSocket&& sock, uint32 threadIndex) { - Instance().OnSocketOpen(std::forward(sock), threadIndex); + Instance().OnSocketOpen(std::move(sock), threadIndex); } private: diff --git a/src/server/shared/Network/AsyncAcceptor.h b/src/server/shared/Network/AsyncAcceptor.h index 71c58ed93..1de085aaa 100644 --- a/src/server/shared/Network/AsyncAcceptor.h +++ b/src/server/shared/Network/AsyncAcceptor.h @@ -20,6 +20,7 @@ #include "IpAddress.h" #include "Log.h" +#include "Socket.h" #include "Systemd.h" #include #include @@ -32,7 +33,7 @@ constexpr auto ACORE_MAX_LISTEN_CONNECTIONS = boost::asio::socket_base::max_list class AsyncAcceptor { public: - typedef void(*AcceptCallback)(tcp::socket&& newSocket, uint32 threadIndex); + typedef void(*AcceptCallback)(IoContextTcpSocket&& newSocket, uint32 threadIndex); AsyncAcceptor(Acore::Asio::IoContext& ioContext, std::string const& bindIp, uint16 port, bool supportSocketActivation = false) : _acceptor(ioContext), _endpoint(Acore::Net::make_address(bindIp), port), @@ -56,7 +57,7 @@ public: template void AsyncAcceptWithCallback() { - tcp::socket* socket; + IoContextTcpSocket* socket; uint32 threadIndex; std::tie(socket, threadIndex) = _socketFactory(); _acceptor.async_accept(*socket, [this, socket, threadIndex](boost::system::error_code error) @@ -129,16 +130,16 @@ public: _acceptor.close(err); } - void SetSocketFactory(std::function()> func) { _socketFactory = func; } + void SetSocketFactory(std::function()> func) { _socketFactory = std::move(func); } private: - std::pair DefaultSocketFactory() { return std::make_pair(&_socket, 0); } + std::pair DefaultSocketFactory() { return std::make_pair(&_socket, 0); } - tcp::acceptor _acceptor; - tcp::endpoint _endpoint; - tcp::socket _socket; + boost::asio::basic_socket_acceptor _acceptor; + boost::asio::ip::tcp::endpoint _endpoint; + IoContextTcpSocket _socket; std::atomic _closed; - std::function()> _socketFactory; + std::function()> _socketFactory; bool _supportSocketActivation; }; diff --git a/src/server/shared/Network/NetworkThread.h b/src/server/shared/Network/NetworkThread.h index 674ad60ec..df39cd7f9 100644 --- a/src/server/shared/Network/NetworkThread.h +++ b/src/server/shared/Network/NetworkThread.h @@ -91,13 +91,13 @@ public: SocketAdded(sock); } - tcp::socket* GetSocketForAccept() { return &_acceptSocket; } + IoContextTcpSocket* GetSocketForAccept() { return &_acceptSocket; } void EnableProxyProtocol() { _proxyHeaderReadingEnabled = true; } protected: - virtual void SocketAdded(std::shared_ptr /*sock*/) { } - virtual void SocketRemoved(std::shared_ptr /*sock*/) { } + virtual void SocketAdded(std::shared_ptr const& /*sock*/) { } + virtual void SocketRemoved(std::shared_ptr const& /*sock*/) { } void AddNewSockets() { @@ -229,7 +229,7 @@ private: SocketContainer _newSockets; Acore::Asio::IoContext _ioContext; - tcp::socket _acceptSocket; + IoContextTcpSocket _acceptSocket; boost::asio::steady_timer _updateTimer; bool _proxyHeaderReadingEnabled; diff --git a/src/server/shared/Network/Socket.h b/src/server/shared/Network/Socket.h index c4ff85c30..4d551b19c 100644 --- a/src/server/shared/Network/Socket.h +++ b/src/server/shared/Network/Socket.h @@ -21,9 +21,8 @@ #include "Log.h" #include "MessageBuffer.h" #include -#include +#include #include -#include #include #include #include @@ -35,6 +34,23 @@ using boost::asio::ip::tcp; #define AC_SOCKET_USE_IOCP #endif +// Specialize boost socket for io_context executor instead of type-erased any_io_executor +// This avoids the type-erasure overhead of any_io_executor +using IoContextTcpSocket = boost::asio::basic_stream_socket; + +enum class SocketReadCallbackResult +{ + KeepReading, + Stop +}; + +enum class SocketState : uint8 +{ + Open = 0, + Closing = 1, + Closed = 2 +}; + enum ProxyHeaderReadingState { PROXY_HEADER_READING_STATE_NOT_STARTED, PROXY_HEADER_READING_STATE_STARTED, @@ -51,8 +67,8 @@ template class Socket : public std::enable_shared_from_this { public: - explicit Socket(tcp::socket&& socket) : _socket(std::move(socket)), _remoteAddress(_socket.remote_endpoint().address()), - _remotePort(_socket.remote_endpoint().port()), _readBuffer(), _closed(false), _closing(false), _isWritingAsync(false), + explicit Socket(IoContextTcpSocket&& socket) : _socket(std::move(socket)), _remoteAddress(_socket.remote_endpoint().address()), + _remotePort(_socket.remote_endpoint().port()), _readBuffer(), _state(SocketState::Open), _isWritingAsync(false), _proxyHeaderReadingState(PROXY_HEADER_READING_STATE_NOT_STARTED) { _readBuffer.Resize(READ_BLOCK_SIZE); @@ -60,7 +76,7 @@ public: virtual ~Socket() { - _closed = true; + _state = SocketState::Closed; boost::system::error_code error; _socket.close(error); } @@ -69,13 +85,14 @@ public: virtual bool Update() { - if (_closed) + SocketState state = _state.load(); + if (state == SocketState::Closed) { return false; } #ifndef AC_SOCKET_USE_IOCP - if (_isWritingAsync || (_writeQueue.empty() && !_closing)) + if (_isWritingAsync || (_writeQueue.empty() && state != SocketState::Closing)) { return true; } @@ -150,12 +167,18 @@ public: [[nodiscard]] ProxyHeaderReadingState GetProxyHeaderReadingState() const { return _proxyHeaderReadingState; } - [[nodiscard]] bool IsOpen() const { return !_closed && !_closing; } + [[nodiscard]] bool IsOpen() const { return _state.load() == SocketState::Open; } void CloseSocket() { - if (_closed.exchange(true)) - return; + SocketState expected = SocketState::Open; + if (!_state.compare_exchange_strong(expected, SocketState::Closed)) + { + // If it was Closing, try to transition to Closed + expected = SocketState::Closing; + if (!_state.compare_exchange_strong(expected, SocketState::Closed)) + return; // Already closed + } boost::system::error_code shutdownError; _socket.shutdown(boost::asio::socket_base::shutdown_send, shutdownError); @@ -168,13 +191,17 @@ public: } /// Marks the socket for closing after write buffer becomes empty - void DelayedCloseSocket() { _closing = true; } + void DelayedCloseSocket() + { + SocketState expected = SocketState::Open; + _state.compare_exchange_strong(expected, SocketState::Closing); + } MessageBuffer& GetReadBuffer() { return _readBuffer; } protected: virtual void OnClose() { } - virtual void ReadHandler() = 0; + virtual SocketReadCallbackResult ReadHandler() = 0; bool AsyncProcessQueue() { @@ -188,7 +215,7 @@ protected: _socket.async_write_some(boost::asio::buffer(buffer.GetReadPointer(), buffer.GetActiveSize()), std::bind(&Socket::WriteHandler, this->shared_from_this(), std::placeholders::_1, std::placeholders::_2)); #else - _socket.async_wait(tcp::socket::wait_write, [self = this->shared_from_this()](boost::system::error_code error) + _socket.async_wait(boost::asio::socket_base::wait_write, [self = this->shared_from_this()](boost::system::error_code error) { self->WriteHandlerWrapper(error, 0); }); @@ -216,7 +243,8 @@ private: } _readBuffer.WriteCompleted(transferredBytes); - ReadHandler(); + if (ReadHandler() == SocketReadCallbackResult::KeepReading) + AsyncRead(); } // ProxyReadHeaderHandler reads Proxy Protocol v2 header (v1 is not supported). @@ -344,7 +372,7 @@ private: if (!_writeQueue.empty()) AsyncProcessQueue(); - else if (_closing) + else if (_state.load() == SocketState::Closing) CloseSocket(); } else @@ -380,7 +408,7 @@ private: _writeQueue.pop(); - if (_closing && _writeQueue.empty()) + if (_state.load() == SocketState::Closing && _writeQueue.empty()) { CloseSocket(); } @@ -391,7 +419,7 @@ private: { _writeQueue.pop(); - if (_closing && _writeQueue.empty()) + if (_state.load() == SocketState::Closing && _writeQueue.empty()) { CloseSocket(); } @@ -406,7 +434,7 @@ private: _writeQueue.pop(); - if (_closing && _writeQueue.empty()) + if (_state.load() == SocketState::Closing && _writeQueue.empty()) { CloseSocket(); } @@ -415,7 +443,7 @@ private: } #endif - tcp::socket _socket; + IoContextTcpSocket _socket; boost::asio::ip::address _remoteAddress; uint16 _remotePort; @@ -423,8 +451,7 @@ private: MessageBuffer _readBuffer; std::queue _writeQueue; - std::atomic _closed; - std::atomic _closing; + std::atomic _state; bool _isWritingAsync; diff --git a/src/server/shared/Network/SocketMgr.h b/src/server/shared/Network/SocketMgr.h index 085e4e238..56d0ed9a7 100644 --- a/src/server/shared/Network/SocketMgr.h +++ b/src/server/shared/Network/SocketMgr.h @@ -91,7 +91,7 @@ public: _threads[i].Wait(); } - virtual void OnSocketOpen(tcp::socket&& sock, uint32 threadIndex) + virtual void OnSocketOpen(IoContextTcpSocket&& sock, uint32 threadIndex) { try { @@ -117,7 +117,7 @@ public: return min; } - std::pair GetSocketForAccept() + std::pair GetSocketForAccept() { uint32 threadIndex = SelectThreadWithMinConnections(); return { _threads[threadIndex].GetSocketForAccept(), threadIndex }; diff --git a/tools/socket_stress_heavy.py b/tools/socket_stress_heavy.py new file mode 100644 index 000000000..8352a8bd8 --- /dev/null +++ b/tools/socket_stress_heavy.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +""" +Socket Stress Test for AzerothCore +Tests authserver and worldserver connection handling under heavy load. + +Usage: + python3 socket_stress_heavy.py [duration_seconds] [auth_threads] [world_threads] + +Defaults: + duration: 300 seconds (5 minutes) + auth_threads: 100 + world_threads: 150 +""" + +import socket +import time +import threading +import sys + +AUTH_PORT = 3724 +WORLD_PORT = 8085 +HOST = '127.0.0.1' + +stats = {'auth_ok': 0, 'auth_fail': 0, 'world_ok': 0, 'world_fail': 0} +running = True + + +def stress_auth(): + """Flood authserver with login challenge packets.""" + global stats + while running: + try: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.settimeout(1) + s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + s.connect((HOST, AUTH_PORT)) + # AUTH_LOGON_CHALLENGE packet + packet = bytes([ + 0x00, # cmd: AUTH_LOGON_CHALLENGE + 0x00, # error + 0x24, 0x00, # size (36) + 0x57, 0x6F, 0x57, 0x00, # 'WoW\0' + 0x03, 0x03, 0x05, # version 3.3.5 + 0x30, 0x30, # build 12340 + 0x78, 0x38, 0x36, 0x00, # 'x86\0' + 0x6E, 0x69, 0x57, 0x00, # 'niW\0' (Win reversed) + 0x53, 0x55, 0x6E, 0x65, # 'SUne' (enUS reversed) + 0x3C, 0x00, 0x00, 0x00, # timezone + 0x7F, 0x00, 0x00, 0x01, # IP 127.0.0.1 + 0x04, # account name length + 0x54, 0x45, 0x53, 0x54 # 'TEST' + ]) + s.sendall(packet) + s.close() + stats['auth_ok'] += 1 + except Exception: + stats['auth_fail'] += 1 + + +def stress_world(): + """Flood worldserver with connection attempts.""" + global stats + while running: + try: + s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + s.settimeout(1) + s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + s.connect((HOST, WORLD_PORT)) + # Wait for SMSG_AUTH_CHALLENGE + s.recv(64) + s.close() + stats['world_ok'] += 1 + except Exception: + stats['world_fail'] += 1 + + +def main(): + global running + + # Parse arguments + duration = int(sys.argv[1]) if len(sys.argv) > 1 else 300 + auth_threads = int(sys.argv[2]) if len(sys.argv) > 2 else 100 + world_threads = int(sys.argv[3]) if len(sys.argv) > 3 else 150 + + print("=" * 60) + print("SOCKET STRESS TEST") + print("=" * 60) + print(f"Duration: {duration} seconds") + print(f"Auth threads: {auth_threads} -> {HOST}:{AUTH_PORT}") + print(f"World threads: {world_threads} -> {HOST}:{WORLD_PORT}") + print("-" * 60) + + threads = [] + + for _ in range(auth_threads): + t = threading.Thread(target=stress_auth, daemon=True) + t.start() + threads.append(t) + + for _ in range(world_threads): + t = threading.Thread(target=stress_world, daemon=True) + t.start() + threads.append(t) + + print(f"Started {len(threads)} threads") + print("-" * 60) + + start = time.time() + try: + while time.time() - start < duration: + elapsed = int(time.time() - start) + total = stats['auth_ok'] + stats['world_ok'] + rate = total / max(elapsed, 1) + print(f"\r[{elapsed:3d}s] Auth: {stats['auth_ok']:7d} ok {stats['auth_fail']:5d} fail | " + f"World: {stats['world_ok']:7d} ok {stats['world_fail']:5d} fail | " + f"Rate: {rate:,.0f}/s ", end='', flush=True) + time.sleep(1) + except KeyboardInterrupt: + print("\n\nInterrupted by user") + + running = False + time.sleep(0.5) + + total_ok = stats['auth_ok'] + stats['world_ok'] + total_fail = stats['auth_fail'] + stats['world_fail'] + elapsed = time.time() - start + + print("\n" + "=" * 60) + print("RESULTS:") + print(f" Duration: {elapsed:.1f} seconds") + print(f" Auth: {stats['auth_ok']:,} ok, {stats['auth_fail']:,} failed") + print(f" World: {stats['world_ok']:,} ok, {stats['world_fail']:,} failed") + print(f" Total: {total_ok:,} ok, {total_fail:,} failed") + print(f" Rate: {total_ok / elapsed:,.0f} connections/sec average") + if total_fail > 0: + print(f" Failure: {total_fail / (total_ok + total_fail) * 100:.2f}%") + print("=" * 60) + + +if __name__ == '__main__': + main()