diff --git a/Hazel.UnitTests/Dtls/AesGcmRecordProtectedTests.cs b/Hazel.UnitTests/Dtls/AesGcmRecordProtectedTests.cs index 8e71f2d..8215686 100644 --- a/Hazel.UnitTests/Dtls/AesGcmRecordProtectedTests.cs +++ b/Hazel.UnitTests/Dtls/AesGcmRecordProtectedTests.cs @@ -53,6 +53,42 @@ public void ServerCanEncryptAndDecryptData() } } + + [TestMethod] + public void ServerCanEncryptWithDuplicateInstances() + { + using (Aes128GcmRecordProtection recordProtection = new Aes128GcmRecordProtection(this.masterSecret, this.serverRandom, this.clientRandom)) + using (IRecordProtection duplicateProtection = recordProtection.Duplicate()) + { + byte[] messageAsBytes = Encoding.UTF8.GetBytes(TestMessage); + + // I want to see that if instances are used at different rates, the outputs are deterministic + Record unusedRecord = new Record(); + unusedRecord.Length = (ushort)recordProtection.GetEncryptedSize(messageAsBytes.Length); + ByteSpan unused = new byte[unusedRecord.Length]; + recordProtection.EncryptClientPlaintext(unused, messageAsBytes, ref unusedRecord); + + Record record = new Record(); + record.ContentType = ContentType.ApplicationData; + record.ProtocolVersion = ProtocolVersion.DTLS1_2; + record.Epoch = 1; + record.SequenceNumber = 124; + record.Length = (ushort)recordProtection.GetEncryptedSize(messageAsBytes.Length); + + ByteSpan encrypted1 = new byte[record.Length]; + recordProtection.EncryptClientPlaintext(encrypted1, messageAsBytes, ref record); + + ByteSpan encrypted2 = new byte[record.Length]; + duplicateProtection.EncryptClientPlaintext(encrypted2, messageAsBytes, ref record); + + Assert.AreEqual(encrypted1.Length, encrypted2.Length); + for (int i = 0; i < encrypted1.Length; i++) + { + Assert.AreEqual(encrypted1[i], encrypted2[i]); + } + } + } + [TestMethod] public void ClientCanEncryptAndDecryptData() { diff --git a/Hazel.UnitTests/Dtls/DtlsConnectionTests.cs b/Hazel.UnitTests/Dtls/DtlsConnectionTests.cs index 55e0c91..dbf587d 100644 --- a/Hazel.UnitTests/Dtls/DtlsConnectionTests.cs +++ b/Hazel.UnitTests/Dtls/DtlsConnectionTests.cs @@ -87,7 +87,7 @@ private static X509Certificate2Collection GetCertificateForClient() protected DtlsConnectionListener CreateListener(int numWorkers, IPEndPoint endPoint, ILogger logger, IPMode ipMode = IPMode.IPv4) { - DtlsConnectionListener listener = new DtlsConnectionListener(2, endPoint, logger, ipMode); + DtlsConnectionListener listener = new DtlsConnectionListener(numWorkers, endPoint, logger, ipMode); listener.SetCertificate(GetCertificateForServer()); return listener; @@ -105,42 +105,38 @@ public void DtlsServerDisposeDisconnectsTest() { IPEndPoint ep = new IPEndPoint(IPAddress.Loopback, 27510); - bool serverConnected = false; + ManualResetEventSlim serverConnected = new ManualResetEventSlim(); bool serverDisconnected = false; - bool clientDisconnected = false; - - Semaphore signal = new Semaphore(0, int.MaxValue); + ManualResetEventSlim clientDisconnected = new ManualResetEventSlim(); - using (var listener = (DtlsConnectionListener)CreateListener(2, new IPEndPoint(IPAddress.Any, ep.Port), new TestLogger())) + using (var listener = CreateListener(2, new IPEndPoint(IPAddress.Any, ep.Port), new TestLogger())) using (var connection = CreateConnection(ep, new TestLogger())) { listener.NewConnection += (evt) => { - serverConnected = true; - signal.Release(); - evt.Connection.Disconnected += (o, et) => { + serverConnected.Set(); + evt.Connection.Disconnected += (o, et) => + { serverDisconnected = true; }; }; - connection.Disconnected += (o, evt) => { - clientDisconnected = true; - signal.Release(); + + connection.Disconnected += (o, evt) => + { + clientDisconnected.Set(); }; listener.Start(); connection.Connect(); // wait for the client to connect - signal.WaitOne(10); + Assert.IsTrue(serverConnected.Wait(1000), "Server never connected"); listener.Dispose(); - // wait for the client to disconnect - signal.WaitOne(100); + Assert.IsTrue(clientDisconnected.Wait(5000), "Client never disconnected"); + Assert.IsFalse(serverDisconnected, "Server shouldn't see a disconnect event when disposing."); - Assert.IsTrue(serverConnected); - Assert.IsTrue(clientDisconnected); - Assert.IsFalse(serverDisconnected); Assert.AreEqual(0, listener.PeerCount); } } @@ -221,7 +217,7 @@ public void TestMalformedApplicationData() IPEndPoint ep = new IPEndPoint(IPAddress.Loopback, 27510); IPEndPoint connectionEndPoint = ep; - DtlsConnectionListener.ConnectionId connectionId = new ThreadLimitedUdpConnectionListener.ConnectionId(); + ConnectionId connectionId = new ConnectionId(); Semaphore signal = new Semaphore(0, int.MaxValue); @@ -277,7 +273,7 @@ public void TestMalformedConnectionData() IPEndPoint ep = new IPEndPoint(IPAddress.Loopback, 27510); IPEndPoint connectionEndPoint = ep; - DtlsConnectionListener.ConnectionId connectionId = new ThreadLimitedUdpConnectionListener.ConnectionId(); + ConnectionId connectionId = new ConnectionId(); Semaphore signal = new Semaphore(0, int.MaxValue); @@ -963,8 +959,8 @@ public void DtlsIPv6ConnectionTest() [TestMethod] public void DtlsUnreliableServerToClientTest() { - using (ThreadLimitedUdpConnectionListener listener = this.CreateListener(2, new IPEndPoint(IPAddress.Any, 4296), new TestLogger())) - using (UdpConnection connection = this.CreateConnection(new IPEndPoint(IPAddress.Loopback, 4296), new TestLogger())) + using (ThreadLimitedUdpConnectionListener listener = this.CreateListener(2, new IPEndPoint(IPAddress.Any, 4296), new TestLogger("Server"))) + using (UdpConnection connection = this.CreateConnection(new IPEndPoint(IPAddress.Loopback, 4296), new TestLogger("Client"))) { TestHelper.RunServerToClientTest(listener, connection, 10, SendOption.None); } @@ -1081,17 +1077,17 @@ public void ClientDisconnectTest() using (var listener = this.CreateListener(2, new IPEndPoint(IPAddress.Any, 4296), new TestLogger("Server"))) using (var connection = this.CreateConnection(new IPEndPoint(IPAddress.Loopback, 4296), new TestLogger("Client"))) { - ManualResetEvent mutex = new ManualResetEvent(false); - ManualResetEvent mutex2 = new ManualResetEvent(false); + ManualResetEvent clientConnected = new ManualResetEvent(false); + ManualResetEvent clientDisconnected = new ManualResetEvent(false); listener.NewConnection += delegate (NewConnectionEventArgs args) { args.Connection.Disconnected += delegate (object sender2, DisconnectedEventArgs args2) { - mutex2.Set(); + clientDisconnected.Set(); }; - mutex.Set(); + clientConnected.Set(); }; listener.Start(); @@ -1099,12 +1095,12 @@ public void ClientDisconnectTest() connection.Connect(); Assert.AreEqual(ConnectionState.Connected, connection.State); - mutex.WaitOne(1000); + Assert.IsTrue(clientConnected.WaitOne(1000), "Client never connected"); Assert.AreEqual(ConnectionState.Connected, connection.State); connection.Disconnect("Testing"); - mutex2.WaitOne(1000); + Assert.IsTrue(clientDisconnected.WaitOne(1000), "Client never disconnected"); Assert.AreEqual(ConnectionState.NotConnected, connection.State); } } diff --git a/Hazel.UnitTests/TestHelper.cs b/Hazel.UnitTests/TestHelper.cs index f3e9dfb..4e3385f 100644 --- a/Hazel.UnitTests/TestHelper.cs +++ b/Hazel.UnitTests/TestHelper.cs @@ -1,11 +1,8 @@ -using System; +using Hazel.Udp.FewerThreads; using Microsoft.VisualStudio.TestTools.UnitTesting; - -using Hazel; -using System.Net; -using System.Threading; +using System; using System.Diagnostics; -using Hazel.Udp.FewerThreads; +using System.Threading; namespace Hazel.UnitTests { @@ -17,90 +14,41 @@ public static class TestHelper /// /// The listener to test. /// The connection to test. - internal static void RunServerToClientTest(ThreadLimitedUdpConnectionListener listener, Connection connection, int dataSize, SendOption sendOption) + internal static void RunServerToClientTest(NetworkConnectionListener listener, Connection connection, int dataSize, SendOption sendOption) { //Setup meta stuff MessageWriter data = BuildData(sendOption, dataSize); - ManualResetEvent mutex = new ManualResetEvent(false); + ManualResetEvent clientConnected = new ManualResetEvent(false); + ManualResetEvent dataReceived = new ManualResetEvent(false); + Connection conn = null; //Setup listener listener.NewConnection += delegate (NewConnectionEventArgs ncArgs) { - ncArgs.Connection.Send(data); + conn = ncArgs.Connection; + clientConnected.Set(); }; - listener.Start(); - + //Setup connection DataReceivedEventArgs? result = null; - //Setup conneciton connection.DataReceived += delegate (DataReceivedEventArgs a) { - Trace.WriteLine("Data was received correctly."); - - try - { - result = a; - } - finally - { - mutex.Set(); - } - }; - - connection.Connect(); - - //Wait until data is received - mutex.WaitOne(); - - var dataReader = ConvertToMessageReader(data); - Assert.AreEqual(dataReader.Length, result.Value.Message.Length); - for (int i = 0; i < dataReader.Length; i++) - { - Assert.AreEqual(dataReader.ReadByte(), result.Value.Message.ReadByte()); - } - - Assert.AreEqual(sendOption, result.Value.SendOption); - } - - /// - /// Runs a general test on the given listener and connection. - /// - /// The listener to test. - /// The connection to test. - internal static void RunServerToClientTest(NetworkConnectionListener listener, Connection connection, int dataSize, SendOption sendOption) - { - //Setup meta stuff - MessageWriter data = BuildData(sendOption, dataSize); - ManualResetEvent mutex = new ManualResetEvent(false); - - //Setup listener - listener.NewConnection += delegate (NewConnectionEventArgs ncArgs) - { - ncArgs.Connection.Send(data); + Console.WriteLine("Data was received correctly."); + result = a; + dataReceived.Set(); }; listener.Start(); + connection.Connect(); - DataReceivedEventArgs? result = null; - //Setup conneciton - connection.DataReceived += delegate (DataReceivedEventArgs a) - { - Trace.WriteLine("Data was received correctly."); + // Wait until data is received + Assert.IsTrue(clientConnected.WaitOne(1000), "Client never connected"); - try - { - result = a; - } - finally - { - mutex.Set(); - } - }; + Assert.AreEqual(SendErrors.None, conn.Send(data)); + Console.WriteLine("Data was sent correctly."); - connection.Connect(); - - //Wait until data is received - mutex.WaitOne(); + // Wait until data is received + Assert.IsTrue(dataReceived.WaitOne(1000), "Data was never received"); var dataReader = ConvertToMessageReader(data); Assert.AreEqual(dataReader.Length, result.Value.Message.Length); @@ -130,10 +78,8 @@ internal static void RunClientToServerTest(NetworkConnectionListener listener, C { args.Connection.DataReceived += delegate (DataReceivedEventArgs innerArgs) { - Trace.WriteLine("Data was received correctly."); - + Console.WriteLine("Data was received correctly."); result = innerArgs; - mutex2.Set(); }; @@ -141,16 +87,13 @@ internal static void RunClientToServerTest(NetworkConnectionListener listener, C }; listener.Start(); - - //Connect connection.Connect(); - mutex.WaitOne(); + Assert.IsTrue(mutex.WaitOne(1000), "Client never connected"); - connection.Send(data); + Assert.AreEqual(SendErrors.None, connection.Send(data)); - //Wait until data is received - mutex2.WaitOne(); + Assert.IsTrue(mutex2.WaitOne(1000), "Data was never received"); var dataReader = ConvertToMessageReader(data); Assert.AreEqual(dataReader.Length, result.Value.Message.Length); @@ -162,57 +105,6 @@ internal static void RunClientToServerTest(NetworkConnectionListener listener, C Assert.AreEqual(sendOption, result.Value.SendOption); } - - /// - /// Runs a general test on the given listener and connection. - /// - /// The listener to test. - /// The connection to test. - internal static void RunClientToServerTest(ThreadLimitedUdpConnectionListener listener, Connection connection, int dataSize, SendOption sendOption) - { - //Setup meta stuff - MessageWriter data = BuildData(sendOption, dataSize); - ManualResetEvent mutex = new ManualResetEvent(false); - ManualResetEvent mutex2 = new ManualResetEvent(false); - - //Setup listener - DataReceivedEventArgs? result = null; - listener.NewConnection += delegate (NewConnectionEventArgs args) - { - args.Connection.DataReceived += delegate (DataReceivedEventArgs innerArgs) - { - Trace.WriteLine("Data was received correctly."); - - result = innerArgs; - - mutex2.Set(); - }; - - mutex.Set(); - }; - - listener.Start(); - - //Connect - connection.Connect(); - - Assert.IsTrue(mutex.WaitOne(100), "Timeout while connecting"); - - connection.Send(data); - - //Wait until data is received - Assert.IsTrue(mutex2.WaitOne(100), "Timeout while sending data"); - - var dataReader = ConvertToMessageReader(data); - Assert.AreEqual(dataReader.Length, result.Value.Message.Length); - for (int i = 0; i < dataReader.Length; i++) - { - Assert.AreEqual(dataReader.ReadByte(), result.Value.Message.ReadByte()); - } - - Assert.AreEqual(sendOption, result.Value.SendOption); - } - /// /// Runs a server disconnect test on the given listener and connection. /// diff --git a/Hazel.UnitTests/ThreadLimitedUdpConnectionTests.cs b/Hazel.UnitTests/ThreadLimitedUdpConnectionTests.cs index 0fe5cfd..f4ba8b6 100644 --- a/Hazel.UnitTests/ThreadLimitedUdpConnectionTests.cs +++ b/Hazel.UnitTests/ThreadLimitedUdpConnectionTests.cs @@ -30,7 +30,8 @@ public void ServerDisposeDisconnectsTest() bool serverConnected = false; bool serverDisconnected = false; - bool clientDisconnected = false; + + ManualResetEvent clientDisconnected = new ManualResetEvent(false); using (ThreadLimitedUdpConnectionListener listener = this.CreateListener(2, new IPEndPoint(IPAddress.Any, 4296), new TestLogger("SERVER"))) using (UdpConnection connection = this.CreateConnection(ep, new TestLogger("CLIENT"))) @@ -40,7 +41,7 @@ public void ServerDisposeDisconnectsTest() serverConnected = true; evt.Connection.Disconnected += (o, et) => serverDisconnected = true; }; - connection.Disconnected += (o, evt) => clientDisconnected = true; + connection.Disconnected += (o, evt) => clientDisconnected.Set(); listener.Start(); connection.Connect(); @@ -50,7 +51,7 @@ public void ServerDisposeDisconnectsTest() Thread.Sleep(100); Assert.IsTrue(serverConnected); - Assert.IsTrue(clientDisconnected); + Assert.IsTrue(clientDisconnected.WaitOne(5000)); Assert.IsFalse(serverDisconnected); } } @@ -378,7 +379,7 @@ public void KeepAliveServerTest() connection.Connect(); - mutex.WaitOne(); + mutex.WaitOne(5000); Assert.AreEqual(ConnectionState.Connected, client.State); diff --git a/Hazel/Dtls/AesGcmRecordProtection.cs b/Hazel/Dtls/AesGcmRecordProtection.cs index 65df39e..5397c11 100644 --- a/Hazel/Dtls/AesGcmRecordProtection.cs +++ b/Hazel/Dtls/AesGcmRecordProtection.cs @@ -1,20 +1,27 @@ using Hazel.Crypto; using System; using System.Diagnostics; +using System.Threading; namespace Hazel.Dtls { /// /// *_AES_128_GCM_* cipher suite /// - public class Aes128GcmRecordProtection: IRecordProtection + public class Aes128GcmRecordProtection : IRecordProtection { + private static int InstanceCount; + public int Id { get; } + private const int ImplicitNonceSize = 4; private const int ExplicitNonceSize = 8; private readonly Aes128Gcm serverWriteCipher; private readonly Aes128Gcm clientWriteCipher; + private readonly ByteSpan clientWriteKey; + private readonly ByteSpan serverWriteKey; + private readonly ByteSpan serverWriteIV; private readonly ByteSpan clientWriteIV; @@ -26,6 +33,8 @@ public class Aes128GcmRecordProtection: IRecordProtection /// Client random data public Aes128GcmRecordProtection(ByteSpan masterSecret, ByteSpan serverRandom, ByteSpan clientRandom) { + this.Id = Interlocked.Increment(ref InstanceCount); + ByteSpan combinedRandom = new byte[serverRandom.Length + clientRandom.Length]; serverRandom.CopyTo(combinedRandom); clientRandom.CopyTo(combinedRandom.Slice(serverRandom.Length)); @@ -43,11 +52,23 @@ public Aes128GcmRecordProtection(ByteSpan masterSecret, ByteSpan serverRandom, B ByteSpan expandedKey = new byte[ExpandedSize]; PrfSha256.ExpandSecret(expandedKey, masterSecret, PrfLabel.KEY_EXPANSION, combinedRandom); - ByteSpan clientWriteKey = expandedKey.Slice(0, Aes128Gcm.KeySize); - ByteSpan serverWriteKey = expandedKey.Slice(Aes128Gcm.KeySize, Aes128Gcm.KeySize); + this.clientWriteKey = expandedKey.Slice(0, Aes128Gcm.KeySize); + this.serverWriteKey = expandedKey.Slice(Aes128Gcm.KeySize, Aes128Gcm.KeySize); this.clientWriteIV = expandedKey.Slice(2 * Aes128Gcm.KeySize, ImplicitNonceSize); this.serverWriteIV = expandedKey.Slice(2 * Aes128Gcm.KeySize + ImplicitNonceSize, ImplicitNonceSize); + this.serverWriteCipher = new Aes128Gcm(this.serverWriteKey); + this.clientWriteCipher = new Aes128Gcm(this.clientWriteKey); + } + + private Aes128GcmRecordProtection(int id, ByteSpan clientWriteKey, ByteSpan serverWriteKey, ByteSpan clientWriteIV, ByteSpan serverWriteIV) + { + this.Id = id; + this.clientWriteKey = clientWriteKey; + this.serverWriteKey = serverWriteKey; + this.clientWriteIV = clientWriteIV; + this.serverWriteIV = serverWriteIV; + this.serverWriteCipher = new Aes128Gcm(serverWriteKey); this.clientWriteCipher = new Aes128Gcm(clientWriteKey); } @@ -125,6 +146,11 @@ public bool DecryptCiphertextFromClient(ByteSpan output, ByteSpan input, ref Rec return DecryptCiphertext(output, input, ref record, this.clientWriteCipher, this.clientWriteIV); } + public IRecordProtection Duplicate() + { + return new Aes128GcmRecordProtection(this.Id, this.clientWriteKey, this.serverWriteKey, this.clientWriteIV, this.serverWriteIV); + } + private static bool DecryptCiphertext(ByteSpan output, ByteSpan input, ref Record record, Aes128Gcm cipher, ByteSpan writeIV) { Debug.Assert(output.Length >= GetDecryptedSizeImpl(input.Length)); diff --git a/Hazel/Dtls/DtlsConnectionListener.cs b/Hazel/Dtls/DtlsConnectionListener.cs index 8cbf038..8bb859d 100644 --- a/Hazel/Dtls/DtlsConnectionListener.cs +++ b/Hazel/Dtls/DtlsConnectionListener.cs @@ -22,143 +22,6 @@ public class DtlsConnectionListener : ThreadLimitedUdpConnectionListener // Min MTU - UDP+IP header - 1 (for good measure. :)) private const int MaxCertFragmentSizeV1 = 576 - 32 - 1; - /// - /// Current state of handshake sequence - /// - enum HandshakeState - { - ExpectingHello, - ExpectingClientKeyExchange, - ExpectingChangeCipherSpec, - ExpectingFinish - } - - /// - /// State to manage the current epoch `N` - /// - struct CurrentEpoch - { - public ulong NextOutgoingSequence; - - public ulong NextExpectedSequence; - public ulong PreviousSequenceWindowBitmask; - - public IRecordProtection RecordProtection; - public IRecordProtection PreviousRecordProtection; - - // Need to keep these around so we can re-transmit our - // last handshake record flight - public ByteSpan ExpectedClientFinishedVerification; - public ByteSpan ServerFinishedVerification; - public ulong NextOutgoingSequenceForPreviousEpoch; - } - - /// - /// State to manage the transition from the current - /// epoch `N` to epoch `N+1` - /// - struct NextEpoch - { - public ushort Epoch; - - public HandshakeState State; - public CipherSuite SelectedCipherSuite; - - public ulong NextOutgoingSequence; - - public IHandshakeCipherSuite Handshake; - public IRecordProtection RecordProtection; - - public ByteSpan ClientRandom; - public ByteSpan ServerRandom; - - public Sha256Stream VerificationStream; - - public ByteSpan ClientVerification; - public ByteSpan ServerVerification; - - } - - /// - /// Per-peer state - /// - sealed class PeerData : IDisposable - { - public ushort Epoch; - public bool CanHandleApplicationData; - - public HazelDtlsSessionInfo Session; - - public CurrentEpoch CurrentEpoch; - public NextEpoch NextEpoch; - - public ConnectionId ConnectionId; - - public readonly List QueuedApplicationDataMessage = new List(); - public readonly ConcurrentBag ApplicationData = new ConcurrentBag(); - public readonly ProtocolVersion ProtocolVersion; - - public DateTime StartOfNegotiation; - - public PeerData(ConnectionId connectionId, ulong nextExpectedSequenceNumber, ProtocolVersion protocolVersion) - { - ByteSpan block = new byte[2 * Finished.Size]; - this.CurrentEpoch.ServerFinishedVerification = block.Slice(0, Finished.Size); - this.CurrentEpoch.ExpectedClientFinishedVerification = block.Slice(Finished.Size, Finished.Size); - this.ProtocolVersion = protocolVersion; - - ResetPeer(connectionId, nextExpectedSequenceNumber); - } - - public void ResetPeer(ConnectionId connectionId, ulong nextExpectedSequenceNumber) - { - Dispose(); - - this.Epoch = 0; - this.CanHandleApplicationData = false; - this.QueuedApplicationDataMessage.Clear(); - - this.CurrentEpoch.NextOutgoingSequence = 2; // Account for our ClientHelloVerify - this.CurrentEpoch.NextExpectedSequence = nextExpectedSequenceNumber; - this.CurrentEpoch.PreviousSequenceWindowBitmask = 0; - this.CurrentEpoch.RecordProtection = NullRecordProtection.Instance; - this.CurrentEpoch.PreviousRecordProtection = null; - this.CurrentEpoch.ServerFinishedVerification.SecureClear(); - this.CurrentEpoch.ExpectedClientFinishedVerification.SecureClear(); - - this.NextEpoch.State = HandshakeState.ExpectingHello; - this.NextEpoch.RecordProtection = null; - this.NextEpoch.Handshake = null; - this.NextEpoch.ClientRandom = new byte[Random.Size]; - this.NextEpoch.ServerRandom = new byte[Random.Size]; - this.NextEpoch.VerificationStream = new Sha256Stream(); - this.NextEpoch.ClientVerification = new byte[Finished.Size]; - this.NextEpoch.ServerVerification = new byte[Finished.Size]; - - this.ConnectionId = connectionId; - - this.StartOfNegotiation = DateTime.UtcNow; - } - - public void Dispose() - { - this.CurrentEpoch.RecordProtection?.Dispose(); - this.CurrentEpoch.PreviousRecordProtection?.Dispose(); - this.NextEpoch.RecordProtection?.Dispose(); - this.NextEpoch.Handshake?.Dispose(); - this.NextEpoch.VerificationStream?.Dispose(); - - while (this.ApplicationData.TryTake(out var msg)) - { - try - { - msg.Recycle(); - } - catch { } - } - } - } - private RandomNumberGenerator random; // Private key component of certificate's public key @@ -387,7 +250,7 @@ private void ProcessIncomingMessage(ByteSpan message, IPEndPoint peerAddress) } // Validate record authenticity - int decryptedSize = peer.CurrentEpoch.RecordProtection.GetDecryptedSize(recordPayload.Length); + int decryptedSize = peer.CurrentEpoch.MasterRecordProtection.GetDecryptedSize(recordPayload.Length); if (decryptedSize < 0) { this.Logger.WriteInfo($"Dropping malformed record: Length {recordPayload.Length} Decrypted length: {decryptedSize}"); @@ -397,7 +260,7 @@ private void ProcessIncomingMessage(ByteSpan message, IPEndPoint peerAddress) ByteSpan decryptedPayload = recordPayload.ReuseSpanIfPossible(decryptedSize); ProtocolVersion protocolVersion = peer.ProtocolVersion; - if (!peer.CurrentEpoch.RecordProtection.DecryptCiphertextFromClient(decryptedPayload, recordPayload, ref record)) + if (!peer.CurrentEpoch.MasterRecordProtection.DecryptCiphertextFromClient(decryptedPayload, recordPayload, ref record)) { this.Logger.WriteVerbose($"Dropping non-authentic {record.ContentType} record from `{peerAddress}`"); return; @@ -446,10 +309,9 @@ private void ProcessIncomingMessage(ByteSpan message, IPEndPoint peerAddress) // Migrate to the next epoch peer.Epoch = peer.NextEpoch.Epoch; peer.CanHandleApplicationData = false; // Need a Finished message - peer.CurrentEpoch.NextOutgoingSequenceForPreviousEpoch = peer.CurrentEpoch.NextOutgoingSequence; + peer.CurrentEpoch.NextOutgoingSequenceForPreviousEpoch = (ulong)peer.CurrentEpoch.NextOutgoingSequence; peer.CurrentEpoch.PreviousRecordProtection?.Dispose(); - peer.CurrentEpoch.PreviousRecordProtection = peer.CurrentEpoch.RecordProtection; - peer.CurrentEpoch.RecordProtection = peer.NextEpoch.RecordProtection; + peer.CurrentEpoch.SetRecordProtection(peer.NextEpoch.RecordProtection); peer.CurrentEpoch.NextOutgoingSequence = 1; peer.CurrentEpoch.NextExpectedSequence = 1; peer.CurrentEpoch.PreviousSequenceWindowBitmask = 0; @@ -703,7 +565,7 @@ private bool ProcessHandshake(PeerData peer, IPEndPoint peerAddress, ref Record // Either way, there is not a feasible // way to progress the connection. MarkConnectionAsStale(peer.ConnectionId); - this.existingPeers.TryRemove(peerAddress, out _); + this.RemovePeerRecord(peer.ConnectionId); return false; } @@ -731,8 +593,8 @@ private bool ProcessHandshake(PeerData peer, IPEndPoint peerAddress, ref Record finishedRecord.ContentType = ContentType.Handshake; finishedRecord.ProtocolVersion = protocolVersion; finishedRecord.Epoch = peer.Epoch; - finishedRecord.SequenceNumber = peer.CurrentEpoch.NextOutgoingSequence; - finishedRecord.Length = (ushort)peer.CurrentEpoch.RecordProtection.GetEncryptedSize(plaintextFinishedPayloadSize); + finishedRecord.SequenceNumber = (ulong)peer.CurrentEpoch.NextOutgoingSequence; + finishedRecord.Length = (ushort)peer.CurrentEpoch.MasterRecordProtection.GetEncryptedSize(plaintextFinishedPayloadSize); ++peer.CurrentEpoch.NextOutgoingSequence; // Encode the flight into wire format @@ -758,7 +620,7 @@ ref changeCipherSpecRecord ); // Protect the Finished Handshake record - peer.CurrentEpoch.RecordProtection.EncryptServerPlaintext( + peer.CurrentEpoch.MasterRecordProtection.EncryptServerPlaintext( startOfFinishedRecord.Slice(Record.Size, finishedRecord.Length), startOfFinishedRecord.Slice(Record.Size, plaintextFinishedPayloadSize), ref finishedRecord @@ -766,8 +628,7 @@ ref finishedRecord // Current epoch can now handle application data peer.CanHandleApplicationData = true; - - base.QueueRawData(packet, peerAddress); + this.sendQueue.Add(new SendMessageInfo(packet, peerAddress)); break; // Drop messages that we do not support @@ -847,7 +708,7 @@ private bool HandleClientHello(PeerData peer, IPEndPoint peerAddress, ref Record outgoingSequence = peer.CurrentEpoch.NextExpectedSequence; ++peer.CurrentEpoch.NextOutgoingSequenceForPreviousEpoch; - recordProtection = peer.CurrentEpoch.RecordProtection; + recordProtection = peer.CurrentEpoch.MasterRecordProtection; } #if DEBUG @@ -990,8 +851,8 @@ private bool HandleClientHello(PeerData peer, IPEndPoint peerAddress, ref Record initialRecord.ContentType = ContentType.Handshake; initialRecord.ProtocolVersion = protocolVersion; initialRecord.Epoch = peer.Epoch; - initialRecord.SequenceNumber = peer.CurrentEpoch.NextOutgoingSequence; - initialRecord.Length = (ushort)peer.CurrentEpoch.RecordProtection.GetEncryptedSize(initialRecordPayloadSize); + initialRecord.SequenceNumber = (ulong)peer.CurrentEpoch.NextOutgoingSequence; + initialRecord.Length = (ushort)peer.CurrentEpoch.MasterRecordProtection.GetEncryptedSize(initialRecordPayloadSize); ++peer.CurrentEpoch.NextOutgoingSequence; // Convert initial record of the flight to @@ -1010,13 +871,13 @@ private bool HandleClientHello(PeerData peer, IPEndPoint peerAddress, ref Record certificateData = certificateData.Slice(certInitialFragmentSize); // Protect initial record of the flight - peer.CurrentEpoch.RecordProtection.EncryptServerPlaintext( + peer.CurrentEpoch.MasterRecordProtection.EncryptServerPlaintext( packet.Slice(Record.Size, initialRecord.Length), packet.Slice(Record.Size, initialRecordPayloadSize), ref initialRecord ); - base.QueueRawData(packet, peerAddress); + this.sendQueue.Add(new SendMessageInfo(packet, peerAddress)); // Record record payload for verification if (recordMessagesForVerifyData) @@ -1054,8 +915,8 @@ ref initialRecord additionalRecord.ContentType = ContentType.Handshake; additionalRecord.ProtocolVersion = protocolVersion; additionalRecord.Epoch = peer.Epoch; - additionalRecord.SequenceNumber = peer.CurrentEpoch.NextOutgoingSequence; - additionalRecord.Length = (ushort)peer.CurrentEpoch.RecordProtection.GetEncryptedSize(additionalRecordPayloadSize); + additionalRecord.SequenceNumber = (ulong)peer.CurrentEpoch.NextOutgoingSequence; + additionalRecord.Length = (ushort)peer.CurrentEpoch.MasterRecordProtection.GetEncryptedSize(additionalRecordPayloadSize); ++peer.CurrentEpoch.NextOutgoingSequence; // Convert record to wire format @@ -1070,13 +931,13 @@ ref initialRecord certificateData = certificateData.Slice(certFragmentSize); // Protect record - peer.CurrentEpoch.RecordProtection.EncryptServerPlaintext( + peer.CurrentEpoch.MasterRecordProtection.EncryptServerPlaintext( packet.Slice(Record.Size, additionalRecord.Length), packet.Slice(Record.Size, additionalRecordPayloadSize), ref additionalRecord ); - base.QueueRawData(packet, peerAddress); + this.sendQueue.Add(new SendMessageInfo(packet, peerAddress)); } // Describe final record of the flight @@ -1102,8 +963,8 @@ ref additionalRecord finalRecord.ContentType = ContentType.Handshake; finalRecord.ProtocolVersion = protocolVersion; finalRecord.Epoch = peer.Epoch; - finalRecord.SequenceNumber = peer.CurrentEpoch.NextOutgoingSequence; - finalRecord.Length = (ushort)peer.CurrentEpoch.RecordProtection.GetEncryptedSize(finalRecordPayloadSize); + finalRecord.SequenceNumber = (ulong)peer.CurrentEpoch.NextOutgoingSequence; + finalRecord.Length = (ushort)peer.CurrentEpoch.MasterRecordProtection.GetEncryptedSize(finalRecordPayloadSize); ++peer.CurrentEpoch.NextOutgoingSequence; // Convert final record of the flight to wire @@ -1130,13 +991,13 @@ ref additionalRecord } // Protect final record of the flight - peer.CurrentEpoch.RecordProtection.EncryptServerPlaintext( + peer.CurrentEpoch.MasterRecordProtection.EncryptServerPlaintext( packet.Slice(Record.Size, finalRecord.Length), packet.Slice(Record.Size, finalRecordPayloadSize), ref finalRecord ); - base.QueueRawData(packet, peerAddress); + this.sendQueue.Add(new SendMessageInfo(packet, peerAddress)); return true; } @@ -1187,7 +1048,7 @@ private void HandleNonPeerRecord(ByteSpan message, IPEndPoint peerAddress) // We only accept Handshake protocol messages from non-peers if (record.ContentType != ContentType.Handshake) { - this.Logger.WriteError($"Dropping non-handhsake message from non-peer `{peerAddress}`"); + this.Logger.WriteError($"Dropping non-handshake message from non-peer `{peerAddress}`: {record.ContentType}"); return; } @@ -1275,92 +1136,61 @@ private void SendHelloVerifyRequest(IPEndPoint peerAddress, ulong recordSequence ref record ); - base.QueueRawData(packet, peerAddress); + this.sendQueue.TryAdd(new SendMessageInfo(packet, peerAddress)); } /// /// Handle a requrest to send a datagram to the network /// - protected override void QueueRawData(ByteSpan span, IPEndPoint remoteEndPoint) + protected override void ProcessQueuedSend(SendMessageInfo msg) { PeerData peer; - if (!this.existingPeers.TryGetValue(remoteEndPoint, out peer)) + if (!this.existingPeers.TryGetValue(msg.Recipient, out peer)) { // Drop messages if we don't know how to send them return; } - lock (peer) + // If we're negotiating a new epoch, queue data + if (peer.Epoch == 0 || peer.NextEpoch.State != HandshakeState.ExpectingHello) { - // If we're negotiating a new epoch, queue data - if (peer.Epoch == 0 || peer.NextEpoch.State != HandshakeState.ExpectingHello) - { - ByteSpan copyOfSpan = new byte[span.Length]; - span.CopyTo(copyOfSpan); + peer.QueuedApplicationDataMessage.Enqueue(msg.Span); + return; + } - peer.QueuedApplicationDataMessage.Add(copyOfSpan); - return; - } + // Send any queued application data now + while (peer.QueuedApplicationDataMessage.TryDequeue(out var queuedSpan)) + { + EncryptEnqueueApplicationDataRecord(msg, peer, queuedSpan); + } - ProtocolVersion protocolVersion = peer.ProtocolVersion; + EncryptEnqueueApplicationDataRecord(msg, peer, msg.Span); + } - // Send any queued application data now - for (int ii = 0, nn = peer.QueuedApplicationDataMessage.Count; ii != nn; ++ii) - { - ByteSpan queuedSpan = peer.QueuedApplicationDataMessage[ii]; - - Record outgoingRecord = new Record(); - outgoingRecord.ContentType = ContentType.ApplicationData; - outgoingRecord.ProtocolVersion = protocolVersion; - outgoingRecord.Epoch = peer.Epoch; - outgoingRecord.SequenceNumber = peer.CurrentEpoch.NextOutgoingSequence; - outgoingRecord.Length = (ushort)peer.CurrentEpoch.RecordProtection.GetEncryptedSize(queuedSpan.Length); - ++peer.CurrentEpoch.NextOutgoingSequence; - - // Encode the record to wire format - ByteSpan packet = new byte[Record.Size + outgoingRecord.Length]; - ByteSpan writer = packet; - outgoingRecord.Encode(writer); - writer = writer.Slice(Record.Size); - queuedSpan.CopyTo(writer); - - // Protect the record - peer.CurrentEpoch.RecordProtection.EncryptServerPlaintext( - packet.Slice(Record.Size, outgoingRecord.Length), - packet.Slice(Record.Size, queuedSpan.Length), - ref outgoingRecord - ); - - base.QueueRawData(packet, remoteEndPoint); - } - peer.QueuedApplicationDataMessage.Clear(); + private void EncryptEnqueueApplicationDataRecord(SendMessageInfo msg, PeerData peer, ByteSpan queuedSpan) + { + Record outgoingRecord = new Record(); + outgoingRecord.ContentType = ContentType.ApplicationData; + outgoingRecord.ProtocolVersion = peer.ProtocolVersion; + outgoingRecord.Epoch = peer.Epoch; + outgoingRecord.SequenceNumber = (ulong)Interlocked.Increment(ref peer.CurrentEpoch.NextOutgoingSequence); + outgoingRecord.Length = (ushort)peer.CurrentEpoch.MasterRecordProtection.GetEncryptedSize(queuedSpan.Length); + + // Encode the record to wire format + ByteSpan packet = new byte[Record.Size + outgoingRecord.Length]; + ByteSpan writer = packet; + outgoingRecord.Encode(writer); + writer = writer.Slice(Record.Size); + queuedSpan.CopyTo(writer); - { - Record outgoingRecord = new Record(); - outgoingRecord.ContentType = ContentType.ApplicationData; - outgoingRecord.ProtocolVersion = protocolVersion; - outgoingRecord.Epoch = peer.Epoch; - outgoingRecord.SequenceNumber = peer.CurrentEpoch.NextOutgoingSequence; - outgoingRecord.Length = (ushort)peer.CurrentEpoch.RecordProtection.GetEncryptedSize(span.Length); - ++peer.CurrentEpoch.NextOutgoingSequence; - - // Encode the record to wire format - ByteSpan packet = new byte[Record.Size + outgoingRecord.Length]; - ByteSpan writer = packet; - outgoingRecord.Encode(writer); - writer = writer.Slice(Record.Size); - span.CopyTo(writer); - - // Protect the record - peer.CurrentEpoch.RecordProtection.EncryptServerPlaintext( - packet.Slice(Record.Size, outgoingRecord.Length), - packet.Slice(Record.Size, span.Length), - ref outgoingRecord - ); - - base.QueueRawData(packet, remoteEndPoint); - } - } + // Protect the record + peer.CurrentEpoch.EncryptServerPlaintext_ThreadSafe( + packet.Slice(Record.Size, outgoingRecord.Length), + packet.Slice(Record.Size, queuedSpan.Length), + ref outgoingRecord + ); + + this.sendQueue.TryAdd(new SendMessageInfo(packet, msg.Recipient)); } private void HandleStaleConnections(object _) @@ -1405,7 +1235,10 @@ protected void MarkConnectionAsStale(ConnectionId connectionId) /// internal override void RemovePeerRecord(ConnectionId connectionId) { - this.existingPeers.TryRemove(connectionId.EndPoint, out _); + if (this.existingPeers.TryRemove(connectionId.EndPoint, out var peer)) + { + peer.Dispose(); + } } /// @@ -1416,6 +1249,5 @@ private ConnectionId AllocateConnectionId(IPEndPoint endPoint) int rawSerialId = Interlocked.Increment(ref this.connectionSerial_unsafe); return ConnectionId.Create(endPoint, rawSerialId); } - } } diff --git a/Hazel/Dtls/DtlsMiscStructs.cs b/Hazel/Dtls/DtlsMiscStructs.cs new file mode 100644 index 0000000..227d887 --- /dev/null +++ b/Hazel/Dtls/DtlsMiscStructs.cs @@ -0,0 +1,127 @@ +using Hazel.Crypto; +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; +using System.Threading; + +namespace Hazel.Dtls +{ + /// + /// Current state of handshake sequence + /// + internal enum HandshakeState + { + ExpectingHello, + ExpectingClientKeyExchange, + ExpectingChangeCipherSpec, + ExpectingFinish + } + + /// + /// State to manage the current epoch `N` + /// + internal struct CurrentEpoch + { + public Udp.FewerThreads.ConnectionId ConnectionId; + public long NextOutgoingSequence; + + public ulong NextExpectedSequence; + public ulong PreviousSequenceWindowBitmask; + + public IRecordProtection MasterRecordProtection { get; private set; } + public IRecordProtection PreviousRecordProtection { get; private set; } + + private ConcurrentBag allRecordProtections; + + // Need to keep these around so we can re-transmit our + // last handshake record flight + public ByteSpan ExpectedClientFinishedVerification; + public ByteSpan ServerFinishedVerification; + public ulong NextOutgoingSequenceForPreviousEpoch; + + public void Init() + { + ByteSpan block = new byte[2 * Finished.Size]; + this.ServerFinishedVerification = block.Slice(0, Finished.Size); + this.ExpectedClientFinishedVerification = block.Slice(Finished.Size, Finished.Size); + this.allRecordProtections = new ConcurrentBag(); + } + + public void SetRecordProtection(IRecordProtection newProtection) + { + this.PreviousRecordProtection = this.MasterRecordProtection; + this.MasterRecordProtection = newProtection; + while (this.allRecordProtections.TryTake(out var i)) + { + i.Dispose(); + } + + if (newProtection == NullRecordProtection.Instance) + { + this.PreviousRecordProtection = null; + return; + } + + for (int i = 0; i < 4; ++i) + { + this.allRecordProtections.Add(newProtection.Duplicate()); + } + } + + public void EncryptServerPlaintext_ThreadSafe(ByteSpan output, ByteSpan input, ref Record record) + { + tryagain: + var master = this.MasterRecordProtection; + if (!this.allRecordProtections.TryTake(out var local)) + { + local = master.Duplicate(); + } + + if (local.Id != master.Id) + { + local.Dispose(); + goto tryagain; + } + + local.EncryptServerPlaintext(output, input, ref record); + + if (local != NullRecordProtection.Instance) + { + this.allRecordProtections.Add(local); + } + } + + public void DisposeThreadStatics() + { + while (this.allRecordProtections.TryTake(out var i)) + { + i.Dispose(); + } + } + } + + /// + /// State to manage the transition from the current + /// epoch `N` to epoch `N+1` + /// + internal struct NextEpoch + { + public ushort Epoch; + + public HandshakeState State; + public CipherSuite SelectedCipherSuite; + + public ulong NextOutgoingSequence; + + public IHandshakeCipherSuite Handshake; + public IRecordProtection RecordProtection; + + public ByteSpan ClientRandom; + public ByteSpan ServerRandom; + + public Sha256Stream VerificationStream; + + public ByteSpan ClientVerification; + public ByteSpan ServerVerification; + } +} diff --git a/Hazel/Dtls/DtlsPeerData.cs b/Hazel/Dtls/DtlsPeerData.cs new file mode 100644 index 0000000..4b4bc68 --- /dev/null +++ b/Hazel/Dtls/DtlsPeerData.cs @@ -0,0 +1,88 @@ +using Hazel.Crypto; +using Hazel.Udp.FewerThreads; +using System; +using System.Collections.Concurrent; +using System.Collections.Generic; + +namespace Hazel.Dtls +{ + /// + /// Per-peer state + /// + internal sealed class PeerData : IDisposable + { + public ushort Epoch; + public bool CanHandleApplicationData; + + public HazelDtlsSessionInfo Session; + + public CurrentEpoch CurrentEpoch; + public NextEpoch NextEpoch; + + public ConnectionId ConnectionId; + + public readonly ConcurrentQueue QueuedApplicationDataMessage = new ConcurrentQueue(); + public readonly ConcurrentBag ApplicationData = new ConcurrentBag(); + public readonly ProtocolVersion ProtocolVersion; + + public DateTime StartOfNegotiation; + + public PeerData(ConnectionId connectionId, ulong nextExpectedSequenceNumber, ProtocolVersion protocolVersion) + { + this.CurrentEpoch.Init(); + this.ProtocolVersion = protocolVersion; + + ResetPeer(connectionId, nextExpectedSequenceNumber); + } + + public void ResetPeer(ConnectionId connectionId, ulong nextExpectedSequenceNumber) + { + Dispose(); + + this.ConnectionId = connectionId; + this.Epoch = 0; + this.CanHandleApplicationData = false; + while (this.QueuedApplicationDataMessage.TryDequeue(out _)); + + this.CurrentEpoch.ConnectionId = this.ConnectionId; + this.CurrentEpoch.NextOutgoingSequence = 2; // Account for our ClientHelloVerify + this.CurrentEpoch.NextExpectedSequence = nextExpectedSequenceNumber; + this.CurrentEpoch.PreviousSequenceWindowBitmask = 0; + this.CurrentEpoch.SetRecordProtection(NullRecordProtection.Instance); + this.CurrentEpoch.ServerFinishedVerification.SecureClear(); + this.CurrentEpoch.ExpectedClientFinishedVerification.SecureClear(); + + this.NextEpoch.State = HandshakeState.ExpectingHello; + this.NextEpoch.RecordProtection = null; + this.NextEpoch.Handshake = null; + this.NextEpoch.ClientRandom = new byte[Random.Size]; + this.NextEpoch.ServerRandom = new byte[Random.Size]; + this.NextEpoch.VerificationStream = new Sha256Stream(); + this.NextEpoch.ClientVerification = new byte[Finished.Size]; + this.NextEpoch.ServerVerification = new byte[Finished.Size]; + + this.StartOfNegotiation = DateTime.UtcNow; + } + + public void Dispose() + { + this.CurrentEpoch.MasterRecordProtection?.Dispose(); + this.CurrentEpoch.PreviousRecordProtection?.Dispose(); + this.CurrentEpoch.DisposeThreadStatics(); + + this.NextEpoch.RecordProtection?.Dispose(); + this.NextEpoch.Handshake?.Dispose(); + this.NextEpoch.VerificationStream?.Dispose(); + + while (this.ApplicationData.TryTake(out var msg)) + { + try + { + msg.Recycle(); + } + catch { } + } + } + } + +} diff --git a/Hazel/Dtls/IRecordProtection.cs b/Hazel/Dtls/IRecordProtection.cs index cbee1b0..83752f2 100644 --- a/Hazel/Dtls/IRecordProtection.cs +++ b/Hazel/Dtls/IRecordProtection.cs @@ -7,6 +7,11 @@ namespace Hazel.Dtls /// public interface IRecordProtection : IDisposable { + /// + /// An id shared by this record protection and all its copies + /// + int Id { get; } + /// /// Calculate the size of an encrypted plaintext /// @@ -62,6 +67,11 @@ public interface IRecordProtection : IDisposable /// Parent DTLS record /// True if the input was authenticated and decrypted. Otherwise false bool DecryptCiphertextFromClient(ByteSpan output, ByteSpan input, ref Record record); + + /// + /// Creates an exact duplicate of this instance. + /// + IRecordProtection Duplicate(); } /// diff --git a/Hazel/Dtls/NullRecordProtection.cs b/Hazel/Dtls/NullRecordProtection.cs index 76fa132..6659373 100644 --- a/Hazel/Dtls/NullRecordProtection.cs +++ b/Hazel/Dtls/NullRecordProtection.cs @@ -11,6 +11,8 @@ namespace Hazel.Dtls /// public class NullRecordProtection : IRecordProtection { + public int Id => -1; + public readonly static NullRecordProtection Instance = new NullRecordProtection(); public void Dispose() @@ -49,6 +51,11 @@ public bool DecryptCiphertextFromClient(ByteSpan output, ByteSpan input, ref Rec return true; } + public IRecordProtection Duplicate() + { + return this; + } + private static void CopyMaybeOverlappingSpans(ByteSpan output, ByteSpan input) { // Early out if the ranges `output` is equal to `input` diff --git a/Hazel/FewerThreads/ConnectionId.cs b/Hazel/FewerThreads/ConnectionId.cs new file mode 100644 index 0000000..8ae9033 --- /dev/null +++ b/Hazel/FewerThreads/ConnectionId.cs @@ -0,0 +1,50 @@ +using System; +using System.Net; + +namespace Hazel.Udp.FewerThreads +{ + public struct ConnectionId : IEquatable + { + public IPEndPoint EndPoint; + public int Serial; + + public static ConnectionId Create(IPEndPoint endPoint, int serial) + { + return new ConnectionId + { + EndPoint = endPoint, + Serial = serial, + }; + } + + public override string ToString() + { + return this.Serial.ToString(); + } + + public bool Equals(ConnectionId other) + { + return this.Serial == other.Serial + && this.EndPoint.Equals(other.EndPoint) + ; + } + + public override bool Equals(object obj) + { + if (obj is ConnectionId) + { + return this.Equals((ConnectionId)obj); + } + + return false; + } + + public override int GetHashCode() + { + ///NOTE(mendsley): We're only hashing the endpoint + /// here, as the common case will have one + /// connection per address+port tuple. + return this.EndPoint.GetHashCode(); + } + } +} diff --git a/Hazel/FewerThreads/HazelPipelineSegment.cs b/Hazel/FewerThreads/HazelPipelineSegment.cs new file mode 100644 index 0000000..b948d77 --- /dev/null +++ b/Hazel/FewerThreads/HazelPipelineSegment.cs @@ -0,0 +1,56 @@ +using System; +using System.Collections.Concurrent; +using System.Threading; + +namespace Hazel +{ + public class HazelPipelineSegment : IDisposable + { + private readonly HazelThreadPool threads; + private readonly BlockingCollection inputs = new BlockingCollection(); + private readonly Action factory; + + public int Count => this.inputs.Count; + + public HazelPipelineSegment(int numThreads, Action factory) + { + this.threads = new HazelThreadPool(numThreads, RunProcessing); + this.factory = factory; + } + + private void RunProcessing() + { + foreach (var item in this.inputs.GetConsumingEnumerable()) + { + try + { + this.factory(item); + } + catch + { + } + } + } + + public void Start() + { + this.threads.Start(); + } + + public void Join() + { + this.inputs.CompleteAdding(); + this.threads.Join(); + } + + public void AddInput(InputType item) + { + this.inputs.TryAdd(item); + } + + public void Dispose() + { + this.inputs.Dispose(); + } + } +} \ No newline at end of file diff --git a/Hazel/FewerThreads/HazelThreadPool.cs b/Hazel/FewerThreads/HazelThreadPool.cs index fb36b00..c18cb1c 100644 --- a/Hazel/FewerThreads/HazelThreadPool.cs +++ b/Hazel/FewerThreads/HazelThreadPool.cs @@ -7,7 +7,7 @@ namespace Hazel { - internal class HazelThreadPool + public class HazelThreadPool { private Thread[] threads; diff --git a/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs b/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs index deaef31..3578117 100644 --- a/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs +++ b/Hazel/FewerThreads/ThreadLimitedUdpConnectionListener.cs @@ -12,10 +12,16 @@ namespace Hazel.Udp.FewerThreads /// public class ThreadLimitedUdpConnectionListener : NetworkConnectionListener { - private struct SendMessageInfo + protected struct SendMessageInfo { public ByteSpan Span; public IPEndPoint Recipient; + + public SendMessageInfo(ByteSpan span, IPEndPoint recipient) + { + this.Span = span; + this.Recipient = recipient; + } } private struct ReceiveMessageInfo @@ -46,49 +52,11 @@ private struct ReceiveMessageInfo public bool ReceiveThreadRunning => this.receiveThread.ThreadState == ThreadState.Running; - public struct ConnectionId : IEquatable - { - public IPEndPoint EndPoint; - public int Serial; - - public static ConnectionId Create(IPEndPoint endPoint, int serial) - { - return new ConnectionId{ - EndPoint = endPoint, - Serial = serial, - }; - } - - public bool Equals(ConnectionId other) - { - return this.Serial == other.Serial - && this.EndPoint.Equals(other.EndPoint) - ; - } - - public override bool Equals(object obj) - { - if (obj is ConnectionId) - { - return this.Equals((ConnectionId)obj); - } - - return false; - } - - public override int GetHashCode() - { - ///NOTE(mendsley): We're only hashing the endpoint - /// here, as the common case will have one - /// connection per address+port tuple. - return this.EndPoint.GetHashCode(); - } - } - protected ConcurrentDictionary allConnections = new ConcurrentDictionary(); private BlockingCollection receiveQueue; - private BlockingCollection sendQueue = new BlockingCollection(); + protected BlockingCollection sendQueue = new BlockingCollection(); + private HazelPipelineSegment appDataQueue; public int MaxAge { @@ -107,7 +75,7 @@ public int MaxAge } public int ConnectionCount { get { return this.allConnections.Count; } } - public int SendQueueLength { get { return this.sendQueue.Count; } } + public int SendQueueLength { get { return this.appDataQueue.Count; } } public int ReceiveQueueLength { get { return this.receiveQueue.Count; } } private bool isActive; @@ -119,6 +87,8 @@ public ThreadLimitedUdpConnectionListener(int numWorkers, IPEndPoint endPoint, I this.IPMode = ipMode; this.receiveQueue = new BlockingCollection(10000); + this.appDataQueue = new HazelPipelineSegment(numWorkers, ProcessQueuedSend); + this.appDataQueue.Start(); this.socket = UdpConnection.CreateSocket(this.IPMode); this.socket.ExclusiveAddressUse = true; @@ -230,9 +200,9 @@ private void ProcessingLoop() { this.ReadCallback(msg.Message, msg.Sender, msg.ConnectionId); } - catch + catch (Exception ex) { - + this.Logger.WriteError("Uncaught exception from read callback: " + ex); } } } @@ -260,7 +230,7 @@ private void SendLoop() } catch (Exception e) { - this.Logger.WriteError("Error in loop while sending: " + e.Message); + this.Logger.WriteError("Error in loop while sending: " + e); Thread.Sleep(1); } } @@ -340,9 +310,24 @@ internal void SendDataRaw(byte[] response, IPEndPoint remoteEndPoint) QueueRawData(response, remoteEndPoint); } - protected virtual void QueueRawData(ByteSpan span, IPEndPoint remoteEndPoint) + internal void SendDisconnect(byte[] buffer, IPEndPoint remoteEndPoint) { - this.sendQueue.TryAdd(new SendMessageInfo() { Span = span, Recipient = remoteEndPoint }); + this.ProcessQueuedSend(new SendMessageInfo(buffer, remoteEndPoint)); + } + + private void QueueRawData(ByteSpan span, IPEndPoint remoteEndPoint) + { + this.appDataQueue.AddInput(new SendMessageInfo(span, remoteEndPoint)); + } + + protected virtual void ProcessQueuedSend(SendMessageInfo info) + { + if (info.Span.GetUnderlyingArray() == null) + { + return; + } + + this.sendQueue.TryAdd(info); } /// @@ -372,7 +357,8 @@ protected override void Dispose(bool disposing) this.isActive = false; // Flush outgoing packets - this.sendQueue?.CompleteAdding(); + this.appDataQueue?.Join(); + this.sendQueue.CompleteAdding(); if (wasActive) { @@ -394,8 +380,8 @@ protected override void Dispose(bool disposing) this.receiveQueue?.Dispose(); this.receiveQueue = null; - this.sendQueue?.Dispose(); - this.sendQueue = null; + this.appDataQueue?.Dispose(); + this.appDataQueue = null; base.Dispose(disposing); } diff --git a/Hazel/FewerThreads/ThreadLimitedUdpServerConnection.cs b/Hazel/FewerThreads/ThreadLimitedUdpServerConnection.cs index bb139c7..7569852 100644 --- a/Hazel/FewerThreads/ThreadLimitedUdpServerConnection.cs +++ b/Hazel/FewerThreads/ThreadLimitedUdpServerConnection.cs @@ -20,7 +20,7 @@ public sealed class ThreadLimitedUdpServerConnection : UdpConnection /// public ThreadLimitedUdpConnectionListener Listener { get; private set; } - public ThreadLimitedUdpConnectionListener.ConnectionId ConnectionId { get; private set; } + public ConnectionId ConnectionId { get; private set; } /// /// Creates a UdpConnection for the virtual connection to the endpoint. @@ -28,7 +28,7 @@ public sealed class ThreadLimitedUdpServerConnection : UdpConnection /// The listener that created this connection. /// The endpoint that we are connected to. /// The IPMode we are connected using. - internal ThreadLimitedUdpServerConnection(ThreadLimitedUdpConnectionListener listener, ThreadLimitedUdpConnectionListener.ConnectionId connectionId, IPEndPoint endPoint, IPMode IPMode, ILogger logger) + internal ThreadLimitedUdpServerConnection(ThreadLimitedUdpConnectionListener listener, ConnectionId connectionId, IPEndPoint endPoint, IPMode IPMode, ILogger logger) : base(logger) { this.Listener = listener; @@ -49,7 +49,7 @@ protected override void WriteBytesToConnection(byte[] bytes, int length) // but I don't want to have a bunch of client references in the send queue... // Does this perhaps mean the encryption is being done in the wrong class? this.Statistics.LogPacketSend(length); - Listener.SendDataRaw(bytes, EndPoint); + this.Listener.SendDataRaw(bytes, EndPoint); } /// @@ -89,7 +89,7 @@ protected override bool SendDisconnect(MessageWriter data = null) try { - this.WriteBytesToConnection(bytes, bytes.Length); + this.Listener.SendDisconnect(bytes, this.EndPoint); } catch { } @@ -103,7 +103,7 @@ protected override void Dispose(bool disposing) SendDisconnect(); } - Listener.RemovePeerRecord(this.ConnectionId); + this.Listener.RemovePeerRecord(this.ConnectionId); base.Dispose(disposing); } } diff --git a/Hazel/MessageWriter.cs b/Hazel/MessageWriter.cs index 7d4e050..9738101 100644 --- a/Hazel/MessageWriter.cs +++ b/Hazel/MessageWriter.cs @@ -135,6 +135,13 @@ public void Recycle() #region WriteMethods + public void CopyFrom(MessageWriter target) + { + target.SendOption = target.SendOption; + System.Buffer.BlockCopy(target.Buffer, 0, this.Buffer, 0, target.Length); + this.Position = this.Length = target.Length; + } + public void CopyFrom(MessageReader target) { int offset, length; diff --git a/Hazel/Udp/UdpConnection.Reliable.cs b/Hazel/Udp/UdpConnection.Reliable.cs index 7494913..cb0f501 100644 --- a/Hazel/Udp/UdpConnection.Reliable.cs +++ b/Hazel/Udp/UdpConnection.Reliable.cs @@ -72,6 +72,10 @@ partial class UdpConnection /// private float _pingMs = 500; + public int MissingAcks => this.reliableDataPacketsSent.Count; + public int MissingPings => this.pingsSinceAck; + public int LastPacketAcked; + /// /// The maximum times a message should be resent before marking the endpoint as disconnected. /// @@ -123,13 +127,14 @@ internal void Set(ushort id, byte[] data, int length, int timeout, Action ackCal // Packets resent public int Resend() { + Packet self; var connection = this.Connection; if (!this.Acknowledged && connection != null) { long lifetimeMs = this.Stopwatch.ElapsedMilliseconds; if (lifetimeMs >= connection.DisconnectTimeoutMs) { - if (connection.reliableDataPacketsSent.TryRemove(this.Id, out Packet self)) + if (connection.reliableDataPacketsSent.TryRemove(this.Id, out self)) { connection.DisconnectInternal(HazelInternalErrors.ReliablePacketWithoutResponse, $"Reliable packet {self.Id} (size={this.Length}) was not ack'd after {lifetimeMs}ms ({self.Retransmissions} resends)"); @@ -145,7 +150,7 @@ public int Resend() if (connection.ResendLimit != 0 && this.Retransmissions > connection.ResendLimit) { - if (connection.reliableDataPacketsSent.TryRemove(this.Id, out Packet self)) + if (connection.reliableDataPacketsSent.TryRemove(this.Id, out self)) { connection.DisconnectInternal(HazelInternalErrors.ReliablePacketWithoutResponse, $"Reliable packet {self.Id} (size={this.Length}) was not ack'd after {self.Retransmissions} resends ({lifetimeMs}ms)"); @@ -438,6 +443,14 @@ private void AcknowledgeMessageId(ushort id) this._pingMs = this._pingMs * .7f + rt * .3f; } } + + lock (PingLock) + { + if (id > this.LastPacketAcked || this.LastPacketAcked - id > 32000) + { + this.LastPacketAcked = id; + } + } } ///