From 78af568603cb52f6a83bb6f65c1300999cc8f499 Mon Sep 17 00:00:00 2001
From: Sanae <atmatm6000@outlook.com>
Date: Wed, 16 Feb 2022 15:20:03 -0600
Subject: [PATCH] Add RentZero extension as rented memory can be non-zero

---
 Server/Client.cs      |  2 +-
 Server/Server.cs      |  8 ++++----
 Shared/Extensions.cs  | 12 +++++++++++-
 TestClient/Program.cs |  2 +-
 4 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/Server/Client.cs b/Server/Client.cs
index 509d986..cb80ecc 100644
--- a/Server/Client.cs
+++ b/Server/Client.cs
@@ -32,7 +32,7 @@ public class Client : IDisposable {
     }
 
     public async Task Send<T>(T packet, Client? sender = null) where T : struct, IPacket {
-        IMemoryOwner<byte> memory = MemoryPool<byte>.Shared.Rent(Constants.MaxPacketSize);
+        IMemoryOwner<byte> memory = MemoryPool<byte>.Shared.RentZero(Constants.MaxPacketSize);
 
         PacketHeader header = new PacketHeader {
             Id = sender?.Id ?? Guid.Empty,
diff --git a/Server/Server.cs b/Server/Server.cs
index 290897e..0ebc965 100644
--- a/Server/Server.cs
+++ b/Server/Server.cs
@@ -58,7 +58,7 @@ public class Server {
         }
     }
     public async Task Broadcast<T>(T packet, Client sender) where T : struct, IPacket {
-        IMemoryOwner<byte> memory = memoryPool.Rent(Constants.MaxPacketSize);
+        IMemoryOwner<byte> memory = memoryPool.RentZero(Constants.MaxPacketSize);
 
         PacketHeader header = new PacketHeader {
             Id = sender?.Id ?? Guid.Empty,
@@ -168,7 +168,7 @@ public class Server {
 
                     List<Client> otherConnectedPlayers = Clients.FindAll(c => c.Id != header.Id && c.Connected && c.Socket != null);
                     await Parallel.ForEachAsync(otherConnectedPlayers, async (other, _) => {
-                        IMemoryOwner<byte> tempBuffer = MemoryPool<byte>.Shared.Rent(Constants.MaxPacketSize);
+                        IMemoryOwner<byte> tempBuffer = MemoryPool<byte>.Shared.RentZero(Constants.MaxPacketSize);
                         PacketHeader connectHeader = new PacketHeader {
                             Id = other.Id,
                             Type = PacketType.Connect
@@ -197,14 +197,14 @@ public class Server {
                     CostumePacket costumePacket = new CostumePacket {
                         BodyName = ""
                     };
-                    costumePacket.Deserialize(memory.Memory.Span[Constants.HeaderSize..]);
+                    costumePacket.Deserialize(memory.Memory.Span[Constants.HeaderSize..Constants.MaxPacketSize]);
                     client.CurrentCostume = costumePacket;
                 }
 
                 try {
                     // if (header.Type is not PacketType.Cap and not PacketType.Player) client.Logger.Warn($"lol {header.Type}");
                     IPacket packet = (IPacket) Activator.CreateInstance(Constants.PacketIdMap[header.Type])!;
-                    packet.Deserialize(memory.Memory.Span[Constants.HeaderSize..]);
+                    packet.Deserialize(memory.Memory.Span[Constants.HeaderSize..Constants.MaxPacketSize]);
                     if (PacketHandler?.Invoke(client, packet) is false) {
                         memory.Dispose();
                         continue;
diff --git a/Shared/Extensions.cs b/Shared/Extensions.cs
index 4129327..345e627 100644
--- a/Shared/Extensions.cs
+++ b/Shared/Extensions.cs
@@ -1,4 +1,8 @@
-namespace Shared;
+using System.Buffers;
+using System.Runtime.InteropServices;
+using System.Security.Cryptography;
+
+namespace Shared;
 
 public static class Extensions {
     public static string Hex(this Span<byte> span) {
@@ -18,4 +22,10 @@ public static class Extensions {
     public static string TrimNullTerm(this string text) {
         return text.Split('\0').FirstOrDefault() ?? "";
     }
+
+    public static IMemoryOwner<byte> RentZero(this MemoryPool<byte> pool, int minSize) {
+        IMemoryOwner<byte> owner = pool.Rent(minSize);
+        CryptographicOperations.ZeroMemory(owner.Memory.Span);
+        return owner;
+    }
 }
\ No newline at end of file
diff --git a/TestClient/Program.cs b/TestClient/Program.cs
index b4a1be3..4dccda3 100644
--- a/TestClient/Program.cs
+++ b/TestClient/Program.cs
@@ -61,7 +61,7 @@ PacketHeader coolHeader = new PacketHeader {
     Type = PacketType.Connect,
     Id = ownId
 };
-IMemoryOwner<byte> owner = MemoryPool<byte>.Shared.Rent(Constants.MaxPacketSize);
+IMemoryOwner<byte> owner = MemoryPool<byte>.Shared.RentZero(Constants.MaxPacketSize);
 MemoryMarshal.Write(owner.Memory.Span[..], ref coolHeader);
 ConnectPacket connect = new ConnectPacket {
     ConnectionType = ConnectionTypes.Reconnecting,