Tweaks.

tmds · Aug 28, 2024 · 8e08dc3 · 8e08dc3
1 parent ac46f41
commit 8e08dc3
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 6 deletions.
diff --git a/src/Tmds.Ssh/ECDHKeyExchange.cs b/src/Tmds.Ssh/ECDHKeyExchange.cs
@@ -35,8 +35,8 @@ public async Task<KeyExchangeOutput> TryExchangeAsync(KeyExchangeContext context
         await context.SendPacketAsync(CreateEcdhInitMessage(sequencePool, q_c), ct).ConfigureAwait(false);
 
         // Receive ECDH_REPLY.
-        using Packet exchangeInitMsg = await context.ReceivePacketAsync(MessageId.SSH_MSG_KEX_ECDH_REPLY, firstPacket.Move(), ct).ConfigureAwait(false);
-        var ecdhReply = ParceEcdhReply(exchangeInitMsg);
+        using Packet ecdhReplyMsg = await context.ReceivePacketAsync(MessageId.SSH_MSG_KEX_ECDH_REPLY, firstPacket.Move(), ct).ConfigureAwait(false);
+        var ecdhReply = ParceEcdhReply(ecdhReplyMsg);
 
         // Verify received key is valid.
         connectionInfo.ServerKey = ecdhReply.public_host_key;

diff --git a/src/Tmds.Ssh/KeyExchangeContext.cs b/src/Tmds.Ssh/KeyExchangeContext.cs
@@ -22,8 +22,15 @@ public KeyExchangeContext(SshConnection connection, SshSession session, bool isI
 
     public SequencePool SequencePool => _connection.SequencePool;
 
-    public ValueTask<Packet> ReceivePacketAsync(CancellationToken ct)
-        => _connection.ReceivePacketAsync(ct);
+    public async ValueTask<Packet> ReceivePacketAsync(CancellationToken ct)
+    {
+        var packet = await _connection.ReceivePacketAsync(ct).ConfigureAwait(false);
+        if (packet.IsEmpty)
+        {
+            ThrowHelper.ThrowProtocolUnexpectedPeerClose();
+        }
+        return packet;
+    }
 
     public async ValueTask<Packet> ReceivePacketAsync(MessageId expected, CancellationToken ct)
     {

diff --git a/src/Tmds.Ssh/SshSession.cs b/src/Tmds.Ssh/SshSession.cs
@@ -398,7 +398,7 @@ private async Task ReceiveLoopAsync(SshConnection connection, SshConnectionInfo
                     // The send loop waits for us to signal kex completion.
                     try
                     {
-                        await PerformKeyExchangeAsync(context, packet, clientKexInitMsg, abortToken).ConfigureAwait(false);
+                        await PerformKeyExchangeAsync(context, serverKexInitMsg: packet, clientKexInitMsg, abortToken).ConfigureAwait(false);
                     }
                     finally
                     {
@@ -419,7 +419,7 @@ private async Task ReceiveLoopAsync(SshConnection connection, SshConnectionInfo
 
     internal void HandleNonKexPacket(MessageId msgId, Packet _p)
     {
-        using Packet packet = _p; // own the dispose.
+        using Packet packet = _p; // Ensure dispose
 
         // Connection Protocol: https://tools.ietf.org/html/rfc4254.
         switch (msgId)