Adds support for custom transport protection algorithms

Motivation:
In some cases AES GCM might not be supported by a server/client, and
clients might require custom algorithms to be implemented in their
projects.

Modifications:
 - Makes NIOSSHTransportProtection, NIOSSHSessionKeys and ExpectedKeySizes public
 - Makes transport algorithms list configurable
 - Adds sequence number of incoming and outgoing messages to parser and
   serializer
 - Tests

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/SSHConnectionStateMachine.swift

@@ -60,11 +60,7 @@ struct SSHConnectionStateMachine {
     /// The state of this state machine.
     private var state: State
 
-    private static let defaultTransportProtectionSchemes: [NIOSSHTransportProtection.Type] = [
-        AES256GCMOpenSSHTransportProtection.self, AES128GCMOpenSSHTransportProtection.self,
-    ]
-
-    init(role: SSHConnectionRole, protectionSchemes: [NIOSSHTransportProtection.Type] = Self.defaultTransportProtectionSchemes) {
+    init(role: SSHConnectionRole, protectionSchemes: [NIOSSHTransportProtection.Type] = SSHConnectionRole.bundledTransportProtectionSchemes) {
         self.state = .idle(IdleState(role: role, protectionSchemes: protectionSchemes))
     }
 

Sources/NIOSSH/Key Exchange/SSHKeyExchangeResult.swift

@@ -46,18 +46,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 var initialInboundIV: [UInt8]
 
-    var initialOutboundIV: [UInt8]
+    public var initialOutboundIV: [UInt8]
 
-    var inboundEncryptionKey: SymmetricKey
+    public var inboundEncryptionKey: SymmetricKey
 
-    var outboundEncryptionKey: SymmetricKey
+    public var outboundEncryptionKey: SymmetricKey
 
-    var inboundMACKey: SymmetricKey
+    public var inboundMACKey: SymmetricKey
 
-    var outboundMACKey: SymmetricKey
+    public 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 +77,18 @@ 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 var ivSize: Int
+
+    public var encryptionKeySize: Int
 
-    var encryptionKeySize: Int
+    public var macKeySize: Int
 
-    var macKeySize: Int
+    public init(ivSize: Int, encryptionKeySize: Int, macKeySize: Int) {
+        self.ivSize = ivSize
+        self.encryptionKeySize = encryptionKeySize
+        self.macKeySize = macKeySize
+    }
 }
+
+extension ExpectedKeySizes: Hashable {}

Sources/NIOSSH/NIOSSHHandler.swift

@@ -57,8 +57,9 @@ public final class NIOSSHHandler {
 
     private var pendingGlobalRequestResponses: CircularBuffer<PendingGlobalRequestResponse?>
 
-    public init(role: SSHConnectionRole, allocator: ByteBufferAllocator, inboundChildChannelInitializer: ((Channel, SSHChannelType) -> EventLoopFuture<Void>)?) {
-        self.stateMachine = SSHConnectionStateMachine(role: role)
+    public init(role: SSHConnectionRole, allocator: ByteBufferAllocator,
+                inboundChildChannelInitializer: ((Channel, SSHChannelType) -> EventLoopFuture<Void>)?) {
+        self.stateMachine = SSHConnectionStateMachine(role: role, protectionSchemes: role.transportProtectionSchemes)
         self.pendingWrite = false
         self.outboundFrameBuffer = allocator.buffer(capacity: 1024)
         self.pendingChannelInitializations = CircularBuffer(initialCapacity: 4)

Sources/NIOSSH/Role.swift

@@ -14,6 +14,10 @@
 
 /// The role of a given party in an SSH connection.
 public enum SSHConnectionRole {
+    public static let bundledTransportProtectionSchemes: [NIOSSHTransportProtection.Type] = [
+        AES256GCMOpenSSHTransportProtection.self, AES128GCMOpenSSHTransportProtection.self,
+    ]
+
     case client(SSHClientConfiguration)
     case server(SSHServerConfiguration)
 
@@ -34,4 +38,14 @@ public enum SSHConnectionRole {
             return true
         }
     }
+
+    internal var transportProtectionSchemes: [NIOSSHTransportProtection.Type] {
+        switch self {
+        case .client(let configuration):
+            return configuration.transportProtectionSchemes
+        case .server(let configuration):
+            return configuration.transportProtectionSchemes
+        }
+    }
 }
+

Sources/NIOSSH/SSHClientConfiguration.swift

@@ -23,11 +23,22 @@ public struct SSHClientConfiguration {
     /// The global request delegate to be used with this client.
     public var globalRequestDelegate: GlobalRequestDelegate
 
+    /// Supported data encryption algorithms
+    public var transportProtectionSchemes: [NIOSSHTransportProtection.Type]
+
     public init(userAuthDelegate: NIOSSHClientUserAuthenticationDelegate,
                 serverAuthDelegate: NIOSSHClientServerAuthenticationDelegate,
                 globalRequestDelegate: GlobalRequestDelegate? = nil) {
+        self.init(userAuthDelegate: userAuthDelegate, serverAuthDelegate: serverAuthDelegate, globalRequestDelegate: globalRequestDelegate, transportProtectionSchemes: SSHConnectionRole.bundledTransportProtectionSchemes)
+    }
+
+    public init(userAuthDelegate: NIOSSHClientUserAuthenticationDelegate,
+                serverAuthDelegate: NIOSSHClientServerAuthenticationDelegate,
+                globalRequestDelegate: GlobalRequestDelegate? = nil,
+                transportProtectionSchemes: [NIOSSHTransportProtection.Type]) {
         self.userAuthDelegate = userAuthDelegate
         self.serverAuthDelegate = serverAuthDelegate
         self.globalRequestDelegate = globalRequestDelegate ?? DefaultGlobalRequestDelegate()
+        self.transportProtectionSchemes = transportProtectionSchemes
     }
 }

Sources/NIOSSH/SSHPacketParser.swift

@@ -25,6 +25,7 @@ struct SSHPacketParser {
 
     private var buffer: ByteBuffer
     private var state: State
+    private(set) var sequenceNumber: UInt32
 
     /// Testing only: the number of bytes we can discard from this buffer.
     internal var _discardableBytes: Int {
@@ -34,6 +35,7 @@ struct SSHPacketParser {
     init(allocator: ByteBufferAllocator) {
         self.buffer = allocator.buffer(capacity: 0)
         self.state = .initialized
+        self.sequenceNumber = 0
     }
 
     mutating func append(bytes: inout ByteBuffer) {
@@ -73,6 +75,7 @@ struct SSHPacketParser {
             if let length = self.buffer.getInteger(at: self.buffer.readerIndex, as: UInt32.self) {
                 if let message = try self.parsePlaintext(length: length) {
                     self.state = .cleartextWaitingForLength
+                    self.sequenceNumber &+= 1
                     return message
                 }
                 self.state = .cleartextWaitingForBytes(length)
@@ -82,6 +85,7 @@ struct SSHPacketParser {
         case .cleartextWaitingForBytes(let length):
             if let message = try self.parsePlaintext(length: length) {
                 self.state = .cleartextWaitingForLength
+                self.sequenceNumber &+= 1
                 return message
             }
             return nil
@@ -92,13 +96,15 @@ struct SSHPacketParser {
 
             if let message = try self.parseCiphertext(length: length, protection: protection) {
                 self.state = .encryptedWaitingForLength(protection)
+                self.sequenceNumber &+= 1
                 return message
             }
             self.state = .encryptedWaitingForBytes(length, protection)
             return nil
         case .encryptedWaitingForBytes(let length, let protection):
             if let message = try self.parseCiphertext(length: length, protection: protection) {
                 self.state = .encryptedWaitingForLength(protection)
+                self.sequenceNumber &+= 1
                 return message
             }
             return nil
@@ -169,7 +175,7 @@ struct SSHPacketParser {
                 return nil
             }
 
-            var content = try protection.decryptAndVerifyRemainingPacket(&buffer)
+            var content = try protection.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: self.sequenceNumber)
             guard let message = try content.readSSHMessage(), content.readableBytes == 0, buffer.readableBytes == 0 else {
                 // Throw this error if the content wasn't exactly the right length for the message.
                 throw NIOSSHError.invalidPacketFormat

Sources/NIOSSH/SSHPacketSerializer.swift

@@ -22,6 +22,7 @@ struct SSHPacketSerializer {
     }
 
     private var state: State = .initialized
+    private(set) var sequenceNumber: UInt32 = 0
 
     /// Encryption schemes can be added to a packet serializer whenever encryption is negotiated.
     mutating func addEncryption(_ protection: NIOSSHTransportProtection) {
@@ -75,9 +76,11 @@ struct SSHPacketSerializer {
             buffer.setInteger(UInt8(paddingLength), at: index + 4)
             /// random padding
             buffer.writeSSHPaddingBytes(count: paddingLength)
+            self.sequenceNumber &+= 1
         case .encrypted(let protection):
             let payload = NIOSSHEncryptablePayload(message: message)
-            try protection.encryptPacket(payload, to: &buffer)
+            try protection.encryptPacket(payload, sequenceNumber: self.sequenceNumber, to: &buffer)
+            self.sequenceNumber &+= 1
         }
     }
 }

Sources/NIOSSH/SSHServerConfiguration.swift

@@ -26,16 +26,28 @@ public struct SSHServerConfiguration {
     /// The ssh banner to display to clients upon authentication
     public var banner: UserAuthBanner?
 
+    /// Supported data encryption algorithms
+    public var transportProtectionSchemes: [NIOSSHTransportProtection.Type]
+
     public init(hostKeys: [NIOSSHPrivateKey], userAuthDelegate: NIOSSHServerUserAuthenticationDelegate, globalRequestDelegate: GlobalRequestDelegate? = nil, banner: UserAuthBanner? = nil) {
-        self.hostKeys = hostKeys
-        self.userAuthDelegate = userAuthDelegate
-        self.globalRequestDelegate = globalRequestDelegate ?? DefaultGlobalRequestDelegate()
-        self.banner = banner
+        self.init(hostKeys: hostKeys, userAuthDelegate: userAuthDelegate, globalRequestDelegate: globalRequestDelegate, banner: banner, transportProtectionSchemes: SSHConnectionRole.bundledTransportProtectionSchemes)
     }
 
     public init(hostKeys: [NIOSSHPrivateKey], userAuthDelegate: NIOSSHServerUserAuthenticationDelegate, globalRequestDelegate: GlobalRequestDelegate? = nil) {
         self.init(hostKeys: hostKeys, userAuthDelegate: userAuthDelegate, globalRequestDelegate: globalRequestDelegate, banner: nil)
     }
+
+    public init(hostKeys: [NIOSSHPrivateKey],
+                userAuthDelegate: NIOSSHServerUserAuthenticationDelegate,
+                globalRequestDelegate: GlobalRequestDelegate? = nil,
+                banner: UserAuthBanner? = nil,
+                transportProtectionSchemes: [NIOSSHTransportProtection.Type]) {
+        self.hostKeys = hostKeys
+        self.userAuthDelegate = userAuthDelegate
+        self.globalRequestDelegate = globalRequestDelegate ?? DefaultGlobalRequestDelegate()
+        self.banner = banner
+        self.transportProtectionSchemes = transportProtectionSchemes
+    }
 }
 
 // MARK: - UserAuthBanner

Sources/NIOSSH/TransportProtection/AESGCM.swift

@@ -79,7 +79,7 @@ extension AESGCMTransportProtection: NIOSSHTransportProtection {
         // unencrypted!
     }
 
-    func decryptAndVerifyRemainingPacket(_ source: inout ByteBuffer) throws -> ByteBuffer {
+    func decryptAndVerifyRemainingPacket(_ source: inout ByteBuffer, sequenceNumber _: UInt32) throws -> ByteBuffer {
         var plaintext: Data
 
         // Establish a nested scope here to avoid the byte buffer views causing an accidental CoW.
@@ -117,7 +117,7 @@ extension AESGCMTransportProtection: NIOSSHTransportProtection {
         return source.readSlice(length: plaintext.count)!
     }
 
-    func encryptPacket(_ packet: NIOSSHEncryptablePayload, to outboundBuffer: inout ByteBuffer) throws {
+    func encryptPacket(_ packet: NIOSSHEncryptablePayload, sequenceNumber _: UInt32, to outboundBuffer: inout ByteBuffer) throws {
         // Keep track of where the length is going to be written.
         let packetLengthIndex = outboundBuffer.writerIndex
         let packetLengthLength = MemoryLayout<UInt32>.size

Sources/NIOSSH/TransportProtection/SSHTransportProtection.swift

@@ -44,7 +44,7 @@ import NIOCore
 /// Implementers of this protocol **must not** expose unauthenticated plaintext, except for the length field. This
 /// is required by the SSH protocol, and swift-nio-ssh does its best to treat the length field as fundamentally
 /// untrusted information.
-protocol NIOSSHTransportProtection: AnyObject {
+public protocol NIOSSHTransportProtection: AnyObject {
     /// The name of the cipher portion of this transport protection scheme as negotiated on the wire.
     static var cipherName: String { get }
 
@@ -87,10 +87,10 @@ protocol NIOSSHTransportProtection: AnyObject {
     /// length, the padding, or the MAC), and update source to indicate the consumed bytes.
     /// It must also perform any integrity checking that
     /// is required and throw if the integrity check fails.
-    func decryptAndVerifyRemainingPacket(_ source: inout ByteBuffer) throws -> ByteBuffer
+    func decryptAndVerifyRemainingPacket(_ source: inout ByteBuffer, sequenceNumber: UInt32) throws -> ByteBuffer
 
     /// Encrypt an entire outbound packet
-    func encryptPacket(_ packet: NIOSSHEncryptablePayload, to outboundBuffer: inout ByteBuffer) throws
+    func encryptPacket(_ packet: NIOSSHEncryptablePayload, sequenceNumber: UInt32, to outboundBuffer: inout ByteBuffer) throws
 }
 
 extension NIOSSHTransportProtection {

Tests/NIOSSHTests/AESGCMTests.swift

@@ -42,7 +42,7 @@ final class AESGCMTests: XCTestCase {
         let initialKeys = self.generateKeys(keySize: .bits128)
 
         let aes128Encryptor = try assertNoThrowWithValue(AES128GCMOpenSSHTransportProtection(initialKeys: initialKeys))
-        XCTAssertNoThrow(try aes128Encryptor.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), to: &self.buffer))
+        XCTAssertNoThrow(try aes128Encryptor.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 0, to: &self.buffer))
 
         // The newKeys message is very straightforward: a single byte. Because of that, we expect that we will need
         // 14 padding bytes: one byte for the padding length, then 14 more to get out to one block size. Thus, the total
@@ -59,7 +59,7 @@ final class AESGCMTests: XCTestCase {
         XCTAssertEqual(bufferCopy, self.buffer)
 
         /// After decryption the plaintext should be a newKeys message.
-        var plaintext = try assertNoThrowWithValue(aes128Decryptor.decryptAndVerifyRemainingPacket(&bufferCopy))
+        var plaintext = try assertNoThrowWithValue(aes128Decryptor.decryptAndVerifyRemainingPacket(&bufferCopy, sequenceNumber: 0))
         XCTAssertEqual(bufferCopy.readableBytes, 0)
         XCTAssertNotEqual(plaintext, self.buffer)
         XCTAssertEqual(plaintext.readableBytes, 1)
@@ -77,7 +77,7 @@ final class AESGCMTests: XCTestCase {
         let initialKeys = self.generateKeys(keySize: .bits256)
 
         let aes256Encryptor = try assertNoThrowWithValue(AES256GCMOpenSSHTransportProtection(initialKeys: initialKeys))
-        XCTAssertNoThrow(try aes256Encryptor.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), to: &self.buffer))
+        XCTAssertNoThrow(try aes256Encryptor.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 0, to: &self.buffer))
 
         // The newKeys message is very straightforward: a single byte. Because of that, we expect that we will need
         // 14 padding bytes: one byte for the padding length, then 14 more to get out to one block size. Thus, the total
@@ -94,7 +94,7 @@ final class AESGCMTests: XCTestCase {
         XCTAssertEqual(bufferCopy, self.buffer)
 
         /// After decryption the plaintext should be a newKeys message.
-        var plaintext = try assertNoThrowWithValue(aes256Decryptor.decryptAndVerifyRemainingPacket(&bufferCopy))
+        var plaintext = try assertNoThrowWithValue(aes256Decryptor.decryptAndVerifyRemainingPacket(&bufferCopy, sequenceNumber: 0))
         XCTAssertEqual(bufferCopy.readableBytes, 0)
         XCTAssertNotEqual(plaintext, self.buffer)
         XCTAssertEqual(plaintext.readableBytes, 1)
@@ -300,7 +300,7 @@ final class AESGCMTests: XCTestCase {
             buffer.clear()
             buffer.writeRepeatingByte(42, count: ciphertextSize)
 
-            XCTAssertThrowsError(try aes128.decryptAndVerifyRemainingPacket(&buffer)) { error in
+            XCTAssertThrowsError(try aes128.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)) { error in
                 XCTAssertEqual((error as? NIOSSHError)?.type, .invalidEncryptedPacketLength)
             }
         }
@@ -320,7 +320,7 @@ final class AESGCMTests: XCTestCase {
             buffer.clear()
             buffer.writeRepeatingByte(42, count: ciphertextSize)
 
-            XCTAssertThrowsError(try aes256.decryptAndVerifyRemainingPacket(&buffer)) { error in
+            XCTAssertThrowsError(try aes256.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)) { error in
                 XCTAssertEqual((error as? NIOSSHError)?.type, .invalidEncryptedPacketLength)
             }
         }
@@ -350,7 +350,7 @@ final class AESGCMTests: XCTestCase {
 
         // We can now attempt to decrypt this packet.
         let aes128 = try assertNoThrowWithValue(AES128GCMOpenSSHTransportProtection(initialKeys: keys))
-        XCTAssertThrowsError(try aes128.decryptAndVerifyRemainingPacket(&buffer)) { error in
+        XCTAssertThrowsError(try aes128.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)) { error in
             XCTAssertEqual((error as? NIOSSHError)?.type, .excessPadding)
         }
     }
@@ -379,7 +379,7 @@ final class AESGCMTests: XCTestCase {
 
         // We can now attempt to decrypt this packet.
         let aes256 = try assertNoThrowWithValue(AES256GCMOpenSSHTransportProtection(initialKeys: keys))
-        XCTAssertThrowsError(try aes256.decryptAndVerifyRemainingPacket(&buffer)) { error in
+        XCTAssertThrowsError(try aes256.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)) { error in
             XCTAssertEqual((error as? NIOSSHError)?.type, .excessPadding)
         }
     }
@@ -408,7 +408,7 @@ final class AESGCMTests: XCTestCase {
 
         // We can now attempt to decrypt this packet.
         let aes128 = try assertNoThrowWithValue(AES128GCMOpenSSHTransportProtection(initialKeys: keys))
-        XCTAssertThrowsError(try aes128.decryptAndVerifyRemainingPacket(&buffer)) { error in
+        XCTAssertThrowsError(try aes128.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)) { error in
             XCTAssertEqual((error as? NIOSSHError)?.type, .insufficientPadding)
         }
     }
@@ -437,7 +437,7 @@ final class AESGCMTests: XCTestCase {
 
         // We can now attempt to decrypt this packet.
         let aes256 = try assertNoThrowWithValue(AES256GCMOpenSSHTransportProtection(initialKeys: keys))
-        XCTAssertThrowsError(try aes256.decryptAndVerifyRemainingPacket(&buffer)) { error in
+        XCTAssertThrowsError(try aes256.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)) { error in
             XCTAssertEqual((error as? NIOSSHError)?.type, .insufficientPadding)
         }
     }