Add extra protection to removing messages

diff --git a/src/Impostor.Api/Exceptions/ImpostorProtocolException.cs b/src/Impostor.Api/Exceptions/ImpostorProtocolException.cs
new file mode 100644 (file)
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 d4668db771d3e6778395b03c1d463d085f90dec1..97996c67ccf14e7669e57fdea1a0fbf00ec6c852 100644 (file)
--- 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 2241fb90f65e16d14d6fed548473527681f4148b..e7fa0dff84fd0ed7e02196e5184fcf5777cdb448 100644 (file)
--- 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]