From a0d686109c8959aed7ae67f0010686f0322f5ebc Mon Sep 17 00:00:00 2001
From: Oliver Booth <me@olivr.me>
Date: Sat, 9 Jul 2022 10:47:39 +0100
Subject: [PATCH] Switch to AES encryption prior to sending session token

---
 .../PacketHandlers/EncryptionResponsePacketHandler.cs |  3 +--
 TcpDotNet/ProtocolClient.cs                           | 11 +++++++----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/TcpDotNet/Protocol/PacketHandlers/EncryptionResponsePacketHandler.cs b/TcpDotNet/Protocol/PacketHandlers/EncryptionResponsePacketHandler.cs
index b80c1fe..c79914a 100644
--- a/TcpDotNet/Protocol/PacketHandlers/EncryptionResponsePacketHandler.cs
+++ b/TcpDotNet/Protocol/PacketHandlers/EncryptionResponsePacketHandler.cs
@@ -32,10 +32,9 @@ internal sealed class EncryptionResponsePacketHandler : PacketHandler<Encryption
         client.Aes = CryptographyUtils.GenerateAes(key);
         client.State = ClientState.Connected;
         client.ParentListener.OnClientConnect(client);
+        client.UseEncryption = true;
 
         var sessionPacket = new SessionExchangePacket(client.SessionId);
         await client.SendPacketAsync(sessionPacket, cancellationToken);
-
-        client.UseEncryption = true;
     }
 }
diff --git a/TcpDotNet/ProtocolClient.cs b/TcpDotNet/ProtocolClient.cs
index f5a62bf..79e823a 100644
--- a/TcpDotNet/ProtocolClient.cs
+++ b/TcpDotNet/ProtocolClient.cs
@@ -114,13 +114,16 @@ public sealed class ProtocolClient : BaseClientNode
         var key = new byte[128];
         using var rng = new RNGCryptoServiceProvider();
         rng.GetBytes(key);
- 
+
         Aes = CryptographyUtils.GenerateAes(key);
-        var encryptionResponse = new EncryptionResponsePacket(encryptedPayload, rsa.Encrypt(key, true));
-        var sessionPacket = await SendAndReceiveAsync<EncryptionResponsePacket, SessionExchangePacket>(encryptionResponse, cancellationToken);
+        byte[] aesKey = rsa.Encrypt(key, true);
+        var encryptionResponse = new EncryptionResponsePacket(encryptedPayload, aesKey);
+        await SendPacketAsync(encryptionResponse, cancellationToken);
+
+        UseEncryption = true;
+        var sessionPacket = await WaitForPacketAsync<SessionExchangePacket>(cancellationToken);
 
         SessionId = sessionPacket.Session;
-        UseEncryption = true;
         State = ClientState.Connected;
     }