HawkeyeVision/framework/Inspectron.HawkEye/UDPB/UDPBSocket.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Sockets;
using System.Threading;

namespace Inspectron.HawkEye.UDPB
{
    public class UDPBSocket : IDisposable
    {
        public int LossSimulation { get; set; }

        public EndPoint LastConnection => _lastConnection;
        private readonly UDPBQueue _queue;

        private readonly UDPBReceiveBuffer _receiveBuffer;

        private readonly byte[] _receivePacketBuffer = new byte[10 * 1024 * 1024];
        private readonly Random _rnd = new Random();
        private Socket _localSocket;
        private uint _currentSquenceId = 1;

        private EndPoint _lastConnection = new IPEndPoint(IPAddress.Any, 27001);
        private bool _isServer;
        private readonly ConcurrentQueue<byte[]> _receiveQueue = new ConcurrentQueue<byte[]>();
        private bool _isReceiveing = true;
        private readonly Thread _receiveThread;
        private readonly List<UDPBPacket> _packetCache = new List<UDPBPacket>();

        public UDPBSocket()
        {
            _queue = new UDPBQueue();
            _receiveBuffer = new UDPBReceiveBuffer(this, _receiveQueue);

            _receiveThread = new Thread(ReceiveLoop);
        }

        public void Dispose()
        {
            if (!_isServer) _localSocket.Disconnect(false);
            _isReceiveing = false;
        }

        public static int FindFreePort(IPAddress adapter)
        {
            var port = 0;
            var socket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
            try
            {
                var localEP = new IPEndPoint(adapter, 0);
                socket.Bind(localEP);
                localEP = (IPEndPoint) socket.LocalEndPoint;
                port = localEP.Port;
            }
            finally
            {
                socket.Close();
            }

            return port;
        }

        public void Listen(IPAddress adapter, int port)
        {
            Console.WriteLine($"Listen on {adapter}");
            _isServer = true;
            _localSocket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
            _localSocket.SetSocketOption(SocketOptionLevel.IP, SocketOptionName.ReuseAddress, true);
            //_localSocket.SetSocketOption(SocketOptionLevel.Udp, SocketOptionName.DontFragment, true);
            _localSocket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReceiveBuffer, 10 * 1024 * 1024);
            _localSocket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.SendBuffer, 10 * 1024 * 1024);

            _localSocket.Bind(new IPEndPoint(adapter, port));
            _receiveThread.Start();
        }

        public void Connect(IPEndPoint endpoint, IPAddress adapter)
        {
            Console.WriteLine($"Connect to {adapter}");
            _localSocket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
            //_localSocket.SetSocketOption(SocketOptionLevel.Udp, SocketOptionName.DontFragment, true);
            _localSocket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReceiveBuffer, 10 * 1024 * 1024);
            _localSocket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.SendBuffer, 10 * 1024 * 1024);

            Console.WriteLine($"Binding on {adapter?.MapToIPv4()}");

            if (adapter != null)
                _localSocket.Bind(new IPEndPoint(adapter.MapToIPv4(), 0));
            _localSocket.Connect(endpoint);
            _receiveThread.Start();
        }


        public void SendData(byte[] data)
        {
            _packetCache.Clear();
            IncrementSequence();
            var begin = new UDPBPacket {Type = UDPBPacket.EPacketType.BeginSequence, SequenceId = _currentSquenceId};
            SendPacket(begin, true);


            _queue.Reset();
            _queue.AddBuffer(data, 0, data.Length);
            var lastId = Flush();

            
            var finish = new UDPBPacket
                {Type = UDPBPacket.EPacketType.CloseSequence, MessageId = lastId + 1, SequenceId = _currentSquenceId};
            SendPacket(finish, true);


            do
            {
                byte[] dataOK;
                  
                    dataOK = Receive();
                

                _localSocket.ReceiveTimeout = 0;

                var packet = new UDPBPacket();
                packet.Deserialize(dataOK, 0, dataOK.Length);

                if (packet.Type == UDPBPacket.EPacketType.Ok)
                {
                    Console.WriteLine("OK");
                    
                    break;
                }

                if (packet.Type == UDPBPacket.EPacketType.Nak)
                {
                    var ms = new MemoryStream(packet.Payload);
                    var br = new BinaryReader(ms);
                    var packets = br.ReadInt32();
                    Console.WriteLine("NAK " + packets);
                    for (var i = 0; i < packets; i++)
                    {
                        var id = br.ReadUInt32();
                        var p = _packetCache.First(x => x.MessageId == id);
                        SendPacket(p, true);
                    }

                    SendPacket(finish, true);
                }
            } while (true);
        }


        public byte[] Receive()
        {
            byte[] res;
            while (!_receiveQueue.TryDequeue(out res)) Thread.Sleep(1);


            return res;
        }

        internal void SendPacket(UDPBPacket packet, bool noCache = false)
        {
            if (!noCache) _packetCache.Add(packet);

            if (LossSimulation > 0)
                if (_rnd.Next(LossSimulation) == LossSimulation - 1)
                    return;

            if (_isServer)
                _localSocket.SendTo(packet.Serialize(), _lastConnection);
            else
                _localSocket.Send(packet.Serialize());
        }

        private void ReceiveLoop()
        {
            while (_isReceiveing)
            {
                var received = ReceiveFrom(_receivePacketBuffer, 0);
                _receiveBuffer.AddPacketBytes(_receivePacketBuffer, 0, received);
            }
        }

        private int ReceiveFrom(byte[] buffer, int bufferOffset)
        {
            return _localSocket.ReceiveFrom(buffer, bufferOffset,
                UDPBQueue.PACKET_SIZE + UDPBPacket.packetHeaderSize,
                SocketFlags.None, ref _lastConnection);
        }


        private void IncrementSequence()
        {
            _currentSquenceId++;
            if (_currentSquenceId == uint.MaxValue) _currentSquenceId = 1;
        }

        private uint Flush()
        {
            uint lastPacketId = 0;
            while (_queue.PacketsToSend() != 0)
            {
                var packet = new UDPBPacket();
                byte[] data = null;
                uint msgNo = 0;
                var payload = _queue.ReadData(ref data, ref msgNo);
                packet.Type = UDPBPacket.EPacketType.Data;
                packet.Length = payload;
                packet.Payload = data;
                packet.MessageId = msgNo;
                packet.SequenceId = _currentSquenceId;
                SendPacket(packet);
                lastPacketId = msgNo;
            }

            return lastPacketId;
        }
    }
}