Socket client

RconClient.cs

using System;
using System.Collections.Generic;
using System.Net;
using System.Net.Sockets;
using FragLabs.Adjutant.Network.Exceptions;
using FragLabs.Adjutant.Network.Interfaces;

namespace FragLabs.Adjutant.Network.Sockets
{
    /// <summary>
    /// Socket client that communicates in a binary rcon protocol.
    /// </summary>
    public class RconClient : IClient
    {
        /// <summary>
        /// Default size of the read buffer.
        /// </summary>
        private const int DefaultReadBufferSize = 1024;

        /// <summary>
        /// Socket connected to client.
        /// </summary>
        private Socket _socket;

        /// <summary>
        /// Socket args used to receive data.
        /// </summary>
        private SocketAsyncEventArgs _readArgs;

        /// <summary>
        /// Backing field for ReadBufferSize.
        /// </summary>
        private int _readBufferSize;

        /// <summary>
        /// Buffer containing bytes read from the socket waiting for be processed.
        /// </summary>
        private readonly List<byte> _messageBuffer = new List<byte>();

        /// <summary>
        /// Message delegates registered per message id.
        /// </summary>
        private readonly Dictionary<uint,Action<IClient,Message>> _messageDelegates = new Dictionary<uint, Action<IClient, Message>>();

        /// <summary>
        /// Sequential message id.
        /// </summary>
        private uint _messageId = 1;

        /// <summary>
        /// Gets the remote endpoint the client is connected to.
        /// </summary>
        public IPEndPoint RemoteEndPoint { get; private set; }

        public IPAddress RemoteIP
        {
            get { return RemoteEndPoint.Address; }
        }

        public IPAddress LocalIP { get
        {
            if (_socket == null)
                return null;
            return ((IPEndPoint) _socket.LocalEndPoint).Address;
        }}

        /// <summary>
        /// Gets or sets the read buffer size for the socket.
        /// Must be set before connecting.
        /// </summary>
        public int ReadBufferSize
        {
            get { return _readBufferSize; }
            set
            {
                if (IsConnected) throw new AlreadyConnectedException();
                _readBufferSize = value;
            }
        }

        /// <summary>
        /// Creates a client from an already connected socket.
        /// </summary>
        /// <param name="clientSocket"></param>
        internal RconClient(Socket clientSocket)
        {
            ReadBufferSize = DefaultReadBufferSize;
            _socket = clientSocket;
            RemoteEndPoint = (IPEndPoint)clientSocket.RemoteEndPoint;
            SetupSocketArgs();
            Read();
        }

        /// <summary>
        /// Creates a client ready to connect to a server.
        /// </summary>
        /// <param name="serverEndPoint">Server EndPoint.</param>
        public RconClient(IPEndPoint serverEndPoint)
        {
            RemoteEndPoint = serverEndPoint;
            ReadBufferSize = DefaultReadBufferSize;
        }

        /// <summary>
        /// Disconnect the client.
        /// </summary>
        public void Disconnect()
        {
            var raiseDisconnectEvent = !(_socket == null);

            if (_socket != null)
            {
                _socket.Close();
                _socket.Dispose();
                _socket = null;
            }

            if (_readArgs != null)
            {
                _readArgs.Dispose();
                _readArgs = null;
            }

            if (raiseDisconnectEvent)
                OnDisconnected(new ClientEventArgs(this,true));
        }

        /// <summary>
        /// Gets if the client is currently connected or not.
        /// </summary>
        public bool IsConnected
        {
            get
            {
                return _socket != null && _socket.Connected;
            }
        }

        /// <summary>
        /// Event raised when a connection attempt has completed.
        /// </summary>
        public event EventHandler<ClientEventArgs> ConnectComplete;

        protected virtual void OnConnectComplete(ClientEventArgs e)
        {
            var handler = ConnectComplete;
            if (handler != null) handler(this, e);
        }

        /// <summary>
        /// Event raised when the client is disconnected.
        /// </summary>
        public event EventHandler<ClientEventArgs> Disconnected;

        protected virtual void OnDisconnected(ClientEventArgs e)
        {
            var handler = Disconnected;
            if (handler != null) handler(this, e);
        }

        /// <summary>
        /// Attempt to connect to the server.
        /// </summary>
        public void Connect()
        {
            Disconnect();
            _socket = new Socket(RemoteEndPoint.AddressFamily, SocketType.Stream, ProtocolType.Tcp);
            var saea = new SocketAsyncEventArgs()
                {
                    RemoteEndPoint = RemoteEndPoint
                };
            saea.Completed += ConnectCompleted;
            if (!_socket.ConnectAsync(saea))
                ConnectCompleted(_socket, saea);
        }

        private void ConnectCompleted(object sender, SocketAsyncEventArgs socketAsyncEventArgs)
        {
            if (socketAsyncEventArgs.SocketError == SocketError.Success)
            {
                SetupSocketArgs();
                Read();
                OnConnectComplete(new ClientEventArgs(this, true));
            }
            else
            {
                OnConnectComplete(new ClientEventArgs(this, false));
                Disconnect();
            }
            socketAsyncEventArgs.Dispose();
        }

        /// <summary>
        /// Prepares the socket args for use.
        /// </summary>
        private void SetupSocketArgs()
        {
            _readArgs = new SocketAsyncEventArgs();
            _readArgs.SetBuffer(new byte[ReadBufferSize], 0, ReadBufferSize);
            _readArgs.Completed += ReadCompleted;
        }

        /// <summary>
        /// Read data from the socket asynchronously.
        /// </summary>
        private void Read()
        {
            if (_readArgs == null || _socket == null) return;
            _readArgs.SetBuffer(0,ReadBufferSize);
            if (!_socket.ReceiveAsync(_readArgs))
                ReadCompleted(_socket, _readArgs);
        }

        private void ReadCompleted(object sender, SocketAsyncEventArgs socketAsyncEventArgs)
        {
            if (socketAsyncEventArgs.SocketError != SocketError.Success || socketAsyncEventArgs.BytesTransferred < 1)
            {
                Disconnect();
                return;
            }
            for (var i = 0; i < socketAsyncEventArgs.BytesTransferred; i++)
                _messageBuffer.Add(socketAsyncEventArgs.Buffer[i]);
            var messages = new List<Message>();
            while (BufferContainsMessage())
            {
                messages.Add(ReadMessage());
            }
            Read();
            foreach (var message in messages)
            {
                DispatchMessage(message);
            }
        }

        /// <summary>
        /// Determins if the buffer contains a full message or not.
        /// </summary>
        /// <returns></returns>
        private bool BufferContainsMessage()
        {
            if (_messageBuffer.Count < 4)
                return false;
            var header = new[]
                {
                    _messageBuffer[0], _messageBuffer[1],
                    _messageBuffer[2], _messageBuffer[3]
                };
            var len = BitConverter.ToUInt32(header, 0);
            return _messageBuffer.Count >= len + 4;
        }

        /// <summary>
        /// Reads a message from the buffer and removes the message from the buffer.
        /// </summary>
        /// <returns></returns>
        private Message ReadMessage()
        {
            if (_messageBuffer.Count < 4)
                return null;
            var header = new[]
                {
                    _messageBuffer[0], _messageBuffer[1],
                    _messageBuffer[2], _messageBuffer[3]
                };
            var len = BitConverter.ToUInt32(header, 0);
            if (len + 4 > _messageBuffer.Count)
                return null;
            var rawData = new byte[len];
            for (var i = 0; i < len; i++)
                rawData[i] = _messageBuffer[i + 4];
            _messageBuffer.RemoveRange(0, (int)(len + 4));
            return Message.Decode(rawData);
        }

        /// <summary>
        /// Dispatches a message to the correct delegate and events.
        /// </summary>
        /// <param name="message"></param>
        private void DispatchMessage(Message message)
        {
            if (message.IsResponse && _messageDelegates.ContainsKey(message.MessageId))
            {
                var callable = _messageDelegates[message.MessageId];
                _messageDelegates.Remove(message.MessageId);
                callable(this, message);
            }
            else
            {
                OnMessageReceived(new ClientEventArgs(this, true, message));
            }
            OnMessageReceivedRaw(new ClientEventArgs(this, true, message));
        }

        /// <summary>
        /// Gets the next message id in the sequence to be used.
        /// </summary>
        /// <returns></returns>
        private uint NextMessageId()
        {
            var ret = _messageId;
            if (_messageId == uint.MaxValue)
                _messageId = 1;
            else _messageId++;
            return ret;
        }

        /// <summary>
        /// Event raised when a message with no registered delegate is received.
        /// </summary>
        public event EventHandler<ClientEventArgs> MessageReceived;

        protected virtual void OnMessageReceived(ClientEventArgs e)
        {
            var handler = MessageReceived;
            if (handler != null) handler(this, e);
        }

        /// <summary>
        /// Event raised when a message is received even if a delegate is registered.
        /// </summary>
        public event EventHandler<ClientEventArgs> MessageReceivedRaw;

        protected virtual void OnMessageReceivedRaw(ClientEventArgs e)
        {
            var handler = MessageReceivedRaw;
            if (handler != null) handler(this, e);
        }


        /// <summary>
        /// Send a message.
        /// </summary>
        /// <param name="message"></param>
        public void Send(Message message)
        {
            if (message == null) throw new ArgumentNullException("message");
            if (!IsConnected)
                throw new Exception("Socket not connected");
            var rawMessage = message.Encode();
            var fullMessage = new byte[rawMessage.Length + 4];
            var len = BitConverter.GetBytes((uint)rawMessage.Length);
            Buffer.BlockCopy(len, 0, fullMessage, 0, 4);
            Buffer.BlockCopy(rawMessage, 0, fullMessage, 4, rawMessage.Length);

            var saea = new SocketAsyncEventArgs();
            saea.SetBuffer(fullMessage, 0, fullMessage.Length);
            saea.Completed += SendComplete;
            if (!_socket.SendAsync(saea))
                SendComplete(_socket, saea);
        }

        private void SendComplete(object sender, SocketAsyncEventArgs socketAsyncEventArgs)
        {
            if (socketAsyncEventArgs.SocketError != SocketError.Success)
                Disconnect();
            socketAsyncEventArgs.Dispose();
        }

        /// <summary>
        /// Sends a message and registers a callback to be executed when a response is received.
        /// </summary>
        /// <param name="message"></param>
        /// <param name="responseCallback"></param>
        public void Send(Message message, Action<IClient, Message> responseCallback)
        {
            if (message == null) throw new ArgumentNullException("message");
            if (message.MessageId == 0)
                message.MessageId = NextMessageId();
            _messageDelegates[message.MessageId] = responseCallback;
            Send(message);
        }

        public void Dispose()
        {
            Disconnect();
        }
    }
}