Add support for custom cryptography

Sources/NIOSSH/ByteBuffer+SSH.swift

@@ -169,7 +169,7 @@ extension ByteBuffer {
 
     /// Writes a given number of SSH-acceptable padding bytes to this buffer.
     @discardableResult
-    mutating func writeSSHPaddingBytes(count: Int) -> Int {
+    public mutating func writeSSHPaddingBytes(count: Int) -> Int {
         // Annoyingly, the system random number generator can only give bytes to us 8 bytes at a time.
         precondition(count >= 0, "Cannot write negative number of padding bytes: \(count)")
 

Sources/NIOSSH/Connection State Machine/Operations/AcceptsKeyExchangeMessages.swift

@@ -30,7 +30,7 @@ extension AcceptsKeyExchangeMessages {
     }
 
     mutating func receiveKeyExchangeInitMessage(_ message: SSHMessage.KeyExchangeECDHInitMessage) throws -> SSHConnectionStateMachine.StateMachineInboundProcessResult {
-        let message = try self.keyExchangeStateMachine.handle(keyExchangeInit: message)
+        let message = try self.keyExchangeStateMachine.handle(keyExchangeInit: message.publicKey)
 
         if let message = message {
             return .emitMessage(message)

Sources/NIOSSH/Connection State Machine/SSHConnectionStateMachine.swift

@@ -60,7 +60,7 @@ struct SSHConnectionStateMachine {
     /// The state of this state machine.
     private var state: State
 
-    private static let defaultTransportProtectionSchemes: [NIOSSHTransportProtection.Type] = [
+    internal static var defaultTransportProtectionSchemes: [NIOSSHTransportProtection.Type] = [
         AES256GCMOpenSSHTransportProtection.self, AES128GCMOpenSSHTransportProtection.self,
     ]
 
@@ -181,6 +181,7 @@ struct SSHConnectionStateMachine {
                 return .noMessage
             case .unimplemented(let unimplemented):
                 throw NIOSSHError.remotePeerDoesNotSupportMessage(unimplemented)
+
             default:
                 // TODO: enforce RFC 4253:
                 //
@@ -305,7 +306,7 @@ struct SSHConnectionStateMachine {
                 let result = try state.receiveUserAuthRequest(message)
                 self.state = .userAuthentication(state)
                 return result
-
+                
             case .userAuthSuccess:
                 let result = try state.receiveUserAuthSuccess()
                 // Hey, auth succeeded!
@@ -818,7 +819,7 @@ struct SSHConnectionStateMachine {
             case .userAuthRequest(let message):
                 try state.writeUserAuthRequest(message, into: &buffer)
                 self.state = .userAuthentication(state)
-
+                
             case .userAuthSuccess:
                 try state.writeUserAuthSuccess(into: &buffer)
                 // Ok we're good to go!

Sources/NIOSSH/Key Exchange/EllipticCurveKeyExchange.swift

@@ -16,33 +16,49 @@ import Crypto
 import NIO
 import NIOFoundationCompat
 
+public struct NIOSSHKeyExchangeServerReply {
+    let hostKey: NIOSSHPublicKey
+    let publicKey: ByteBuffer
+    let signature: NIOSSHSignature
+}
+
 /// This protocol defines a container used by the key exchange state machine to manage key exchange.
 /// This type erases the specific key exchanger.
-protocol EllipticCurveKeyExchangeProtocol {
+public protocol NIOSSHKeyExchangeAlgorithmProtocol {
+    static var keyExchangeInitMessageId: UInt8 { get }
+    static var keyExchangeReplyMessageId: UInt8 { get }
+
     init(ourRole: SSHConnectionRole, previousSessionIdentifier: ByteBuffer?)
 
-    func initiateKeyExchangeClientSide(allocator: ByteBufferAllocator) -> SSHMessage.KeyExchangeECDHInitMessage
+    func initiateKeyExchangeClientSide(allocator: ByteBufferAllocator) -> ByteBuffer
 
-    mutating func completeKeyExchangeServerSide(clientKeyExchangeMessage message: SSHMessage.KeyExchangeECDHInitMessage,
-                                                serverHostKey: NIOSSHPrivateKey,
-                                                initialExchangeBytes: inout ByteBuffer,
-                                                allocator: ByteBufferAllocator,
-                                                expectedKeySizes: ExpectedKeySizes) throws -> (KeyExchangeResult, SSHMessage.KeyExchangeECDHReplyMessage)
+    mutating func completeKeyExchangeServerSide(
+        clientKeyExchangeMessage message: ByteBuffer,
+        serverHostKey: NIOSSHPrivateKey,
+        initialExchangeBytes: inout ByteBuffer,
+        allocator: ByteBufferAllocator,
+        expectedKeySizes: ExpectedKeySizes
+    ) throws -> (KeyExchangeResult, NIOSSHKeyExchangeServerReply)
 
-    mutating func receiveServerKeyExchangePayload(serverKeyExchangeMessage message: SSHMessage.KeyExchangeECDHReplyMessage,
-                                                  initialExchangeBytes: inout ByteBuffer,
-                                                  allocator: ByteBufferAllocator,
-                                                  expectedKeySizes: ExpectedKeySizes) throws -> KeyExchangeResult
+    mutating func receiveServerKeyExchangePayload(
+        serverKeyExchangeMessage: NIOSSHKeyExchangeServerReply,
+        initialExchangeBytes: inout ByteBuffer,
+        allocator: ByteBufferAllocator,
+        expectedKeySizes: ExpectedKeySizes
+    ) throws -> KeyExchangeResult
 
     static var keyExchangeAlgorithmNames: [Substring] { get }
 }
 
-struct EllipticCurveKeyExchange<PrivateKey: ECDHCompatiblePrivateKey>: EllipticCurveKeyExchangeProtocol {
+struct EllipticCurveKeyExchange<PrivateKey: ECDHCompatiblePrivateKey>:  NIOSSHKeyExchangeAlgorithmProtocol {
     private var previousSessionIdentifier: ByteBuffer?
     private var ourKey: PrivateKey
     private var theirKey: PrivateKey.PublicKey?
     private var ourRole: SSHConnectionRole
     private var sharedSecret: SharedSecret?
+
+    static var keyExchangeInitMessageId: UInt8 { 30 }
+    static var keyExchangeReplyMessageId: UInt8 { 31 }
 
     init(ourRole: SSHConnectionRole, previousSessionIdentifier: ByteBuffer?) {
         self.ourRole = ourRole
@@ -59,14 +75,13 @@ extension EllipticCurveKeyExchange {
     /// Initiates key exchange by producing an SSH message.
     ///
     /// For now, we just return the ByteBuffer containing the SSH string.
-    func initiateKeyExchangeClientSide(allocator: ByteBufferAllocator) -> SSHMessage.KeyExchangeECDHInitMessage {
+    func initiateKeyExchangeClientSide(allocator: ByteBufferAllocator) -> ByteBuffer {
         precondition(self.ourRole.isClient, "Only clients may initiate the client side key exchange!")
 
         // The largest key we're likely to end up with here is 256 bytes.
         var buffer = allocator.buffer(capacity: 256)
         self.ourKey.publicKey.write(to: &buffer)
-
-        return .init(publicKey: buffer)
+        return buffer
     }
 
     /// Handles receiving the client key exchange payload on the server side.
@@ -77,15 +92,15 @@ extension EllipticCurveKeyExchange {
     ///     - initialExchangeBytes: The initial bytes of the exchange, suitable for writing into the exchange hash.
     ///     - allocator: A `ByteBufferAllocator` suitable for this connection.
     ///     - expectedKeySizes: The sizes of the keys we need to generate.
-    mutating func completeKeyExchangeServerSide(clientKeyExchangeMessage message: SSHMessage.KeyExchangeECDHInitMessage,
+    mutating func completeKeyExchangeServerSide(clientKeyExchangeMessage message: ByteBuffer,
                                                 serverHostKey: NIOSSHPrivateKey,
                                                 initialExchangeBytes: inout ByteBuffer,
                                                 allocator: ByteBufferAllocator,
-                                                expectedKeySizes: ExpectedKeySizes) throws -> (KeyExchangeResult, SSHMessage.KeyExchangeECDHReplyMessage) {
+                                                expectedKeySizes: ExpectedKeySizes) throws -> (KeyExchangeResult, NIOSSHKeyExchangeServerReply) {
         precondition(self.ourRole.isServer, "Only servers may receive a client key exchange packet!")
 
         // With that, we have enough to finalize the key exchange.
-        let kexResult = try self.finalizeKeyExchange(theirKeyBytes: message.publicKey,
+        let kexResult = try self.finalizeKeyExchange(theirKeyBytes: message,
                                                      initialExchangeBytes: &initialExchangeBytes,
                                                      serverHostKey: serverHostKey.publicKey,
                                                      allocator: allocator,
@@ -100,9 +115,9 @@ extension EllipticCurveKeyExchange {
         self.ourKey.publicKey.write(to: &publicKeyBytes)
 
         // Now we have all we need.
-        let responseMessage = SSHMessage.KeyExchangeECDHReplyMessage(hostKey: serverHostKey.publicKey,
-                                                                     publicKey: publicKeyBytes,
-                                                                     signature: exchangeHashSignature)
+        let responseMessage = NIOSSHKeyExchangeServerReply(hostKey: serverHostKey.publicKey,
+                                                           publicKey: publicKeyBytes,
+                                                           signature: exchangeHashSignature)
 
         return (KeyExchangeResult(kexResult), responseMessage)
     }
@@ -116,12 +131,14 @@ extension EllipticCurveKeyExchange {
     ///     - initialExchangeBytes: The initial bytes of the exchange, suitable for writing into the exchange hash.
     ///     - allocator: A `ByteBufferAllocator` suitable for this connection.
     ///     - expectedKeySizes: The sizes of the keys we need to generate.
-    mutating func receiveServerKeyExchangePayload(serverKeyExchangeMessage message: SSHMessage.KeyExchangeECDHReplyMessage,
-                                                  initialExchangeBytes: inout ByteBuffer,
-                                                  allocator: ByteBufferAllocator,
-                                                  expectedKeySizes: ExpectedKeySizes) throws -> KeyExchangeResult {
+    mutating func receiveServerKeyExchangePayload(
+        serverKeyExchangeMessage: NIOSSHKeyExchangeServerReply,
+        initialExchangeBytes: inout ByteBuffer,
+        allocator: ByteBufferAllocator,
+        expectedKeySizes: ExpectedKeySizes
+    ) throws -> KeyExchangeResult {
         precondition(self.ourRole.isClient, "Only clients may receive a server key exchange packet!")
-
+        
         // Ok, we have a few steps here. Firstly, we need to extract the server's public key and generate our shared
         // secret. Then we need to validate that we didn't generate a weak shared secret (possible under some cases),
         // as this must fail the key exchange process.
@@ -131,14 +148,14 @@ extension EllipticCurveKeyExchange {
         //
         // Finally, we return our generated keys to the state machine.
 
-        let kexResult = try self.finalizeKeyExchange(theirKeyBytes: message.publicKey,
+        let kexResult = try self.finalizeKeyExchange(theirKeyBytes: serverKeyExchangeMessage.publicKey,
                                                      initialExchangeBytes: &initialExchangeBytes,
-                                                     serverHostKey: message.hostKey,
+                                                     serverHostKey: serverKeyExchangeMessage.hostKey,
                                                      allocator: allocator,
                                                      expectedKeySizes: expectedKeySizes)
 
         // We can now verify signature over the exchange hash.
-        guard message.hostKey.isValidSignature(message.signature, for: kexResult.exchangeHash) else {
+        guard serverKeyExchangeMessage.hostKey.isValidSignature(serverKeyExchangeMessage.signature, for: kexResult.exchangeHash) else {
             throw NIOSSHError.invalidExchangeHashSignature
         }
 

Sources/NIOSSH/Key Exchange/SSHKeyExchangeResult.swift

@@ -19,11 +19,16 @@ import NIO
 ///
 /// A round of key exchange generates a number of keys and also generates an exchange hash.
 /// This exchange hash is used for a number of purposes.
-struct KeyExchangeResult {
+public struct KeyExchangeResult {
     /// The session ID to use for this connection. Will be static across the lifetime of a connection.
     var sessionID: ByteBuffer
 
     var keys: NIOSSHSessionKeys
+
+    public init(sessionID: ByteBuffer, keys: NIOSSHSessionKeys) {
+        self.sessionID = sessionID
+        self.keys = keys
+    }
 }
 
 extension KeyExchangeResult: Equatable {}
@@ -46,18 +51,27 @@ extension KeyExchangeResult: Equatable {}
 /// Of these types, the encryption keys and the MAC keys are intended to be secret, and so
 /// we store them in the `SymmetricKey` types. The IVs do not need to be secret, and so are
 /// stored in regular heap buffers.
-struct NIOSSHSessionKeys {
-    var initialInboundIV: [UInt8]
+public struct NIOSSHSessionKeys {
+    public internal(set) var initialInboundIV: [UInt8]
 
-    var initialOutboundIV: [UInt8]
+    public internal(set) var initialOutboundIV: [UInt8]
 
-    var inboundEncryptionKey: SymmetricKey
+    public internal(set) var inboundEncryptionKey: SymmetricKey
 
-    var outboundEncryptionKey: SymmetricKey
+    public internal(set) var outboundEncryptionKey: SymmetricKey
 
-    var inboundMACKey: SymmetricKey
+    public internal(set) var inboundMACKey: SymmetricKey
 
-    var outboundMACKey: SymmetricKey
+    public internal(set) var outboundMACKey: SymmetricKey
+
+    public init(initialInboundIV: [UInt8], initialOutboundIV: [UInt8], inboundEncryptionKey: SymmetricKey, outboundEncryptionKey: SymmetricKey, inboundMACKey: SymmetricKey, outboundMACKey: SymmetricKey) {
+        self.initialInboundIV = initialInboundIV
+        self.initialOutboundIV = initialOutboundIV
+        self.inboundEncryptionKey = inboundEncryptionKey
+        self.outboundEncryptionKey = outboundEncryptionKey
+        self.inboundMACKey = inboundMACKey
+        self.outboundMACKey = outboundMACKey
+    }
 }
 
 extension NIOSSHSessionKeys: Equatable {}
@@ -68,10 +82,16 @@ extension NIOSSHSessionKeys: Equatable {}
 /// hash function invocations. The output of these hash functions is truncated to an appropriate
 /// length as needed, which means we need to ensure the code doing the calculation knows how
 /// to truncate appropriately.
-struct ExpectedKeySizes {
-    var ivSize: Int
+public struct ExpectedKeySizes {
+    public internal(set) var ivSize: Int
 
-    var encryptionKeySize: Int
+    public internal(set) var encryptionKeySize: Int
 
-    var macKeySize: Int
+    public internal(set) var macKeySize: Int
+
+    public init(ivSize: Int, encryptionKeySize: Int, macKeySize: Int) {
+        self.ivSize = ivSize
+        self.encryptionKeySize = encryptionKeySize
+        self.macKeySize = macKeySize
+    }
 }