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

1 files changed, 109 insertions, 0 deletions

diff --git a/src/server/shared/Networking/Socket.h b/src/server/shared/Networking/Socket.h
new file mode 100644
index 00000000000..daefa0d4ad5
--- /dev/null
+++ b/src/server/shared/Networking/Socket.h
@@ -0,0 +1,109 @@
+/*
+ * Copyright (C) 2008-2014 TrinityCore <http://www.trinitycore.org/>
+ * Copyright (C) 2005-2009 MaNGOS <http://getmangos.com/>
+ *
+ * This program is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License as published by the
+ * Free Software Foundation; either version 2 of the License, or (at your
+ * option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
+ * more details.
+ *
+ * You should have received a copy of the GNU General Public License along
+ * with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#ifndef __SOCKET_H__
+#define __SOCKET_H__
+
+#include "Define.h"
+#include "Log.h"
+#include <vector>
+#include <mutex>
+#include <queue>
+#include <memory>
+#include <boost/asio/ip/tcp.hpp>
+#include <boost/asio/placeholders.hpp>
+#include <boost/bind.hpp>
+
+using boost::asio::ip::tcp;
+
+template<class T>
+class Socket : public std::enable_shared_from_this<T>
+{
+public:
+    Socket(tcp::socket&& socket, std::size_t headerSize) : _socket(std::move(socket)), _headerSize(headerSize) { }
+
+    virtual void Start() = 0;
+
+    boost::asio::ip::address GetRemoteIpAddress() const { return _socket.remote_endpoint().address(); };
+    uint16 GetRemotePort() const { return _socket.remote_endpoint().port(); }
+
+    void AsyncReadHeader()
+    {
+        _socket.async_read_some(boost::asio::buffer(_readBuffer, _headerSize), boost::bind(&Socket::ReadHeaderHandlerInternal, shared_from_this(), boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
+    }
+
+    void AsyncReadData(std::size_t size, std::size_t bufferOffset)
+    {
+        _socket.async_read_some(boost::asio::buffer(&_readBuffer[bufferOffset], size), boost::bind(&Socket::ReadDataHandlerInternal, shared_from_this(), boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
+    }
+
+    void ReadData(std::size_t size, std::size_t bufferOffset)
+    {
+        _socket.read_some(boost::asio::buffer(&_readBuffer[bufferOffset], size));
+    }
+
+    void AsyncWrite(std::vector<uint8> const& data)
+    {
+        boost::asio::async_write(_socket, boost::asio::buffer(data), boost::bind(&Socket::WriteHandler, shared_from_this(), boost::asio::placeholders::error, boost::asio::placeholders::bytes_transferred));
+    }
+
+    bool IsOpen() const { return _socket.is_open(); }
+    void CloseSocket()
+    {
+        boost::system::error_code socketError;
+        _socket.close(socketError);
+        if (socketError)
+            TC_LOG_DEBUG("network", "Socket::CloseSocket: %s errored when closing socket: %i (%s)", GetRemoteIpAddress().to_string().c_str(), socketError.value(), socketError.message().c_str());
+    }
+
+    uint8* GetReadBuffer() { return _readBuffer; }
+
+protected:
+    virtual void ReadHeaderHandler(boost::system::error_code error, size_t transferedBytes) = 0;
+    virtual void ReadDataHandler(boost::system::error_code error, size_t transferedBytes) = 0;
+
+    std::mutex _writeLock;
+    std::queue<std::vector<uint8> > _writeQueue;
+
+private:
+    void ReadHeaderHandlerInternal(boost::system::error_code error, size_t transferedBytes) { ReadHeaderHandler(error, transferedBytes); }
+    void ReadDataHandlerInternal(boost::system::error_code error, size_t transferedBytes) { ReadDataHandler(error, transferedBytes); }
+
+    void WriteHandler(boost::system::error_code error, size_t /*transferedBytes*/)
+    {
+        if (!error)
+        {
+            std::lock_guard<std::mutex> deleteGuard(_writeLock);
+
+            _writeQueue.pop();
+
+            if (!_writeQueue.empty())
+                AsyncWrite(_writeQueue.front());
+        }
+        else
+            CloseSocket();
+    }
+
+    tcp::socket _socket;
+
+    uint8 _readBuffer[4096];
+
+    std::size_t _headerSize;
+};
+
+#endif // __SOCKET_H__