From 14fc2bc5039e0cb40cb256bf4687605b9071e4b3 Mon Sep 17 00:00:00 2001
From: AeonLucid <aeonlucid@outlook.com>
Date: Sun, 1 Nov 2020 18:50:14 +0100
Subject: [PATCH] RemoveMessage progress

---
 .../Net/Messages/IMessageReader.cs            |  2 +
 src/Impostor.Hazel/MessageReader.cs           | 57 +++++++++++++++++--
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/src/Impostor.Api/Net/Messages/IMessageReader.cs b/src/Impostor.Api/Net/Messages/IMessageReader.cs
index d3320f3..87c06c4 100644
--- a/src/Impostor.Api/Net/Messages/IMessageReader.cs
+++ b/src/Impostor.Api/Net/Messages/IMessageReader.cs
@@ -61,6 +61,8 @@ namespace Impostor.Api.Net.Messages
 
         void Seek(int position);
 
+        void RemoveMessage(IMessageReader message);
+
         IMessageReader Copy(int offset = 0);
     }
 }
diff --git a/src/Impostor.Hazel/MessageReader.cs b/src/Impostor.Hazel/MessageReader.cs
index 82c998e..e615503 100644
--- a/src/Impostor.Hazel/MessageReader.cs
+++ b/src/Impostor.Hazel/MessageReader.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Buffers;
 using System.Buffers.Binary;
 using System.Runtime.CompilerServices;
 using System.Text;
@@ -9,11 +10,11 @@ namespace Impostor.Hazel
 {
     public class MessageReader : IMessageReader
     {
+        private static readonly ArrayPool<byte> ArrayPool = ArrayPool<byte>.Shared;
+
         private readonly ObjectPool<MessageReader> _pool;
         private bool _inUse;
 
-        private byte _tag;
-
         internal MessageReader(ObjectPool<MessageReader> pool)
         {
             _pool = pool;
@@ -29,9 +30,11 @@ namespace Impostor.Hazel
 
         public byte Tag { get; private set; }
 
+        public MessageReader Parent { get; private set; }
+
         private int ReadPosition => Offset + Position;
 
-        public void Update(byte[] buffer, int offset = 0, int position = 0, int? length = null, byte tag = byte.MaxValue)
+        public void Update(byte[] buffer, int offset = 0, int position = 0, int? length = null, byte tag = byte.MaxValue, MessageReader parent = null)
         {
             _inUse = true;
 
@@ -40,6 +43,7 @@ namespace Impostor.Hazel
             Position = position;
             Length = length ?? buffer.Length;
             Tag = tag;
+            Parent = parent;
         }
 
         internal void Reset()
@@ -51,6 +55,7 @@ namespace Impostor.Hazel
             Offset = 0;
             Position = 0;
             Length = 0;
+            Parent = null;
         }
 
         public IMessageReader ReadMessage()
@@ -62,7 +67,7 @@ namespace Impostor.Hazel
             Position += length;
 
             var reader = _pool.Get();
-            reader.Update(Buffer, pos, 0, length, tag);
+            reader.Update(Buffer, pos, 0, length, tag, this);
             return reader;
         }
 
@@ -181,6 +186,50 @@ namespace Impostor.Hazel
             Position = position;
         }
 
+        public void RemoveMessage(IMessageReader message)
+        {
+            var pool = ArrayPool.Rent(message.Buffer.Length);
+
+            try
+            {
+                var offsetHeader = message.Offset - 3;
+                var offsetEnd = message.Offset + message.Length;
+                var len = message.Buffer.Length - offsetEnd;
+
+                Array.Copy(message.Buffer, offsetEnd, pool, 0, len);
+                Array.Copy(pool, 0, this.Buffer, offsetHeader, len);
+
+                AdjustLength(message.Offset, message.Length + 3);
+            }
+            finally
+            {
+                ArrayPool.Return(pool);
+            }
+        }
+
+        private void AdjustLength(int offset, int amount)
+        {
+            if (this.ReadPosition > offset)
+            {
+                this.Position -= amount;
+            }
+
+            if (Parent != null)
+            {
+                var lengthOffset = this.Offset - 3;
+                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);
+
+                Parent.AdjustLength(offset, amount);
+            }
+        }
+
         public IMessageReader Copy(int offset = 0)
         {
             var reader = _pool.Get();
-- 
2.39.5