From 4cb2813116be3279f9abb1b56881f2f1bd15d8a9 Mon Sep 17 00:00:00 2001
From: Chris Dimitrov <ChrisDimitrov95@gmail.com>
Date: Sun, 13 Jan 2019 21:23:06 +0200
Subject: [PATCH] Added functionality to buffer packets before the initial
 packet arrives

---
 .../Settings/QuicSettings.cs                  |  5 +++
 QuicNet/Buffers/ConnectionBuffer.cs           | 38 +++++++++++++++++
 QuicNet/QuicListener.cs                       | 42 ++++++++++++++++++-
 QuicNet/QuicNet.csproj                        |  1 +
 4 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 QuicNet/Buffers/ConnectionBuffer.cs

diff --git a/QuicNet.Infrastructure/Settings/QuicSettings.cs b/QuicNet.Infrastructure/Settings/QuicSettings.cs
index f52fe23..8b257aa 100644
--- a/QuicNet.Infrastructure/Settings/QuicSettings.cs
+++ b/QuicNet.Infrastructure/Settings/QuicSettings.cs
@@ -37,6 +37,11 @@ public class QuicSettings
         /// </summary>
         public const bool ShouldBufferPacketsBeforeConnection = false;
 
+        /// <summary>
+        /// Time before the buffer expires. 30 seconds.
+        /// </summary>
+        public const double BufferTimeInverval = 30000;
+
         /// <summary>
         /// Limit the maximum number of frames a packet can carry.
         /// </summary>
diff --git a/QuicNet/Buffers/ConnectionBuffer.cs b/QuicNet/Buffers/ConnectionBuffer.cs
new file mode 100644
index 0000000..b0022ed
--- /dev/null
+++ b/QuicNet/Buffers/ConnectionBuffer.cs
@@ -0,0 +1,38 @@
+using QuicNet.Infrastructure.Packets;
+using QuicNet.Infrastructure.Settings;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using System.Timers;
+
+namespace QuicNet.Buffers
+{
+    public class ConnectionBuffer
+    {
+        private Queue<Packet> _queue;
+        private Timer _expiry;
+
+        public ConnectionBuffer()
+        {
+            _queue = new Queue<Packet>();
+            _expiry = new Timer(QuicSettings.BufferTimeInverval);
+
+            _expiry.Start();
+        }
+
+        public void Push(Packet packet)
+        {
+            _queue.Enqueue(packet);
+        }
+
+        public List<Packet> ReadAll()
+        {
+            List<Packet> result = new List<Packet>();
+            result = _queue.ToList();
+
+            return result;
+        }
+    }
+}
diff --git a/QuicNet/QuicListener.cs b/QuicNet/QuicListener.cs
index c7d233c..c2c7a21 100644
--- a/QuicNet/QuicListener.cs
+++ b/QuicNet/QuicListener.cs
@@ -1,4 +1,5 @@
 using QuickNet.Utilities;
+using QuicNet.Buffers;
 using QuicNet.Connections;
 using QuicNet.Context;
 using QuicNet.Exceptions;
@@ -27,6 +28,7 @@ public class QuicListener
 
         private Unpacker _unpacker;
         private InitialPacketCreator _packetCreator;
+        private Dictionary<UInt64, ConnectionBuffer> _buffers;
 
         private PacketWireTransfer _pwt;
 
@@ -44,6 +46,7 @@ public QuicListener(int port)
             
             _unpacker = new Unpacker();
             _packetCreator = new InitialPacketCreator();
+            _buffers = new Dictionary<ulong, ConnectionBuffer>();
         }
 
         /// <summary>
@@ -83,6 +86,11 @@ public QuicConnection AcceptQuicClient()
                 if (packet is InitialPacket)
                 {
                     QuicConnection connection = ProcessInitialPacket(packet, _pwt.LastTransferEndpoint());
+
+                    // If the connection has been successfully created, replay buffer.
+                    if (connection != null && _buffers.ContainsKey(connection.ConnectionId))
+                        ReplayBuffer(connection.ConnectionId);
+
                     return connection;
                 }
 
@@ -186,11 +194,43 @@ private void ProcessShortHeaderPacket(Packet packet)
 
             QuicConnection connection = ConnectionPool.Find(shp.DestinationConnectionId);
 
-            // No suitable connection found. Discard the packet.
+            // No suitable connection found. Buffer the packet.
             if (connection == null)
+            {
+                BufferPacket(shp.DestinationConnectionId, packet);
+
                 return;
+            }
 
             connection.ProcessFrames(shp.GetFrames());
         }
+
+        private void BufferPacket(UInt64 connectionId, Packet packet)
+        {
+            if (_buffers.ContainsKey(connectionId) == false)
+            {
+                ConnectionBuffer buffer = new ConnectionBuffer();
+                buffer.Push(packet);
+
+                _buffers.Add(connectionId, buffer);
+            }
+            else
+            {
+                ConnectionBuffer buffer = _buffers[connectionId];
+                buffer.Push(packet);
+            }
+        }
+
+        private void ReplayBuffer(UInt64 connectionId)
+        {
+            ConnectionBuffer buffer = _buffers[connectionId];
+            List<Packet> packets = buffer.ReadAll();
+
+            // Replay
+            foreach (Packet bufferedPacket in packets)
+            {
+                OrchestratePacket(bufferedPacket);
+            }
+        }
     }
 }
diff --git a/QuicNet/QuicNet.csproj b/QuicNet/QuicNet.csproj
index 7b0431c..6c70e75 100644
--- a/QuicNet/QuicNet.csproj
+++ b/QuicNet/QuicNet.csproj
@@ -40,6 +40,7 @@
     <Reference Include="System.Xml" />
   </ItemGroup>
   <ItemGroup>
+    <Compile Include="Buffers\ConnectionBuffer.cs" />
     <Compile Include="Connections\ConnectionPool.cs" />
     <Compile Include="Connections\ConnectionState.cs" />
     <Compile Include="Connections\QuicConnection.cs" />