refactoring of the encryptPacket method (#123)

Motivation:
I think there is a bit of a mix of responsibilities between
TransportProtection and PacketSerializer. PacketSerializer is
responsible for writing out plaintext packet structure, but encrypted
packet structure is written out by the TransportProtection
implementations. This means that we will have duplicate implementations
of basic packet writing logic, every implementation of TP will have to
write length, padding length, payload and padding. Also, right now this
API uses EncryptablePayload, which we cannot make fully public and
accessible for writing tests as it will mean making all SSHMessages
public as well.

Modifications:
 - Extract packet writing logic to a separate function that will be
   used by both plaintext and ciphertext modes
 - Removes usage of EncryptablePayload
 - Write plaintext packet structure (including payload) in the parser
 - TransportProtection implementations now only have to encrypt and tag

Sources/NIOSSH/SSHPacketSerializer.swift

@@ -47,40 +47,55 @@ struct SSHPacketSerializer {
                 preconditionFailure("only .version message is allowed at this point")
             }
         case .cleartext:
-            let index = buffer.writerIndex
+            buffer.writeSSHPacket(message: message, lengthEncrypted: true, blockSize: 8)
+            self.sequenceNumber &+= 1
+        case .encrypted(let protection):
+            let index = buffer.readerIndex
+            buffer.moveReaderIndex(to: buffer.writerIndex)
+            buffer.writeSSHPacket(message: message, lengthEncrypted: protection.lengthEncrypted, blockSize: protection.cipherBlockSize)
+            try protection.encryptPacket(&buffer, sequenceNumber: self.sequenceNumber)
+            buffer.moveReaderIndex(to: index)
+            self.sequenceNumber &+= 1
+        }
+    }
+}
 
-            /// Each packet is in the following format:
-            ///
-            ///   uint32        packet_length
-            ///   byte           padding_length
-            ///   byte[n1]  payload; n1 = packet_length - padding_length - 1
-            ///   byte[n2]  random padding; n2 = padding_length
-            ///   byte[m]   mac (Message Authentication Code - MAC); m = mac_length
+extension ByteBuffer {
+    mutating func writeSSHPacket(message: SSHMessage, lengthEncrypted: Bool, blockSize: Int) {
+        let index = self.writerIndex
 
-            /// payload
-            buffer.moveWriterIndex(forwardBy: 5)
-            let messageLength = buffer.writeSSHMessage(message)
+        /// Each packet is in the following format:
+        ///
+        ///   uint32        packet_length
+        ///   byte           padding_length
+        ///   byte[n1]  payload; n1 = packet_length - padding_length - 1
+        ///   byte[n2]  random padding; n2 = padding_length
+        ///   byte[m]   mac (Message Authentication Code - MAC); m = mac_length
 
-            /// RFC 4253 § 6:
-            /// random padding
-            ///   Arbitrary-length padding, such that the total length of (packet_length || padding_length || payload || random padding)
-            ///   is a multiple of the cipher block size or 8, whichever is larger.  There MUST be at least four bytes of padding.  The
-            ///   padding SHOULD consist of random bytes.  The maximum amount of padding is 255 bytes.
-            let blockSize = 8
-            let paddingLength = 3 + blockSize - ((messageLength + blockSize) % blockSize)
+        /// payload
+        self.moveWriterIndex(forwardBy: 5)
+        let messageLength = self.writeSSHMessage(message)
 
-            /// packet_length
-            ///   The length of the packet in bytes, not including 'mac' or the 'packet_length' field itself.
-            buffer.setInteger(UInt32(1 + messageLength + paddingLength), at: index)
-            /// padding_length
-            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, sequenceNumber: self.sequenceNumber, to: &buffer)
-            self.sequenceNumber &+= 1
+        // Depending on on whether packet length is encrypted, padding should reflect that
+        let payloadLength = lengthEncrypted ? messageLength + 5 : messageLength + 1
+
+        /// RFC 4253 § 6:
+        /// random padding
+        ///   Arbitrary-length padding, such that the total length of (packet_length || padding_length || payload || random padding)
+        ///   is a multiple of the cipher block size or 8, whichever is larger.  There MUST be at least four bytes of padding.  The
+        ///   padding SHOULD consist of random bytes.  The maximum amount of padding is 255 bytes.
+        var paddingLength = blockSize - (payloadLength % blockSize)
+        if paddingLength < 4 {
+            paddingLength += blockSize
         }
+
+        /// packet_length
+        ///   The length of the packet in bytes, not including 'mac' or the 'packet_length' field itself.
+        let packetLength = 1 + messageLength + paddingLength
+        self.setInteger(UInt32(packetLength), at: index)
+        /// padding_length
+        self.setInteger(UInt8(paddingLength), at: index + 4)
+        /// random padding
+        self.writeSSHPaddingBytes(count: paddingLength)
     }
 }

Sources/NIOSSH/TransportProtection/AESGCM.swift

@@ -62,6 +62,10 @@ extension AESGCMTransportProtection: NIOSSHTransportProtection {
         16
     }
 
+    var lengthEncrypted: Bool {
+        false
+    }
+
     func updateKeys(_ newKeys: NIOSSHSessionKeys) throws {
         guard newKeys.outboundEncryptionKey.bitCount == Self.keySizes.encryptionKeySize * 8,
             newKeys.inboundEncryptionKey.bitCount == Self.keySizes.encryptionKeySize * 8 else {
@@ -117,61 +121,29 @@ extension AESGCMTransportProtection: NIOSSHTransportProtection {
         return source.readSlice(length: plaintext.count)!
     }
 
-    func encryptPacket(_ packet: NIOSSHEncryptablePayload, sequenceNumber _: UInt32, to outboundBuffer: inout ByteBuffer) throws {
+    func encryptPacket(_ destination: inout ByteBuffer, sequenceNumber: UInt32) throws {
         // Keep track of where the length is going to be written.
-        let packetLengthIndex = outboundBuffer.writerIndex
+        let packetLengthIndex = destination.readerIndex
         let packetLengthLength = MemoryLayout<UInt32>.size
-        let packetPaddingIndex = outboundBuffer.writerIndex + packetLengthLength
-        let packetPaddingLength = MemoryLayout<UInt8>.size
-
-        outboundBuffer.moveWriterIndex(forwardBy: packetLengthLength + packetPaddingLength)
-
-        // First, we write the packet.
-        let payloadBytes = outboundBuffer.writeEncryptablePayload(packet)
-
-        // Ok, now we need to pad. The rules for padding for AES GCM are:
-        //
-        // 1. We must pad out such that the total encrypted content (padding length byte,
-        //     plus content bytes, plus padding bytes) is a multiple of the block size.
-        // 2. At least 4 bytes of padding MUST be added.
-        // 3. This padding SHOULD be random.
-        //
-        // Note that, unlike other protection modes, the length is not encrypted, and so we
-        // must exclude it from the padding calculation.
-        //
-        // So we check how many bytes we've already written, use modular arithmetic to work out
-        // how many more bytes we need, and then if that's fewer than 4 we add a block size to it
-        // to fill it out.
-        var encryptedBufferSize = payloadBytes + packetPaddingLength
-        var necessaryPaddingBytes = Self.cipherBlockSize - (encryptedBufferSize % Self.cipherBlockSize)
-        if necessaryPaddingBytes < 4 {
-            necessaryPaddingBytes += Self.cipherBlockSize
-        }
-
-        // We now want to write that many padding bytes to the end of the buffer. These are supposed to be
-        // random bytes. We're going to get those from the system random number generator.
-        encryptedBufferSize += outboundBuffer.writeSSHPaddingBytes(count: necessaryPaddingBytes)
-        precondition(encryptedBufferSize % Self.cipherBlockSize == 0, "Incorrectly counted buffer size; got \(encryptedBufferSize)")
+        let packetPaddingIndex = destination.readerIndex + packetLengthLength
 
-        // We now know the length: it's going to be "encrypted buffer size". The length does not include the tag, so don't add it.
-        // Let's write that in. We also need to write the number of padding bytes in.
-        outboundBuffer.setInteger(UInt32(encryptedBufferSize), at: packetLengthIndex)
-        outboundBuffer.setInteger(UInt8(necessaryPaddingBytes), at: packetPaddingIndex)
+        // We encrypte everything, except the length
+        let encryptedBufferSize = destination.readableBytes - packetLengthLength
 
-        // Ok, nice! Now we need to encrypt the data. We pass the length field as additional authenticated data, and the encrypted
+        // Now we need to encrypt the data. We pass the length field as additional authenticated data, and the encrypted
         // payload portion as the data to encrypt. We know these views will be valid, so we forcibly unwrap them: if they're invalid,
         // our math was wrong and we cannot recover.
-        let sealedBox = try AES.GCM.seal(outboundBuffer.viewBytes(at: packetPaddingIndex, length: encryptedBufferSize)!,
+        let sealedBox = try AES.GCM.seal(destination.viewBytes(at: packetPaddingIndex, length: encryptedBufferSize)!,
                                          using: self.outboundEncryptionKey,
                                          nonce: try AES.GCM.Nonce(data: self.outboundNonce),
-                                         authenticating: outboundBuffer.viewBytes(at: packetLengthIndex, length: packetLengthLength)!)
+                                         authenticating: destination.viewBytes(at: packetLengthIndex, length: packetLengthLength)!)
 
         assert(sealedBox.ciphertext.count == encryptedBufferSize)
 
-        // We now want to overwrite the portion of the bytebuffer that contains the plaintext with the ciphertext, and then append the
-        // tag.
-        outboundBuffer.setContiguousBytes(sealedBox.ciphertext, at: packetPaddingIndex)
-        let tagLength = outboundBuffer.writeContiguousBytes(sealedBox.tag)
+        // We now want to overwrite the portion of the bytebuffer that contains the plaintext with the ciphertext,
+        // and then append the tag.
+        destination.setContiguousBytes(sealedBox.ciphertext, at: packetPaddingIndex)
+        let tagLength = destination.writeContiguousBytes(sealedBox.tag)
         precondition(tagLength == self.macBytes, "Unexpected short tag")
 
         // Now we increment the Nonce for the next use, and then we're done!

Sources/NIOSSH/TransportProtection/SSHTransportProtection.swift

@@ -61,6 +61,10 @@ public protocol NIOSSHTransportProtection: AnyObject {
     /// The number of bytes consumed by the MAC
     var macBytes: Int { get }
 
+    /// Whether legnth of the packet will be encrypted. If length is not encrypted, it should be counted
+    /// when padding size is calculated.
+    var lengthEncrypted: Bool { get }
+
     /// Create a new instance of this transport protection scheme with the given keys.
     init(initialKeys: NIOSSHSessionKeys) throws
 
@@ -90,7 +94,7 @@ public protocol NIOSSHTransportProtection: AnyObject {
     func decryptAndVerifyRemainingPacket(_ source: inout ByteBuffer, sequenceNumber: UInt32) throws -> ByteBuffer
 
     /// Encrypt an entire outbound packet
-    func encryptPacket(_ packet: NIOSSHEncryptablePayload, sequenceNumber: UInt32, to outboundBuffer: inout ByteBuffer) throws
+    func encryptPacket(_ destination: inout ByteBuffer, sequenceNumber: UInt32) throws
 }
 
 extension NIOSSHTransportProtection {

Tests/NIOSSHTests/AESGCMTests.swift

@@ -42,7 +42,9 @@ final class AESGCMTests: XCTestCase {
         let initialKeys = self.generateKeys(keySize: .bits128)
 
         let aes128Encryptor = try assertNoThrowWithValue(AES128GCMOpenSSHTransportProtection(initialKeys: initialKeys))
-        XCTAssertNoThrow(try aes128Encryptor.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 0, to: &self.buffer))
+
+        self.buffer.writeSSHPacket(message: .newKeys, lengthEncrypted: aes128Encryptor.lengthEncrypted, blockSize: aes128Encryptor.cipherBlockSize)
+        XCTAssertNoThrow(try aes128Encryptor.encryptPacket(&self.buffer, sequenceNumber: 0))
 
         // 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
@@ -77,7 +79,9 @@ final class AESGCMTests: XCTestCase {
         let initialKeys = self.generateKeys(keySize: .bits256)
 
         let aes256Encryptor = try assertNoThrowWithValue(AES256GCMOpenSSHTransportProtection(initialKeys: initialKeys))
-        XCTAssertNoThrow(try aes256Encryptor.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 0, to: &self.buffer))
+
+        self.buffer.writeSSHPacket(message: .newKeys, lengthEncrypted: aes256Encryptor.lengthEncrypted, blockSize: aes256Encryptor.cipherBlockSize)
+        XCTAssertNoThrow(try aes256Encryptor.encryptPacket(&self.buffer, sequenceNumber: 0))
 
         // 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

Tests/NIOSSHTests/SSHKeyExchangeStateMachineTests.swift

@@ -144,7 +144,8 @@ final class SSHKeyExchangeStateMachineTests: XCTestCase {
         var buffer = ByteBufferAllocator().buffer(capacity: 1024)
 
         do {
-            try client.encryptPacket(.init(message: message), sequenceNumber: 0, to: &buffer)
+            buffer.writeSSHPacket(message: message, lengthEncrypted: client.lengthEncrypted, blockSize: client.cipherBlockSize)
+            try client.encryptPacket(&buffer, sequenceNumber: 0)
             try server.decryptFirstBlock(&buffer)
             var messageBuffer = try server.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)
             let decrypted = try messageBuffer.readSSHMessage()
@@ -158,7 +159,8 @@ final class SSHKeyExchangeStateMachineTests: XCTestCase {
         buffer.clear()
 
         do {
-            try server.encryptPacket(.init(message: message), sequenceNumber: 0, to: &buffer)
+            buffer.writeSSHPacket(message: message, lengthEncrypted: server.lengthEncrypted, blockSize: server.cipherBlockSize)
+            try server.encryptPacket(&buffer, sequenceNumber: 0)
             try client.decryptFirstBlock(&buffer)
             var messageBuffer = try client.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)
             let decrypted = try messageBuffer.readSSHMessage()

Tests/NIOSSHTests/SSHPacketParserTests.swift

@@ -251,7 +251,8 @@ final class SSHPacketParserTests: XCTestCase {
         parser.addEncryption(protection)
 
         part = allocator.buffer(capacity: 1024)
-        XCTAssertNoThrow(try protection.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 2, to: &part))
+        part.writeSSHPacket(message: .newKeys, lengthEncrypted: protection.lengthEncrypted, blockSize: protection.cipherBlockSize)
+        XCTAssertNoThrow(try protection.encryptPacket(&part, sequenceNumber: 2))
         var subpart = part.readSlice(length: 2)!
         parser.append(bytes: &subpart)
 
@@ -268,7 +269,8 @@ final class SSHPacketParserTests: XCTestCase {
         }
 
         part = allocator.buffer(capacity: 1024)
-        XCTAssertNoThrow(try protection.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 2, to: &part))
+        part.writeSSHPacket(message: .newKeys, lengthEncrypted: protection.lengthEncrypted, blockSize: protection.cipherBlockSize)
+        XCTAssertNoThrow(try protection.encryptPacket(&part, sequenceNumber: 2))
         parser.append(bytes: &part)
 
         switch try parser.nextPacket() {

Tests/NIOSSHTests/Utilities.swift

@@ -111,6 +111,10 @@ class TestTransportProtection: NIOSSHTransportProtection {
         32
     }
 
+    var lengthEncrypted: Bool {
+        true
+    }
+
     static var cipherName: String {
         "insecure-tiny-encription-cipher"
     }
@@ -211,41 +215,20 @@ class TestTransportProtection: NIOSSHTransportProtection {
         return plaintext.readSlice(length: plaintext.readableBytes - Int(paddingLength))!
     }
 
-    func encryptPacket(_ packet: NIOSSHEncryptablePayload, sequenceNumber: UInt32, to outboundBuffer: inout ByteBuffer) throws {
-        let packetLengthIndex = outboundBuffer.writerIndex
-        let packetLengthLength = MemoryLayout<UInt32>.size
-        let packetPaddingIndex = outboundBuffer.writerIndex + packetLengthLength
-        let packetPaddingLength = MemoryLayout<UInt8>.size
-
-        outboundBuffer.moveWriterIndex(forwardBy: packetLengthLength + packetPaddingLength)
-
-        let payloadBytes = outboundBuffer.writeEncryptablePayload(packet)
-
-        var encryptedBufferSize = payloadBytes + packetPaddingLength + packetLengthLength
-        var necessaryPaddingBytes = Self.cipherBlockSize - (encryptedBufferSize % Self.cipherBlockSize)
-        if necessaryPaddingBytes < 4 {
-            necessaryPaddingBytes += Self.cipherBlockSize
-        }
-
-        // We now want to write that many padding bytes to the end of the buffer. These are supposed to be
-        // random bytes. We're going to get those from the system random number generator.
-        encryptedBufferSize += outboundBuffer.writeSSHPaddingBytes(count: necessaryPaddingBytes)
-        precondition(encryptedBufferSize % Self.cipherBlockSize == 0, "Incorrectly counted buffer size; got \(encryptedBufferSize)")
+    func encryptPacket(_ destination: inout ByteBuffer, sequenceNumber: UInt32) throws {
+        let packetLengthIndex = destination.readerIndex
+        let encryptedBufferSize = destination.readableBytes
 
-        // We now know the length: it's going to be "encrypted buffer size". The length does not include the tag, so don't add it.
-        // Let's write that in. We also need to write the number of padding bytes in.
-        outboundBuffer.setInteger(UInt32(encryptedBufferSize - packetLengthLength), at: packetLengthIndex)
-        outboundBuffer.setInteger(UInt8(necessaryPaddingBytes), at: packetPaddingIndex)
-        let plaintextView = outboundBuffer.viewBytes(at: packetLengthIndex, length: encryptedBufferSize)!
+        let plaintextView = destination.viewBytes(at: packetLengthIndex, length: encryptedBufferSize)!
         let ciphertext = InsecureEncryptionAlgorithm.encrypt(key: self.outboundEncryptionKey, plaintext: ByteBuffer(bytes: plaintextView))
         assert(ciphertext.readableBytes == encryptedBufferSize)
 
         var hmac = HMAC<SHA256>.init(key: self.outboundMACKey)
         hmac.update(data: plaintextView)
 
         // We now want to overwrite the portion of the bytebuffer that contains the plaintext with the ciphertext, and then append the tag.
-        outboundBuffer.setBytes(ciphertext.readableBytesView, at: packetLengthIndex)
-        let tagLength = outboundBuffer.writeBytes(hmac.finalize())
+        destination.setBytes(ciphertext.readableBytesView, at: packetLengthIndex)
+        let tagLength = destination.writeBytes(hmac.finalize())
         precondition(tagLength == self.macBytes, "Unexpected short tag")
     }
 }

Tests/NIOSSHTests/UtilitiesTests.swift

@@ -52,7 +52,10 @@ final class UtilitiesTests: XCTestCase {
         let message = SSHMessage.channelRequest(.init(recipientChannel: 1, type: .exec("uname"), wantReply: false))
         let allocator = ByteBufferAllocator()
         var buffer = allocator.buffer(capacity: 1024)
-        XCTAssertNoThrow(try client.encryptPacket(.init(message: message), sequenceNumber: 0, to: &buffer))
+
+        buffer.writeSSHPacket(message: message, lengthEncrypted: client.lengthEncrypted, blockSize: client.cipherBlockSize)
+
+        XCTAssertNoThrow(try client.encryptPacket(&buffer, sequenceNumber: 0))
         XCTAssertNoThrow(try server.decryptFirstBlock(&buffer))
         var decoded = try server.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)
         XCTAssertEqual(message, try decoded.readSSHMessage())