maki/Maki/Gateway/GatewayShard.cs

using Maki.Structures.Gateway;
using Newtonsoft.Json;
using System;
using WebSocketSharp;

namespace Maki.Gateway
{
    /// <summary>
    /// Gateway connection shard
    /// </summary>
    class GatewayShard : IDisposable
    {
        private const string GATEWAY_URL = "{0}?v={1}&encoding=json";

        private readonly string GatewayUrl;

        /// <summary>
        /// Session key for continuing a resuming after disconnecting
        /// </summary>
        private string Session;

        /// <summary>
        /// Websocket container
        /// </summary>
        private WebSocket WebSocket;
        
        /// <summary>
        /// Interval at which heartbeats are sent
        /// </summary>
        public TimeSpan HeartbeatInterval { get; private set; }

        /// <summary>
        /// Last sequence received from the gateway, primarily used for resuming connections and getting the data that was missed
        /// </summary>
        public int? LastSequence { get; private set; } = null;

        /// <summary>
        /// Identifier for this Shard
        /// </summary>
        public readonly int Id;

        /// <summary>
        /// Parent DiscordClient instance
        /// </summary>
        private readonly Discord Client;

        /// <summary>
        /// Heartbeat handler
        /// </summary>
        private readonly GatewayHeartbeatManager HeartbeatHandler;
        
        /// <summary>
        /// Fires when a payload is received.
        /// </summary>
        public event Action<GatewayShard, GatewayPayload> OnSocketPayload;

        /// <summary>
        /// Constructor
        /// </summary>
        /// <param name="id">Shard Id</param>
        /// <param name="client">Parent DiscordClient instance</param>
        /// <param name="connect">Whether to immediately call Connect()</param>
        public GatewayShard(int id, Discord client, bool connect = true)
        {
            Id = id;
            Client = client;

            HeartbeatHandler = new GatewayHeartbeatManager(this);
            GatewayUrl = string.Format(GATEWAY_URL, Client.Gateway, Discord.GATEWAY_VERSION);

            if (connect)
                Connect();
        }
        
        /// <summary>
        /// Event handler for WebSocketSharp's OnClose event, stops heartbeats and reconnects if the close wasn't clean
        /// </summary>
        /// <param name="sender">Sender object</param>
        /// <param name="e">Event arguments</param>
        private void WebSocket_OnClose(object sender, CloseEventArgs e)
        {
            HeartbeatHandler.Stop();

            if (!e.WasClean || e.Code != 1000)
                Connect();
        }

        /// <summary>
        /// Event handler for WebSocketSharp's OnOpen event, handles gateway instructions and forwards to our events
        /// </summary>
        /// <param name="sender">Sender object</param>
        /// <param name="e">Event arguments</param>
        private void WebSocket_OnMessage(object sender, MessageEventArgs e)
        {
            // ignore non-text messages
            if (!e.IsText)
                return;

            Console.WriteLine(e.Data.Replace(Client.Token, new string('*', Client.Token.Length)));

            GatewayPayload payload = JsonConvert.DeserializeObject<GatewayPayload>(e.Data);
            OnSocketPayload?.Invoke(this, payload);

            switch (payload.OPCode) {
                case GatewayOPCode.Dispatch:
                    LastSequence = payload.Sequence;

                    if (payload.Name == "READY")
                    {
                        GatewayReady ready = payload.DataAs<GatewayReady>();
                        Session = ready.Session;
                    }
                    break;

                case GatewayOPCode.Hello:
                    GatewayHello hello = payload.DataAs<GatewayHello>();

                    HeartbeatInterval = TimeSpan.FromMilliseconds(hello.HeartbeatInterval);
                    HeartbeatHandler.Start();

                    if (string.IsNullOrEmpty(Session))
                        Send(GatewayOPCode.Identify, new GatewayIdentification
                        {
                            Token = Client.Token,
                            Compress = false,
                            LargeThreshold = 250,
                            Shard = new int[] { Id, Client.ShardClient.ShardCount },
                            Properties = new GatewayIdentificationProperties
                            {
                                OperatingSystem = "windows",
                                Browser = "Maki",
                                Device = "Maki",
                                Referrer = string.Empty,
                                ReferringDomain = string.Empty,
                            }
                        });
                    else
                        Send(GatewayOPCode.Resume, new GatewayResume
                        {
                            LastSequence = LastSequence ?? default(int),
                            Token = Client.Token,
                            Session = Session,
                        });
                    break;
            }
        }

        #region IDisposable

        private bool IsDisposed = false;

        private void Dispose(bool disposing)
        {
            if (IsDisposed)
                return;

            IsDisposed = true;
            Disconnect();

            if (disposing)
                GC.SuppressFinalize(this);
        }

        ~GatewayShard()
            => Dispose(false);

        public void Dispose()
            => Dispose(true);

        #endregion

        /// <summary>
        /// Connects to the gateway
        /// </summary>
        public void Connect()
        {
            WebSocket = new WebSocket(GatewayUrl);

            // make wss not log anything on its own
            WebSocket.Log.Output = (LogData logData, string path) => { };

            WebSocket.OnClose += WebSocket_OnClose;
            WebSocket.OnMessage += WebSocket_OnMessage;
        
            WebSocket.Connect();
        }

        /// <summary>
        /// Stops heartbeats and closes gateway connection for this shard
        /// </summary>
        public void Disconnect()
        {
            HeartbeatHandler.Stop();

            if (WebSocket.ReadyState != WebSocketState.Closed)
                WebSocket?.Close(CloseStatusCode.Normal);
        }

        /// <summary>
        /// Serialises and sends a json object
        /// </summary>
        /// <typeparam name="T">Type to serialise as</typeparam>
        /// <param name="data">Data to serialise and send</param>
        public void Send<T>(T data)
        {
            string json = JsonConvert.SerializeObject(data, typeof(T), new JsonSerializerSettings());

            Console.WriteLine(json.Replace(Client.Token, new string('*', Client.Token.Length)));

            WebSocket.Send(json);
        }

        /// <summary>
        /// Serialises and sends a Gateway Payload
        /// </summary>
        /// <param name="opcode">Opcode to use</param>
        /// <param name="data">Data to send</param>
        public void Send(GatewayOPCode opcode, object data)
            => Send(new GatewayPayload
            {
                OPCode = opcode,
                Data = data
            });
    }
}