inabox/InABox.Client.RPC/Transports/Socket/RPCClientSocketTransport.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.WebSockets;
using System.Threading;
using System.Threading.Tasks;
using InABox.Clients;
using InABox.Core;

namespace InABox.Rpc
{
    public class RpcClientSocketTransport : RpcClientTransport, IDisposable
    {
        private ClientWebSocket? _socket;
        private Task? readTask;

        private string? _host;
        private object _sendLock = new object();
        private CancellationTokenSource _tokenSource = new CancellationTokenSource();
        
        private string[] _urls;

        public RpcClientSocketTransport(string[] urls)
        {
            _urls = urls;
        }

        // Returns true if we are to continue the receive loop.
        private bool DoReceive()
        {
            if(_socket != null)
            {
                try
                {
                    var buffer = new ArraySegment<byte>(new byte[1024]);
                    using (var ms = new MemoryStream())
                    {
                        WebSocketReceiveResult result;
                        do
                        {
                            var task = _socket.ReceiveAsync(buffer, _tokenSource.Token);
                            task.Wait();
                            result = task.Result;

                            ms.Write(buffer.Array, buffer.Offset, result.Count);
                        } while (!result.EndOfMessage);

                        ms.Seek(0, SeekOrigin.Begin);

                        if (result.MessageType == WebSocketMessageType.Close)
                        {
                            if (result.CloseStatus == WebSocketCloseStatus.NormalClosure)
                            {
                                Task.Run(() =>
                                {
                                    DoClose(RpcTransportCloseEventType.Closed);
                                });
                                return false;
                            }
                            else
                            {
                                DoException(new Exception(result.CloseStatusDescription));
                                Task.Run(() =>
                                {
                                    DoClose(RpcTransportCloseEventType.Error);
                                });
                                return false;
                            }
                        }
                        else
                        {
                            RpcMessage? rpcMessage = null;
                            if (result.MessageType == WebSocketMessageType.Binary)
                            {
                                rpcMessage = Serialization.ReadBinary<RpcMessage>(ms, BinarySerializationSettings.Latest);
                            }
                            else if (result.MessageType == WebSocketMessageType.Text)
                            {
                                rpcMessage = Serialization.Deserialize<RpcMessage>(ms);
                            }
                            Accept(rpcMessage);
                        }
                    }
                }
                catch(Exception e)
                {
                    DoException(e);
                    if (!IsConnected())
                    {
                        Task.Run(() =>
                        {
                            DoClose(RpcTransportCloseEventType.Error);
                        });
                    }
                    return false;
                }
            }
            return true;
        }
        
        private ClientWebSocket? CreateSocket(string url, bool secure)
        {
            var client = new ClientWebSocket();

            //WebsocketClient client;

            //WebSocket socket = null;
            var address = $"{(secure ? "wss" : "ws")}://{url}";
            try
            {
                client.ConnectAsync(new Uri(address), _tokenSource.Token).Wait();
                //socket = new WebSocket(address);
            }
            catch (Exception e)
            {
                client.Dispose();
                return null;
            }

            DoOpen();
            _host = url;

            // Time to wait before disconnect - the default meant that the client disconnected during debugging, since the ping would fail
            /*socket.WaitTime = TimeSpan.FromSeconds(20);
            socket.OnOpen -= Socket_OnOpen;
            socket.OnError -= Socket_OnError;
            socket.OnClose -= Socket_OnClose;
            socket.OnMessage -= Socket_OnMessage;
            socket.Connect();
            if (socket.ReadyState == WebSocketState.Open)
            {
                DoOpen();

                socket.OnOpen += Socket_OnOpen;
                socket.OnError += Socket_OnError;
                socket.OnClose += Socket_OnClose;
                socket.OnMessage += Socket_OnMessage;

                return socket;
            }*/
            return client;
        }

        public override void Connect()
        {
            _socket?.Dispose();

            var tasks = new List<Task<ClientWebSocket?>>();
            foreach (var url in _urls)
            {
                tasks.Add(Task.Run(() => CreateSocket(url, true)));
                tasks.Add(Task.Run(() => CreateSocket(url, false)));
            }
            while (tasks.Count > 0)
            {
                var result = Task.WhenAny(tasks).Result;
                if (result.Result == null)
                    tasks.Remove(result);
                else
                {
                    _socket = result.Result;

                    Task.Run(() =>
                    {
                        while (IsConnected())
                        {
                            if (!DoReceive())
                            {
                                break;
                            }
                        }
                    });

                    return;
                }
            }
        }

        public override bool IsConnected() => _socket?.State == WebSocketState.Open;
        public override bool IsSecure() => false;

        public override String? ServerName() => _host;
        
        public override void Disconnect()
        {
            _socket?.CloseAsync(WebSocketCloseStatus.NormalClosure, "", _tokenSource.Token).Wait();
        }

        public override void Send(RpcMessage message)
        {
            lock (_sendLock)
            {
                var buffer = message.WriteBinary(BinarySerializationSettings.Latest);
                _socket?.SendAsync(buffer, WebSocketMessageType.Binary, true, _tokenSource.Token)?.Wait();
            }
        }

        protected override RpcClientTransport Clone() => new RpcClientSocketTransport(_urls);

        public void Dispose()
        {
            if (IsConnected())
                Disconnect();
            _tokenSource.Cancel();
            _socket?.Dispose();
        }
    }
}