From 4b2bec46c202e904ff9814585b97a5bd8d3d4422 Mon Sep 17 00:00:00 2001
From: Matthew Endsley <mendsley@gmail.com>
Date: Fri, 5 Mar 2021 17:53:23 -0800
Subject: [PATCH] Remove DTLS peers that fail to negotiate a valid epoch

---
 Hazel/Dtls/DtlsConnectionListener.cs          | 27 +++++++++++++++++++
 .../ThreadLimitedUdpConnectionListener.cs     |  2 +-
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/Hazel/Dtls/DtlsConnectionListener.cs b/Hazel/Dtls/DtlsConnectionListener.cs
index 25794eb..32c7835 100644
--- a/Hazel/Dtls/DtlsConnectionListener.cs
+++ b/Hazel/Dtls/DtlsConnectionListener.cs
@@ -92,6 +92,8 @@ namespace Hazel.Dtls
 
             public readonly List<ByteSpan> QueuedApplicationDataMessage = new List<ByteSpan>();
 
+            public DateTime StartOfNegotiation;
+
             public PeerData()
             {
                 ByteSpan block = new byte[2 * Finished.Size];
@@ -126,6 +128,8 @@ namespace Hazel.Dtls
                 this.NextEpoch.ServerVerification = new byte[Finished.Size];
 
                 this.ConnectionId = connectionId;
+
+                this.StartOfNegotiation = DateTime.UtcNow;
             }
 
             public void Dispose()
@@ -1283,6 +1287,29 @@ namespace Hazel.Dtls
             }
         }
 
+        /// <inheritdoc />
+        public override void DisconnectOldConnections(TimeSpan maxAge, MessageWriter disconnectMessage)
+        {
+            DateTime now = DateTime.UtcNow;
+            foreach (KeyValuePair<IPEndPoint, PeerData> kvp in this.existingPeers)
+            {
+                PeerData peer = kvp.Value;
+                lock(peer)
+                {
+                    if (peer.Epoch == 0 || peer.NextEpoch.State != HandshakeState.ExpectingHello)
+                    {
+                        TimeSpan negotiationAge = now - peer.StartOfNegotiation;
+                        if (negotiationAge > maxAge)
+                        {
+                            base.MarkConnectionAsStale(peer.ConnectionId);
+                        }
+                    }
+                }
+            }
+
+            base.DisconnectOldConnections(maxAge, disconnectMessage);
+        }
+
         /// <summary>
         /// Allocate a new connection id
         /// </summary>
diff --git a/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs b/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs
index 84ac5a4..deda91d 100644
--- a/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs
+++ b/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs
@@ -151,7 +151,7 @@ namespace Hazel.Udp.FewerThreads
             }
         }
 
-        public void DisconnectOldConnections(TimeSpan maxAge, MessageWriter disconnectMessage)
+        public virtual void DisconnectOldConnections(TimeSpan maxAge, MessageWriter disconnectMessage)
         {
             var now = DateTime.UtcNow;
             foreach (var conn in this.allConnections.Values)
-- 
2.39.5