From 795c19871c5ee50d4a03efba96106bd9647a8dbe Mon Sep 17 00:00:00 2001 From: Max Radermacher Date: Wed, 15 Oct 2025 12:20:15 -0500 Subject: [PATCH] Handle overflow in PaddingBucket --- .../Attachments/PaddingBucket.swift | 50 ++++++++++++------- .../Attachments/PaddingBucketTest.swift | 12 +++-- .../Cryptography/Cryptography.swift | 26 +++++----- .../Cryptography/CryptographyTests.swift | 2 +- .../StreamTransform/GzipStreamTransform.swift | 2 +- 5 files changed, 56 insertions(+), 36 deletions(-) diff --git a/SignalServiceKit/Attachments/PaddingBucket.swift b/SignalServiceKit/Attachments/PaddingBucket.swift index 742c04c78a..5642273080 100644 --- a/SignalServiceKit/Attachments/PaddingBucket.swift +++ b/SignalServiceKit/Attachments/PaddingBucket.swift @@ -28,29 +28,40 @@ struct PaddingBucket { let bucketNumber: Int - init(bucketNumber: Int) { - self.bucketNumber = max(bucketNumber, Constants.smallestBucketNumber) - } - /// The plaintext size with padding. - var plaintextSize: UInt64 { - return UInt64(floor(pow(Constants.paddingMultiplier, Double(bucketNumber)))) - } + let plaintextSize: UInt64 /// The encrypted size with padding & encryption overhead. - var encryptedSize: UInt64 { - return Self.addingEncryptionOverhead(to: self.plaintextSize) + let encryptedSize: UInt64 + + init?(bucketNumber: Int) { + self.bucketNumber = max(bucketNumber, Constants.smallestBucketNumber) + let plaintextSize = UInt64(exactly: floor(pow(Constants.paddingMultiplier, Double(self.bucketNumber)))) + guard let plaintextSize else { + return nil + } + self.plaintextSize = plaintextSize + let encryptedSize = Self.addingEncryptionOverhead(to: plaintextSize) + guard let encryptedSize else { + return nil + } + self.encryptedSize = encryptedSize } - static func addingEncryptionOverhead(to paddedValue: UInt64) -> UInt64 { - var result = paddedValue - result += Constants.blockLength - result % Constants.blockLength - result += Constants.ivLength - result += Constants.hmacLength - return result + static func addingEncryptionOverhead(to paddedValue: UInt64) -> UInt64? { + let result = paddedValue.addingReportingOverflow( + Constants.ivLength + + Constants.blockLength + - paddedValue % Constants.blockLength + + Constants.hmacLength + ) + if result.overflow { + return nil + } + return result.partialValue } - static func forUnpaddedPlaintextSize(_ unpaddedPlaintextSize: UInt64) -> PaddingBucket { + static func forUnpaddedPlaintextSize(_ unpaddedPlaintextSize: UInt64) -> PaddingBucket? { let bucketNumber: Int if unpaddedPlaintextSize == 0 { bucketNumber = 0 @@ -70,7 +81,7 @@ struct PaddingBucket { + Constants.hmacLength ) if worstCasePlaintextLimit.overflow || worstCasePlaintextLimit.partialValue == 0 { - return PaddingBucket(bucketNumber: 0) + return PaddingBucket(bucketNumber: 0)! } // Taking the `floor(...)` here may cause us to pick a bucket one smaller // than we should when `encryptedSize` is exactly the size of a bucket. @@ -81,9 +92,10 @@ struct PaddingBucket { // We check one optimistic bucket because the minimum spacing is 27 bytes // (which is larger than the 15 + 1 worst-case bytes mentioned above). let optimisticPaddingBucket = PaddingBucket(bucketNumber: Int(worstCaseBucketNumber) + 1) - if optimisticPaddingBucket.encryptedSize <= encryptedSize { + if let optimisticPaddingBucket, optimisticPaddingBucket.encryptedSize <= encryptedSize { return optimisticPaddingBucket } - return PaddingBucket(bucketNumber: Int(worstCaseBucketNumber)) + // By definition, this bucket can't overflow the encrypted size limit. + return PaddingBucket(bucketNumber: Int(worstCaseBucketNumber))! } } diff --git a/SignalServiceKit/Attachments/PaddingBucketTest.swift b/SignalServiceKit/Attachments/PaddingBucketTest.swift index 25b319a1fc..b29a5fe5a2 100644 --- a/SignalServiceKit/Attachments/PaddingBucketTest.swift +++ b/SignalServiceKit/Attachments/PaddingBucketTest.swift @@ -25,7 +25,7 @@ struct PaddingBucketTest { (1_000_000_000, 1_012_633_066), ]) func testPaddedSize(testCase: (unpaddedSize: UInt64, paddedSize: UInt64)) { - #expect(PaddingBucket.forUnpaddedPlaintextSize(testCase.unpaddedSize).plaintextSize == testCase.paddedSize) + #expect(PaddingBucket.forUnpaddedPlaintextSize(testCase.unpaddedSize)?.plaintextSize == testCase.paddedSize) } @Test(arguments: [ @@ -44,7 +44,7 @@ struct PaddingBucketTest { (1_000_000_000, 1_012_633_120), ]) func testEncryptedSize(testCase: (unpaddedSize: UInt64, encryptedSize: UInt64)) { - #expect(PaddingBucket.forUnpaddedPlaintextSize(testCase.unpaddedSize).encryptedSize == testCase.encryptedSize) + #expect(PaddingBucket.forUnpaddedPlaintextSize(testCase.unpaddedSize)?.encryptedSize == testCase.encryptedSize) } @Test(arguments: [ @@ -57,8 +57,14 @@ struct PaddingBucketTest { @Test(arguments: 130...483) func testAllInterestingLimits(bucketNumber: Int) { - let encryptedSize = PaddingBucket(bucketNumber: bucketNumber).encryptedSize + let encryptedSize = PaddingBucket(bucketNumber: bucketNumber)!.encryptedSize #expect(PaddingBucket.forEncryptedSizeLimit(encryptedSize).bucketNumber == bucketNumber) #expect(PaddingBucket.forEncryptedSizeLimit(encryptedSize - 1).bucketNumber == bucketNumber - 1) } + + @Test + func testOverflow() { + let largestBucket = PaddingBucket.forEncryptedSizeLimit(.max) + #expect(largestBucket.bucketNumber == 909) + } } diff --git a/SignalServiceKit/Cryptography/Cryptography.swift b/SignalServiceKit/Cryptography/Cryptography.swift index f4489d49d8..c70270b34b 100644 --- a/SignalServiceKit/Cryptography/Cryptography.swift +++ b/SignalServiceKit/Cryptography/Cryptography.swift @@ -98,8 +98,8 @@ public extension Cryptography { static let diskPageSize = 8192 } - static func paddedSize(unpaddedSize: UInt) -> UInt { - return UInt(PaddingBucket.forUnpaddedPlaintextSize(UInt64(unpaddedSize)).plaintextSize) + static func paddedSize(unpaddedSize: UInt) -> UInt? { + return PaddingBucket.forUnpaddedPlaintextSize(UInt64(unpaddedSize)).flatMap { UInt(exactly: $0.plaintextSize) } } /// Given an unencrypted, unpadded byte count, returns the *estimated* byte count of the final padded, encrypted blob @@ -111,7 +111,7 @@ public extension Cryptography { /// 2. Our padding scheme may change between when this is checked and when we upload(ed). static func estimatedMediaTierCDNSize(unencryptedSize: UInt32) -> UInt32 { let transitTierSize = UInt64(estimatedTransitTierCDNSize(unencryptedSize: unencryptedSize)) - return UInt32(clamping: PaddingBucket.addingEncryptionOverhead(to: transitTierSize)) + return UInt32(clamping: PaddingBucket.addingEncryptionOverhead(to: transitTierSize) ?? .max) } /// Given an unencrypted, unpadded byte count, returns the *estimated* byte count of the final padded, encrypted blob @@ -122,7 +122,7 @@ public extension Cryptography { /// 1. It may be a different client uploading with a differing padding scheme (or a bug with its padding scheme) /// 2. Our padding scheme may change between when this is checked and when we upload(ed). static func estimatedTransitTierCDNSize(unencryptedSize: UInt32) -> UInt32 { - return UInt32(clamping: PaddingBucket.forUnpaddedPlaintextSize(UInt64(unencryptedSize)).encryptedSize) + return UInt32(clamping: PaddingBucket.forUnpaddedPlaintextSize(UInt64(unencryptedSize))?.encryptedSize ?? .max) } static func randomAttachmentEncryptionKey() -> Data { @@ -387,15 +387,17 @@ public extension Cryptography { } // Add zero padding to the plaintext attachment data if necessary. - let paddedPlaintextLength = paddedSize(unpaddedSize: unpaddedPlaintextLength) - if applyExtraPadding, paddedPlaintextLength > unpaddedPlaintextLength { - let ciphertextBlock = try cipherContext.update( - Data(repeating: 0, count: Int(paddedPlaintextLength - unpaddedPlaintextLength)) - ) + if applyExtraPadding { + let paddedPlaintextLength = paddedSize(unpaddedSize: unpaddedPlaintextLength)! + if paddedPlaintextLength > unpaddedPlaintextLength { + let ciphertextBlock = try cipherContext.update( + Data(repeating: 0, count: Int(paddedPlaintextLength - unpaddedPlaintextLength)) + ) - hmac.update(data: ciphertextBlock) - sha256.update(data: ciphertextBlock) - output(ciphertextBlock) + hmac.update(data: ciphertextBlock) + sha256.update(data: ciphertextBlock) + output(ciphertextBlock) + } } // Finalize the encryption and write out the last block. diff --git a/SignalServiceKit/Cryptography/CryptographyTests.swift b/SignalServiceKit/Cryptography/CryptographyTests.swift index fe64be45ed..3434c81af4 100644 --- a/SignalServiceKit/Cryptography/CryptographyTests.swift +++ b/SignalServiceKit/Cryptography/CryptographyTests.swift @@ -364,7 +364,7 @@ class CryptographyTestsSwift: XCTestCase { // When we encrypt, we add custom padding 0s to a determined length. // Normally these get truncated in the final output using the hint of plaintextLength; // since we are omitting that we need to expect them in the final output. - let customPaddedLength = UInt32(Cryptography.paddedSize(unpaddedSize: UInt(plaintextLength))) + let customPaddedLength = UInt32(Cryptography.paddedSize(unpaddedSize: UInt(plaintextLength))!) let customPaddingLength = customPaddedLength - plaintextLength let expectedPlaintextOutput = plaintextData + Data(repeating: 0, count: Int(customPaddingLength)) diff --git a/SignalServiceKit/Util/StreamTransform/GzipStreamTransform.swift b/SignalServiceKit/Util/StreamTransform/GzipStreamTransform.swift index ff264c41dc..d6d9113958 100644 --- a/SignalServiceKit/Util/StreamTransform/GzipStreamTransform.swift +++ b/SignalServiceKit/Util/StreamTransform/GzipStreamTransform.swift @@ -216,7 +216,7 @@ public class GzipStreamTransform: StreamTransform, FinalizableStreamTransform { // Pad the gzip similar to how attachments are padded. // gzip will ignore this trailing data during decompression. let unpaddedSize = UInt(bitPattern: outputCount) - let paddedSize = Cryptography.paddedSize(unpaddedSize: unpaddedSize) + let paddedSize = Cryptography.paddedSize(unpaddedSize: unpaddedSize)! if paddedSize > unpaddedSize { finalData.append(Data(repeating: 0, count: Int(paddedSize - unpaddedSize))) }