Rework handshake

Reactor.Example/ExampleRpc.cs

@@ -11,7 +11,7 @@ namespace Reactor.Example
 
         public readonly struct Data
         {
-            public string Message { get; }
+            public readonly string Message;
 
             public Data(string message)
             {

Reactor/CustomRpc.cs

@@ -1,11 +1,14 @@
 using System;
 using System.Collections.Generic;
+using System.IO;
 using System.Linq;
 using System.Reflection;
+using System.Text;
 using BepInEx;
 using BepInEx.IL2CPP;
 using HarmonyLib;
 using Hazel;
+using Reactor.Net;
 
 namespace Reactor
 {
@@ -38,7 +41,7 @@ namespace Reactor
         public abstract void UnsafeHandle(InnerNetObject innerNetObject, object data);
     }
 
-    public abstract class CustomRpc<TPlugin, TInnerNetObject, TData> : UnsafeCustomRpc where TPlugin : BasePlugin where TInnerNetObject : InnerNetObject where TData : struct
+    public abstract class CustomRpc<TPlugin, TInnerNetObject, TData> : UnsafeCustomRpc where TPlugin : BasePlugin where TInnerNetObject : InnerNetObject
     {
         protected CustomRpc(TPlugin plugin) : base(plugin)
         {
@@ -80,6 +83,9 @@ namespace Reactor
             _map.SelectMany(pair => pair.Value, (pair, value) => new { pair.Key, Value = value })
                 .ToLookup(pair => pair.Key, pair => pair.Value);
 
+        public Dictionary<string, int> PluginIdMap { get; private set; }
+        public Dictionary<int, string> PluginIdMapReversed { get; private set; }
+
         public CustomRpcManager()
         {
             foreach (var type in HandleRpcPatch.InnerNetObjectTypes)
@@ -88,6 +94,24 @@ namespace Reactor
             }
         }
 
+        public void ReloadPluginIdMap()
+        {
+            PluginIdMap = new Dictionary<string, int>();
+            PluginIdMapReversed = new Dictionary<int, string>();
+
+            var i = -1;
+
+            foreach (var mod in ModList.GetCurrent())
+            {
+                if (mod.Side == PluginSide.Both)
+                {
+                    PluginIdMap[mod.Id] = i;
+                    PluginIdMapReversed[i] = mod.Id;
+                    i--;
+                }
+            }
+        }
+
         public UnsafeCustomRpc Register(UnsafeCustomRpc customRpc)
         {
             customRpc.Id = List.Count(x => x.UnsafePlugin == customRpc.UnsafePlugin);
@@ -124,8 +148,16 @@ namespace Reactor
                 true => AmongUsClient.Instance.StartRpcImmediately(netObject.NetId, CallId, SendOption.Reliable, -1)
             };
 
-            writer.Write(customRpc.PluginId);
-            writer.Write(customRpc.Id!.Value);
+            if (PluginIdMap.TryGetValue(customRpc.PluginId, out var pluginId))
+            {
+                writer.WritePacked(pluginId);
+            }
+            else
+            {
+                writer.Write(customRpc.PluginId);
+            }
+
+            writer.WritePacked(customRpc.Id!.Value);
             customRpc.UnsafeWrite(writer, data);
 
             if (immediately)
@@ -153,14 +185,29 @@ namespace Reactor
                 return InnerNetObjectTypes.Select(x => x.GetMethod(nameof(InnerNetObject.HandleRpc), AccessTools.all));
             }
 
+            private static string ReadString(MessageReader reader, int len)
+            {
+                if (reader.BytesRemaining < len) throw new InvalidDataException($"Read length is longer than message length: {len} of {reader.BytesRemaining}");
+
+                var output = Encoding.UTF8.GetString(reader.Buffer, reader.readHead, len);
+                reader.Position += len;
+                return output;
+            }
+
             public static bool Prefix(InnerNetObject __instance, [HarmonyArgument(0)] byte callId, [HarmonyArgument(1)] MessageReader reader)
             {
                 if (callId == CallId)
                 {
-                    var customRpcs = PluginSingleton<ReactorPlugin>.Instance.CustomRpcManager.Map[__instance.GetType()];
+                    var manager = PluginSingleton<ReactorPlugin>.Instance.CustomRpcManager;
+                    var customRpcs = manager.Map[__instance.GetType()];
 
-                    var pluginId = reader.ReadString();
-                    var id = reader.ReadInt32();
+                    var lengthOrShortId = reader.ReadPackedInt32();
+
+                    var pluginId = lengthOrShortId < 0
+                        ? manager.PluginIdMapReversed[lengthOrShortId]
+                        : ReadString(reader, lengthOrShortId);
+
+                    var id = reader.ReadPackedInt32();
 
                     var customRpc = customRpcs.Single(x => x.PluginId == pluginId && x.Id == id);
 

Reactor/Net/Mod.cs

@@ -0,0 +1,21 @@
+namespace Reactor.Net
+{
+    public readonly struct Mod
+    {
+        public readonly string Id;
+        public readonly string Version;
+        public readonly PluginSide Side;
+
+        public Mod(string id, string version, PluginSide side)
+        {
+            Id = id;
+            Version = version;
+            Side = side;
+        }
+
+        public override string ToString()
+        {
+            return $"{Id} ({Version})";
+        }
+    }
+}

Reactor/Net/ModList.cs

@@ -0,0 +1,34 @@
+using System.Collections.Generic;
+using System.Linq;
+using System.Reflection;
+using BepInEx.IL2CPP;
+using Hazel;
+
+namespace Reactor.Net
+{
+    public static class ModList
+    {
+        public static void Serialize(MessageWriter writer, ISet<Mod> mods)
+        {
+            writer.WritePacked(mods.Count);
+
+            foreach (var mod in mods)
+            {
+                writer.Write(mod.Id);
+                writer.Write(mod.Version);
+                writer.Write((byte) mod.Side);
+            }
+        }
+
+        public static ISet<Mod> GetCurrent()
+        {
+            return IL2CPPChainloader.Instance.Plugins
+                .Select(plugin => new Mod(
+                    plugin.Key,
+                    plugin.Value.Metadata.Version.ToString(),
+                    plugin.Value.Instance.GetType().GetCustomAttribute<ReactorPluginSideAttribute>()?.Side ?? PluginSide.Both
+                ))
+                .ToHashSet();
+        }
+    }
+}

Reactor/Net/ModdedHandshakeC2S.cs

@@ -0,0 +1,20 @@
+using System.Collections.Generic;
+using Hazel;
+using UnhollowerBaseLib;
+
+namespace Reactor.Net
+{
+    public static class ModdedHandshakeC2S
+    {
+        public static void Serialize(MessageWriter writer, Il2CppStructArray<byte> clientVersion, string name, ISet<Mod> mods = null)
+        {
+            writer.Write(clientVersion);
+            writer.Write(name);
+
+            if (mods != null)
+            {
+                ModList.Serialize(writer, mods);
+            }
+        }
+    }
+}

Reactor/Net/ModdedHandshakeS2C.cs

@@ -0,0 +1,14 @@
+using Hazel;
+
+namespace Reactor.Net
+{
+    public static class ModdedHandshakeS2C
+    {
+        public static void Serialize(MessageWriter writer, string serverBrand)
+        {
+            writer.StartMessage(byte.MaxValue);
+            writer.Write(serverBrand);
+            writer.EndMessage();
+        }
+    }
+}

Reactor/Patches/ModdedHandshake.cs

@@ -1,80 +0,0 @@
-using System.Linq;
-using System.Reflection;
-using BepInEx.IL2CPP;
-using HarmonyLib;
-using Hazel;
-using UnhollowerBaseLib;
-
-namespace Reactor.Patches
-{
-    internal static class ModdedHandshake
-    {
-        // TODO add back when vent issue is fixed
-        // [HarmonyPatch(typeof(InnerNetServer), nameof(InnerNetServer.OnServerConnect))]
-        // public static class FreeplayPatch
-        // {
-        //     public static void Prefix([HarmonyArgument(0)] NewConnectionEventArgs evt)
-        //     {
-        //         var handshakeData = evt.HandshakeData;
-        //
-        //         var protocolVersion = handshakeData.ReadInt32();
-        //         if (protocolVersion != -1)
-        //         {
-        //             // reset position if vanilla
-        //             handshakeData._position -= sizeof(int);
-        //             handshakeData.readHead -= sizeof(int);
-        //         }
-        //     }
-        // }
-
-        // workaround for ^
-        [HarmonyPatch(typeof(InnerNetServer), nameof(InnerNetServer.IsCompatibleVersion))]
-        public static class FreeplayPatch
-        {
-            public static bool Prefix([HarmonyArgument(0)] int version, ref bool __result)
-            {
-                if (version == -1)
-                {
-                    __result = true;
-                    return false;
-                }
-
-                return true;
-            }
-        }
-
-        [HarmonyPatch(typeof(AmongUsClient), nameof(AmongUsClient.GetConnectionData))]
-        public static class HandshakePatch
-        {
-            public static bool Prefix(out Il2CppStructArray<byte> __result)
-            {
-                var handshake = MessageWriter.Get(SendOption.Reliable);
-
-                var plugins = IL2CPPChainloader.Instance.Plugins;
-
-                if (plugins.Values
-                    .Select(pluginInfo => pluginInfo.Instance.GetType().GetCustomAttribute<ReactorPluginSideAttribute>())
-                    .Any(attribute => attribute == null || attribute.Side == PluginSide.Both))
-                {
-                    handshake.Write(-1);
-                }
-
-                handshake.Write(Constants.GetBroadcastVersion());
-                handshake.Write(SaveManager.PlayerName);
-
-                handshake.WritePacked(plugins.Count);
-
-                foreach (var (pluginId, plugin) in plugins)
-                {
-                    handshake.Write(pluginId);
-                    handshake.Write(plugin.Metadata.Version.ToString());
-                }
-
-                __result = handshake.ToByteArray(false);
-                handshake.Recycle();
-
-                return false;
-            }
-        }
-    }
-}

Reactor/Patches/ModdedHandshakePatches.cs

@@ -0,0 +1,127 @@
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using System.Reflection;
+using HarmonyLib;
+using Hazel;
+using Hazel.Udp;
+using Reactor.Extensions;
+using Reactor.Net;
+using UnhollowerBaseLib;
+using UnityEngine;
+
+namespace Reactor.Patches
+{
+    internal static class ModdedHandshakePatches
+    {
+        // [HarmonyPatch(typeof(InnerNetServer), nameof(InnerNetServer.OnServerConnect))]
+        [HarmonyPatch]
+        public static class ServerPatch
+        {
+            public static IEnumerable<MethodBase> TargetMethods()
+            {
+                return typeof(InnerNetServer).GetMethods(typeof(void), typeof(NewConnectionEventArgs));
+            }
+
+            private static readonly string ServerBrand = $"Among Us {Application.version}";
+
+            public static void Postfix([HarmonyArgument(0)] NewConnectionEventArgs evt)
+            {
+                var writer = MessageWriter.Get(SendOption.Reliable);
+                ModdedHandshakeS2C.Serialize(writer, ServerBrand);
+                evt.Connection.Send(writer);
+                writer.Recycle();
+            }
+        }
+
+        // [HarmonyPatch(typeof(InnerNetClient), nameof(InnerNetClient.HandleMessage))]
+        [HarmonyPatch]
+        public static class HandleMessagePatch
+        {
+            public static IEnumerable<MethodBase> TargetMethods()
+            {
+                return typeof(InnerNetClient).GetMethods(typeof(void), typeof(MessageReader), typeof(SendOption));
+            }
+
+            public static bool Prefix(InnerNetClient __instance, [HarmonyArgument(0)] MessageReader reader, [HarmonyArgument(1)] SendOption sendOption)
+            {
+                if (reader.Tag == byte.MaxValue)
+                {
+                    var serverBrand = reader.ReadString();
+
+                    __instance.Field_23.State = ConnectionState.Connected;
+
+                    PluginSingleton<ReactorPlugin>.Instance.Log.LogDebug($"Connected to modded server ({serverBrand})");
+
+                    return false;
+                }
+
+                return true;
+            }
+        }
+
+        [HarmonyPatch(typeof(UdpConnection), nameof(UdpConnection.HandleSend))]
+        public static class SendHelloPatch
+        {
+            public static void Prefix(UdpConnection __instance, byte sendOption, ref Il2CppSystem.Action ackCallback)
+            {
+                if (sendOption != 8 || AmongUsClient.Instance == null || AmongUsClient.Instance.Field_23 == null || !AmongUsClient.Instance.Field_23.Equals(__instance))
+                {
+                    return;
+                }
+
+                if (!PluginSingleton<ReactorPlugin>.Instance.ModdedHandshake.Value)
+                {
+                    return;
+                }
+
+                ackCallback = (Action) (() =>
+                {
+                    PluginSingleton<ReactorPlugin>.Instance.Log.LogDebug("Hello was acked, waiting for modded handshake response");
+
+                    Coroutines.Start(Coroutine());
+
+                    __instance.add_Disconnected((Action<Il2CppSystem.Object, DisconnectedEventArgs>) ((_, _) =>
+                    {
+                        Coroutines.Stop(Coroutine());
+                    }));
+                });
+            }
+        }
+
+        private static IEnumerator Coroutine()
+        {
+            yield return new WaitForSeconds(3);
+
+            var client = AmongUsClient.Instance;
+
+            if (client != null && client.Field_23 != null && client.Field_23.State == ConnectionState.Connecting)
+            {
+                client.LastDisconnectReason = DisconnectReasons.Custom;
+                client.LastCustomDisconnect = "Server didn't respond to modded handshake";
+                client.HandleDisconnect(client.LastDisconnectReason, client.LastCustomDisconnect);
+            }
+        }
+
+        [HarmonyPatch(typeof(AmongUsClient), nameof(AmongUsClient.GetConnectionData))]
+        public static class HandshakePatch
+        {
+            public static bool Prefix(out Il2CppStructArray<byte> __result)
+            {
+                var handshake = MessageWriter.Get(SendOption.Reliable);
+
+                ModdedHandshakeC2S.Serialize(
+                    handshake,
+                    Constants.GetBroadcastVersion(),
+                    SaveManager.PlayerName,
+                    ModList.GetCurrent()
+                );
+
+                __result = handshake.ToByteArray(false);
+                handshake.Recycle();
+
+                return false;
+            }
+        }
+    }
+}

Reactor/ReactorPlugin.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Linq;
 using BepInEx;
+using BepInEx.Configuration;
 using BepInEx.IL2CPP;
 using HarmonyLib;
 using Reactor.Extensions;
@@ -21,12 +22,16 @@ namespace Reactor
         public Harmony Harmony { get; } = new Harmony(Id);
         public CustomRpcManager CustomRpcManager { get; } = new CustomRpcManager();
 
+        public ConfigEntry<bool> ModdedHandshake;
+
         private RegionInfoWatcher RegionInfoWatcher { get; } = new RegionInfoWatcher();
 
         public override void Load()
         {
             RegisterInIl2CppAttribute.Register();
 
+            ModdedHandshake = Config.Bind("Features", "Modded handshake", true);
+
             var gameObject = new GameObject(nameof(ReactorPlugin)).DontDestroy();
             gameObject.AddComponent<ReactorComponent>().Plugin = this;
 
@@ -54,6 +59,8 @@ namespace Reactor
 
             public void Start()
             {
+                Plugin.CustomRpcManager.ReloadPluginIdMap();
+
                 Camera.onPostRender = Camera.onPostRender == null
                     ? new Action<Camera>(OnPostRenderM)
                     : Il2CppSystem.Delegate.Combine(Camera.onPostRender, Il2CppSystem.Delegate.CreateDelegate(GetIl2CppType(), GetIl2CppType().GetMethod(nameof(OnPostRenderM), Il2CppSystem.Reflection.BindingFlags.Static | Il2CppSystem.Reflection.BindingFlags.Public))).Cast<Camera.CameraCallback>();

Reactor/ReactorPluginSideAttribute.cs

@@ -24,7 +24,7 @@ namespace Reactor
     /// <summary>
     /// Plugin side used in modded handshake
     /// </summary>
-    public enum PluginSide
+    public enum PluginSide : byte
     {
         /// <summary>
         /// Required by both sides, reject connection if missing on the other side