Handle overflow in PaddingBucket

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))!
     }
 }

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)
+    }
 }

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.

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))
 

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)))
             }