From 272230e4fffeed2f65c34f42a2e3d69b61100d5b Mon Sep 17 00:00:00 2001
From: AeonLucid <aeonlucid@outlook.com>
Date: Sun, 1 Nov 2020 23:57:50 +0100
Subject: [PATCH] Add extra protection to removing messages

---
 .../Exceptions/ImpostorProtocolException.cs   | 24 +++++++++++++++++++
 src/Impostor.Hazel/MessageReader.cs           |  6 +++++
 .../Hazel/MessageReaderTests.cs               | 11 +++++++++
 3 files changed, 41 insertions(+)
 create mode 100644 src/Impostor.Api/Exceptions/ImpostorProtocolException.cs

diff --git a/src/Impostor.Api/Exceptions/ImpostorProtocolException.cs b/src/Impostor.Api/Exceptions/ImpostorProtocolException.cs
new file mode 100644
index 0000000..864602d
--- /dev/null
+++ b/src/Impostor.Api/Exceptions/ImpostorProtocolException.cs
@@ -0,0 +1,24 @@
+using System;
+using System.Runtime.Serialization;
+
+namespace Impostor.Api
+{
+    public class ImpostorProtocolException : ImpostorException
+    {
+        public ImpostorProtocolException()
+        {
+        }
+
+        protected ImpostorProtocolException(SerializationInfo info, StreamingContext context) : base(info, context)
+        {
+        }
+
+        public ImpostorProtocolException(string? message) : base(message)
+        {
+        }
+
+        public ImpostorProtocolException(string? message, Exception? innerException) : base(message, innerException)
+        {
+        }
+    }
+}
diff --git a/src/Impostor.Hazel/MessageReader.cs b/src/Impostor.Hazel/MessageReader.cs
index d4668db..97996c6 100644
--- a/src/Impostor.Hazel/MessageReader.cs
+++ b/src/Impostor.Hazel/MessageReader.cs
@@ -3,6 +3,7 @@ using System.Buffers;
 using System.Buffers.Binary;
 using System.Runtime.CompilerServices;
 using System.Text;
+using Impostor.Api;
 using Impostor.Api.Net.Messages;
 using Microsoft.Extensions.ObjectPool;
 
@@ -188,6 +189,11 @@ namespace Impostor.Hazel
 
         public void RemoveMessage(IMessageReader message)
         {
+            if (message.Buffer != Buffer)
+            {
+                throw new ImpostorProtocolException("Tried to remove message from a message that does not have the same buffer.");
+            }
+
             // Offset of where to start removing.
             var offsetStart = message.Offset - 3;
 
diff --git a/src/Impostor.Tests/Hazel/MessageReaderTests.cs b/src/Impostor.Tests/Hazel/MessageReaderTests.cs
index 2241fb9..e7fa0df 100644
--- a/src/Impostor.Tests/Hazel/MessageReaderTests.cs
+++ b/src/Impostor.Tests/Hazel/MessageReaderTests.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Linq;
+using Impostor.Api;
 using Impostor.Hazel;
 using Impostor.Hazel.Extensions;
 using Microsoft.Extensions.DependencyInjection;
@@ -315,6 +316,9 @@ namespace Impostor.Tests.Hazel
             var bufferCopy = new byte[messageWriter.Length];
             Buffer.BlockCopy(messageWriter.Buffer, 0, bufferCopy, 0, bufferCopy.Length);
 
+            var bufferCopyTwo = new byte[messageWriter.Length];
+            Buffer.BlockCopy(messageWriter.Buffer, 0, bufferCopyTwo, 0, bufferCopyTwo.Length);
+
             // Do the magic.
             var readerPool = CreateReaderPool();
             var reader = readerPool.Get();
@@ -350,6 +354,13 @@ namespace Impostor.Tests.Hazel
             // Check if the magic was successful.
             Assert.Equal(messageExpected.Length, reader.Length);
             Assert.Equal(messageExpected.ToByteArray(true), reader.Buffer.Take(reader.Length).ToArray());
+
+            // Test ownership.
+            var readerTwo = readerPool.Get();
+
+            readerTwo.Update(bufferCopyTwo);
+
+            Assert.Throws<ImpostorProtocolException>(() => reader.RemoveMessage(readerTwo.ReadMessage()));
         }
 
         [Fact]
-- 
2.39.5