New network handling

Managers/NetworkManager.cs

@@ -0,0 +1,320 @@
+using BepInEx.Unity.IL2CPP.Utils;
+using BetterAmongUs.Helpers;
+using BetterAmongUs.Modules;
+using BetterAmongUs.Modules.AntiCheat;
+using HarmonyLib;
+using Hazel;
+using InnerNet;
+using System.Collections;
+using UnityEngine;
+
+namespace BetterAmongUs.Managers;
+
+internal class NetworkManager
+{
+    internal static InnerNetClient? InnerNetClient => AmongUsClient.Instance;
+
+    public static void HandleGameData(MessageReader parentReader)
+    {
+        try
+        {
+            while (parentReader.Position < parentReader.Length)
+            {
+                MessageReader messageReader = parentReader.ReadMessageAsNewBuffer();
+                int currentMessageNumber = InnerNetClient.msgNum++;
+                InnerNetClient.StartCoroutine(HandleGameDataInner(messageReader, currentMessageNumber));
+                RPCHandler.HandleRPC(parentReader.Tag, null, MessageReader.Get(parentReader), HandlerFlag.HandleGameDataTag);
+            }
+        }
+        finally
+        {
+            parentReader.Recycle();
+        }
+    }
+
+    private static IEnumerator HandleGameDataInner(MessageReader reader, int msgNum)
+    {
+        int attemptCount = 0;
+        reader.Position = 0;
+        byte tag = reader.Tag;
+
+        switch (tag)
+        {
+            case 1: // Object deserialization
+                yield return HandleObjectDeserialization(reader, msgNum, attemptCount);
+                break;
+
+            case 2: // RPC handling
+                yield return HandleRpcCall(reader, msgNum, attemptCount);
+                break;
+
+            case 4: // Spawn handling
+                InnerNetClient.StartCoroutine(InnerNetClient.CoHandleSpawn(reader));
+                break;
+
+            case 5: // Object destruction
+                yield return HandleObjectDestruction(reader);
+                break;
+
+            case 7: // Client ready status
+                yield return HandleClientReady(reader);
+                break;
+
+            case 6: // Scene change
+                HandleSceneChange(reader);
+                break;
+
+            case 207: // Special case (ulong parsing)
+                yield return HandleSpecialCase(reader);
+                break;
+
+            default: // Invalid tags
+                HandleInvalidTag(reader);
+                break;
+        }
+
+        yield break;
+    }
+
+    private static IEnumerator HandleObjectDeserialization(MessageReader reader, int msgNum, int initialAttemptCount)
+    {
+        int attemptCount = initialAttemptCount;
+        try
+        {
+            InnerNetObject innerNetObject;
+            while (true)
+            {
+                uint netId = reader.ReadPackedUInt32();
+
+                if (InnerNetClient.allObjectsFast.TryGetValue(netId, out innerNetObject))
+                {
+                    innerNetObject.Deserialize(reader, false);
+                    break;
+                }
+
+                if (InnerNetClient.DestroyedObjects.Contains(netId))
+                {
+                    break;
+                }
+
+                Debug.LogWarning("Stored data for " + netId.ToString());
+                attemptCount++;
+
+                if (attemptCount > 10)
+                {
+                    yield break;
+                }
+
+                reader.Position = 0;
+                yield return Effects.Wait(0.1f);
+            }
+        }
+        finally
+        {
+            reader.Recycle();
+        }
+    }
+
+    private static IEnumerator HandleRpcCall(MessageReader reader, int msgNum, int initialAttemptCount)
+    {
+        int attemptCount = initialAttemptCount;
+        try
+        {
+            while (true)
+            {
+                uint netId;
+                byte rpcCall;
+
+                try
+                {
+                    netId = reader.ReadPackedUInt32();
+                    rpcCall = reader.ReadByte();
+                }
+                catch
+                {
+                    throw;
+                }
+
+                if (InnerNetClient.allObjectsFast.TryGetValue(netId, out InnerNetObject innerNetObject))
+                {
+                    if (!HandleInnerNetObject(innerNetObject, rpcCall, reader))
+                    {
+                        break;
+                    }
+
+                    if (Enum.IsDefined(typeof(RpcCalls), rpcCall))
+                    {
+                        innerNetObject.HandleRpc(rpcCall, reader);
+                    }
+                    else
+                    {
+                        if (innerNetObject is PlayerControl player)
+                        {
+                            RPC.HandleCustomRPC(player, rpcCall, reader);
+                        }
+                    }
+
+                    break;
+                }
+
+                if (netId == 4294967295U || InnerNetClient.DestroyedObjects.Contains(netId))
+                {
+                    break;
+                }
+
+                Debug.LogWarning($"Stored Msg {msgNum} RPC {(RpcCalls)rpcCall} for {netId}");
+                attemptCount++;
+
+                if (attemptCount > 10)
+                {
+                    yield break;
+                }
+
+                reader.Position = 0;
+                yield return Effects.Wait(0.1f);
+            }
+        }
+        finally
+        {
+            reader.Recycle();
+        }
+    }
+
+    private static IEnumerator HandleObjectDestruction(MessageReader reader)
+    {
+        try
+        {
+            uint netId = reader.ReadPackedUInt32();
+            InnerNetClient.DestroyedObjects.Add(netId);
+
+            InnerNetObject innerNetObject = InnerNetClient.FindObjectByNetId<InnerNetObject>(netId);
+            if (innerNetObject && !innerNetObject.AmOwner)
+            {
+                InnerNetClient.RemoveNetObject(innerNetObject);
+                innerNetObject.gameObject.DestroyObj();
+            }
+        }
+        finally
+        {
+            reader.Recycle();
+        }
+        yield break;
+    }
+
+    private static IEnumerator HandleClientReady(MessageReader reader)
+    {
+        try
+        {
+            ClientData clientData = InnerNetClient.FindClientById(reader.ReadPackedInt32());
+            if (clientData != null)
+            {
+                clientData.IsReady = true;
+            }
+        }
+        finally
+        {
+            reader.Recycle();
+        }
+        yield break;
+    }
+
+    private static void HandleSceneChange(MessageReader reader)
+    {
+        int clientId = reader.ReadPackedInt32();
+        ClientData clientData = InnerNetClient.FindClientById(clientId);
+        string sceneName = reader.ReadString();
+
+        if (clientData != null && !string.IsNullOrWhiteSpace(sceneName))
+        {
+            InnerNetClient.StartCoroutine(InnerNetClient.CoOnPlayerChangedScene(clientData, sceneName));
+        }
+        else
+        {
+            Debug.Log($"Couldn't find client {clientId} to change scene to {sceneName}");
+            reader.Recycle();
+        }
+    }
+
+    private static IEnumerator HandleSpecialCase(MessageReader reader)
+    {
+        try
+        {
+            ulong.Parse(reader.ReadString());
+        }
+        finally
+        {
+            reader.Recycle();
+        }
+        yield break;
+    }
+
+    private static void HandleInvalidTag(MessageReader reader)
+    {
+        Debug.Log($"Bad tag {reader.Tag} at {reader.Offset}+{reader.Position}={reader.Length}: " +
+                  string.Join(" ", reader.Buffer.Take(128)));
+        reader.Recycle();
+    }
+
+    private static bool HandleInnerNetObject(InnerNetObject netObj, byte callId, MessageReader reader)
+    {
+        if (netObj is PlayerControl player)
+        {
+            if (!PlayerRpc(player, callId, reader))
+            {
+                return false;
+            }
+        }
+        else if (netObj is PlayerPhysics physics)
+        {
+            if (!PlayerRpc(physics.myPlayer, callId, reader))
+            {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    private static bool PlayerRpc(PlayerControl player, byte callId, MessageReader reader)
+    {
+        player.BetterData().AntiCheatInfo.RPCSentPS++;
+        if (player.BetterData().AntiCheatInfo.RPCSentPS >= ExtendedAntiCheatInfo.MaxRPCSent)
+        {
+            return false;
+        }
+
+        BAUAntiCheat.HandleCheatRPCBeforeCheck(player, callId, reader);
+
+        if (BAUAntiCheat.CheckCancelRPC(player, callId, reader) != true)
+        {
+            Logger.LogCheat($"RPC canceled by Anti-Cheat: {Enum.GetName((RpcCalls)callId)}{Enum.GetName((CustomRPC)callId)} - {callId}");
+            return false;
+        }
+
+        BAUAntiCheat.CheckRPC(player, callId, reader);
+        RPC.HandleRPC(player, callId, reader);
+
+        return true;
+    }
+
+    [HarmonyPatch(typeof(ShipStatus), nameof(ShipStatus.UpdateSystem), typeof(SystemTypes), typeof(PlayerControl), typeof(MessageReader))]
+    internal static class MessageReaderUpdateSystemPatch
+    {
+        internal static bool Prefix(/*ShipStatus __instance,*/ [HarmonyArgument(0)] SystemTypes systemType, [HarmonyArgument(1)] PlayerControl player, [HarmonyArgument(2)] MessageReader reader)
+        {
+            player.BetterData().AntiCheatInfo.RPCSentPS++;
+            if (player.BetterData().AntiCheatInfo.RPCSentPS >= ExtendedAntiCheatInfo.MaxRPCSent)
+            {
+                return false;
+            }
+
+            if (BAUAntiCheat.RpcUpdateSystemCheck(player, systemType, reader) != true)
+            {
+                Logger.LogCheat($"RPC canceled by Anti-Cheat: {Enum.GetName(typeof(SystemTypes), (int)systemType)} - {MessageReader.Get(reader).ReadByte()}");
+                return false;
+            }
+
+            return true;
+        }
+    }
+}

Modules/RPC.cs

@@ -1,7 +1,6 @@
 using BetterAmongUs.Helpers;
 using BetterAmongUs.Managers;
 using BetterAmongUs.Modules.AntiCheat;
-using HarmonyLib;
 using Hazel;
 
 namespace BetterAmongUs.Modules;
@@ -30,84 +29,6 @@ enum HandleGameDataTags : byte
     NetObjectSpawn = 4,
     NetObjectDespawn = 5,
     ClientDataReady = 7,
-}
-
-[HarmonyPatch(typeof(PlayerControl), nameof(PlayerControl.HandleRpc))]
-internal class PlayerControlRPCHandlerPatch
-{
-    internal static bool Prefix(PlayerControl __instance, [HarmonyArgument(0)] byte callId, [HarmonyArgument(1)] MessageReader reader)
-    {
-        __instance.BetterData().AntiCheatInfo.RPCSentPS++;
-        if (__instance.BetterData().AntiCheatInfo.RPCSentPS >= ExtendedAntiCheatInfo.MaxRPCSent)
-        {
-            return false;
-        }
-
-        BAUAntiCheat.HandleCheatRPCBeforeCheck(__instance, callId, reader);
-
-        if (BAUAntiCheat.CheckCancelRPC(__instance, callId, reader) != true)
-        {
-            Logger.LogCheat($"RPC canceled by Anti-Cheat: {Enum.GetName((RpcCalls)callId)}{Enum.GetName((CustomRPC)callId)} - {callId}");
-            return false;
-        }
-
-        BAUAntiCheat.CheckRPC(__instance, callId, reader);
-        RPC.HandleRPC(__instance, callId, reader);
-
-        return true;
-    }
-
-    internal static void Postfix(PlayerControl __instance, [HarmonyArgument(0)] byte callId, [HarmonyArgument(1)] MessageReader reader)
-    {
-        if (!Enum.IsDefined(typeof(RpcCalls), callId))
-        {
-            RPC.HandleCustomRPC(__instance, callId, reader);
-        }
-    }
-}
-
-[HarmonyPatch(typeof(PlayerPhysics), nameof(PlayerPhysics.HandleRpc))]
-internal class PlayerPhysicsRPCHandlerPatch
-{
-    internal static bool Prefix(PlayerPhysics __instance, [HarmonyArgument(0)] byte callId, [HarmonyArgument(1)] MessageReader reader)
-    {
-        __instance.myPlayer.BetterData().AntiCheatInfo.RPCSentPS++;
-        if (__instance.myPlayer.BetterData().AntiCheatInfo.RPCSentPS >= ExtendedAntiCheatInfo.MaxRPCSent)
-        {
-            return false;
-        }
-
-        if (BAUAntiCheat.CheckCancelRPC(__instance.myPlayer, callId, reader) != true)
-        {
-            Logger.LogCheat($"RPC canceled by Anti-Cheat: {Enum.GetName((RpcCalls)callId)}{Enum.GetName((CustomRPC)callId)} - {callId}");
-        }
-
-        BAUAntiCheat.CheckRPC(__instance.myPlayer, callId, reader);
-        RPC.HandleRPC(__instance.myPlayer, callId, reader);
-
-        return true;
-    }
-}
-
-[HarmonyPatch(typeof(ShipStatus), nameof(ShipStatus.UpdateSystem), typeof(SystemTypes), typeof(PlayerControl), typeof(MessageReader))]
-internal static class MessageReaderUpdateSystemPatch
-{
-    internal static bool Prefix(/*ShipStatus __instance,*/ [HarmonyArgument(0)] SystemTypes systemType, [HarmonyArgument(1)] PlayerControl player, [HarmonyArgument(2)] MessageReader reader)
-    {
-        player.BetterData().AntiCheatInfo.RPCSentPS++;
-        if (player.BetterData().AntiCheatInfo.RPCSentPS >= ExtendedAntiCheatInfo.MaxRPCSent)
-        {
-            return false;
-        }
-
-        if (BAUAntiCheat.RpcUpdateSystemCheck(player, systemType, reader) != true)
-        {
-            Logger.LogCheat($"RPC canceled by Anti-Cheat: {Enum.GetName(typeof(SystemTypes), (int)systemType)} - {MessageReader.Get(reader).ReadByte()}");
-            return false;
-        }
-
-        return true;
-    }
 }
 
 internal static class RPC

Patches/ClientPatch.cs

@@ -2,7 +2,6 @@
 using BetterAmongUs.Items;
 using BetterAmongUs.Managers;
 using BetterAmongUs.Modules;
-using BetterAmongUs.Modules.AntiCheat;
 using HarmonyLib;
 using Hazel;
 using InnerNet;
@@ -133,16 +132,12 @@ internal class ClientPatch
     {
         [HarmonyPatch(nameof(InnerNetClient.HandleGameData))]
         [HarmonyPrefix]
-        internal static void HandleGameDataInner_Prefix(/*InnerNetClient __instance,*/ [HarmonyArgument(0)] MessageReader oldReader)
+        internal static bool HandleGameDataInner_Prefix([HarmonyArgument(0)] MessageReader oldReader)
         {
-            var parentReader = MessageReader.Get(oldReader);
-            while (parentReader.Position < parentReader.Length)
-            {
-                MessageReader messageReader = parentReader.ReadMessageAsNewBuffer();
-                MessageReader reader = messageReader;
-                RPCHandler.HandleRPC(reader.Tag, null, MessageReader.Get(reader), HandlerFlag.HandleGameDataTag);
-            }
+            NetworkManager.HandleGameData(oldReader);
+            return false;
         }
+
         [HarmonyPatch(nameof(InnerNetClient.CanBan))]
         [HarmonyPrefix]
         internal static bool CanBan_Prefix(ref bool __result)
@@ -150,6 +145,7 @@ internal class ClientPatch
             __result = GameState.IsHost;
             return false;
         }
+
         [HarmonyPatch(nameof(InnerNetClient.CanKick))]
         [HarmonyPrefix]
         internal static bool CanKick_Prefix(ref bool __result)
@@ -157,6 +153,7 @@ internal class ClientPatch
             __result = GameState.IsHost || (GameState.IsInGamePlay && (GameState.IsMeeting || GameState.IsExilling));
             return false;
         }
+
         [HarmonyPatch(nameof(InnerNetClient.KickPlayer))]
         [HarmonyPrefix]
         internal static void KickPlayer_Prefix(ref int clientId, ref bool ban)

Patches/PrivateLobbyPatch.cs

@@ -10,6 +10,7 @@ namespace BetterAmongUs.Patches;
 internal class PrivateLobbyPatch
 {
     private static AprilFoolsModeToggleButton? toggle;
+    private static TextMeshPro? toggleText;
 
     [HarmonyPatch(typeof(CreateGameOptions))]
     [HarmonyPatch(nameof(CreateGameOptions.Show))]
@@ -44,6 +45,7 @@ internal class PrivateLobbyPatch
             {
                 text.DestroyTextTranslator();
                 text.text = "Private Only Lobby";
+                toggleText = text;
             }
 
             var banner = toggle.transform.Find("Banner")?.GetComponent<SpriteRenderer>();
@@ -74,6 +76,7 @@ internal class PrivateLobbyPatch
                 __instance.onButton.OutColor = Color.white;
                 __instance.onText.color = Color.white;
                 __instance.onButtonSprite.color = Color.green;
+                if (toggleText != null) toggleText.color = new(1, 1, 1, 1);
             }
             else
             {
@@ -83,6 +86,7 @@ internal class PrivateLobbyPatch
                 __instance.offButton.OutColor = Color.white;
                 __instance.offText.color = Color.white;
                 __instance.offButtonSprite.color = Color.green;
+                if (toggleText != null) toggleText.color = new(1, 1, 1, 0.5f);
             }
 
             return false;