From 66ec3847125bd0db64ea48d13300c59d12e46b47 Mon Sep 17 00:00:00 2001
From: AeonLucid <aeonlucid@outlook.com>
Date: Sun, 1 Nov 2020 20:26:19 +0100
Subject: [PATCH] Fix stack overflow on CTRL+C

---
 .../Tests/MessageReaderBenchmark.cs                    |  4 ++--
 src/Impostor.Hazel/MessageReader.cs                    | 10 +++++++++-
 src/Impostor.Hazel/MessageReaderPolicy.cs              |  2 +-
 src/Impostor.Hazel/Udp/UdpConnection.cs                |  5 +++--
 4 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/Impostor.Benchmarks/Tests/MessageReaderBenchmark.cs b/src/Impostor.Benchmarks/Tests/MessageReaderBenchmark.cs
index c8ad228..abf1642 100644
--- a/src/Impostor.Benchmarks/Tests/MessageReaderBenchmark.cs
+++ b/src/Impostor.Benchmarks/Tests/MessageReaderBenchmark.cs
@@ -14,7 +14,7 @@ namespace Impostor.Benchmarks.Tests
     public class MessageReaderBenchmark
     {
         private byte[] _data;
-        private Microsoft.Extensions.ObjectPool.ObjectPool<MessageReader_Bytes_Pooled_Improved> _pool;
+        private ObjectPool<MessageReader_Bytes_Pooled_Improved> _pool;
 
         [GlobalSetup]
         public void Setup()
@@ -45,7 +45,7 @@ namespace Impostor.Benchmarks.Tests
 
             _pool = services
                 .BuildServiceProvider()
-                .GetRequiredService<Microsoft.Extensions.ObjectPool.ObjectPool<MessageReader_Bytes_Pooled_Improved>>();
+                .GetRequiredService<ObjectPool<MessageReader_Bytes_Pooled_Improved>>();
         }
 
         [Benchmark]
diff --git a/src/Impostor.Hazel/MessageReader.cs b/src/Impostor.Hazel/MessageReader.cs
index df6d014..ab6aea8 100644
--- a/src/Impostor.Hazel/MessageReader.cs
+++ b/src/Impostor.Hazel/MessageReader.cs
@@ -10,6 +10,7 @@ namespace Impostor.Hazel
     public class MessageReader : IMessageReader
     {
         private readonly ObjectPool<MessageReader> _pool;
+        private bool _inUse;
 
         public byte Tag { get; private set; }
         public ReadOnlyMemory<byte> Buffer { get; private set; }
@@ -28,6 +29,8 @@ namespace Impostor.Hazel
 
         public void Update(byte tag, ReadOnlyMemory<byte> buffer)
         {
+            _inUse = true;
+
             Tag = tag;
             Buffer = buffer;
             Position = 0;
@@ -35,6 +38,8 @@ namespace Impostor.Hazel
 
         internal void Reset()
         {
+            _inUse = false;
+
             Tag = byte.MaxValue;
             Buffer = null;
             Position = 0;
@@ -185,7 +190,10 @@ namespace Impostor.Hazel
 
         public void Dispose()
         {
-            _pool.Return(this);
+            if (_inUse)
+            {
+                _pool.Return(this);
+            }
         }
     }
 }
diff --git a/src/Impostor.Hazel/MessageReaderPolicy.cs b/src/Impostor.Hazel/MessageReaderPolicy.cs
index 995ad78..ef3939a 100644
--- a/src/Impostor.Hazel/MessageReaderPolicy.cs
+++ b/src/Impostor.Hazel/MessageReaderPolicy.cs
@@ -15,7 +15,7 @@ namespace Impostor.Hazel
 
         public MessageReader Create()
         {
-            return new MessageReader(_serviceProvider.GetRequiredService<Microsoft.Extensions.ObjectPool.ObjectPool<MessageReader>>());
+            return new MessageReader(_serviceProvider.GetRequiredService<ObjectPool<MessageReader>>());
         }
 
         public bool Return(MessageReader obj)
diff --git a/src/Impostor.Hazel/Udp/UdpConnection.cs b/src/Impostor.Hazel/Udp/UdpConnection.cs
index 4fffb30..46ad376 100644
--- a/src/Impostor.Hazel/Udp/UdpConnection.cs
+++ b/src/Impostor.Hazel/Udp/UdpConnection.cs
@@ -3,6 +3,7 @@ using System.Threading;
 using System.Threading.Channels;
 using System.Threading.Tasks;
 using Impostor.Api.Net.Messages;
+using Microsoft.Extensions.ObjectPool;
 using Serilog;
 
 namespace Impostor.Hazel.Udp
@@ -17,14 +18,14 @@ namespace Impostor.Hazel.Udp
 
         private static readonly ILogger Logger = Log.ForContext<UdpConnection>();
         private readonly ConnectionListener _listener;
-        private readonly Microsoft.Extensions.ObjectPool.ObjectPool<MessageReader> _readerPool;
+        private readonly ObjectPool<MessageReader> _readerPool;
         private readonly CancellationTokenSource _stoppingCts;
 
         private bool _isDisposing;
         private bool _isFirst = true;
         private Task _executingTask;
 
-        protected UdpConnection(ConnectionListener listener, Microsoft.Extensions.ObjectPool.ObjectPool<MessageReader> readerPool)
+        protected UdpConnection(ConnectionListener listener, ObjectPool<MessageReader> readerPool)
         {
             _listener = listener;
             _readerPool = readerPool;
-- 
2.39.5