Skip to content

Commit a48586f

Browse files
artemredkinartemredkin
authored
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
1 parent c905128 commit a48586f

File tree

8 files changed

+94
-109
lines changed

8 files changed

+94
-109
lines changed

‎Sources/NIOSSH/SSHPacketSerializer.swift

Lines changed: 45 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -47,40 +47,55 @@ struct SSHPacketSerializer {
4747
preconditionFailure("only .version message is allowed at this point")
4848
}
4949
case .cleartext:
50-
let index = buffer.writerIndex
50+
buffer.writeSSHPacket(message: message, lengthEncrypted: true, blockSize: 8)
51+
self.sequenceNumber &+= 1
52+
case .encrypted(let protection):
53+
let index = buffer.readerIndex
54+
buffer.moveReaderIndex(to: buffer.writerIndex)
55+
buffer.writeSSHPacket(message: message, lengthEncrypted: protection.lengthEncrypted, blockSize: protection.cipherBlockSize)
56+
try protection.encryptPacket(&buffer, sequenceNumber: self.sequenceNumber)
57+
buffer.moveReaderIndex(to: index)
58+
self.sequenceNumber &+= 1
59+
}
60+
}
61+
}
5162

52-
/// Each packet is in the following format:
53-
///
54-
/// uint32 packet_length
55-
/// byte padding_length
56-
/// byte[n1] payload; n1 = packet_length - padding_length - 1
57-
/// byte[n2] random padding; n2 = padding_length
58-
/// byte[m] mac (Message Authentication Code - MAC); m = mac_length
63+
extension ByteBuffer {
64+
mutating func writeSSHPacket(message: SSHMessage, lengthEncrypted: Bool, blockSize: Int) {
65+
let index = self.writerIndex
5966

60-
/// payload
61-
buffer.moveWriterIndex(forwardBy: 5)
62-
let messageLength = buffer.writeSSHMessage(message)
67+
/// Each packet is in the following format:
68+
///
69+
/// uint32 packet_length
70+
/// byte padding_length
71+
/// byte[n1] payload; n1 = packet_length - padding_length - 1
72+
/// byte[n2] random padding; n2 = padding_length
73+
/// byte[m] mac (Message Authentication Code - MAC); m = mac_length
6374

64-
/// RFC 4253 § 6:
65-
/// random padding
66-
/// Arbitrary-length padding, such that the total length of (packet_length || padding_length || payload || random padding)
67-
/// is a multiple of the cipher block size or 8, whichever is larger. There MUST be at least four bytes of padding. The
68-
/// padding SHOULD consist of random bytes. The maximum amount of padding is 255 bytes.
69-
let blockSize = 8
70-
let paddingLength = 3 + blockSize - ((messageLength + blockSize) % blockSize)
75+
/// payload
76+
self.moveWriterIndex(forwardBy: 5)
77+
let messageLength = self.writeSSHMessage(message)
7178

72-
/// packet_length
73-
/// The length of the packet in bytes, not including 'mac' or the 'packet_length' field itself.
74-
buffer.setInteger(UInt32(1 + messageLength + paddingLength), at: index)
75-
/// padding_length
76-
buffer.setInteger(UInt8(paddingLength), at: index + 4)
77-
/// random padding
78-
buffer.writeSSHPaddingBytes(count: paddingLength)
79-
self.sequenceNumber &+= 1
80-
case .encrypted(let protection):
81-
let payload = NIOSSHEncryptablePayload(message: message)
82-
try protection.encryptPacket(payload, sequenceNumber: self.sequenceNumber, to: &buffer)
83-
self.sequenceNumber &+= 1
79+
// Depending on on whether packet length is encrypted, padding should reflect that
80+
let payloadLength = lengthEncrypted ? messageLength + 5 : messageLength + 1
81+
82+
/// RFC 4253 § 6:
83+
/// random padding
84+
/// Arbitrary-length padding, such that the total length of (packet_length || padding_length || payload || random padding)
85+
/// is a multiple of the cipher block size or 8, whichever is larger. There MUST be at least four bytes of padding. The
86+
/// padding SHOULD consist of random bytes. The maximum amount of padding is 255 bytes.
87+
var paddingLength = blockSize - (payloadLength % blockSize)
88+
if paddingLength < 4 {
89+
paddingLength += blockSize
8490
}
91+
92+
/// packet_length
93+
/// The length of the packet in bytes, not including 'mac' or the 'packet_length' field itself.
94+
let packetLength = 1 + messageLength + paddingLength
95+
self.setInteger(UInt32(packetLength), at: index)
96+
/// padding_length
97+
self.setInteger(UInt8(paddingLength), at: index + 4)
98+
/// random padding
99+
self.writeSSHPaddingBytes(count: paddingLength)
85100
}
86101
}

‎Sources/NIOSSH/TransportProtection/AESGCM.swift

Lines changed: 16 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,10 @@ extension AESGCMTransportProtection: NIOSSHTransportProtection {
6262
16
6363
}
6464

65+
var lengthEncrypted: Bool {
66+
false
67+
}
68+
6569
func updateKeys(_ newKeys: NIOSSHSessionKeys) throws {
6670
guard newKeys.outboundEncryptionKey.bitCount == Self.keySizes.encryptionKeySize * 8,
6771
newKeys.inboundEncryptionKey.bitCount == Self.keySizes.encryptionKeySize * 8 else {
@@ -117,61 +121,29 @@ extension AESGCMTransportProtection: NIOSSHTransportProtection {
117121
return source.readSlice(length: plaintext.count)!
118122
}
119123

120-
func encryptPacket(_ packet: NIOSSHEncryptablePayload, sequenceNumber _: UInt32, to outboundBuffer: inout ByteBuffer) throws {
124+
func encryptPacket(_ destination: inout ByteBuffer, sequenceNumber: UInt32) throws {
121125
// Keep track of where the length is going to be written.
122-
let packetLengthIndex = outboundBuffer.writerIndex
126+
let packetLengthIndex = destination.readerIndex
123127
let packetLengthLength = MemoryLayout<UInt32>.size
124-
let packetPaddingIndex = outboundBuffer.writerIndex + packetLengthLength
125-
let packetPaddingLength = MemoryLayout<UInt8>.size
126-
127-
outboundBuffer.moveWriterIndex(forwardBy: packetLengthLength + packetPaddingLength)
128-
129-
// First, we write the packet.
130-
let payloadBytes = outboundBuffer.writeEncryptablePayload(packet)
131-
132-
// Ok, now we need to pad. The rules for padding for AES GCM are:
133-
//
134-
// 1. We must pad out such that the total encrypted content (padding length byte,
135-
// plus content bytes, plus padding bytes) is a multiple of the block size.
136-
// 2. At least 4 bytes of padding MUST be added.
137-
// 3. This padding SHOULD be random.
138-
//
139-
// Note that, unlike other protection modes, the length is not encrypted, and so we
140-
// must exclude it from the padding calculation.
141-
//
142-
// So we check how many bytes we've already written, use modular arithmetic to work out
143-
// how many more bytes we need, and then if that's fewer than 4 we add a block size to it
144-
// to fill it out.
145-
var encryptedBufferSize = payloadBytes + packetPaddingLength
146-
var necessaryPaddingBytes = Self.cipherBlockSize - (encryptedBufferSize % Self.cipherBlockSize)
147-
if necessaryPaddingBytes < 4 {
148-
necessaryPaddingBytes += Self.cipherBlockSize
149-
}
150-
151-
// We now want to write that many padding bytes to the end of the buffer. These are supposed to be
152-
// random bytes. We're going to get those from the system random number generator.
153-
encryptedBufferSize += outboundBuffer.writeSSHPaddingBytes(count: necessaryPaddingBytes)
154-
precondition(encryptedBufferSize % Self.cipherBlockSize == 0, "Incorrectly counted buffer size; got \(encryptedBufferSize)")
128+
let packetPaddingIndex = destination.readerIndex + packetLengthLength
155129

156-
// 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.
157-
// Let's write that in. We also need to write the number of padding bytes in.
158-
outboundBuffer.setInteger(UInt32(encryptedBufferSize), at: packetLengthIndex)
159-
outboundBuffer.setInteger(UInt8(necessaryPaddingBytes), at: packetPaddingIndex)
130+
// We encrypte everything, except the length
131+
let encryptedBufferSize = destination.readableBytes - packetLengthLength
160132

161-
// Ok, nice! Now we need to encrypt the data. We pass the length field as additional authenticated data, and the encrypted
133+
// Now we need to encrypt the data. We pass the length field as additional authenticated data, and the encrypted
162134
// payload portion as the data to encrypt. We know these views will be valid, so we forcibly unwrap them: if they're invalid,
163135
// our math was wrong and we cannot recover.
164-
let sealedBox = try AES.GCM.seal(outboundBuffer.viewBytes(at: packetPaddingIndex, length: encryptedBufferSize)!,
136+
let sealedBox = try AES.GCM.seal(destination.viewBytes(at: packetPaddingIndex, length: encryptedBufferSize)!,
165137
using: self.outboundEncryptionKey,
166138
nonce: try AES.GCM.Nonce(data: self.outboundNonce),
167-
authenticating: outboundBuffer.viewBytes(at: packetLengthIndex, length: packetLengthLength)!)
139+
authenticating: destination.viewBytes(at: packetLengthIndex, length: packetLengthLength)!)
168140

169141
assert(sealedBox.ciphertext.count == encryptedBufferSize)
170142

171-
// We now want to overwrite the portion of the bytebuffer that contains the plaintext with the ciphertext, and then append the
172-
// tag.
173-
outboundBuffer.setContiguousBytes(sealedBox.ciphertext, at: packetPaddingIndex)
174-
let tagLength = outboundBuffer.writeContiguousBytes(sealedBox.tag)
143+
// We now want to overwrite the portion of the bytebuffer that contains the plaintext with the ciphertext,
144+
// and then append the tag.
145+
destination.setContiguousBytes(sealedBox.ciphertext, at: packetPaddingIndex)
146+
let tagLength = destination.writeContiguousBytes(sealedBox.tag)
175147
precondition(tagLength == self.macBytes, "Unexpected short tag")
176148

177149
// Now we increment the Nonce for the next use, and then we're done!

‎Sources/NIOSSH/TransportProtection/SSHTransportProtection.swift

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,10 @@ public protocol NIOSSHTransportProtection: AnyObject {
6161
/// The number of bytes consumed by the MAC
6262
var macBytes: Int { get }
6363

64+
/// Whether legnth of the packet will be encrypted. If length is not encrypted, it should be counted
65+
/// when padding size is calculated.
66+
var lengthEncrypted: Bool { get }
67+
6468
/// Create a new instance of this transport protection scheme with the given keys.
6569
init(initialKeys: NIOSSHSessionKeys) throws
6670

@@ -90,7 +94,7 @@ public protocol NIOSSHTransportProtection: AnyObject {
9094
func decryptAndVerifyRemainingPacket(_ source: inout ByteBuffer, sequenceNumber: UInt32) throws -> ByteBuffer
9195

9296
/// Encrypt an entire outbound packet
93-
func encryptPacket(_ packet: NIOSSHEncryptablePayload, sequenceNumber: UInt32, to outboundBuffer: inout ByteBuffer) throws
97+
func encryptPacket(_ destination: inout ByteBuffer, sequenceNumber: UInt32) throws
9498
}
9599

96100
extension NIOSSHTransportProtection {

‎Tests/NIOSSHTests/AESGCMTests.swift

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,9 @@ final class AESGCMTests: XCTestCase {
4242
let initialKeys = self.generateKeys(keySize: .bits128)
4343

4444
let aes128Encryptor = try assertNoThrowWithValue(AES128GCMOpenSSHTransportProtection(initialKeys: initialKeys))
45-
XCTAssertNoThrow(try aes128Encryptor.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 0, to: &self.buffer))
45+
46+
self.buffer.writeSSHPacket(message: .newKeys, lengthEncrypted: aes128Encryptor.lengthEncrypted, blockSize: aes128Encryptor.cipherBlockSize)
47+
XCTAssertNoThrow(try aes128Encryptor.encryptPacket(&self.buffer, sequenceNumber: 0))
4648

4749
// The newKeys message is very straightforward: a single byte. Because of that, we expect that we will need
4850
// 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 {
7779
let initialKeys = self.generateKeys(keySize: .bits256)
7880

7981
let aes256Encryptor = try assertNoThrowWithValue(AES256GCMOpenSSHTransportProtection(initialKeys: initialKeys))
80-
XCTAssertNoThrow(try aes256Encryptor.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 0, to: &self.buffer))
82+
83+
self.buffer.writeSSHPacket(message: .newKeys, lengthEncrypted: aes256Encryptor.lengthEncrypted, blockSize: aes256Encryptor.cipherBlockSize)
84+
XCTAssertNoThrow(try aes256Encryptor.encryptPacket(&self.buffer, sequenceNumber: 0))
8185

8286
// The newKeys message is very straightforward: a single byte. Because of that, we expect that we will need
8387
// 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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,8 @@ final class SSHKeyExchangeStateMachineTests: XCTestCase {
144144
var buffer = ByteBufferAllocator().buffer(capacity: 1024)
145145

146146
do {
147-
try client.encryptPacket(.init(message: message), sequenceNumber: 0, to: &buffer)
147+
buffer.writeSSHPacket(message: message, lengthEncrypted: client.lengthEncrypted, blockSize: client.cipherBlockSize)
148+
try client.encryptPacket(&buffer, sequenceNumber: 0)
148149
try server.decryptFirstBlock(&buffer)
149150
var messageBuffer = try server.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)
150151
let decrypted = try messageBuffer.readSSHMessage()
@@ -158,7 +159,8 @@ final class SSHKeyExchangeStateMachineTests: XCTestCase {
158159
buffer.clear()
159160

160161
do {
161-
try server.encryptPacket(.init(message: message), sequenceNumber: 0, to: &buffer)
162+
buffer.writeSSHPacket(message: message, lengthEncrypted: server.lengthEncrypted, blockSize: server.cipherBlockSize)
163+
try server.encryptPacket(&buffer, sequenceNumber: 0)
162164
try client.decryptFirstBlock(&buffer)
163165
var messageBuffer = try client.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)
164166
let decrypted = try messageBuffer.readSSHMessage()

‎Tests/NIOSSHTests/SSHPacketParserTests.swift

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,8 @@ final class SSHPacketParserTests: XCTestCase {
251251
parser.addEncryption(protection)
252252

253253
part = allocator.buffer(capacity: 1024)
254-
XCTAssertNoThrow(try protection.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 2, to: &part))
254+
part.writeSSHPacket(message: .newKeys, lengthEncrypted: protection.lengthEncrypted, blockSize: protection.cipherBlockSize)
255+
XCTAssertNoThrow(try protection.encryptPacket(&part, sequenceNumber: 2))
255256
var subpart = part.readSlice(length: 2)!
256257
parser.append(bytes: &subpart)
257258

@@ -268,7 +269,8 @@ final class SSHPacketParserTests: XCTestCase {
268269
}
269270

270271
part = allocator.buffer(capacity: 1024)
271-
XCTAssertNoThrow(try protection.encryptPacket(NIOSSHEncryptablePayload(message: .newKeys), sequenceNumber: 2, to: &part))
272+
part.writeSSHPacket(message: .newKeys, lengthEncrypted: protection.lengthEncrypted, blockSize: protection.cipherBlockSize)
273+
XCTAssertNoThrow(try protection.encryptPacket(&part, sequenceNumber: 2))
272274
parser.append(bytes: &part)
273275

274276
switch try parser.nextPacket() {

‎Tests/NIOSSHTests/Utilities.swift

Lines changed: 10 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,10 @@ class TestTransportProtection: NIOSSHTransportProtection {
111111
32
112112
}
113113

114+
var lengthEncrypted: Bool {
115+
true
116+
}
117+
114118
static var cipherName: String {
115119
"insecure-tiny-encription-cipher"
116120
}
@@ -211,41 +215,20 @@ class TestTransportProtection: NIOSSHTransportProtection {
211215
return plaintext.readSlice(length: plaintext.readableBytes - Int(paddingLength))!
212216
}
213217

214-
func encryptPacket(_ packet: NIOSSHEncryptablePayload, sequenceNumber: UInt32, to outboundBuffer: inout ByteBuffer) throws {
215-
let packetLengthIndex = outboundBuffer.writerIndex
216-
let packetLengthLength = MemoryLayout<UInt32>.size
217-
let packetPaddingIndex = outboundBuffer.writerIndex + packetLengthLength
218-
let packetPaddingLength = MemoryLayout<UInt8>.size
219-
220-
outboundBuffer.moveWriterIndex(forwardBy: packetLengthLength + packetPaddingLength)
221-
222-
let payloadBytes = outboundBuffer.writeEncryptablePayload(packet)
223-
224-
var encryptedBufferSize = payloadBytes + packetPaddingLength + packetLengthLength
225-
var necessaryPaddingBytes = Self.cipherBlockSize - (encryptedBufferSize % Self.cipherBlockSize)
226-
if necessaryPaddingBytes < 4 {
227-
necessaryPaddingBytes += Self.cipherBlockSize
228-
}
229-
230-
// We now want to write that many padding bytes to the end of the buffer. These are supposed to be
231-
// random bytes. We're going to get those from the system random number generator.
232-
encryptedBufferSize += outboundBuffer.writeSSHPaddingBytes(count: necessaryPaddingBytes)
233-
precondition(encryptedBufferSize % Self.cipherBlockSize == 0, "Incorrectly counted buffer size; got \(encryptedBufferSize)")
218+
func encryptPacket(_ destination: inout ByteBuffer, sequenceNumber: UInt32) throws {
219+
let packetLengthIndex = destination.readerIndex
220+
let encryptedBufferSize = destination.readableBytes
234221

235-
// 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.
236-
// Let's write that in. We also need to write the number of padding bytes in.
237-
outboundBuffer.setInteger(UInt32(encryptedBufferSize - packetLengthLength), at: packetLengthIndex)
238-
outboundBuffer.setInteger(UInt8(necessaryPaddingBytes), at: packetPaddingIndex)
239-
let plaintextView = outboundBuffer.viewBytes(at: packetLengthIndex, length: encryptedBufferSize)!
222+
let plaintextView = destination.viewBytes(at: packetLengthIndex, length: encryptedBufferSize)!
240223
let ciphertext = InsecureEncryptionAlgorithm.encrypt(key: self.outboundEncryptionKey, plaintext: ByteBuffer(bytes: plaintextView))
241224
assert(ciphertext.readableBytes == encryptedBufferSize)
242225

243226
var hmac = HMAC<SHA256>.init(key: self.outboundMACKey)
244227
hmac.update(data: plaintextView)
245228

246229
// We now want to overwrite the portion of the bytebuffer that contains the plaintext with the ciphertext, and then append the tag.
247-
outboundBuffer.setBytes(ciphertext.readableBytesView, at: packetLengthIndex)
248-
let tagLength = outboundBuffer.writeBytes(hmac.finalize())
230+
destination.setBytes(ciphertext.readableBytesView, at: packetLengthIndex)
231+
let tagLength = destination.writeBytes(hmac.finalize())
249232
precondition(tagLength == self.macBytes, "Unexpected short tag")
250233
}
251234
}

‎Tests/NIOSSHTests/UtilitiesTests.swift

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,10 @@ final class UtilitiesTests: XCTestCase {
5252
let message = SSHMessage.channelRequest(.init(recipientChannel: 1, type: .exec("uname"), wantReply: false))
5353
let allocator = ByteBufferAllocator()
5454
var buffer = allocator.buffer(capacity: 1024)
55-
XCTAssertNoThrow(try client.encryptPacket(.init(message: message), sequenceNumber: 0, to: &buffer))
55+
56+
buffer.writeSSHPacket(message: message, lengthEncrypted: client.lengthEncrypted, blockSize: client.cipherBlockSize)
57+
58+
XCTAssertNoThrow(try client.encryptPacket(&buffer, sequenceNumber: 0))
5659
XCTAssertNoThrow(try server.decryptFirstBlock(&buffer))
5760
var decoded = try server.decryptAndVerifyRemainingPacket(&buffer, sequenceNumber: 0)
5861
XCTAssertEqual(message, try decoded.readSSHMessage())

0 commit comments

Comments
 (0)