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,12 +60,12 @@ struct SSHConnectionStateMachine {
     /// The state of this state machine.
     private var state: State
 
-    private static let defaultTransportProtectionSchemes: [NIOSSHTransportProtection.Type] = [
+    static let bundledTransportProtectionSchemes: [NIOSSHTransportProtection.Type] = [
         AES256GCMOpenSSHTransportProtection.self, AES128GCMOpenSSHTransportProtection.self,
     ]
 
-    init(role: SSHConnectionRole, protectionSchemes: [NIOSSHTransportProtection.Type] = Self.defaultTransportProtectionSchemes) {
-        self.state = .idle(IdleState(role: role, protectionSchemes: protectionSchemes))
+    init(role: SSHConnectionRole) {
+        self.state = .idle(IdleState(role: role))
     }
 
     func start() -> SSHMultiMessage? {
@@ -182,6 +182,7 @@ struct SSHConnectionStateMachine {
                 return .noMessage
             case .unimplemented(let unimplemented):
                 throw NIOSSHError.remotePeerDoesNotSupportMessage(unimplemented)
+
             default:
                 // TODO: enforce RFC 4253:
                 //

Sources/NIOSSH/Connection State Machine/States/ActiveState.swift

@@ -27,16 +27,13 @@ extension SSHConnectionStateMachine {
 
         internal var remoteVersion: String
 
-        internal var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         internal var sessionIdentifier: ByteBuffer
 
         init(_ previous: UserAuthenticationState) {
             self.role = previous.role
             self.serializer = previous.serializer
             self.parser = previous.parser
             self.remoteVersion = previous.remoteVersion
-            self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentifier = previous.sessionIdentifier
         }
 
@@ -45,7 +42,6 @@ extension SSHConnectionStateMachine {
             self.serializer = previous.serializer
             self.parser = previous.parser
             self.remoteVersion = previous.remoteVersion
-            self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentifier = previous.sessionIdentifier
         }
 
@@ -54,7 +50,6 @@ extension SSHConnectionStateMachine {
             self.serializer = previous.serializer
             self.parser = previous.parser
             self.remoteVersion = previous.remoteVersion
-            self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentifier = previous.sessionIdentifier
         }
     }

Sources/NIOSSH/Connection State Machine/States/IdleState.swift

@@ -23,10 +23,13 @@ extension SSHConnectionStateMachine {
 
         internal var protectionSchemes: [NIOSSHTransportProtection.Type]
 
-        init(role: SSHConnectionRole, protectionSchemes: [NIOSSHTransportProtection.Type]) {
+        internal var keyExchangeAlgorithms: [NIOSSHKeyExchangeAlgorithmProtocol.Type]
+
+        init(role: SSHConnectionRole) {
             self.role = role
             self.serializer = SSHPacketSerializer()
-            self.protectionSchemes = protectionSchemes
+            self.protectionSchemes = role.transportProtectionSchemes
+            self.keyExchangeAlgorithms = role.keyExchangeAlgorithms
         }
     }
 }

Sources/NIOSSH/Connection State Machine/States/KeyExchangeState.swift

@@ -28,8 +28,6 @@ extension SSHConnectionStateMachine {
 
         var remoteVersion: String
 
-        var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         /// The backing state machine.
         var keyExchangeStateMachine: SSHKeyExchangeStateMachine
 
@@ -38,8 +36,7 @@ extension SSHConnectionStateMachine {
             self.parser = state.parser
             self.serializer = state.serializer
             self.remoteVersion = remoteVersion
-            self.protectionSchemes = state.protectionSchemes
-            self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: state.role, remoteVersion: remoteVersion, protectionSchemes: state.protectionSchemes, previousSessionIdentifier: nil)
+            self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: state.role, remoteVersion: remoteVersion, keyExchangeAlgorithms: state.role.keyExchangeAlgorithms, transportProtectionSchemes: state.role.transportProtectionSchemes, previousSessionIdentifier: nil)
         }
     }
 }

Sources/NIOSSH/Connection State Machine/States/ReceivedKexInitWhenActiveState.swift

@@ -28,8 +28,6 @@ extension SSHConnectionStateMachine {
 
         internal var remoteVersion: String
 
-        internal var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         internal var keyExchangeStateMachine: SSHKeyExchangeStateMachine
 
         internal var sessionIdentifier: ByteBuffer
@@ -39,9 +37,8 @@ extension SSHConnectionStateMachine {
             self.serializer = previous.serializer
             self.parser = previous.parser
             self.remoteVersion = previous.remoteVersion
-            self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentifier = previous.sessionIdentifier
-            self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: previous.role, remoteVersion: previous.remoteVersion, protectionSchemes: previous.protectionSchemes, previousSessionIdentifier: self.sessionIdentifier)
+            self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: previous.role, remoteVersion: previous.remoteVersion, keyExchangeAlgorithms: self.role.keyExchangeAlgorithms, transportProtectionSchemes: self.role.transportProtectionSchemes, previousSessionIdentifier: self.sessionIdentifier)
         }
     }
 }

Sources/NIOSSH/Connection State Machine/States/ReceivedNewKeysState.swift

@@ -29,8 +29,6 @@ extension SSHConnectionStateMachine {
 
         var remoteVersion: String
 
-        var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         var sessionIdentifier: ByteBuffer
 
         /// The backing state machine.
@@ -45,7 +43,6 @@ extension SSHConnectionStateMachine {
             self.parser = state.parser
             self.serializer = state.serializer
             self.remoteVersion = state.remoteVersion
-            self.protectionSchemes = state.protectionSchemes
             self.keyExchangeStateMachine = state.keyExchangeStateMachine
 
             // We force unwrap the session ID because it's programmer error to not have it at this time.

Sources/NIOSSH/Connection State Machine/States/RekeyingReceivedNewKeysState.swift

@@ -29,8 +29,6 @@ extension SSHConnectionStateMachine {
 
         var remoteVersion: String
 
-        var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         var sessionIdentifier: ByteBuffer
 
         /// The backing state machine.
@@ -41,7 +39,6 @@ extension SSHConnectionStateMachine {
             self.parser = previousState.parser
             self.serializer = previousState.serializer
             self.remoteVersion = previousState.remoteVersion
-            self.protectionSchemes = previousState.protectionSchemes
             self.sessionIdentifier = previousState.sessionIdentifier
             self.keyExchangeStateMachine = previousState.keyExchangeStateMachine
         }

Sources/NIOSSH/Connection State Machine/States/RekeyingSentNewKeysState.swift

@@ -29,8 +29,6 @@ extension SSHConnectionStateMachine {
 
         var remoteVersion: String
 
-        var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         var sessionIdentifier: ByteBuffer
 
         /// The backing state machine.
@@ -41,7 +39,6 @@ extension SSHConnectionStateMachine {
             self.parser = previousState.parser
             self.serializer = previousState.serializer
             self.remoteVersion = previousState.remoteVersion
-            self.protectionSchemes = previousState.protectionSchemes
             self.sessionIdentifier = previousState.sessionIdentifier
             self.keyExchangeStateMachine = previousState.keyExchangeStateMachine
         }

Sources/NIOSSH/Connection State Machine/States/RekeyingState.swift

@@ -28,8 +28,6 @@ extension SSHConnectionStateMachine {
 
         var remoteVersion: String
 
-        var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         var sessionIdentifier: ByteBuffer
 
         /// The backing state machine.
@@ -40,7 +38,6 @@ extension SSHConnectionStateMachine {
             self.parser = previousState.parser
             self.serializer = previousState.serializer
             self.remoteVersion = previousState.remoteVersion
-            self.protectionSchemes = previousState.protectionSchemes
             self.sessionIdentifier = previousState.sessionIdentifier
             self.keyExchangeStateMachine = previousState.keyExchangeStateMachine
         }
@@ -50,7 +47,6 @@ extension SSHConnectionStateMachine {
             self.parser = previousState.parser
             self.serializer = previousState.serializer
             self.remoteVersion = previousState.remoteVersion
-            self.protectionSchemes = previousState.protectionSchemes
             self.sessionIdentifier = previousState.sessionIdentitifier
             self.keyExchangeStateMachine = previousState.keyExchangeStateMachine
         }

Sources/NIOSSH/Connection State Machine/States/SentKexInitWhenActiveState.swift

@@ -28,8 +28,6 @@ extension SSHConnectionStateMachine {
 
         internal var remoteVersion: String
 
-        internal var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         internal var sessionIdentitifier: ByteBuffer
 
         internal var keyExchangeStateMachine: SSHKeyExchangeStateMachine
@@ -39,9 +37,8 @@ extension SSHConnectionStateMachine {
             self.serializer = previous.serializer
             self.parser = previous.parser
             self.remoteVersion = previous.remoteVersion
-            self.protectionSchemes = previous.protectionSchemes
             self.sessionIdentitifier = previous.sessionIdentifier
-            self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: self.role, remoteVersion: self.remoteVersion, protectionSchemes: self.protectionSchemes, previousSessionIdentifier: previous.sessionIdentifier)
+            self.keyExchangeStateMachine = SSHKeyExchangeStateMachine(allocator: allocator, loop: loop, role: self.role, remoteVersion: self.remoteVersion, keyExchangeAlgorithms: self.role.keyExchangeAlgorithms, transportProtectionSchemes: self.role.transportProtectionSchemes, previousSessionIdentifier: previous.sessionIdentifier)
         }
     }
 }

Sources/NIOSSH/Connection State Machine/States/SentNewKeysState.swift

@@ -29,8 +29,6 @@ extension SSHConnectionStateMachine {
 
         var remoteVersion: String
 
-        var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         var sessionIdentifier: ByteBuffer
 
         /// The backing state machine.
@@ -46,7 +44,6 @@ extension SSHConnectionStateMachine {
             self.serializer = state.serializer
             self.keyExchangeStateMachine = state.keyExchangeStateMachine
             self.remoteVersion = state.remoteVersion
-            self.protectionSchemes = state.protectionSchemes
 
             // We force unwrap the session ID here because it's programmer error to not have it at this stage.
             self.sessionIdentifier = self.keyExchangeStateMachine.sessionID!

Sources/NIOSSH/Connection State Machine/States/SentVersionState.swift

@@ -26,14 +26,11 @@ extension SSHConnectionStateMachine {
         /// The packet serializer used by this state machine.
         var serializer: SSHPacketSerializer
 
-        var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         private let allocator: ByteBufferAllocator
 
         init(idleState state: IdleState, allocator: ByteBufferAllocator) {
             self.role = state.role
             self.serializer = state.serializer
-            self.protectionSchemes = state.protectionSchemes
 
             self.parser = SSHPacketParser(allocator: allocator)
             self.allocator = allocator

Sources/NIOSSH/Connection State Machine/States/UserAuthenticationState.swift

@@ -28,8 +28,6 @@ extension SSHConnectionStateMachine {
 
         var remoteVersion: String
 
-        var protectionSchemes: [NIOSSHTransportProtection.Type]
-
         var sessionIdentifier: ByteBuffer
 
         /// The backing state machine.
@@ -41,7 +39,6 @@ extension SSHConnectionStateMachine {
             self.serializer = state.serializer
             self.userAuthStateMachine = state.userAuthStateMachine
             self.remoteVersion = state.remoteVersion
-            self.protectionSchemes = state.protectionSchemes
             self.sessionIdentifier = state.sessionIdentifier
         }
 
@@ -51,7 +48,6 @@ extension SSHConnectionStateMachine {
             self.serializer = state.serializer
             self.userAuthStateMachine = state.userAuthStateMachine
             self.remoteVersion = state.remoteVersion
-            self.protectionSchemes = state.protectionSchemes
             self.sessionIdentifier = state.sessionIdentifier
         }
     }