Cleanup around DtlsVerifier

2 files changed, 25 insertions, 15 deletions

diff --git a/crypto/test/src/tls/test/DtlsProtocolTest.cs b/crypto/test/src/tls/test/DtlsProtocolTest.cs
index 7fc49fb51..79b919999 100644
--- a/crypto/test/src/tls/test/DtlsProtocolTest.cs
+++ b/crypto/test/src/tls/test/DtlsProtocolTest.cs
@@ -5,6 +5,8 @@ using System.Threading;
 using NUnit.Framework;
 
 using Org.BouncyCastle.Security;
+using Org.BouncyCastle.Tls.Crypto;
+using Org.BouncyCastle.Tls.Crypto.Impl.BC;
 using Org.BouncyCastle.Utilities;
 
 namespace Org.BouncyCastle.Tls.Tests
@@ -70,47 +72,50 @@ namespace Org.BouncyCastle.Tls.Tests
             {
                 try
                 {
-                    MockDtlsServer server = new MockDtlsServer();
+                    TlsCrypto serverCrypto = new BcTlsCrypto();
 
                     DtlsRequest request = null;
 
                     // Use DtlsVerifier to require a HelloVerifyRequest cookie exchange before accepting
                     {
-                        DtlsVerifier verifier = new DtlsVerifier(server.Crypto);
+                        DtlsVerifier verifier = new DtlsVerifier(serverCrypto);
 
                         // NOTE: Test value only - would typically be the client IP address
                         byte[] clientID = Encoding.UTF8.GetBytes("MockDtlsClient");
 
                         int receiveLimit = m_serverTransport.GetReceiveLimit();
-                        int dummyOffset = server.Crypto.SecureRandom.Next(16) + 1;
-                        byte[] transportBuf = new byte[dummyOffset + m_serverTransport.GetReceiveLimit()];
+                        int dummyOffset = serverCrypto.SecureRandom.Next(16) + 1;
+                        byte[] buf = new byte[dummyOffset + m_serverTransport.GetReceiveLimit()];
 
                         do
                         {
                             if (m_isShutdown)
                                 return;
 
-                            int length = m_serverTransport.Receive(transportBuf, dummyOffset, receiveLimit, 1000);
+                            int length = m_serverTransport.Receive(buf, dummyOffset, receiveLimit, 1000);
                             if (length > 0)
                             {
-                                request = verifier.VerifyRequest(clientID, transportBuf, dummyOffset, length,
-                                    m_serverTransport);
+                                request = verifier.VerifyRequest(clientID, buf, dummyOffset, length, m_serverTransport);
                             }
                         }
                         while (request == null);
                     }
 
-                    DtlsTransport dtlsServer = m_serverProtocol.Accept(server, m_serverTransport, request);
-                    byte[] buf = new byte[dtlsServer.GetReceiveLimit()];
-                    while (!m_isShutdown)
+                    // NOTE: A real server would handle each DtlsRequest in a new task/thread and continue accepting
                     {
-                        int length = dtlsServer.Receive(buf, 0, buf.Length, 1000);
-                        if (length >= 0)
+                        MockDtlsServer server = new MockDtlsServer(serverCrypto);
+                        DtlsTransport dtlsTransport = m_serverProtocol.Accept(server, m_serverTransport, request);
+                        byte[] buf = new byte[dtlsTransport.GetReceiveLimit()];
+                        while (!m_isShutdown)
                         {
-                            dtlsServer.Send(buf, 0, length);
+                            int length = dtlsTransport.Receive(buf, 0, buf.Length, 1000);
+                            if (length >= 0)
+                            {
+                                dtlsTransport.Send(buf, 0, length);
+                            }
                         }
+                        dtlsTransport.Close();
                     }
-                    dtlsServer.Close();
                 }
                 catch (Exception e)
                 {
diff --git a/crypto/test/src/tls/test/MockDtlsServer.cs b/crypto/test/src/tls/test/MockDtlsServer.cs
index f69b6df91..6baae3f49 100644
--- a/crypto/test/src/tls/test/MockDtlsServer.cs
+++ b/crypto/test/src/tls/test/MockDtlsServer.cs
@@ -13,7 +13,12 @@ namespace Org.BouncyCastle.Tls.Tests
         : DefaultTlsServer
     {
         internal MockDtlsServer()
-            : base(new BcTlsCrypto())
+            : this(new BcTlsCrypto())
+        {
+        }
+
+        internal MockDtlsServer(TlsCrypto crypto)
+            : base(crypto)
         {
         }