Add RestrictiveSocketSend

sgkoishi · Feb 14, 2024 · 3711d68 · 3711d68
1 parent 017bb47
commit 3711d68
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 5 deletions.
diff --git a/Core/Config.cs b/Core/Config.cs
@@ -706,6 +706,17 @@ public record class MitigationSettings
         /// </summary>
         public Optional<Dictionary<PacketFilter, LimiterConfig>?> PacketSpamLimit = Optional.Default<Dictionary<PacketFilter, LimiterConfig>?>(null);
 
+        /// <summary>
+        /// <para>
+        /// Restrict all socket send operations to have exactly one message/packet per 
+        /// Send call. Requires AnotherAsyncSocket.
+        /// </para>
+        /// <para>
+        /// If you SendRawData and it contains more than one message/packet, turn this off.
+        /// </para>
+        /// </summary>
+        public Optional<bool> RestrictiveSocketSend = Optional.Default(false);
+
         public enum DisabledDamageAction
         {
             AsIs,

diff --git a/Core/Mitigations.cs b/Core/Mitigations.cs
@@ -80,7 +80,7 @@ private void OTHook_Mitigation_GetData(object? sender, OTAPI.Hooks.MessageBuffer
                 if (mitigation.SwapWhileUsePE)
                 {
                     var existingItem = Terraria.Main.player[index].GetInventory(slot);
-                    if (Terraria.Main.player[index].controlUseItem && slot == Terraria.Main.player[index].selectedItem 
+                    if (Terraria.Main.player[index].controlUseItem && slot == Terraria.Main.player[index].selectedItem
                         && type != existingItem.netID && (type != 0 || stack != 0))
                     {
                         this.Statistics.MitigationRejectedSwapWhileUse++;

diff --git a/Core/SocketProvider.cs b/Core/SocketProvider.cs
@@ -31,12 +31,18 @@ private void OTHook_Socket_OnCreate(object? sender, OTAPI.Hooks.Netplay.CreateTc
                 args.Result = new Socket.HackyAsyncSocket();
                 return;
             case Config.EnhancementsSettings.SocketType.AnotherAsyncSocket:
-                args.Result = new Socket.AnotherAsyncSocket();
+                args.Result = new Socket.AnotherAsyncSocket
+                {
+                    EnforceMessageSize = !this.config.Mitigation.Value.DisableAllMitigation && this.config.Mitigation.Value.RestrictiveSocketSend
+                };
                 return;
             case Config.EnhancementsSettings.SocketType.AnotherAsyncSocketAsFallback:
                 if (args.Result is TShockAPI.Sockets.LinuxTcpSocket)
                 {
-                    args.Result = new Socket.AnotherAsyncSocket();
+                    args.Result = new Socket.AnotherAsyncSocket
+                    {
+                        EnforceMessageSize = !this.config.Mitigation.Value.DisableAllMitigation && this.config.Mitigation.Value.RestrictiveSocketSend
+                    };
                 }
                 return;
         }
@@ -229,6 +235,7 @@ public override void AsyncReceive(byte[] data, int offset, int size, SocketRecei
 
     public class AnotherAsyncSocket : SelfSocket
     {
+        internal bool EnforceMessageSize;
         public AnotherAsyncSocket() : base()
         {
         }
@@ -239,7 +246,10 @@ public AnotherAsyncSocket(TcpClient tcpClient) : base(tcpClient)
 
         internal override ISocket New(TcpClient client)
         {
-            return new AnotherAsyncSocket(client);
+            return new AnotherAsyncSocket(client)
+            {
+                EnforceMessageSize = this.EnforceMessageSize
+            };
         }
 
         public override async void AsyncSend(byte[] data, int offset, int size, SocketSendCallback callback, object? state = null)
@@ -250,7 +260,25 @@ public override async void AsyncSend(byte[] data, int offset, int size, SocketSe
                 // ArrayPool sounds smarter than vanilla's Net.LegacyNetBufferPool
                 var buffer = ArrayPool<byte>.Shared.Rent(size);
                 Buffer.BlockCopy(data, offset, buffer, 0, size);
-                await this._connection.GetStream().WriteAsync(buffer.AsMemory(0, size));
+                if (this.EnforceMessageSize)
+                {
+                    if (size >= 3)
+                    {
+                        var claimedSize = BitConverter.ToInt16(buffer, 0);
+                        if (claimedSize != size)
+                        {
+                            TShockAPI.TShock.Log.ConsoleWarn($"[DbgPkt] Outbound message size mismatch: {size} != {claimedSize} (to {this._remoteAddress})");
+                        }
+                        else
+                        {
+                            await this._connection.GetStream().WriteAsync(buffer.AsMemory(0, size));
+                        }
+                    }
+                }
+                else
+                {
+                    await this._connection.GetStream().WriteAsync(buffer.AsMemory(0, size));
+                }
                 callback(state);
                 ArrayPool<byte>.Shared.Return(buffer);
             }