Add Ack callback support to custom rpcs

Reactor.Example/ExamplePlugin.cs

@@ -1,6 +1,8 @@
 using System;
 using BepInEx;
 using BepInEx.Unity.IL2CPP;
+using Il2CppInterop.Runtime.Attributes;
+using Reactor.Extensions;
 using Reactor.Networking;
 using Reactor.Networking.MethodRpc;
 using UnityEngine;
@@ -21,31 +23,52 @@ public partial class ExamplePlugin : BasePlugin
     [RegisterInIl2Cpp]
     public class ExampleComponent : MonoBehaviour
     {
+        [HideFromIl2Cpp]
+        public DragWindow TestWindow { get; }
+
         public ExampleComponent(IntPtr ptr) : base(ptr)
         {
+            TestWindow = new DragWindow(new Rect(60, 20, 0, 0), "Example", () =>
+            {
+                if (AmongUsClient.Instance && PlayerControl.LocalPlayer)
+                {
+                    if (GUILayout.Button("Send ExampleRpc"))
+                    {
+                        var name = PlayerControl.LocalPlayer.Data.PlayerName;
+                        Rpc<ExampleRpc>.Instance.Send(new ExampleRpc.Data($"Send: from {name}"), ackCallback: () =>
+                        {
+                            Logger<ExamplePlugin>.Info("Got an acknowledgement for example rpc");
+                        });
+
+                        if (!AmongUsClient.Instance.AmHost)
+                        {
+                            Rpc<ExampleRpc>.Instance.SendTo(AmongUsClient.Instance.HostId, new ExampleRpc.Data($"SendTo: from {name} to host"));
+                        }
+                    }
+
+                    if (GUILayout.Button("Send MethodExampleRpc"))
+                    {
+                        RpcSay(PlayerControl.LocalPlayer, "Hello from method rpc", Random.value, PlayerControl.LocalPlayer);
+                    }
+                }
+            })
+            {
+                Enabled = false,
+            };
         }
 
         private void Update()
         {
-            if (AmongUsClient.Instance && PlayerControl.LocalPlayer)
+            if (Input.GetKeyDown(KeyCode.F2))
             {
-                if (Input.GetKeyDown(KeyCode.F3))
-                {
-                    var name = PlayerControl.LocalPlayer.Data.PlayerName;
-                    Rpc<ExampleRpc>.Instance.Send(new ExampleRpc.Data($"Send: from {name}"));
-
-                    if (!AmongUsClient.Instance.AmHost)
-                    {
-                        Rpc<ExampleRpc>.Instance.SendTo(AmongUsClient.Instance.HostId, new ExampleRpc.Data($"SendTo: from {name} to host"));
-                    }
-                }
-
-                if (Input.GetKeyDown(KeyCode.F4))
-                {
-                    RpcSay(PlayerControl.LocalPlayer, "Hello from method rpc", Random.value, PlayerControl.LocalPlayer);
-                }
+                TestWindow.Enabled = !TestWindow.Enabled;
             }
         }
+
+        private void OnGUI()
+        {
+            TestWindow.OnGUI();
+        }
     }
 
     [MethodRpc((uint) CustomRpcCalls.MethodRpcExample)]

Reactor/Networking/CustomRpc.cs

@@ -33,14 +33,14 @@ public abstract class CustomRpc<TPlugin, TInnerNetObject, TData> : UnsafeCustomR
         Handle((TInnerNetObject) innerNetObject, (TData?) data);
     }
 
-    public void Send(InnerNetObject netObject, TData data, bool immediately = false)
+    public void Send(InnerNetObject netObject, TData data, bool immediately = false, Action? ackCallback = null)
     {
-        UnsafeSend(netObject, data, immediately);
+        UnsafeSend(netObject, data, immediately, ackCallback: ackCallback);
     }
 
-    public void SendTo(InnerNetObject netObject, int targetId, TData data)
+    public void SendTo(InnerNetObject netObject, int targetId, TData data, Action? ackCallback = null)
     {
-        UnsafeSend(netObject, data, true, targetId);
+        UnsafeSend(netObject, data, true, targetId, ackCallback);
     }
 }
 
@@ -69,13 +69,13 @@ public abstract class CustomRpc<TPlugin, TInnerNetObject> : UnsafeCustomRpc wher
         Handle((TInnerNetObject) innerNetObject);
     }
 
-    public void Send(InnerNetObject netObject, bool immediately = false)
+    public void Send(InnerNetObject netObject, bool immediately = false, Action? ackCallback = null)
     {
-        UnsafeSend(netObject, immediately);
+        UnsafeSend(netObject, immediately, ackCallback: ackCallback);
     }
 
-    public void SendTo(InnerNetObject netObject, int targetId)
+    public void SendTo(InnerNetObject netObject, int targetId, Action? ackCallback = null)
     {
-        UnsafeSend(netObject, null, true, targetId);
+        UnsafeSend(netObject, null, true, targetId, ackCallback);
     }
 }

Reactor/Networking/PlayerCustomRpc.cs

@@ -1,3 +1,4 @@
+using System;
 using BepInEx.Unity.IL2CPP;
 
 namespace Reactor.Networking;
@@ -8,14 +9,14 @@ public abstract class PlayerCustomRpc<TPlugin, TData> : CustomRpc<TPlugin, Playe
     {
     }
 
-    public void Send(TData data, bool immediately = false)
+    public void Send(TData data, bool immediately = false, Action? ackCallback = null)
     {
-        Send(PlayerControl.LocalPlayer, data, immediately);
+        Send(PlayerControl.LocalPlayer, data, immediately, ackCallback);
     }
 
-    public void SendTo(int targetId, TData data)
+    public void SendTo(int targetId, TData data, Action? ackCallback = null)
     {
-        SendTo(PlayerControl.LocalPlayer, targetId, data);
+        SendTo(PlayerControl.LocalPlayer, targetId, data, ackCallback);
     }
 }
 
@@ -25,13 +26,13 @@ public abstract class PlayerCustomRpc<TPlugin> : CustomRpc<TPlugin, PlayerContro
     {
     }
 
-    public void Send(bool immediately = false)
+    public void Send(bool immediately = false, Action? ackCallback = null)
     {
-        Send(PlayerControl.LocalPlayer, immediately);
+        Send(PlayerControl.LocalPlayer, immediately, ackCallback);
     }
 
-    public void SendTo(int targetId)
+    public void SendTo(int targetId, Action? ackCallback = null)
     {
-        SendTo(PlayerControl.LocalPlayer, targetId);
+        SendTo(PlayerControl.LocalPlayer, targetId, ackCallback);
     }
 }

Reactor/Networking/UnsafeCustomRpc.cs

@@ -1,8 +1,12 @@
 using System;
-using BepInEx;
+using System.Collections.Concurrent;
 using BepInEx.Unity.IL2CPP;
+using HarmonyLib;
 using Hazel;
+using Hazel.Udp;
+using Il2CppInterop.Runtime.InteropTypes.Arrays;
 using InnerNet;
+using Buffer = Il2CppSystem.Buffer;
 
 namespace Reactor.Networking;
 
@@ -31,7 +35,7 @@ public abstract class UnsafeCustomRpc
     public abstract object? UnsafeRead(MessageReader reader);
     public abstract void UnsafeHandle(InnerNetObject innerNetObject, object? data);
 
-    public void UnsafeSend(InnerNetObject innerNetObject, object? data, bool immediately = false, int targetClientId = -1)
+    public void UnsafeSend(InnerNetObject innerNetObject, object? data, bool immediately = false, int targetClientId = -1, Action? ackCallback = null)
     {
         if (innerNetObject == null) throw new ArgumentNullException(nameof(innerNetObject));
 
@@ -59,6 +63,12 @@ public abstract class UnsafeCustomRpc
         UnsafeWrite(writer, data);
         writer.EndMessage();
 
+        if (ackCallback != null)
+        {
+            if (SendOption != SendOption.Reliable) throw new ArgumentException("Can't add an ack callback to unreliable rpc", nameof(ackCallback));
+            AckCallbacks.TryAdd(writer, ackCallback);
+        }
+
         if (immediately)
         {
             AmongUsClient.Instance.FinishRpcImmediately(writer);
@@ -73,4 +83,30 @@ public abstract class UnsafeCustomRpc
             UnsafeHandle(innerNetObject, data);
         }
     }
+
+    private static ConcurrentDictionary<MessageWriter, Action> AckCallbacks { get; } = new(Il2CppEqualityComparer<MessageWriter>.Instance);
+
+    [HarmonyPatch(typeof(UdpConnection), nameof(UdpConnection.Send))]
+    private static class MessageAckPatch
+    {
+        public static bool Prefix(UdpConnection __instance, MessageWriter msg)
+        {
+            if (__instance._state != ConnectionState.Connected) return true;
+
+            if (msg.SendOption == SendOption.Reliable && AckCallbacks.TryRemove(msg, out var ackCallback))
+            {
+                var buffer = new Il2CppStructArray<byte>(msg.Length);
+                Buffer.BlockCopy(new Il2CppSystem.Array(msg.Buffer.Pointer), 0, new Il2CppSystem.Array(buffer.Pointer), 0, msg.Length);
+
+                __instance.ResetKeepAliveTimer();
+
+                __instance.AttachReliableID(buffer, 1, ackCallback);
+                __instance.WriteBytesToConnection(buffer, buffer.Length);
+
+                return false;
+            }
+
+            return true;
+        }
+    }
 }