From db89da2f435896e75de64eb932744c1e37160718 Mon Sep 17 00:00:00 2001
From: Forest <chocozilla@gmail.com>
Date: Mon, 7 Jun 2021 17:05:03 -0700
Subject: [PATCH] Fix some bounds checking while parsing ClientHellos and
 improve or disable some flaky tests

---
 Hazel.UnitTests/UPnPTests.cs               |  3 ++-
 Hazel.UnitTests/UdpConnectionTests.cs      | 13 ++++++++++---
 Hazel.UnitTests/UnityUdpConnectionTests.cs | 13 ++++++++++---
 Hazel/Dtls/Handshake.cs                    | 16 ++++++++++++----
 4 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/Hazel.UnitTests/UPnPTests.cs b/Hazel.UnitTests/UPnPTests.cs
index 44f6aea..657c2e7 100644
--- a/Hazel.UnitTests/UPnPTests.cs
+++ b/Hazel.UnitTests/UPnPTests.cs
@@ -4,7 +4,8 @@ using Microsoft.VisualStudio.TestTools.UnitTesting;
 
 namespace Hazel.UnitTests
 {
-    [TestClass]
+    // [TestClass]
+    // TODO: These tests are super flaky because of hardware differences. Not sure what can be done.
     public class UPnPTests
     {
         [TestMethod]
diff --git a/Hazel.UnitTests/UdpConnectionTests.cs b/Hazel.UnitTests/UdpConnectionTests.cs
index 1b3fb27..79e9ae9 100644
--- a/Hazel.UnitTests/UdpConnectionTests.cs
+++ b/Hazel.UnitTests/UdpConnectionTests.cs
@@ -4,6 +4,7 @@ using System.Net;
 using System.Threading;
 using Hazel.Udp;
 using System.Net.Sockets;
+using System.Threading.Tasks;
 
 namespace Hazel.UnitTests
 {
@@ -489,9 +490,15 @@ namespace Hazel.UnitTests
 
                 listener.NewConnection += delegate (NewConnectionEventArgs args)
                 {
-                    MessageWriter writer = MessageWriter.Get(SendOption.None);
-                    writer.Write("Goodbye");
-                    args.Connection.Disconnect("Testing", writer);
+                    // As it turns out, the UdpConnectionListener can have an issue on loopback where the disconnect can happen before the hello confirm
+                    // Tossing it on a different thread makes this test more reliable. Perhaps something to think about elsewhere though.
+                    Task.Run(async () =>
+                    {
+                        await Task.Delay(1);
+                        MessageWriter writer = MessageWriter.Get(SendOption.None);
+                        writer.Write("Goodbye");
+                        args.Connection.Disconnect("Testing", writer);
+                    });
                 };
 
                 listener.Start();
diff --git a/Hazel.UnitTests/UnityUdpConnectionTests.cs b/Hazel.UnitTests/UnityUdpConnectionTests.cs
index 83f62fc..0597df4 100644
--- a/Hazel.UnitTests/UnityUdpConnectionTests.cs
+++ b/Hazel.UnitTests/UnityUdpConnectionTests.cs
@@ -4,6 +4,7 @@ using System.Net;
 using System.Threading;
 using Hazel.Udp;
 using System.Net.Sockets;
+using System.Threading.Tasks;
 
 namespace Hazel.UnitTests
 {
@@ -459,9 +460,15 @@ namespace Hazel.UnitTests
 
                 listener.NewConnection += delegate (NewConnectionEventArgs args)
                 {
-                    MessageWriter writer = MessageWriter.Get(SendOption.None);
-                    writer.Write("Goodbye");
-                    args.Connection.Disconnect("Testing", writer);
+                    // As it turns out, the UdpConnectionListener can have an issue on loopback where the disconnect can happen before the hello confirm
+                    // Tossing it on a different thread makes this test more reliable. Perhaps something to think about elsewhere though.
+                    Task.Run(async () =>
+                    {
+                        await Task.Delay(1);
+                        MessageWriter writer = MessageWriter.Get(SendOption.None);
+                        writer.Write("Goodbye");
+                        args.Connection.Disconnect("Testing", writer);
+                    });
                 };
 
                 listener.Start();
diff --git a/Hazel/Dtls/Handshake.cs b/Hazel/Dtls/Handshake.cs
index d8bde25..b0fe376 100644
--- a/Hazel/Dtls/Handshake.cs
+++ b/Hazel/Dtls/Handshake.cs
@@ -246,16 +246,23 @@ namespace Hazel.Dtls
                     break;
                 }
             }
-            span = span.Slice(1 + compressionMethodsSize);
 
-            if (!foundNullCompressionMethod)
+            if (!foundNullCompressionMethod
+                || span.Length < 1 + compressionMethodsSize)
             {
                 return false;
             }
 
+            span = span.Slice(1 + compressionMethodsSize);
+
             // Parse extensions
             if (span.Length > 0)
             {
+                if (span.Length < 2)
+                {
+                    return false;
+                }
+
                 ushort extensionsSize = span.ReadBigEndian16();
                 span = span.Slice(2);
                 if (span.Length != extensionsSize)
@@ -273,12 +280,13 @@ namespace Hazel.Dtls
 
                     ExtensionType extensionType = (ExtensionType)span.ReadBigEndian16(0);
                     ushort extensionLength = span.ReadBigEndian16(2);
-                    ByteSpan extensionData = span.Slice(4, extensionLength);
-                    if (extensionData.Length < extensionLength)
+
+                    if (span.Length < 4 + extensionLength)
                     {
                         return false;
                     }
 
+                    ByteSpan extensionData = span.Slice(4, extensionLength);
                     span = span.Slice(4 + extensionLength);
                     result.ParseExtension(extensionType, extensionData);
                 }
-- 
2.39.5