Harden against "leaking" message registration.

Fixes issue #55.

src/Renci.SshNet/KeyboardInteractiveAuthenticationMethod.cs

@@ -57,15 +57,16 @@ namespace Renci.SshNet
             session.UserAuthenticationFailureReceived += Session_UserAuthenticationFailureReceived;
             session.MessageReceived += Session_MessageReceived;
 
+            session.RegisterMessage("SSH_MSG_USERAUTH_INFO_REQUEST");
+
             try
             {
-                session.RegisterMessage("SSH_MSG_USERAUTH_INFO_REQUEST");
                 session.SendMessage(_requestMessage);
                 session.WaitOnHandle(_authenticationCompleted);
-                session.UnRegisterMessage("SSH_MSG_USERAUTH_INFO_REQUEST");
             }
             finally
             {
+                session.UnRegisterMessage("SSH_MSG_USERAUTH_INFO_REQUEST");
                 session.UserAuthenticationSuccessReceived -= Session_UserAuthenticationSuccessReceived;
                 session.UserAuthenticationFailureReceived -= Session_UserAuthenticationFailureReceived;
                 session.MessageReceived -= Session_MessageReceived;

src/Renci.SshNet/PasswordAuthenticationMethod.cs

@@ -89,6 +89,7 @@ namespace Renci.SshNet
             }
             finally 
             {
+                session.UnRegisterMessage("SSH_MSG_USERAUTH_PASSWD_CHANGEREQ");
                 session.UserAuthenticationSuccessReceived -= Session_UserAuthenticationSuccessReceived;
                 session.UserAuthenticationFailureReceived -= Session_UserAuthenticationFailureReceived;
                 session.MessageReceived -= Session_MessageReceived;

src/Renci.SshNet/PrivateKeyAuthenticationMethod.cs

@@ -60,55 +60,65 @@ namespace Renci.SshNet
 
             session.RegisterMessage("SSH_MSG_USERAUTH_PK_OK");
 
-            foreach (var keyFile in KeyFiles)
+            try
             {
-                _authenticationCompleted.Reset();
-                _isSignatureRequired = false;
+                foreach (var keyFile in KeyFiles)
+                {
+                    _authenticationCompleted.Reset();
+                    _isSignatureRequired = false;
 
-                var message = new RequestMessagePublicKey(ServiceName.Connection, Username, keyFile.HostKey.Name, keyFile.HostKey.Data);
+                    var message = new RequestMessagePublicKey(ServiceName.Connection,
+                                                              Username,
+                                                              keyFile.HostKey.Name,
+                                                              keyFile.HostKey.Data);
 
-                if (KeyFiles.Count < 2)
-                {
-                    //  If only one key file provided then send signature for very first request
-                    var signatureData = new SignatureData(message, session.SessionId).GetBytes();
+                    if (KeyFiles.Count < 2)
+                    {
+                        //  If only one key file provided then send signature for very first request
+                        var signatureData = new SignatureData(message, session.SessionId).GetBytes();
 
-                    message.Signature = keyFile.HostKey.Sign(signatureData);
-                }
+                        message.Signature = keyFile.HostKey.Sign(signatureData);
+                    }
 
-                //  Send public key authentication request
-                session.SendMessage(message);
+                    //  Send public key authentication request
+                    session.SendMessage(message);
 
-                session.WaitOnHandle(_authenticationCompleted);
+                    session.WaitOnHandle(_authenticationCompleted);
 
-                if (_isSignatureRequired)
-                {
-                    _authenticationCompleted.Reset();
+                    if (_isSignatureRequired)
+                    {
+                        _authenticationCompleted.Reset();
 
-                    var signatureMessage = new RequestMessagePublicKey(ServiceName.Connection, Username, keyFile.HostKey.Name, keyFile.HostKey.Data);
+                        var signatureMessage = new RequestMessagePublicKey(ServiceName.Connection,
+                                                                           Username,
+                                                                           keyFile.HostKey.Name,
+                                                                           keyFile.HostKey.Data);
 
-                    var signatureData = new SignatureData(message, session.SessionId).GetBytes();
+                        var signatureData = new SignatureData(message, session.SessionId).GetBytes();
 
-                    signatureMessage.Signature = keyFile.HostKey.Sign(signatureData);
+                        signatureMessage.Signature = keyFile.HostKey.Sign(signatureData);
 
-                    //  Send public key authentication request with signature
-                    session.SendMessage(signatureMessage);
-                }
+                        //  Send public key authentication request with signature
+                        session.SendMessage(signatureMessage);
+                    }
 
-                session.WaitOnHandle(_authenticationCompleted);
+                    session.WaitOnHandle(_authenticationCompleted);
 
-                if (_authenticationResult == AuthenticationResult.Success)
-                {
-                    break;
+                    if (_authenticationResult == AuthenticationResult.Success)
+                    {
+                        break;
+                    }
                 }
-            }
-            
-            session.UserAuthenticationSuccessReceived -= Session_UserAuthenticationSuccessReceived;
-            session.UserAuthenticationFailureReceived -= Session_UserAuthenticationFailureReceived;
-            session.MessageReceived -= Session_MessageReceived;
-
-            session.UnRegisterMessage("SSH_MSG_USERAUTH_PK_OK");
 
-            return _authenticationResult;
+                return _authenticationResult;
+            }
+            finally
+            {
+                session.UserAuthenticationSuccessReceived -= Session_UserAuthenticationSuccessReceived;
+                session.UserAuthenticationFailureReceived -= Session_UserAuthenticationFailureReceived;
+                session.MessageReceived -= Session_MessageReceived;
+                session.UnRegisterMessage("SSH_MSG_USERAUTH_PK_OK");
+            }
         }
 
         private void Session_UserAuthenticationSuccessReceived(object sender, MessageEventArgs<SuccessMessage> e)
@@ -168,8 +178,8 @@ namespace Renci.SshNet
                 var authenticationCompleted = _authenticationCompleted;
                 if (authenticationCompleted != null)
                 {
-                    authenticationCompleted.Dispose();
                     _authenticationCompleted = null;
+                    authenticationCompleted.Dispose();
                 }
 
                 _isDisposed = true;
@@ -243,6 +253,5 @@ namespace Renci.SshNet
                 WriteBinaryString(_message.PublicKeyData);
             }
         }
-
     }
 }

src/Renci.SshNet/Security/KeyExchangeDiffieHellmanGroupExchangeShaBase.cs

@@ -50,11 +50,10 @@ namespace Renci.SshNet.Security
             base.Start(session, message);
 
             Session.RegisterMessage("SSH_MSG_KEX_DH_GEX_GROUP");
-            Session.RegisterMessage("SSH_MSG_KEX_DH_GEX_REPLY");
 
             Session.MessageReceived += Session_MessageReceived;
 
-            //  1. send SSH_MSG_KEY_DH_GEX_REQUEST
+            // 1. client sends SSH_MSG_KEY_DH_GEX_REQUEST
             SendMessage(new KeyExchangeDhGroupExchangeRequest(MinimumGroupSize, PreferredGroupSize,
                 MaximumProupSize));
         }
@@ -71,34 +70,38 @@ namespace Renci.SshNet.Security
 
         private void Session_MessageReceived(object sender, MessageEventArgs<Message> e)
         {
+            // 2. server sends SSH_MSG_KEX_DH_GEX_GROUP
             var groupMessage = e.Message as KeyExchangeDhGroupExchangeGroup;
             if (groupMessage != null)
             {
-                //  Unregister message once received
+                // Unregister SSH_MSG_KEX_DH_GEX_GROUP message once received
                 Session.UnRegisterMessage("SSH_MSG_KEX_DH_GEX_GROUP");
+                // Register in order to be able to receive SSH_MSG_KEX_DH_GEX_REPLY message
+                Session.RegisterMessage("SSH_MSG_KEX_DH_GEX_REPLY");
 
-                //  2. Receive SSH_MSG_KEX_DH_GEX_GROUP
                 _prime = groupMessage.SafePrime;
                 _group = groupMessage.SubGroup;
 
                 PopulateClientExchangeValue();
 
-                //  3. Send SSH_MSG_KEX_DH_GEX_INIT
+                // 3. client sends SSH_MSG_KEX_DH_GEX_INIT
                 SendMessage(new KeyExchangeDhGroupExchangeInit(_clientExchangeValue));
+
+                // Skip further execution as we'll be waiting for the SSH_MSG_KEX_DH_GEX_REPLY message
+                return;
             }
-            else
+
+            // 4. server sends SSH_MSG_KEX_DH_GEX_REPLY
+            var replyMessage = e.Message as KeyExchangeDhGroupExchangeReply;
+            if (replyMessage != null)
             {
-                var replyMessage = e.Message as KeyExchangeDhGroupExchangeReply;
-                if (replyMessage != null)
-                {
-                    //  Unregister message once received
-                    Session.UnRegisterMessage("SSH_MSG_KEX_DH_GEX_REPLY");
+                // Unregister SSH_MSG_KEX_DH_GEX_REPLY message once received
+                Session.UnRegisterMessage("SSH_MSG_KEX_DH_GEX_REPLY");
 
-                    HandleServerDhReply(replyMessage.HostKey, replyMessage.F, replyMessage.Signature);
+                HandleServerDhReply(replyMessage.HostKey, replyMessage.F, replyMessage.Signature);
 
-                    //  When SSH_MSG_KEX_DH_GEX_REPLY received key exchange is completed
-                    Finish();
-                }
+                // When SSH_MSG_KEX_DH_GEX_REPLY received key exchange is completed
+                Finish();
             }
         }
     }