PackageSocket.cs

using UnityEngine;
using System;
using System.Collections;
using System.Net;
using System.Net.Sockets;
using System.Collections.Generic;

public enum PackageSocketError
{
    None,
    SocketShutdown,
    RecviveBufferNotEnough,
    RecviveTimeout,
}

public class PackageSocket
{
    enum State
    {
        Closed,
        Connecting,
        Connected,
        UserClosed,
        Error,
    }

    public delegate void ConnectCallback(bool connected);

    public delegate void RecviveCallback(byte[] data, int start, int length);

    public delegate void DisconnectCallback(SocketError socketError, PackageSocketError packageSocketError);

    public ConnectCallback OnConnect;
    public RecviveCallback OnRecvive;
    public DisconnectCallback OnDisconnect;
    public int ConnectTimeoutSetting = 3000;
    public int ReceiveTimeoutSetting = 10000;
    public int AutoSendPingSetting = 10000 / 3;

    private State CurState;
    private DateTime CheckTimeout;

    private Socket CurSocket;
    private SocketError LastSystemSocketError;
    private PackageSocketError LastPackageSocketError;

    private Queue<byte[]> SendQueue;
    private int SendQueueStartIndex;

    private byte[] RecvBuffer;
    private int RecvBufferBeginIndex;
    private int RecvBufferEndIndex;

    private byte[] PingBuffer;
    private DateTime NextSendPingTime;

    const int MaxSizePerSend = 1024 * 4;

    public PackageSocket()
    {
        SendQueue = new Queue<byte[]>();
        RecvBuffer = new byte[(int) ushort.MaxValue + 2];
        Reset();
    }

    public void Connect(string ip, int port)
    {
        Reset();

        CurSocket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        CurSocket.Blocking = false;

        try
        {
            var ipAddress = IPAddress.Parse(ip);
            CurSocket.Connect(ipAddress, port);
        }
        catch (SocketException e)
        {
            if (e.SocketErrorCode != SocketError.WouldBlock && e.SocketErrorCode != SocketError.InProgress)
                throw e;
        }

        CurState = State.Connecting;
        CheckTimeout = DateTime.Now.AddMilliseconds(ConnectTimeoutSetting);
    }

    public void Close()
    {
        if (CurState != State.Closed)
        {
            CurState = State.UserClosed;
        }
    }

    public void Update()
    {
        switch (CurState)
        {
            case State.Closed:
                return;

            case State.Connecting:
                ProcessConnect();
                break;

            case State.Connected:
                ProcessSend();
                ProcessReceive();
                break;

            case State.Error:
                Reset();
                OnDisconnect(LastSystemSocketError, LastPackageSocketError);
                break;

            case State.UserClosed:
                Reset();
                break;
        }
    }

    public void Send(Byte[] buffer, int start, int length)
    {
        var headerLen = GetHeaderLength();
        var data = new byte[headerLen + length];

        EncodeHeader(length, data, 0);
        Array.Copy(buffer, start, data, headerLen, length);
        SendQueue.Enqueue(data);

        if (CurState == State.Connected)
            ProcessSend();
    }

    public void SendPing()
    {
        if (PingBuffer == null)
        {
            var headerLen = GetHeaderLength();
            PingBuffer = new byte[headerLen];
            EncodeHeader(0, PingBuffer, 0);
        }

        SendQueue.Enqueue(PingBuffer);
        ProcessSend();
    }

    void ProcessConnect()
    {
        var now = DateTime.Now;

        // http://www.dpull.com/blog/csharp_socket/
        CurSocket.Poll(0, SelectMode.SelectWrite);

        if (!CurSocket.Connected && CheckTimeout < now)
        {
            Reset();
            OnConnect(false);
        }
        else if (CurSocket.Connected)
        {
            OnConnect(true);
            CurState = State.Connected;
            CheckTimeout = now.AddMilliseconds(ReceiveTimeoutSetting);
        }
    }

    void ProcessSend()
    {
        if (!CurSocket.Poll(0, SelectMode.SelectWrite))
            return;

        while (CurState == State.Connected && SendQueue.Count > 0)
        {
            var buffer = SendQueue.Peek();
            int leftLength = buffer.Length - SendQueueStartIndex;

            while (leftLength > 0)
            {
                var trySend = Math.Min(MaxSizePerSend, leftLength);
                SocketError error;

                var send = CurSocket.Send(buffer, SendQueueStartIndex, trySend, SocketFlags.None, out error);
                if (error != SocketError.Success)
                {
                    if (error != SocketError.WouldBlock && error != SocketError.Interrupted &&
                        error != SocketError.TryAgain)
                    {
                        SetError(error, PackageSocketError.None);
                        return;
                    }

                    break;
                }

                SendQueueStartIndex += send;
                leftLength -= send;
            }

            if (leftLength == 0)
            {
                SendQueue.Dequeue();
                SendQueueStartIndex = 0;
            }
        }
    }

    void ProcessReceive()
    {
        var now = DateTime.Now;
        if (!CurSocket.Poll(0, SelectMode.SelectRead))
        {
            if (NextSendPingTime < now)
            {
                SendPing();
                NextSendPingTime = now.AddMilliseconds(AutoSendPingSetting);
            }

            if (CheckTimeout < now)
                SetError(SocketError.Success, PackageSocketError.RecviveTimeout);
            return;
        }

        while (CurState == State.Connected)
        {
            var leftBufferSize = RecvBuffer.Length - RecvBufferEndIndex;
            SocketError error;
            var receive = CurSocket.Receive(RecvBuffer, RecvBufferEndIndex, leftBufferSize, SocketFlags.None,
                out error);
            if (error != SocketError.Success)
            {
                if (error != SocketError.WouldBlock && error != SocketError.Interrupted &&
                    error != SocketError.TryAgain)
                    SetError(error, PackageSocketError.None);
                break;
            }

            RecvBufferEndIndex += receive;
            ProcessPackage();

            if (receive == 0)
            {
                SetError(error, PackageSocketError.SocketShutdown);
                break;
            }

            CheckTimeout = now.AddMilliseconds(ReceiveTimeoutSetting);
            NextSendPingTime = now.AddMilliseconds(AutoSendPingSetting);
        }
    }

    void ProcessPackage()
    {
        while (CurState == State.Connected)
        {
            var bufferLen = RecvBufferEndIndex - RecvBufferBeginIndex;
            var dataLen = 0;
            var headerLen = DecodeHeader(RecvBuffer, RecvBufferBeginIndex, bufferLen, ref dataLen);
            if (headerLen == 0)
            {
                headerLen = GetHeaderLength();
                if (RecvBufferBeginIndex + headerLen > RecvBuffer.Length)
                    MemmoveRecvBuffer();
                break;
            }

            var packageLen = headerLen + dataLen;
            if (packageLen > RecvBuffer.Length)
            {
                SetError(SocketError.Success, PackageSocketError.RecviveBufferNotEnough);
                break;
            }

            if (packageLen <= bufferLen)
            {
                if (dataLen > 0)
                    OnRecvive(RecvBuffer, RecvBufferBeginIndex + headerLen, dataLen);
                RecvBufferBeginIndex += packageLen;
                continue;
            }

            if (RecvBufferBeginIndex + packageLen > RecvBuffer.Length)
                MemmoveRecvBuffer();

            break;
        }
    }

    public void MemmoveRecvBuffer()
    {
        var bufferLen = RecvBufferEndIndex - RecvBufferBeginIndex;
        Array.Copy(RecvBuffer, RecvBufferBeginIndex, RecvBuffer, 0, bufferLen);
        RecvBufferBeginIndex = 0;
        RecvBufferEndIndex = bufferLen;
    }

    void SetError(SocketError socketError, PackageSocketError packageSocketError)
    {
        LastSystemSocketError = socketError;
        LastPackageSocketError = packageSocketError;
        CurState = State.Error;
    }

    void Reset()
    {
        CurState = State.Closed;
        CurSocket = null;
        RecvBufferBeginIndex = 0;
        RecvBufferEndIndex = 0;
        SendQueue.Clear();
        SendQueueStartIndex = 0;
    }

    protected virtual int DecodeHeader(byte[] buffer, int start, int length, ref int dataLength)
    {
        if (length < 2)
            return 0;

        dataLength = buffer[start] << 8 | buffer[start + 1];
        return 2;
    }

    protected virtual void EncodeHeader(int length, byte[] buffer, int start)
    {
        buffer[start + 0] = (byte) ((length >> 8) & 0xff);
        buffer[start + 1] = (byte) (length & 0xff);
    }

    protected virtual int GetHeaderLength()
    {
        return 2;
    }

    public override string ToString()
    {
        return string.Format("State:{0} SendQueue:{1}", CurState, SendQueue.Count);
    }
}