BoneSync/BoneSync/Networking/Packet.cs

using BoneSync.Data;
using BoneSync.Networking.Messages;
using BoneSync.Sync;
using Facepunch.Steamworks;
using MelonLoader;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
#if TEST
using Xunit;
#endif

namespace BoneSync.Networking
{
    
    public enum PacketReliability
    {
        Unreliable = 0,
        Reliable = 1,
        ReliableFast = 2,
    }
    public struct PacketInfo
    {
        public ulong id;
        public ulong senderId;
        public ulong receiverId;
        public byte sceneIndex;
        public PacketType packetType;
        public PacketReliability reliability;

        public PacketInfo(
            ulong senderId, 
            ulong receiverId, 
            PacketType packetType, 
            PacketReliability reliability,
            ulong id,
            byte sceneIndex
        )
        {
            this.senderId = senderId;
            this.receiverId = receiverId;
            this.packetType = packetType;
            this.reliability = reliability;
            this.id = id == 0 ? Packet.GenerateId() : id;
            this.sceneIndex = sceneIndex;
        }
    }
    public class Packet
    {
        private static Dictionary<byte, Queue<NetworkMessage>> _packetQueues = new Dictionary<byte, Queue<NetworkMessage>>();

        private static ulong _packetId = 1;

        private static Dictionary<ulong, Dictionary<PacketType, ulong>> _latestPacketIds = new Dictionary<ulong, Dictionary<PacketType, ulong>>();

        private static void SetLatestPacketId(ulong id, PacketType packetType, ulong value)
        {
            if (!_latestPacketIds.ContainsKey(id))
            {
                _latestPacketIds.Add(id, new Dictionary<PacketType, ulong>());
            }
            _latestPacketIds[id][packetType] = value;
        }

        private static ulong GetLatestPacketId(ulong id, PacketType packetType)
        {
            if (!_latestPacketIds.ContainsKey(id))
            {
                return 0;
            }
            if (!_latestPacketIds[id].ContainsKey(packetType))
            {
                return 0;
            }
            return _latestPacketIds[id][packetType];
        }

        public PacketInfo Info
        {
            get => _packetInfo;
            private set => _packetInfo = value;
        }

        private PacketInfo _packetInfo;
        private byte[] _dataBytes;

        public Packet(PacketInfo packetInfo, byte[] data)
        {
            _packetInfo = packetInfo;
            _dataBytes = data;
        }

        public static Packet FromBytes(byte[] bytes)
        {
            ByteEncoder byteEncoder = new ByteEncoder();
            byteEncoder.WriteBytes(bytes);

            byte packetType = byteEncoder.ReadByte();
            byte reliability = byteEncoder.ReadByte();
            byte sceneIndex = byteEncoder.ReadByte();
            ulong id = byteEncoder.ReadULong();
            ulong senderId = byteEncoder.ReadULong();
            ulong receiverId = byteEncoder.ReadULong();

            PacketInfo packetInfo = new PacketInfo(senderId, receiverId, (PacketType)packetType, (PacketReliability)reliability, id, sceneIndex);
            return new Packet(packetInfo, byteEncoder.ToArray());
        }

        public byte[] ToBytes()
        {
            ByteEncoder byteEncoder = new ByteEncoder();
            byteEncoder.WriteByte((byte)_packetInfo.packetType);
            byteEncoder.WriteByte((byte)_packetInfo.reliability);
            byteEncoder.WriteByte(_packetInfo.sceneIndex);
            byteEncoder.WriteULong(_packetInfo.id);
            byteEncoder.WriteULong(_packetInfo.senderId);
            byteEncoder.WriteULong(_packetInfo.receiverId);
            byteEncoder.WriteBytes(_dataBytes);
            return byteEncoder.ToArray();
        }

        private static void EnsureQueueForScene(byte sceneIndex)
        {
            if (!_packetQueues.ContainsKey(sceneIndex))
            {
                _packetQueues.Add(sceneIndex, new Queue<NetworkMessage>());
            }
        }

        public static bool OnPacketReceived(Packet packet)
        {
            if (packet._packetInfo.reliability == PacketReliability.ReliableFast)
            {
                if (packet.Info.id <= GetLatestPacketId(packet.Info.senderId, packet.Info.packetType))
                {
                    return false;
                }
                SetLatestPacketId(packet.Info.senderId, packet.Info.packetType, packet.Info.id);
            }
           
            NetworkMessage networkMessage = NetworkMessage.ParsePacket(packet);

            if (packet._packetInfo.sceneIndex != SceneSync.CurrentSceneIndex && !networkMessage.GetAlwaysExecute())
            {
                bool addToQueue = networkMessage.GetCatchupOrder() > 0;
                if (!addToQueue) return false;
                SyncLogger.Msg("Adding packet to queue for scene " + packet._packetInfo.sceneIndex);
                EnsureQueueForScene(packet._packetInfo.sceneIndex);
                _packetQueues[packet._packetInfo.sceneIndex].Enqueue(networkMessage);
                return true;
            }

            networkMessage.Execute();
            return true;
        }

        public static void TryCatchup()
        {
            if (!BoneSync.IsConnected)
            {
                if (_packetQueues.Count > 0) _packetQueues.Clear();
                return;
            }
            byte sceneIndex = (byte)SceneSync.CurrentSceneIndex;
            if (!_packetQueues.ContainsKey(sceneIndex))
            {
                return;
            }
            Queue<NetworkMessage> queue = _packetQueues[sceneIndex];
            int processed = 0;
            while (queue.Count > 0 && processed < 10 && SceneSync.TimeSinceLastSceneChange > SceneSync.MAP_LOAD_GRACE_PERIOD)
            {
                processed++;
                NetworkMessage networkMessage = queue.Dequeue();
                networkMessage.Execute();
            }
            if (processed > 0)
            {
                SyncLogger.Msg("Processed " + processed + " queued packets for scene " + sceneIndex);
            }
        }

        public byte[] Data => _dataBytes;

        public static ulong GenerateId()
        {
            return _packetId++;
        }
    }
#if TEST
    public class PacketTests
    {
        [Fact]
        public void TestPacketParsing()
        {
            Random random = new Random();

            byte[] randomDataBytes = new byte[1024];
            random.NextBytes(randomDataBytes);
            PacketInfo originalPacketInfo = new PacketInfo(Packet.GenerateId(), 1, 2, PacketType.LobbyInfo);
            Packet originalPacket = new Packet(originalPacketInfo, randomDataBytes);

            byte[] bytes = originalPacket.ToBytes();

            Packet parsedPacket = Packet.FromBytes(bytes);

            Assert.Equal(originalPacket.Info.id, parsedPacket.Info.id);
            Assert.Equal(originalPacket.Info.senderId, parsedPacket.Info.senderId);
            Assert.Equal(originalPacket.Info.receiverId, parsedPacket.Info.receiverId);
            Assert.Equal(originalPacket.Info.packetType, parsedPacket.Info.packetType);

            Assert.Equal(originalPacket.Data, parsedPacket.Data);


            Assert.Equal(randomDataBytes, parsedPacket.Data);
            Assert.NotEqual(bytes, randomDataBytes);
        }
    }
    #endif
}