From a8d388840967652c06c88694c86e637760515b82 Mon Sep 17 00:00:00 2001 From: Forest Date: Sat, 31 Oct 2020 15:04:38 -0700 Subject: [PATCH] Add functionality to rewrite packets --- Hazel.UnitTests/MessageReaderTests.cs | 86 +++++++++++++++++++ .../ThreadLimitedUdpConnectionTests.cs | 30 ------- .../ThreadLimitedUdpConnectionListener.cs | 29 ++----- Hazel/MessageReader.cs | 85 ++++++++++++++++-- Hazel/MessageWriter.cs | 7 +- 5 files changed, 176 insertions(+), 61 deletions(-) diff --git a/Hazel.UnitTests/MessageReaderTests.cs b/Hazel.UnitTests/MessageReaderTests.cs index ad713ba..ee3b427 100644 --- a/Hazel.UnitTests/MessageReaderTests.cs +++ b/Hazel.UnitTests/MessageReaderTests.cs @@ -88,7 +88,61 @@ namespace Hazel.UnitTests MessageReader reader = MessageReader.Get(msg.Buffer, 0); Assert.AreEqual(Test1, reader.ReadSingle()); + } + + [TestMethod] + public void RemoveMessageWorks() + { + const byte Test0 = 11; + const byte Test3 = 33; + const byte Test4 = 44; + const byte Test5 = 55; + + var msg = new MessageWriter(2048); + msg.StartMessage(0); + msg.Write(Test0); + msg.EndMessage(); + + msg.StartMessage(12); + msg.StartMessage(23); + + msg.StartMessage(34); + msg.Write(Test3); + msg.EndMessage(); + + msg.StartMessage(45); + msg.Write(Test4); + msg.EndMessage(); + msg.EndMessage(); + msg.EndMessage(); + + msg.StartMessage(56); + msg.Write(Test5); + msg.EndMessage(); + + MessageReader reader = MessageReader.Get(msg.Buffer); + reader.Length = msg.Length; + + var zero = reader.ReadMessage(); + + var one = reader.ReadMessage(); + var two = one.ReadMessage(); + var three = two.ReadMessage(); + two.RemoveMessage(three); + + // Reader becomes invalid + Assert.AreNotEqual(Test3, three.ReadByte()); + + // Unrealistic, but nice. Earlier data is not affected + Assert.AreEqual(Test0, zero.ReadByte()); + + // Continuing to read depth-first works + var four = two.ReadMessage(); + Assert.AreEqual(Test4, four.ReadByte()); + + var five = reader.ReadMessage(); + Assert.AreEqual(Test5, five.ReadByte()); } [TestMethod] @@ -163,6 +217,38 @@ namespace Hazel.UnitTests Assert.AreEqual(2, sub.Tag); } + [TestMethod] + public void ReadMessageAsNewBufferLength() + { + var msg = new MessageWriter(2048); + msg.StartMessage(1); + msg.Write(65534); + msg.StartMessage(2); + msg.Write("HO"); + msg.EndMessage(); + msg.StartMessage(232); + msg.EndMessage(); + msg.EndMessage(); + + Assert.AreEqual(msg.Length, msg.Position); + + MessageReader reader = MessageReader.Get(msg.Buffer, 0); + Assert.AreEqual(1, reader.Tag); + Assert.AreEqual(65534, reader.ReadInt32()); // Content + + var sub = reader.ReadMessageAsNewBuffer(); + Assert.AreEqual(3, sub.Length); + Assert.AreEqual(2, sub.Tag); + Assert.AreEqual("HO", sub.ReadString()); + + sub.Recycle(); + + sub = reader.ReadMessageAsNewBuffer(); + Assert.AreEqual(0, sub.Length); + Assert.AreEqual(232, sub.Tag); + sub.Recycle(); + } + [TestMethod] public void ReadStringProtectsAgainstOverrun() { diff --git a/Hazel.UnitTests/ThreadLimitedUdpConnectionTests.cs b/Hazel.UnitTests/ThreadLimitedUdpConnectionTests.cs index 42b7180..b9a3a0a 100644 --- a/Hazel.UnitTests/ThreadLimitedUdpConnectionTests.cs +++ b/Hazel.UnitTests/ThreadLimitedUdpConnectionTests.cs @@ -309,36 +309,6 @@ namespace Hazel.UnitTests } } - /// - /// Tests the keepalive functionality from the client, - /// - [TestMethod] - public void PingDisconnectClientTest() - { -#if DEBUG - using (ThreadLimitedUdpConnectionListener listener = new ThreadLimitedUdpConnectionListener(2, new IPEndPoint(IPAddress.Any, 4296), new NullLogger())) - using (UdpConnection connection = new UdpClientConnection(new IPEndPoint(IPAddress.Loopback, 4296))) - { - listener.Start(); - - connection.Connect(); - - // After connecting, quietly stop responding to all messages to fake connection loss. - Thread.Sleep(10); - // listener.TestDropRate = 1; - - connection.KeepAliveInterval = 100; - - Thread.Sleep(1050); //Enough time for ~10 keep alive packets - - Assert.AreEqual(ConnectionState.NotConnected, connection.State); - Assert.AreEqual(3 * connection.MissingPingsUntilDisconnect + 4, connection.Statistics.TotalBytesSent); // + 4 for connecting overhead - } -#else - Assert.Inconclusive("Only works in DEBUG"); -#endif - } - /// /// Tests the keepalive functionality from the client, /// diff --git a/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs b/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs index 78c061c..9e227fb 100644 --- a/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs +++ b/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs @@ -52,7 +52,7 @@ namespace Hazel.Udp.FewerThreads private ConcurrentDictionary allConnections = new ConcurrentDictionary(); private BlockingCollection receiveQueue; - private Queue sendQueue = new Queue(); + private BlockingCollection sendQueue = new BlockingCollection(); public int MaxAge { @@ -71,7 +71,7 @@ namespace Hazel.Udp.FewerThreads } public int ConnectionCount { get { return this.allConnections.Count; } } - public int SendQueueLength { get { lock(this.sendQueue) return this.sendQueue.Count; } } + public int SendQueueLength { get { return this.sendQueue.Count; } } public int ReceiveQueueLength { get { return this.receiveQueue.Count; } } private bool isActive; @@ -197,21 +197,7 @@ namespace Hazel.Udp.FewerThreads { while (this.isActive) { - SendMessageInfo msg; - lock (this.sendQueue) - { - if (this.sendQueue.Count == 0) - { - Monitor.Wait(this.sendQueue); - - if (this.sendQueue.Count == 0) - { - continue; - } - } - - msg = this.sendQueue.Dequeue(); - } + SendMessageInfo msg = this.sendQueue.Take(); try { @@ -297,11 +283,7 @@ namespace Hazel.Udp.FewerThreads internal void SendDataRaw(byte[] response, EndPoint remoteEndPoint) { - lock (this.sendQueue) - { - this.sendQueue.Enqueue(new SendMessageInfo() { Buffer = response, Recipient = remoteEndPoint }); - Monitor.Pulse(this.sendQueue); - } + this.sendQueue.Add(new SendMessageInfo() { Buffer = response, Recipient = remoteEndPoint }); } /// @@ -326,9 +308,8 @@ namespace Hazel.Udp.FewerThreads this.isActive = false; - lock (this.sendQueue) Monitor.PulseAll(this.sendQueue); - this.receiveQueue.CompleteAdding(); + this.sendQueue.CompleteAdding(); this.reliablePacketThread.Join(); this.sendThread.Join(); diff --git a/Hazel/MessageReader.cs b/Hazel/MessageReader.cs index cc3ef2a..43ed4f0 100644 --- a/Hazel/MessageReader.cs +++ b/Hazel/MessageReader.cs @@ -1,5 +1,6 @@ using System; using System.IO; +using System.Linq; using System.Runtime.CompilerServices; using System.Text; @@ -17,6 +18,8 @@ namespace Hazel public int BytesRemaining => this.Length - this.Position; + private MessageReader Parent; + public int Position { get { return this._position; } @@ -29,7 +32,7 @@ namespace Hazel private int _position; private int readHead; - + public static MessageReader GetSized(int minSize) { var output = ReaderPool.GetObject(); @@ -39,10 +42,11 @@ namespace Hazel } output.Offset = 0; + output.Position = 0; output.Tag = byte.MaxValue; return output; } - + public static MessageReader Get(byte[] buffer) { var output = ReaderPool.GetObject(); @@ -52,7 +56,7 @@ namespace Hazel output.Position = 0; output.Length = buffer.Length; output.Tag = byte.MaxValue; - + return output; } @@ -64,7 +68,7 @@ namespace Hazel output.Offset = 0; output.Position = 0; output.Length = source.Length + 3; - + return output; } @@ -104,7 +108,9 @@ namespace Hazel return output; } - /// + /// + /// Produces a MessageReader using the parent's buffer. This MessageReader should **NOT** be recycled. + /// public MessageReader ReadMessage() { // Ensure there is at least a header @@ -112,6 +118,7 @@ namespace Hazel var output = new MessageReader(); + output.Parent = this; output.Buffer = this.Buffer; output.Offset = this.readHead; output.Position = 0; @@ -128,8 +135,76 @@ namespace Hazel return output; } + /// + /// Produces a MessageReader with a new buffer. This MessageReader should be recycled. + /// + public MessageReader ReadMessageAsNewBuffer() + { + if (this.BytesRemaining < 3) throw new InvalidDataException($"ReadMessage header is longer than message length: 3 of {this.BytesRemaining}"); + + var len = this.ReadUInt16(); + var tag = this.ReadByte(); + + if (this.BytesRemaining < len) throw new InvalidDataException($"Message length is longer than message length: {len} of {this.BytesRemaining}"); + + var output = MessageReader.GetSized(len); + + output.Parent = this; + Array.Copy(this.Buffer, this.readHead, output.Buffer, 0, len); + + output.Length = len; + output.Tag = tag; + + this.Position += output.Length; + return output; + } + + public void RemoveMessage(MessageReader reader) + { + var temp = MessageReader.GetSized(reader.Buffer.Length); + try + { + var headerOffset = reader.Offset - 3; + var endOfMessage = reader.Offset + reader.Length; + var len = reader.Buffer.Length - endOfMessage; + + Array.Copy(reader.Buffer, endOfMessage, temp.Buffer, 0, len); + Array.Copy(temp.Buffer, 0, this.Buffer, headerOffset, len); + + this.AdjustLength(reader.Offset, reader.Length + 3); + } + finally + { + temp.Recycle(); + } + } + + private void AdjustLength(int offset, int amount) + { + if (this.readHead > 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 void Recycle() { + this.Parent = null; ReaderPool.PutObject(this); } diff --git a/Hazel/MessageWriter.cs b/Hazel/MessageWriter.cs index 73a6efd..0c94643 100644 --- a/Hazel/MessageWriter.cs +++ b/Hazel/MessageWriter.cs @@ -84,8 +84,11 @@ namespace Hazel /// public void StartMessage(byte typeFlag) { - messageStarts.Push(this.Position); - this.Position += 2; // Skip for size + var messageStart = this.Position; + messageStarts.Push(messageStart); + this.Buffer[messageStart] = 0; + this.Buffer[messageStart + 1] = 0; + this.Position += 2; this.Write(typeFlag); } -- 2.39.5