From ed055749549f8673eb726eee89e9e92f48f5866b Mon Sep 17 00:00:00 2001
From: AeonLucid <aeonlucid@outlook.com>
Date: Sun, 1 Nov 2020 23:52:56 +0100
Subject: [PATCH] Make RemoveMessage work

---
 src/Impostor.Hazel/MessageReader.cs           | 30 ++++++++-----------
 .../Hazel/MessageReaderTests.cs               | 20 +++++++++----
 2 files changed, 28 insertions(+), 22 deletions(-)

diff --git a/src/Impostor.Hazel/MessageReader.cs b/src/Impostor.Hazel/MessageReader.cs
index e615503..d4668db 100644
--- a/src/Impostor.Hazel/MessageReader.cs
+++ b/src/Impostor.Hazel/MessageReader.cs
@@ -188,27 +188,24 @@ namespace Impostor.Hazel
 
         public void RemoveMessage(IMessageReader message)
         {
-            var pool = ArrayPool.Rent(message.Buffer.Length);
+            // Offset of where to start removing.
+            var offsetStart = message.Offset - 3;
 
-            try
-            {
-                var offsetHeader = message.Offset - 3;
-                var offsetEnd = message.Offset + message.Length;
-                var len = message.Buffer.Length - offsetEnd;
+            // Offset of where to end removing.
+            var offsetEnd = message.Offset + message.Length;
 
-                Array.Copy(message.Buffer, offsetEnd, pool, 0, len);
-                Array.Copy(pool, 0, this.Buffer, offsetHeader, len);
+            // The amount of bytes to copy over ourselves.
+            var lengthToCopy = message.Buffer.Length - offsetEnd;
 
-                AdjustLength(message.Offset, message.Length + 3);
-            }
-            finally
-            {
-                ArrayPool.Return(pool);
-            }
+            System.Buffer.BlockCopy(Buffer, offsetEnd, Buffer, offsetStart, lengthToCopy);
+
+            ((MessageReader) message).Parent.AdjustLength(message.Offset, message.Length + 3);
         }
 
         private void AdjustLength(int offset, int amount)
         {
+            this.Length -= amount;
+
             if (this.ReadPosition > offset)
             {
                 this.Position -= amount;
@@ -217,11 +214,10 @@ namespace Impostor.Hazel
             if (Parent != null)
             {
                 var lengthOffset = this.Offset - 3;
-                var curLen = this.Buffer[lengthOffset]
-                             | (this.Buffer[lengthOffset + 1] << 8);
+                var curLen = this.Buffer[lengthOffset] |
+                             (this.Buffer[lengthOffset + 1] << 8);
 
                 curLen -= amount;
-                this.Length -= amount;
 
                 this.Buffer[lengthOffset] = (byte)curLen;
                 this.Buffer[lengthOffset + 1] = (byte)(this.Buffer[lengthOffset + 1] >> 8);
diff --git a/src/Impostor.Tests/Hazel/MessageReaderTests.cs b/src/Impostor.Tests/Hazel/MessageReaderTests.cs
index a51b8bf..2241fb9 100644
--- a/src/Impostor.Tests/Hazel/MessageReaderTests.cs
+++ b/src/Impostor.Tests/Hazel/MessageReaderTests.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Linq;
 using Impostor.Hazel;
 using Impostor.Hazel.Extensions;
 using Microsoft.Extensions.DependencyInjection;
@@ -7,7 +8,7 @@ using Xunit;
 
 namespace Impostor.Tests.Hazel
 {
-    public class BufferMessageReaderTests
+    public class MessageReaderTests
     {
         private ObjectPool<MessageReader> CreateReaderPool()
         {
@@ -288,10 +289,10 @@ namespace Impostor.Tests.Hazel
 
             messageExpected.StartMessage(0);
             messageExpected.StartMessage(1);
-            messageExpected.Write("HiTest");
+            messageExpected.Write("HiTest1");
             messageExpected.EndMessage();
             messageExpected.StartMessage(2);
-            messageExpected.Write("HiTest");
+            messageExpected.Write("HiTest2");
             messageExpected.EndMessage();
             messageExpected.EndMessage();
 
@@ -310,10 +311,14 @@ namespace Impostor.Tests.Hazel
                 messageWriter.EndMessage();
             messageWriter.EndMessage();
 
+            // Copy buffer.
+            var bufferCopy = new byte[messageWriter.Length];
+            Buffer.BlockCopy(messageWriter.Buffer, 0, bufferCopy, 0, bufferCopy.Length);
+
             // Do the magic.
             var readerPool = CreateReaderPool();
             var reader = readerPool.Get();
-            reader.Update(messageWriter.Buffer);
+            reader.Update(bufferCopy);
             var inner = reader.ReadMessage();
 
             while (inner.Position < inner.Length)
@@ -336,10 +341,15 @@ namespace Impostor.Tests.Hazel
                 {
                     Assert.Equal("HiTest2", message.ReadString());
                 }
+                else
+                {
+                    Assert.True(false, "Invalid tag was read.");
+                }
             }
 
             // Check if the magic was successful.
-            Assert.Equal(messageExpected.ToByteArray(true), messageWriter.ToByteArray(true));
+            Assert.Equal(messageExpected.Length, reader.Length);
+            Assert.Equal(messageExpected.ToByteArray(true), reader.Buffer.Take(reader.Length).ToArray());
         }
 
         [Fact]
-- 
2.39.5