Stack-allocate CBS and make var rawRepresentation non-throwing

apple · Oct 22, 2024 · 5950937 · 5950937
1 parent acd486a
commit 5950937
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 31 deletions.
diff --git a/Sources/_CryptoExtras/MLDSA/MLDSA_boring.swift b/Sources/_CryptoExtras/MLDSA/MLDSA_boring.swift
@@ -132,11 +132,8 @@ extension MLDSA {
                 self.pointer = UnsafeMutablePointer<MLDSA65_private_key>.allocate(capacity: 1)
 
                 try rawRepresentation.regions.flatMap { $0 }.withUnsafeBufferPointer { buffer in
-                    let cbsPointer = UnsafeMutablePointer<CBS>.allocate(capacity: 1)
-                    defer { cbsPointer.deallocate() }
-                    cbsPointer.pointee = CBS(data: buffer.baseAddress, len: buffer.count)
-
-                    guard CCryptoBoringSSL_MLDSA65_parse_private_key(self.pointer, cbsPointer) == 1 else {
+                    var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
+                    guard CCryptoBoringSSL_MLDSA65_parse_private_key(self.pointer, &cbs) == 1 else {
                         throw CryptoKitError.internalBoringSSLError()
                     }
                 }
@@ -205,9 +202,7 @@ extension MLDSA {
 
         /// The raw binary representation of the public key.
         public var rawRepresentation: Data {
-            get throws {
-                try self.backing.rawRepresentation
-            }
+            self.backing.rawRepresentation
         }
 
         /// Verify a signature for the given data.
@@ -248,34 +243,20 @@ extension MLDSA {
                 self.pointer = UnsafeMutablePointer<MLDSA65_public_key>.allocate(capacity: 1)
 
                 try rawRepresentation.regions.flatMap { $0 }.withUnsafeBufferPointer { buffer in
-                    let cbsPointer = UnsafeMutablePointer<CBS>.allocate(capacity: 1)
-                    defer { cbsPointer.deallocate() }
-                    cbsPointer.pointee = CBS(data: buffer.baseAddress, len: buffer.count)
-
-                    guard CCryptoBoringSSL_MLDSA65_parse_public_key(self.pointer, cbsPointer) == 1 else {
+                    var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
+                    guard CCryptoBoringSSL_MLDSA65_parse_public_key(self.pointer, &cbs) == 1 else {
                         throw CryptoKitError.internalBoringSSLError()
                     }
                 }
             }
 
             /// The raw binary representation of the public key.
             var rawRepresentation: Data {
-                get throws {
-                    var cbb = CBB()
-                    // `CBB_init` can only return 0 on allocation failure, which we define as impossible.
-                    CCryptoBoringSSL_CBB_init(&cbb, MLDSA.PublicKey.Backing.bytesCount)
-
-                    guard CCryptoBoringSSL_MLDSA65_marshal_public_key(&cbb, self.pointer) == 1 else {
-                        CCryptoBoringSSL_CBB_cleanup(&cbb)
-                        throw CryptoKitError.internalBoringSSLError()
-                    }
-
-                    guard let data = CCryptoBoringSSL_CBB_data(&cbb) else {
-                        CCryptoBoringSSL_CBB_cleanup(&cbb)
-                        throw CryptoKitError.internalBoringSSLError()
-                    }
-                    return Data(bytes: data, count: CCryptoBoringSSL_CBB_len(&cbb))
-                }
+                var cbb = CBB()
+                // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
+                CCryptoBoringSSL_CBB_init(&cbb, MLDSA.PublicKey.Backing.bytesCount)
+                CCryptoBoringSSL_MLDSA65_marshal_public_key(&cbb, self.pointer)
+                return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
             }
 
             /// Verify a signature for the given data.

diff --git a/Tests/_CryptoExtrasTests/MLDSATests.swift b/Tests/_CryptoExtrasTests/MLDSATests.swift
@@ -97,7 +97,7 @@ final class MLDSATests: XCTestCase {
         let seed: [UInt8] = (0..<32).map { _ in UInt8.random(in: 0...255) }
         let key = try MLDSA.PrivateKey(seed: seed)
         let publicKey = key.publicKey
-        try encodedPublicKey.replaceSubrange(0..<MLDSA.PublicKey.bytesCount, with: publicKey.rawRepresentation)
+        encodedPublicKey.replaceSubrange(0..<MLDSA.PublicKey.bytesCount, with: publicKey.rawRepresentation)
 
         // Public key is 1 byte too short.
         let shortPublicKey = Array(encodedPublicKey.prefix(MLDSA.PublicKey.bytesCount - 1))
@@ -117,7 +117,7 @@ final class MLDSATests: XCTestCase {
             let publicKey = try MLDSA.PublicKey(rawRepresentation: Data(hexString: testVector.pub))
 
             let expectedkey = try MLDSA.PrivateKey(seed: seed).publicKey
-            try XCTAssertEqual(publicKey.rawRepresentation, expectedkey.rawRepresentation)
+            XCTAssertEqual(publicKey.rawRepresentation, expectedkey.rawRepresentation)
         }
     }