Correctly handle host key negotiation.

Sources/NIOSSH/Key Exchange/SSHKeyExchangeStateMachine.swift

@@ -244,6 +244,11 @@ struct SSHKeyExchangeStateMachine {
         case .keyExchangeInitSent(exchange: var exchanger, negotiated: let negotiated):
             switch self.role {
             case .client:
+                guard message.hostKey.keyPrefix.elementsEqual(negotiated.negotiatedHostKeyAlgorithm.utf8) else {
+                    throw NIOSSHError.invalidHostKeyForKeyExchange(expected: negotiated.negotiatedHostKeyAlgorithm,
+                                                                   got: message.hostKey.keyPrefix)
+                }
+
                 let result = try exchanger.receiveServerKeyExchangePayload(
                     serverKeyExchangeMessage: message,
                     initialExchangeBytes: &self.initialExchangeBytes,
@@ -347,15 +352,20 @@ struct SSHKeyExchangeStateMachine {
         // Ok, rephrase as client and server instead of us and them.
         let clientAlgorithms: [Substring]
         let serverAlgorithms: [Substring]
-        let supportedHostKeyAlgorithms = self.supportedHostKeyAlgorithms
+        let clientHostKeyAlgorithms: [Substring]
+        let serverHostKeyAlgorithms: [Substring]
 
         switch self.role {
         case .client:
             clientAlgorithms = Self.supportedKeyExchangeAlgorithms
             serverAlgorithms = peerKeyExchangeAlgorithms
+            clientHostKeyAlgorithms = self.supportedHostKeyAlgorithms
+            serverHostKeyAlgorithms = peerHostKeyAlgorithms
         case .server:
             clientAlgorithms = peerKeyExchangeAlgorithms
             serverAlgorithms = Self.supportedKeyExchangeAlgorithms
+            clientHostKeyAlgorithms = peerHostKeyAlgorithms
+            serverHostKeyAlgorithms = self.supportedHostKeyAlgorithms
         }
 
         // Let's find the first protocol the client supports that the server does too.
@@ -365,9 +375,9 @@ struct SSHKeyExchangeStateMachine {
             }
 
             // Ok, got one. We need a signing capable host key algorithm, which for us is all of them.
-            // We iterate over our list because it's almost certainly smaller than the peer's.
-            for hostKeyAlgorithm in supportedHostKeyAlgorithms {
-                guard peerHostKeyAlgorithms.contains(hostKeyAlgorithm) else {
+            // Again, we prefer the first one the client supports that the server does too.
+            for hostKeyAlgorithm in clientHostKeyAlgorithms {
+                guard serverHostKeyAlgorithms.contains(hostKeyAlgorithm) else {
                     continue
                 }
 
@@ -507,6 +517,23 @@ extension SSHKeyExchangeStateMachine {
             return nil
         }
     }
+
+    var _testOnly_negotiatedHostKeyAlgorithm: Substring? {
+        switch self.state {
+        case .idle, .keyExchangeSent, .complete:
+            return nil
+
+        case .keyExchangeReceived(_, negotiated: let negotiated, _),
+             .awaitingKeyExchangeInitInvalidGuess(_, negotiated: let negotiated),
+             .awaitingKeyExchangeInit(_, negotiated: let negotiated),
+             .keyExchangeInitReceived(_, negotiated: let negotiated),
+             .keyExchangeInitSent(_, negotiated: let negotiated),
+             .keysExchanged(_, _, negotiated: let negotiated),
+             .newKeysReceived(_, _, negotiated: let negotiated),
+             .newKeysSent(_, _, negotiated: let negotiated):
+            return negotiated.negotiatedHostKeyAlgorithm
+        }
+    }
 }
 
 extension CSPRNG {

Sources/NIOSSH/NIOSSHError.swift

@@ -118,6 +118,11 @@ extension NIOSSHError {
     internal static func remotePeerDoesNotSupportMessage(_ message: SSHMessage.UnimplementedMessage) -> NIOSSHError {
         NIOSSHError(type: .remotePeerDoesNotSupportMessage, diagnostics: "Sequence Number: \(message.sequenceNumber)")
     }
+
+    @inline(never)
+    internal static func invalidHostKeyForKeyExchange(expected: Substring, got actual: String.UTF8View) -> NIOSSHError {
+        NIOSSHError(type: .invalidHostKeyForKeyExchange, diagnostics: "Expected \(String(expected)), got \(String(actual))")
+    }
 }
 
 // MARK: - NIOSSHError CustomStringConvertible conformance.
@@ -161,6 +166,7 @@ extension NIOSSHError {
             case globalRequestRefused
             case missingGlobalRequestResponse
             case remotePeerDoesNotSupportMessage
+            case invalidHostKeyForKeyExchange
         }
 
         private var base: Base
@@ -249,6 +255,9 @@ extension NIOSSHError {
 
         /// The remote peer sent an "unimplemented" message, indicating they do not support a message we sent.
         public static let remotePeerDoesNotSupportMessage: ErrorType = .init(.remotePeerDoesNotSupportMessage)
+
+        /// The peer has sent a host key that does not correspond to the one negotiated in key exchange.
+        public static let invalidHostKeyForKeyExchange: ErrorType = .init(.invalidHostKeyForKeyExchange)
     }
 }
 

Tests/NIOSSHTests/SSHKeyExchangeStateMachineTests.swift

@@ -682,6 +682,123 @@ final class SSHKeyExchangeStateMachineTests: XCTestCase {
         self.assertCompatibleProtection(client: clientInboundProtection, server: serverInboundProtection)
         XCTAssertTrue(clientInboundProtection is AES128GCMOpenSSHTransportProtection)
     }
+
+    func testKeyExchangeDifferentPriorityBetweenServerAndClient() throws {
+        // This test runs a full key exchange where the server's host keys are in a different order to the client.
+        // This flushes out a bug where we'd choose the wrong host key algorithm because we preferred our own list in
+        // all cases, rather than iterating the client's.
+        let allocator = ByteBufferAllocator()
+        var keys = [NIOSSHPrivateKey(p256Key: .init()), NIOSSHPrivateKey(ed25519Key: .init())]
+
+        let serverKeyAlgorithms = keys.flatMap { $0.hostKeyAlgorithms }
+        XCTAssertEqual(serverKeyAlgorithms.count, 2)
+        let clientKeyAlgorithms = SSHKeyExchangeStateMachine.supportedServerHostKeyAlgorithms.filter { serverKeyAlgorithms.contains($0) }
+        XCTAssertEqual(clientKeyAlgorithms.count, 2)
+
+        // We force these out of order so that there's no way they agree.
+        if serverKeyAlgorithms.first == clientKeyAlgorithms.first {
+            keys.swapAt(0, 1)
+        }
+
+        // From here on, things should run smoothly.
+        var client = SSHKeyExchangeStateMachine(
+            allocator: allocator,
+            role: .client(.init(userAuthDelegate: ExplodingAuthDelegate())),
+            remoteVersion: Constants.version,
+            protectionSchemes: [AES256GCMOpenSSHTransportProtection.self],
+            previousSessionIdentifier: nil
+        )
+        var server = SSHKeyExchangeStateMachine(
+            allocator: allocator,
+            role: .server(.init(hostKeys: keys, userAuthDelegate: DenyAllServerAuthDelegate())),
+            remoteVersion: Constants.version,
+            protectionSchemes: [AES256GCMOpenSSHTransportProtection.self],
+            previousSessionIdentifier: nil
+        )
+
+        // Both sides begin by generating a key exchange message.
+        let serverMessage = server.createKeyExchangeMessage()
+        let clientMessage = client.createKeyExchangeMessage()
+        server.send(keyExchange: serverMessage)
+        client.send(keyExchange: clientMessage)
+
+        // The server does not generate a response message, but the client does.
+        try self.assertGeneratesNoMessage(server.handle(keyExchange: clientMessage))
+
+        let ecdhInit = try assertGeneratesECDHKeyExchangeInit(client.handle(keyExchange: serverMessage))
+        client.send(keyExchangeInit: ecdhInit)
+
+        // Now the server receives the ECDH init message and generates the reply, as well as the newKeys message.
+        let ecdhReply = try assertGeneratesECDHKeyExchangeReplyAndNewKeys(server.handle(keyExchangeInit: ecdhInit))
+        XCTAssertNoThrow(try server.send(keyExchangeReply: ecdhReply))
+        let serverOutboundProtection = server.sendNewKeys()
+
+        // At this time, both parties have negotiated the algorithm in question, so we check it here. This should be
+        // the one that was _last_ in the server' list.
+        XCTAssertEqual(client._testOnly_negotiatedHostKeyAlgorithm, server._testOnly_negotiatedHostKeyAlgorithm)
+        XCTAssertEqual(client._testOnly_negotiatedHostKeyAlgorithm, keys.last.flatMap { $0.hostKeyAlgorithms.first })
+
+        // Now the client receives the reply and newKeys, and generates a newKeys message.
+        try self.assertGeneratesNewKeys(client.handle(keyExchangeReply: ecdhReply))
+        let clientInboundProtection = try assertNoThrowWithValue(client.handleNewKeys())
+        let clientOutboundProtection = client.sendNewKeys()
+
+        // Server receives the newKeys.
+        let serverInboundProtection = try assertNoThrowWithValue(server.handleNewKeys())
+
+        // Each peer has generated the exact same protection object for both directions.
+        XCTAssertTrue(clientInboundProtection === clientOutboundProtection)
+        XCTAssertTrue(serverInboundProtection === serverOutboundProtection)
+
+        self.assertCompatibleProtection(client: clientInboundProtection, server: serverInboundProtection)
+    }
+
+    func testKeyExchangeServerReturnsWrongHostKeyType() throws {
+        // This test runs a full key exchange where the server returns a host key type it didn't negotiate. This should
+        // error.
+        let allocator = ByteBufferAllocator()
+
+        var client = SSHKeyExchangeStateMachine(
+            allocator: allocator,
+            role: .client(.init(userAuthDelegate: ExplodingAuthDelegate())),
+            remoteVersion: Constants.version,
+            protectionSchemes: [AES256GCMOpenSSHTransportProtection.self],
+            previousSessionIdentifier: nil
+        )
+        var server = SSHKeyExchangeStateMachine(
+            allocator: allocator,
+            role: .server(.init(hostKeys: [NIOSSHPrivateKey(p256Key: .init())], userAuthDelegate: DenyAllServerAuthDelegate())),
+            remoteVersion: Constants.version,
+            protectionSchemes: [AES256GCMOpenSSHTransportProtection.self],
+            previousSessionIdentifier: nil
+        )
+
+        // Both sides begin by generating a key exchange message.
+        let serverMessage = server.createKeyExchangeMessage()
+        let clientMessage = client.createKeyExchangeMessage()
+        server.send(keyExchange: serverMessage)
+        client.send(keyExchange: clientMessage)
+
+        // The server does not generate a response message, but the client does.
+        try self.assertGeneratesNoMessage(server.handle(keyExchange: clientMessage))
+
+        let ecdhInit = try assertGeneratesECDHKeyExchangeInit(client.handle(keyExchange: serverMessage))
+        client.send(keyExchangeInit: ecdhInit)
+
+        // Now the server receives the ECDH init message and generates the reply. We ignore newKeys here as we
+        // won't get that far.
+        var ecdhReply = try assertGeneratesECDHKeyExchangeReplyAndNewKeys(server.handle(keyExchangeInit: ecdhInit))
+        XCTAssertNoThrow(try server.send(keyExchangeReply: ecdhReply))
+
+        // We tweak the reply. This would lead to an invalid signature, but we expect it to be tripped up on
+        // the wrong negotiation first.
+        ecdhReply.hostKey = NIOSSHPrivateKey(ed25519Key: .init()).publicKey
+
+        // Now the client receives the reply and errors.
+        XCTAssertThrowsError(try client.handle(keyExchangeReply: ecdhReply)) { error in
+            XCTAssertEqual((error as? NIOSSHError)?.type, .invalidHostKeyForKeyExchange)
+        }
+    }
 }
 
 extension SSHKeyExchangeStateMachineTests {