From 02f2c89f40ee63e67ffb35525737fe1dee4f966c Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Sat, 23 Nov 2024 20:09:40 +0100
Subject: [PATCH 01/13] Initial commit

---
 .../include/CCryptoBoringSSL.h                |  1 +
 .../_CryptoExtras/MLKEM/MLKEM_boring.swift    | 85 +++++++++++++++++++
 scripts/vendor-boringssl.sh                   |  1 +
 3 files changed, 87 insertions(+)
 create mode 100644 Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift

diff --git a/Sources/CCryptoBoringSSL/include/CCryptoBoringSSL.h b/Sources/CCryptoBoringSSL/include/CCryptoBoringSSL.h
index 3ff3aef5..dae17cff 100644
--- a/Sources/CCryptoBoringSSL/include/CCryptoBoringSSL.h
+++ b/Sources/CCryptoBoringSSL/include/CCryptoBoringSSL.h
@@ -44,6 +44,7 @@
 #include "CCryptoBoringSSL_hrss.h"
 #include "CCryptoBoringSSL_md4.h"
 #include "CCryptoBoringSSL_md5.h"
+#include "CCryptoBoringSSL_mlkem.h"
 #include "CCryptoBoringSSL_obj_mac.h"
 #include "CCryptoBoringSSL_objects.h"
 #include "CCryptoBoringSSL_opensslv.h"
diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
new file mode 100644
index 00000000..010c45d5
--- /dev/null
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
@@ -0,0 +1,85 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftCrypto open source project
+//
+// Copyright (c) 2024 Apple Inc. and the SwiftCrypto project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftCrypto project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import Crypto
+import Foundation
+
+@_implementationOnly import CCryptoBoringSSL
+
+/// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
+public enum MLKEM {}
+
+extension MLKEM {
+    /// A ML-KEM-768 private key.
+    public struct PrivateKey: Sendable {
+        private var backing: Backing
+
+        /// Initialize a ML-KEM-768 private key from a random seed.
+        public init() {
+            self.backing = Backing()
+        }
+
+        /// Initialize a ML-KEM-768 private key from a seed.
+        /// 
+        /// - Parameter seed: The seed to use to generate the private key.
+        /// 
+        /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
+        public init(seed: some DataProtocol) throws {
+            self.backing = try Backing(seed: seed)
+        }
+
+        fileprivate final class Backing {
+            var key: MLKEM768_private_key
+            var seed: Data
+
+            /// Initialize a ML-KEM-768 private key from a random seed.
+            init() {
+                self.key = .init()
+                self.seed = Data()
+
+                self.seed = withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(MLKEM_SEED_BYTES)) { seedPtr in
+                    withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(MLKEM768_PUBLIC_KEY_BYTES)) { publicKeyPtr in
+                        MLKEM768_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
+
+                        return Data(bytes: seedPtr.baseAddress!, count: Int(MLKEM_SEED_BYTES))
+                    }
+                }
+            }
+
+            /// Initialize a ML-KEM-768 private key from a seed.
+            /// 
+            /// - Parameter seed: The seed to use to generate the private key.
+            /// 
+            /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
+            init(seed: some DataProtocol) throws {
+                guard seed.count == MLKEM_SEED_BYTES else {
+                    throw CryptoKitError.incorrectKeySize
+                }
+
+                self.key = .init()
+                self.seed = Data(seed)
+
+                guard self.seed.withUnsafeBytes({ seedPtr in
+                    MLKEM768_private_key_from_seed(
+                        &self.key,
+                        seedPtr.baseAddress,
+                        seedPtr.count
+                    )
+                }) == 1 else {
+                    throw CryptoKitError.internalBoringSSLError()
+                }
+            }
+        }
+    }
+}
\ No newline at end of file
diff --git a/scripts/vendor-boringssl.sh b/scripts/vendor-boringssl.sh
index a3f13fb1..272be30c 100755
--- a/scripts/vendor-boringssl.sh
+++ b/scripts/vendor-boringssl.sh
@@ -344,6 +344,7 @@ cat << EOF > "$DSTROOT/include/CCryptoBoringSSL.h"
 #include "CCryptoBoringSSL_hrss.h"
 #include "CCryptoBoringSSL_md4.h"
 #include "CCryptoBoringSSL_md5.h"
+#include "CCryptoBoringSSL_mlkem.h"
 #include "CCryptoBoringSSL_obj_mac.h"
 #include "CCryptoBoringSSL_objects.h"
 #include "CCryptoBoringSSL_opensslv.h"

From 4cf36daa21afcd04103cf191cfedcff8efb29251 Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Sat, 23 Nov 2024 20:54:52 +0100
Subject: [PATCH 02/13] Add encap/decap for ML-KEM-768

---
 .../_CryptoExtras/MLKEM/MLKEM_boring.swift    | 138 +++++++++++++++++-
 Tests/_CryptoExtrasTests/MLKEMTests.swift     |  28 ++++
 2 files changed, 161 insertions(+), 5 deletions(-)
 create mode 100644 Tests/_CryptoExtrasTests/MLKEMTests.swift

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
index 010c45d5..4648abef 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
@@ -18,11 +18,13 @@ import Foundation
 @_implementationOnly import CCryptoBoringSSL
 
 /// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
+@available(macOS 14.0, *)
 public enum MLKEM {}
 
+@available(macOS 14.0, *)
 extension MLKEM {
     /// A ML-KEM-768 private key.
-    public struct PrivateKey: Sendable {
+    public struct PrivateKey: Sendable, KEMPrivateKey {
         private var backing: Backing
 
         /// Initialize a ML-KEM-768 private key from a random seed.
@@ -30,6 +32,13 @@ extension MLKEM {
             self.backing = Backing()
         }
 
+        /// Generate a ML-KEM-768 private key from a random seed.
+        ///
+        /// - Returns: The generated private key.
+        public static func generate() -> MLKEM.PrivateKey {
+            .init()
+        }
+
         /// Initialize a ML-KEM-768 private key from a seed.
         /// 
         /// - Parameter seed: The seed to use to generate the private key.
@@ -39,6 +48,27 @@ extension MLKEM {
             self.backing = try Backing(seed: seed)
         }
 
+        /// The seed from which this private key was generated.
+        public var seed: Data {
+            self.backing.seed
+        }
+
+        /// The public key associated with this private key.
+        public var publicKey: PublicKey {
+            self.backing.publicKey
+        }
+
+        /// Decapsulate a shared secret and create a symmetric key.
+        ///
+        /// - Parameter encapsulated: The encapsulated shared secret.
+        ///
+        /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
+        ///
+        /// - Returns: The symmetric key.
+        public func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
+            try self.backing.decapsulate(encapsulated)
+        }
+
         fileprivate final class Backing {
             var key: MLKEM768_private_key
             var seed: Data
@@ -48,11 +78,11 @@ extension MLKEM {
                 self.key = .init()
                 self.seed = Data()
 
-                self.seed = withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(MLKEM_SEED_BYTES)) { seedPtr in
-                    withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(MLKEM768_PUBLIC_KEY_BYTES)) { publicKeyPtr in
+                self.seed = withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.seedSizeInBytes) { seedPtr in
+                    withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.PublicKey.bytesCount) { publicKeyPtr in
                         MLKEM768_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
 
-                        return Data(bytes: seedPtr.baseAddress!, count: Int(MLKEM_SEED_BYTES))
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
                     }
                 }
             }
@@ -63,7 +93,7 @@ extension MLKEM {
             /// 
             /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
             init(seed: some DataProtocol) throws {
-                guard seed.count == MLKEM_SEED_BYTES else {
+                guard seed.count == MLKEM.seedSizeInBytes else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -80,6 +110,104 @@ extension MLKEM {
                     throw CryptoKitError.internalBoringSSLError()
                 }
             }
+
+            /// The public key associated with this private key.
+            var publicKey: PublicKey {
+                PublicKey(privateKeyBacking: self)
+            }
+
+            /// Decapsulate a shared secret and create a symmetric key.
+            ///
+            /// - Parameter encapsulated: The encapsulated shared secret.
+            ///
+            /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
+            ///
+            /// - Returns: The symmetric key.
+            func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
+                guard encapsulated.count == Int(MLKEM768_CIPHERTEXT_BYTES) else {
+                    throw CryptoKitError.incorrectParameterSize
+                }
+
+                let output = try Array<UInt8>(unsafeUninitializedCapacity: Int(MLKEM_SHARED_SECRET_BYTES)) { bufferPtr, length in
+                    let bytes: ContiguousBytes = encapsulated.regions.count == 1 ? encapsulated.regions.first! : Array(encapsulated)
+                    let result = bytes.withUnsafeBytes { encapsulatedPtr in
+                        MLKEM768_decap(
+                            bufferPtr.baseAddress,
+                            encapsulatedPtr.baseAddress,
+                            encapsulatedPtr.count,
+                            &self.key
+                        )
+                    }
+
+                    guard result == 1 else {
+                        throw CryptoKitError.internalBoringSSLError()
+                    }
+
+                    length = Int(MLKEM_SHARED_SECRET_BYTES)
+                }
+
+                return SymmetricKey(data: Data(output))
+            }
+        }
+    }
+}
+
+@available(macOS 14.0, *)
+extension MLKEM {
+    /// A ML-KEM-768 public key.
+    public struct PublicKey: Sendable, KEMPublicKey {
+        private var backing: Backing
+
+        fileprivate init(privateKeyBacking: PrivateKey.Backing) {
+            self.backing = Backing(privateKeyBacking: privateKeyBacking)
+        }
+
+        /// Encapsulate a shared secret.
+        ///
+        /// - Returns: The shared secret and its encapsulated version.
+        public func encapsulate() -> KEM.EncapsulationResult {
+            self.backing.encapsulate()
+        }
+
+        /// The size of the public key in bytes.
+        static let bytesCount = Backing.bytesCount
+
+        fileprivate final class Backing {
+            var key: MLKEM768_public_key
+
+            init(privateKeyBacking: PrivateKey.Backing) {
+                self.key = .init()
+                MLKEM768_public_from_private(&self.key, &privateKeyBacking.key)
+            }
+
+            /// The size of the public key in bytes.
+            static let bytesCount = Int(MLKEM768_PUBLIC_KEY_BYTES)
+            
+            /// Encapsulate a shared secret.
+            ///
+            /// - Returns: The shared secret and its encapsulated version.
+            func encapsulate() -> KEM.EncapsulationResult {
+                withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(MLKEM768_CIPHERTEXT_BYTES)) { encapsulatedPtr in
+                    withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(MLKEM_SHARED_SECRET_BYTES)) { secretPtr in
+                        MLKEM768_encap(
+                            encapsulatedPtr.baseAddress,
+                            secretPtr.baseAddress,
+                            &self.key
+                        )
+
+                        return KEM.EncapsulationResult(
+                            sharedSecret: SymmetricKey(data: Data(bytes: secretPtr.baseAddress!, count: Int(MLKEM_SHARED_SECRET_BYTES))),
+                            encapsulated: Data(bytes: encapsulatedPtr.baseAddress!, count: Int(MLKEM768_CIPHERTEXT_BYTES))
+                        )
+                    }
+                }
+            }
         }
     }
+}
+
+@available(macOS 14.0, *)
+extension MLKEM {
+    /// The size of the seed in bytes.
+    private static let seedSizeInBytes = Int(MLKEM_SEED_BYTES)
 }
\ No newline at end of file
diff --git a/Tests/_CryptoExtrasTests/MLKEMTests.swift b/Tests/_CryptoExtrasTests/MLKEMTests.swift
new file mode 100644
index 00000000..b269a037
--- /dev/null
+++ b/Tests/_CryptoExtrasTests/MLKEMTests.swift
@@ -0,0 +1,28 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftCrypto open source project
+//
+// Copyright (c) 2024 Apple Inc. and the SwiftCrypto project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftCrypto project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import XCTest
+
+@testable import _CryptoExtras
+
+@available(macOS 14.0, *)
+final class MLKEMTests: XCTestCase {
+    func testMLKEMEncapDecap() throws {
+        let privateKey = MLKEM.PrivateKey()
+        let publicKey = privateKey.publicKey
+        let encapsulationResult = publicKey.encapsulate()
+        let sharedSecret = try privateKey.decapsulate(encapsulationResult.encapsulated)
+        XCTAssertEqual(sharedSecret, encapsulationResult.sharedSecret)
+    }
+}
\ No newline at end of file

From 05fa85bdb3b9757aa14c1a5cb9caea9fd7508209 Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Wed, 27 Nov 2024 22:54:39 +0100
Subject: [PATCH 03/13] Add `rawRepresentation` for public keys and improve
 tests

---
 .../_CryptoExtras/MLKEM/MLKEM_boring.swift    | 86 +++++++++++++++----
 Tests/_CryptoExtrasTests/MLKEMTests.swift     | 40 ++++++++-
 2 files changed, 108 insertions(+), 18 deletions(-)

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
index 4648abef..e498c3f8 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
@@ -80,7 +80,7 @@ extension MLKEM {
 
                 self.seed = withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.seedSizeInBytes) { seedPtr in
                     withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.PublicKey.bytesCount) { publicKeyPtr in
-                        MLKEM768_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
+                        CCryptoBoringSSL_MLKEM768_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
 
                         return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
                     }
@@ -101,7 +101,7 @@ extension MLKEM {
                 self.seed = Data(seed)
 
                 guard self.seed.withUnsafeBytes({ seedPtr in
-                    MLKEM768_private_key_from_seed(
+                    CCryptoBoringSSL_MLKEM768_private_key_from_seed(
                         &self.key,
                         seedPtr.baseAddress,
                         seedPtr.count
@@ -124,14 +124,14 @@ extension MLKEM {
             ///
             /// - Returns: The symmetric key.
             func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
-                guard encapsulated.count == Int(MLKEM768_CIPHERTEXT_BYTES) else {
+                guard encapsulated.count == MLKEM.ciphertextSizeInBytes else {
                     throw CryptoKitError.incorrectParameterSize
                 }
 
-                let output = try Array<UInt8>(unsafeUninitializedCapacity: Int(MLKEM_SHARED_SECRET_BYTES)) { bufferPtr, length in
+                let output = try Array<UInt8>(unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes) { bufferPtr, length in
                     let bytes: ContiguousBytes = encapsulated.regions.count == 1 ? encapsulated.regions.first! : Array(encapsulated)
                     let result = bytes.withUnsafeBytes { encapsulatedPtr in
-                        MLKEM768_decap(
+                        CCryptoBoringSSL_MLKEM768_decap(
                             bufferPtr.baseAddress,
                             encapsulatedPtr.baseAddress,
                             encapsulatedPtr.count,
@@ -143,7 +143,7 @@ extension MLKEM {
                         throw CryptoKitError.internalBoringSSLError()
                     }
 
-                    length = Int(MLKEM_SHARED_SECRET_BYTES)
+                    length = MLKEM.sharedSecretSizeInBytes
                 }
 
                 return SymmetricKey(data: Data(output))
@@ -162,6 +162,20 @@ extension MLKEM {
             self.backing = Backing(privateKeyBacking: privateKeyBacking)
         }
 
+        /// Initialize a ML-KEM-768 public key from a raw representation.
+        /// 
+        /// - Parameter rawRepresentation: The public key bytes.
+        /// 
+        /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
+        init(rawRepresentation: some DataProtocol) throws {
+            self.backing = try Backing(rawRepresentation: rawRepresentation)
+        }
+
+        /// The raw binary representation of the public key.
+        public var rawRepresentation: Data {
+            self.backing.rawRepresentation
+        }
+
         /// Encapsulate a shared secret.
         ///
         /// - Returns: The shared secret and its encapsulated version.
@@ -177,37 +191,79 @@ extension MLKEM {
 
             init(privateKeyBacking: PrivateKey.Backing) {
                 self.key = .init()
-                MLKEM768_public_from_private(&self.key, &privateKeyBacking.key)
+                CCryptoBoringSSL_MLKEM768_public_from_private(&self.key, &privateKeyBacking.key)
             }
 
-            /// The size of the public key in bytes.
-            static let bytesCount = Int(MLKEM768_PUBLIC_KEY_BYTES)
+            /// Initialize a ML-KEM-768 public key from a raw representation.
+            /// 
+            /// - Parameter rawRepresentation: The public key bytes.
+            /// 
+            /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
+            init(rawRepresentation: some DataProtocol) throws {
+                guard rawRepresentation.count == MLKEM.PublicKey.bytesCount else {
+                    throw CryptoKitError.incorrectKeySize
+                }
+
+                self.key = .init()
+
+                let bytes: ContiguousBytes = rawRepresentation.regions.count == 1 ? rawRepresentation.regions.first! : Array(rawRepresentation)
+                try bytes.withUnsafeBytes { rawBuffer in
+                    try rawBuffer.withMemoryRebound(to: UInt8.self) { buffer in
+                        var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
+                        guard CCryptoBoringSSL_MLKEM768_parse_public_key(&self.key, &cbs) == 1 else {
+                            throw CryptoKitError.internalBoringSSLError()
+                        }
+                    }
+                }
+            }
+
+            /// The raw binary representation of the public key.
+            var rawRepresentation: Data {
+                var cbb = CBB()
+                // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM.PublicKey.Backing.bytesCount)
+                defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
+                CCryptoBoringSSL_MLKEM768_marshal_public_key(&cbb, &self.key)
+                return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
+            }
             
             /// Encapsulate a shared secret.
             ///
             /// - Returns: The shared secret and its encapsulated version.
             func encapsulate() -> KEM.EncapsulationResult {
-                withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(MLKEM768_CIPHERTEXT_BYTES)) { encapsulatedPtr in
-                    withUnsafeTemporaryAllocation(of: UInt8.self, capacity: Int(MLKEM_SHARED_SECRET_BYTES)) { secretPtr in
-                        MLKEM768_encap(
+                withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.ciphertextSizeInBytes) { encapsulatedPtr in
+                    withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.sharedSecretSizeInBytes) { secretPtr in
+                        CCryptoBoringSSL_MLKEM768_encap(
                             encapsulatedPtr.baseAddress,
                             secretPtr.baseAddress,
                             &self.key
                         )
 
                         return KEM.EncapsulationResult(
-                            sharedSecret: SymmetricKey(data: Data(bytes: secretPtr.baseAddress!, count: Int(MLKEM_SHARED_SECRET_BYTES))),
-                            encapsulated: Data(bytes: encapsulatedPtr.baseAddress!, count: Int(MLKEM768_CIPHERTEXT_BYTES))
+                            sharedSecret: SymmetricKey(data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)),
+                            encapsulated: Data(bytes: encapsulatedPtr.baseAddress!, count: MLKEM.ciphertextSizeInBytes)
                         )
                     }
                 }
             }
+
+            /// The size of the public key in bytes.
+            static let bytesCount = 1184
         }
     }
 }
 
+@available(macOS 14.0, *)
+extension MLKEM {
+    /// The size of the encapsulated shared secret in bytes.
+    private static let ciphertextSizeInBytes = 1088
+}
+
 @available(macOS 14.0, *)
 extension MLKEM {
     /// The size of the seed in bytes.
-    private static let seedSizeInBytes = Int(MLKEM_SEED_BYTES)
+    private static let seedSizeInBytes = 64
+
+    // The size of the shared secret in bytes.
+    private static let sharedSecretSizeInBytes = 32
 }
\ No newline at end of file
diff --git a/Tests/_CryptoExtrasTests/MLKEMTests.swift b/Tests/_CryptoExtrasTests/MLKEMTests.swift
index b269a037..dc2b0ab9 100644
--- a/Tests/_CryptoExtrasTests/MLKEMTests.swift
+++ b/Tests/_CryptoExtrasTests/MLKEMTests.swift
@@ -18,11 +18,45 @@ import XCTest
 
 @available(macOS 14.0, *)
 final class MLKEMTests: XCTestCase {
-    func testMLKEMEncapDecap() throws {
+    func testMLKEM_768() throws {
+        // Generate a key pair
         let privateKey = MLKEM.PrivateKey()
         let publicKey = privateKey.publicKey
+        
+        // Serialize and deserialize the private key
+        let seed = privateKey.seed
+        let privateKey2 = try MLKEM.PrivateKey(seed: seed)
+        XCTAssertEqual(privateKey.seed, privateKey2.seed)
+        
+        // Serialize and deserialize the public key
+        let publicKeyBytes = publicKey.rawRepresentation
+        var modifiedPublicKeyBytes = publicKeyBytes
+        modifiedPublicKeyBytes[0] = 0xff
+        modifiedPublicKeyBytes[1] = 0xff
+        // Parsing should fail because the first coefficient is >= kPrime;
+        XCTAssertThrowsError(try MLKEM.PublicKey(rawRepresentation: modifiedPublicKeyBytes))
+        
+        let publicKey2 = try MLKEM.PublicKey(rawRepresentation: publicKeyBytes)
+        XCTAssertEqual(publicKeyBytes, publicKey2.rawRepresentation)
+        
+        // Ensure public key derived from private key matches the original public key
+        let derivedPublicKey = privateKey.publicKey
+        XCTAssertEqual(publicKeyBytes, derivedPublicKey.rawRepresentation)
+        
+        // Serialize and deserialize the private key with modifications
+        var modifiedSeed = privateKey.seed
+        modifiedSeed[0] = 0xff
+        XCTAssertNotEqual(
+            try MLKEM.PrivateKey(seed: modifiedSeed).publicKey.rawRepresentation,
+            publicKeyBytes
+        )
+        
+        // Encapsulation and decapsulation
         let encapsulationResult = publicKey.encapsulate()
-        let sharedSecret = try privateKey.decapsulate(encapsulationResult.encapsulated)
-        XCTAssertEqual(sharedSecret, encapsulationResult.sharedSecret)
+        let sharedSecret1 = encapsulationResult.sharedSecret
+        let ciphertext = encapsulationResult.encapsulated
+        
+        let sharedSecret2 = try privateKey.decapsulate(ciphertext)
+        XCTAssertEqual(sharedSecret1, sharedSecret2)
     }
 }
\ No newline at end of file

From dee35107835fe5d4733c35bba27bf758e7754711 Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Sat, 30 Nov 2024 12:43:58 +0100
Subject: [PATCH 04/13] Swift Format

---
 .../_CryptoExtras/MLKEM/MLKEM_boring.swift    | 69 ++++++++++++-------
 Tests/_CryptoExtrasTests/MLKEMTests.swift     | 16 ++---
 2 files changed, 51 insertions(+), 34 deletions(-)

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
index e498c3f8..dfd0c8c0 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
@@ -12,11 +12,10 @@
 //
 //===----------------------------------------------------------------------===//
 
+@_implementationOnly import CCryptoBoringSSL
 import Crypto
 import Foundation
 
-@_implementationOnly import CCryptoBoringSSL
-
 /// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
 @available(macOS 14.0, *)
 public enum MLKEM {}
@@ -40,9 +39,9 @@ extension MLKEM {
         }
 
         /// Initialize a ML-KEM-768 private key from a seed.
-        /// 
+        ///
         /// - Parameter seed: The seed to use to generate the private key.
-        /// 
+        ///
         /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
         public init(seed: some DataProtocol) throws {
             self.backing = try Backing(seed: seed)
@@ -79,7 +78,9 @@ extension MLKEM {
                 self.seed = Data()
 
                 self.seed = withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.seedSizeInBytes) { seedPtr in
-                    withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.PublicKey.bytesCount) { publicKeyPtr in
+                    withUnsafeTemporaryAllocation(
+                        of: UInt8.self, capacity: MLKEM.PublicKey.bytesCount
+                    ) { publicKeyPtr in
                         CCryptoBoringSSL_MLKEM768_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
 
                         return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
@@ -88,9 +89,9 @@ extension MLKEM {
             }
 
             /// Initialize a ML-KEM-768 private key from a seed.
-            /// 
+            ///
             /// - Parameter seed: The seed to use to generate the private key.
-            /// 
+            ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
             init(seed: some DataProtocol) throws {
                 guard seed.count == MLKEM.seedSizeInBytes else {
@@ -100,13 +101,15 @@ extension MLKEM {
                 self.key = .init()
                 self.seed = Data(seed)
 
-                guard self.seed.withUnsafeBytes({ seedPtr in
-                    CCryptoBoringSSL_MLKEM768_private_key_from_seed(
-                        &self.key,
-                        seedPtr.baseAddress,
-                        seedPtr.count
-                    )
-                }) == 1 else {
+                guard
+                    self.seed.withUnsafeBytes({ seedPtr in
+                        CCryptoBoringSSL_MLKEM768_private_key_from_seed(
+                            &self.key,
+                            seedPtr.baseAddress,
+                            seedPtr.count
+                        )
+                    }) == 1
+                else {
                     throw CryptoKitError.internalBoringSSLError()
                 }
             }
@@ -128,8 +131,13 @@ extension MLKEM {
                     throw CryptoKitError.incorrectParameterSize
                 }
 
-                let output = try Array<UInt8>(unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes) { bufferPtr, length in
-                    let bytes: ContiguousBytes = encapsulated.regions.count == 1 ? encapsulated.regions.first! : Array(encapsulated)
+                let output = try [UInt8](
+                    unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes
+                ) { bufferPtr, length in
+                    let bytes: ContiguousBytes =
+                        encapsulated.regions.count == 1
+                        ? encapsulated.regions.first!
+                        : Array(encapsulated)
                     let result = bytes.withUnsafeBytes { encapsulatedPtr in
                         CCryptoBoringSSL_MLKEM768_decap(
                             bufferPtr.baseAddress,
@@ -163,9 +171,9 @@ extension MLKEM {
         }
 
         /// Initialize a ML-KEM-768 public key from a raw representation.
-        /// 
+        ///
         /// - Parameter rawRepresentation: The public key bytes.
-        /// 
+        ///
         /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
         init(rawRepresentation: some DataProtocol) throws {
             self.backing = try Backing(rawRepresentation: rawRepresentation)
@@ -195,9 +203,9 @@ extension MLKEM {
             }
 
             /// Initialize a ML-KEM-768 public key from a raw representation.
-            /// 
+            ///
             /// - Parameter rawRepresentation: The public key bytes.
-            /// 
+            ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
             init(rawRepresentation: some DataProtocol) throws {
                 guard rawRepresentation.count == MLKEM.PublicKey.bytesCount else {
@@ -206,7 +214,10 @@ extension MLKEM {
 
                 self.key = .init()
 
-                let bytes: ContiguousBytes = rawRepresentation.regions.count == 1 ? rawRepresentation.regions.first! : Array(rawRepresentation)
+                let bytes: ContiguousBytes =
+                    rawRepresentation.regions.count == 1
+                    ? rawRepresentation.regions.first!
+                    : Array(rawRepresentation)
                 try bytes.withUnsafeBytes { rawBuffer in
                     try rawBuffer.withMemoryRebound(to: UInt8.self) { buffer in
                         var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
@@ -226,13 +237,17 @@ extension MLKEM {
                 CCryptoBoringSSL_MLKEM768_marshal_public_key(&cbb, &self.key)
                 return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
             }
-            
+
             /// Encapsulate a shared secret.
             ///
             /// - Returns: The shared secret and its encapsulated version.
             func encapsulate() -> KEM.EncapsulationResult {
-                withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.ciphertextSizeInBytes) { encapsulatedPtr in
-                    withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.sharedSecretSizeInBytes) { secretPtr in
+                withUnsafeTemporaryAllocation(
+                    of: UInt8.self, capacity: MLKEM.ciphertextSizeInBytes
+                ) { encapsulatedPtr in
+                    withUnsafeTemporaryAllocation(
+                        of: UInt8.self, capacity: MLKEM.sharedSecretSizeInBytes
+                    ) { secretPtr in
                         CCryptoBoringSSL_MLKEM768_encap(
                             encapsulatedPtr.baseAddress,
                             secretPtr.baseAddress,
@@ -240,7 +255,9 @@ extension MLKEM {
                         )
 
                         return KEM.EncapsulationResult(
-                            sharedSecret: SymmetricKey(data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)),
+                            sharedSecret: SymmetricKey(
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)
+                            ),
                             encapsulated: Data(bytes: encapsulatedPtr.baseAddress!, count: MLKEM.ciphertextSizeInBytes)
                         )
                     }
@@ -266,4 +283,4 @@ extension MLKEM {
 
     // The size of the shared secret in bytes.
     private static let sharedSecretSizeInBytes = 32
-}
\ No newline at end of file
+}
diff --git a/Tests/_CryptoExtrasTests/MLKEMTests.swift b/Tests/_CryptoExtrasTests/MLKEMTests.swift
index dc2b0ab9..bda3539e 100644
--- a/Tests/_CryptoExtrasTests/MLKEMTests.swift
+++ b/Tests/_CryptoExtrasTests/MLKEMTests.swift
@@ -22,12 +22,12 @@ final class MLKEMTests: XCTestCase {
         // Generate a key pair
         let privateKey = MLKEM.PrivateKey()
         let publicKey = privateKey.publicKey
-        
+
         // Serialize and deserialize the private key
         let seed = privateKey.seed
         let privateKey2 = try MLKEM.PrivateKey(seed: seed)
         XCTAssertEqual(privateKey.seed, privateKey2.seed)
-        
+
         // Serialize and deserialize the public key
         let publicKeyBytes = publicKey.rawRepresentation
         var modifiedPublicKeyBytes = publicKeyBytes
@@ -35,14 +35,14 @@ final class MLKEMTests: XCTestCase {
         modifiedPublicKeyBytes[1] = 0xff
         // Parsing should fail because the first coefficient is >= kPrime;
         XCTAssertThrowsError(try MLKEM.PublicKey(rawRepresentation: modifiedPublicKeyBytes))
-        
+
         let publicKey2 = try MLKEM.PublicKey(rawRepresentation: publicKeyBytes)
         XCTAssertEqual(publicKeyBytes, publicKey2.rawRepresentation)
-        
+
         // Ensure public key derived from private key matches the original public key
         let derivedPublicKey = privateKey.publicKey
         XCTAssertEqual(publicKeyBytes, derivedPublicKey.rawRepresentation)
-        
+
         // Serialize and deserialize the private key with modifications
         var modifiedSeed = privateKey.seed
         modifiedSeed[0] = 0xff
@@ -50,13 +50,13 @@ final class MLKEMTests: XCTestCase {
             try MLKEM.PrivateKey(seed: modifiedSeed).publicKey.rawRepresentation,
             publicKeyBytes
         )
-        
+
         // Encapsulation and decapsulation
         let encapsulationResult = publicKey.encapsulate()
         let sharedSecret1 = encapsulationResult.sharedSecret
         let ciphertext = encapsulationResult.encapsulated
-        
+
         let sharedSecret2 = try privateKey.decapsulate(ciphertext)
         XCTAssertEqual(sharedSecret1, sharedSecret2)
     }
-}
\ No newline at end of file
+}

From bebed707c27f5917f1578770ae0844cca23029bb Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Thu, 12 Dec 2024 17:01:51 +0100
Subject: [PATCH 05/13] Rename `MLKEM` to `MLKEM768`

---
 ...KEM_boring.swift => MLKEM768_boring.swift} | 38 +++++++++----------
 .../{MLKEMTests.swift => MLKEM768Tests.swift} | 14 +++----
 2 files changed, 26 insertions(+), 26 deletions(-)
 rename Sources/_CryptoExtras/MLKEM/{MLKEM_boring.swift => MLKEM768_boring.swift} (90%)
 rename Tests/_CryptoExtrasTests/{MLKEMTests.swift => MLKEM768Tests.swift} (81%)

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
similarity index 90%
rename from Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
rename to Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
index dfd0c8c0..031d6b4f 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
@@ -18,10 +18,10 @@ import Foundation
 
 /// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
 @available(macOS 14.0, *)
-public enum MLKEM {}
+public enum MLKEM768 {}
 
 @available(macOS 14.0, *)
-extension MLKEM {
+extension MLKEM768 {
     /// A ML-KEM-768 private key.
     public struct PrivateKey: Sendable, KEMPrivateKey {
         private var backing: Backing
@@ -34,7 +34,7 @@ extension MLKEM {
         /// Generate a ML-KEM-768 private key from a random seed.
         ///
         /// - Returns: The generated private key.
-        public static func generate() -> MLKEM.PrivateKey {
+        public static func generate() -> MLKEM768.PrivateKey {
             .init()
         }
 
@@ -77,13 +77,13 @@ extension MLKEM {
                 self.key = .init()
                 self.seed = Data()
 
-                self.seed = withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM.seedSizeInBytes) { seedPtr in
+                self.seed = withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM768.seedSizeInBytes) { seedPtr in
                     withUnsafeTemporaryAllocation(
-                        of: UInt8.self, capacity: MLKEM.PublicKey.bytesCount
+                        of: UInt8.self, capacity: MLKEM768.PublicKey.bytesCount
                     ) { publicKeyPtr in
                         CCryptoBoringSSL_MLKEM768_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
 
-                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM768.seedSizeInBytes)
                     }
                 }
             }
@@ -94,7 +94,7 @@ extension MLKEM {
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
             init(seed: some DataProtocol) throws {
-                guard seed.count == MLKEM.seedSizeInBytes else {
+                guard seed.count == MLKEM768.seedSizeInBytes else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -127,12 +127,12 @@ extension MLKEM {
             ///
             /// - Returns: The symmetric key.
             func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
-                guard encapsulated.count == MLKEM.ciphertextSizeInBytes else {
+                guard encapsulated.count == MLKEM768.ciphertextSizeInBytes else {
                     throw CryptoKitError.incorrectParameterSize
                 }
 
                 let output = try [UInt8](
-                    unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes
+                    unsafeUninitializedCapacity: MLKEM768.sharedSecretSizeInBytes
                 ) { bufferPtr, length in
                     let bytes: ContiguousBytes =
                         encapsulated.regions.count == 1
@@ -151,7 +151,7 @@ extension MLKEM {
                         throw CryptoKitError.internalBoringSSLError()
                     }
 
-                    length = MLKEM.sharedSecretSizeInBytes
+                    length = MLKEM768.sharedSecretSizeInBytes
                 }
 
                 return SymmetricKey(data: Data(output))
@@ -161,7 +161,7 @@ extension MLKEM {
 }
 
 @available(macOS 14.0, *)
-extension MLKEM {
+extension MLKEM768 {
     /// A ML-KEM-768 public key.
     public struct PublicKey: Sendable, KEMPublicKey {
         private var backing: Backing
@@ -208,7 +208,7 @@ extension MLKEM {
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
             init(rawRepresentation: some DataProtocol) throws {
-                guard rawRepresentation.count == MLKEM.PublicKey.bytesCount else {
+                guard rawRepresentation.count == MLKEM768.PublicKey.bytesCount else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -232,7 +232,7 @@ extension MLKEM {
             var rawRepresentation: Data {
                 var cbb = CBB()
                 // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
-                CCryptoBoringSSL_CBB_init(&cbb, MLKEM.PublicKey.Backing.bytesCount)
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM768.PublicKey.Backing.bytesCount)
                 defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
                 CCryptoBoringSSL_MLKEM768_marshal_public_key(&cbb, &self.key)
                 return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
@@ -243,10 +243,10 @@ extension MLKEM {
             /// - Returns: The shared secret and its encapsulated version.
             func encapsulate() -> KEM.EncapsulationResult {
                 withUnsafeTemporaryAllocation(
-                    of: UInt8.self, capacity: MLKEM.ciphertextSizeInBytes
+                    of: UInt8.self, capacity: MLKEM768.ciphertextSizeInBytes
                 ) { encapsulatedPtr in
                     withUnsafeTemporaryAllocation(
-                        of: UInt8.self, capacity: MLKEM.sharedSecretSizeInBytes
+                        of: UInt8.self, capacity: MLKEM768.sharedSecretSizeInBytes
                     ) { secretPtr in
                         CCryptoBoringSSL_MLKEM768_encap(
                             encapsulatedPtr.baseAddress,
@@ -256,9 +256,9 @@ extension MLKEM {
 
                         return KEM.EncapsulationResult(
                             sharedSecret: SymmetricKey(
-                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM768.sharedSecretSizeInBytes)
                             ),
-                            encapsulated: Data(bytes: encapsulatedPtr.baseAddress!, count: MLKEM.ciphertextSizeInBytes)
+                            encapsulated: Data(bytes: encapsulatedPtr.baseAddress!, count: MLKEM768.ciphertextSizeInBytes)
                         )
                     }
                 }
@@ -271,13 +271,13 @@ extension MLKEM {
 }
 
 @available(macOS 14.0, *)
-extension MLKEM {
+extension MLKEM768 {
     /// The size of the encapsulated shared secret in bytes.
     private static let ciphertextSizeInBytes = 1088
 }
 
 @available(macOS 14.0, *)
-extension MLKEM {
+extension MLKEM768 {
     /// The size of the seed in bytes.
     private static let seedSizeInBytes = 64
 
diff --git a/Tests/_CryptoExtrasTests/MLKEMTests.swift b/Tests/_CryptoExtrasTests/MLKEM768Tests.swift
similarity index 81%
rename from Tests/_CryptoExtrasTests/MLKEMTests.swift
rename to Tests/_CryptoExtrasTests/MLKEM768Tests.swift
index bda3539e..875ef08f 100644
--- a/Tests/_CryptoExtrasTests/MLKEMTests.swift
+++ b/Tests/_CryptoExtrasTests/MLKEM768Tests.swift
@@ -17,15 +17,15 @@ import XCTest
 @testable import _CryptoExtras
 
 @available(macOS 14.0, *)
-final class MLKEMTests: XCTestCase {
-    func testMLKEM_768() throws {
+final class MLKEM768Tests: XCTestCase {
+    func testMLKEM768() throws {
         // Generate a key pair
-        let privateKey = MLKEM.PrivateKey()
+        let privateKey = MLKEM768.PrivateKey()
         let publicKey = privateKey.publicKey
 
         // Serialize and deserialize the private key
         let seed = privateKey.seed
-        let privateKey2 = try MLKEM.PrivateKey(seed: seed)
+        let privateKey2 = try MLKEM768.PrivateKey(seed: seed)
         XCTAssertEqual(privateKey.seed, privateKey2.seed)
 
         // Serialize and deserialize the public key
@@ -34,9 +34,9 @@ final class MLKEMTests: XCTestCase {
         modifiedPublicKeyBytes[0] = 0xff
         modifiedPublicKeyBytes[1] = 0xff
         // Parsing should fail because the first coefficient is >= kPrime;
-        XCTAssertThrowsError(try MLKEM.PublicKey(rawRepresentation: modifiedPublicKeyBytes))
+        XCTAssertThrowsError(try MLKEM768.PublicKey(rawRepresentation: modifiedPublicKeyBytes))
 
-        let publicKey2 = try MLKEM.PublicKey(rawRepresentation: publicKeyBytes)
+        let publicKey2 = try MLKEM768.PublicKey(rawRepresentation: publicKeyBytes)
         XCTAssertEqual(publicKeyBytes, publicKey2.rawRepresentation)
 
         // Ensure public key derived from private key matches the original public key
@@ -47,7 +47,7 @@ final class MLKEMTests: XCTestCase {
         var modifiedSeed = privateKey.seed
         modifiedSeed[0] = 0xff
         XCTAssertNotEqual(
-            try MLKEM.PrivateKey(seed: modifiedSeed).publicKey.rawRepresentation,
+            try MLKEM768.PrivateKey(seed: modifiedSeed).publicKey.rawRepresentation,
             publicKeyBytes
         )
 

From 4d6751f59a18a04ea024d69330ed58212f2b4f9e Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Thu, 12 Dec 2024 17:13:01 +0100
Subject: [PATCH 06/13] Swift Format

---
 .../_CryptoExtras/MLKEM/MLKEM768_boring.swift | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
index 031d6b4f..c0673bb9 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
@@ -77,9 +77,13 @@ extension MLKEM768 {
                 self.key = .init()
                 self.seed = Data()
 
-                self.seed = withUnsafeTemporaryAllocation(of: UInt8.self, capacity: MLKEM768.seedSizeInBytes) { seedPtr in
+                self.seed = withUnsafeTemporaryAllocation(
+                    of: UInt8.self,
+                    capacity: MLKEM768.seedSizeInBytes
+                ) { seedPtr in
                     withUnsafeTemporaryAllocation(
-                        of: UInt8.self, capacity: MLKEM768.PublicKey.bytesCount
+                        of: UInt8.self,
+                        capacity: MLKEM768.PublicKey.bytesCount
                     ) { publicKeyPtr in
                         CCryptoBoringSSL_MLKEM768_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
 
@@ -243,10 +247,12 @@ extension MLKEM768 {
             /// - Returns: The shared secret and its encapsulated version.
             func encapsulate() -> KEM.EncapsulationResult {
                 withUnsafeTemporaryAllocation(
-                    of: UInt8.self, capacity: MLKEM768.ciphertextSizeInBytes
+                    of: UInt8.self,
+                    capacity: MLKEM768.ciphertextSizeInBytes
                 ) { encapsulatedPtr in
                     withUnsafeTemporaryAllocation(
-                        of: UInt8.self, capacity: MLKEM768.sharedSecretSizeInBytes
+                        of: UInt8.self,
+                        capacity: MLKEM768.sharedSecretSizeInBytes
                     ) { secretPtr in
                         CCryptoBoringSSL_MLKEM768_encap(
                             encapsulatedPtr.baseAddress,
@@ -258,7 +264,10 @@ extension MLKEM768 {
                             sharedSecret: SymmetricKey(
                                 data: Data(bytes: secretPtr.baseAddress!, count: MLKEM768.sharedSecretSizeInBytes)
                             ),
-                            encapsulated: Data(bytes: encapsulatedPtr.baseAddress!, count: MLKEM768.ciphertextSizeInBytes)
+                            encapsulated: Data(
+                                bytes: encapsulatedPtr.baseAddress!,
+                                count: MLKEM768.ciphertextSizeInBytes
+                            )
                         )
                     }
                 }

From 0172516c7686d4e6c82149b93ba483e21ad5d381 Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Thu, 12 Dec 2024 17:21:50 +0100
Subject: [PATCH 07/13] Add `MLKEM1024`

---
 .../MLKEM/MLKEM1024_boring.swift              | 295 ++++++++++++++++++
 Tests/_CryptoExtrasTests/MLKEM1024Tests.swift |  62 ++++
 2 files changed, 357 insertions(+)
 create mode 100644 Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift
 create mode 100644 Tests/_CryptoExtrasTests/MLKEM1024Tests.swift

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift
new file mode 100644
index 00000000..1e286208
--- /dev/null
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift
@@ -0,0 +1,295 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftCrypto open source project
+//
+// Copyright (c) 2024 Apple Inc. and the SwiftCrypto project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftCrypto project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+@_implementationOnly import CCryptoBoringSSL
+import Crypto
+import Foundation
+
+/// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
+@available(macOS 14.0, *)
+public enum MLKEM1024 {}
+
+@available(macOS 14.0, *)
+extension MLKEM1024 {
+    /// A ML-KEM-1024 private key.
+    public struct PrivateKey: Sendable, KEMPrivateKey {
+        private var backing: Backing
+
+        /// Initialize a ML-KEM-1024 private key from a random seed.
+        public init() {
+            self.backing = Backing()
+        }
+
+        /// Generate a ML-KEM-1024 private key from a random seed.
+        ///
+        /// - Returns: The generated private key.
+        public static func generate() -> MLKEM1024.PrivateKey {
+            .init()
+        }
+
+        /// Initialize a ML-KEM-1024 private key from a seed.
+        ///
+        /// - Parameter seed: The seed to use to generate the private key.
+        ///
+        /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
+        public init(seed: some DataProtocol) throws {
+            self.backing = try Backing(seed: seed)
+        }
+
+        /// The seed from which this private key was generated.
+        public var seed: Data {
+            self.backing.seed
+        }
+
+        /// The public key associated with this private key.
+        public var publicKey: PublicKey {
+            self.backing.publicKey
+        }
+
+        /// Decapsulate a shared secret and create a symmetric key.
+        ///
+        /// - Parameter encapsulated: The encapsulated shared secret.
+        ///
+        /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
+        ///
+        /// - Returns: The symmetric key.
+        public func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
+            try self.backing.decapsulate(encapsulated)
+        }
+
+        fileprivate final class Backing {
+            var key: MLKEM1024_private_key
+            var seed: Data
+
+            /// Initialize a ML-KEM-1024 private key from a random seed.
+            init() {
+                self.key = .init()
+                self.seed = Data()
+
+                self.seed = withUnsafeTemporaryAllocation(
+                    of: UInt8.self,
+                    capacity: MLKEM1024.seedSizeInBytes
+                ) { seedPtr in
+                    withUnsafeTemporaryAllocation(
+                        of: UInt8.self,
+                        capacity: MLKEM1024.PublicKey.bytesCount
+                    ) { publicKeyPtr in
+                        CCryptoBoringSSL_MLKEM1024_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
+
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM1024.seedSizeInBytes)
+                    }
+                }
+            }
+
+            /// Initialize a ML-KEM-1024 private key from a seed.
+            ///
+            /// - Parameter seed: The seed to use to generate the private key.
+            ///
+            /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
+            init(seed: some DataProtocol) throws {
+                guard seed.count == MLKEM1024.seedSizeInBytes else {
+                    throw CryptoKitError.incorrectKeySize
+                }
+
+                self.key = .init()
+                self.seed = Data(seed)
+
+                guard
+                    self.seed.withUnsafeBytes({ seedPtr in
+                        CCryptoBoringSSL_MLKEM1024_private_key_from_seed(
+                            &self.key,
+                            seedPtr.baseAddress,
+                            seedPtr.count
+                        )
+                    }) == 1
+                else {
+                    throw CryptoKitError.internalBoringSSLError()
+                }
+            }
+
+            /// The public key associated with this private key.
+            var publicKey: PublicKey {
+                PublicKey(privateKeyBacking: self)
+            }
+
+            /// Decapsulate a shared secret and create a symmetric key.
+            ///
+            /// - Parameter encapsulated: The encapsulated shared secret.
+            ///
+            /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
+            ///
+            /// - Returns: The symmetric key.
+            func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
+                guard encapsulated.count == MLKEM1024.ciphertextSizeInBytes else {
+                    throw CryptoKitError.incorrectParameterSize
+                }
+
+                let output = try [UInt8](
+                    unsafeUninitializedCapacity: MLKEM1024.sharedSecretSizeInBytes
+                ) { bufferPtr, length in
+                    let bytes: ContiguousBytes =
+                        encapsulated.regions.count == 1
+                        ? encapsulated.regions.first!
+                        : Array(encapsulated)
+                    let result = bytes.withUnsafeBytes { encapsulatedPtr in
+                        CCryptoBoringSSL_MLKEM1024_decap(
+                            bufferPtr.baseAddress,
+                            encapsulatedPtr.baseAddress,
+                            encapsulatedPtr.count,
+                            &self.key
+                        )
+                    }
+
+                    guard result == 1 else {
+                        throw CryptoKitError.internalBoringSSLError()
+                    }
+
+                    length = MLKEM1024.sharedSecretSizeInBytes
+                }
+
+                return SymmetricKey(data: Data(output))
+            }
+        }
+    }
+}
+
+@available(macOS 14.0, *)
+extension MLKEM1024 {
+    /// A ML-KEM-1024 public key.
+    public struct PublicKey: Sendable, KEMPublicKey {
+        private var backing: Backing
+
+        fileprivate init(privateKeyBacking: PrivateKey.Backing) {
+            self.backing = Backing(privateKeyBacking: privateKeyBacking)
+        }
+
+        /// Initialize a ML-KEM-1024 public key from a raw representation.
+        ///
+        /// - Parameter rawRepresentation: The public key bytes.
+        ///
+        /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
+        init(rawRepresentation: some DataProtocol) throws {
+            self.backing = try Backing(rawRepresentation: rawRepresentation)
+        }
+
+        /// The raw binary representation of the public key.
+        public var rawRepresentation: Data {
+            self.backing.rawRepresentation
+        }
+
+        /// Encapsulate a shared secret.
+        ///
+        /// - Returns: The shared secret and its encapsulated version.
+        public func encapsulate() -> KEM.EncapsulationResult {
+            self.backing.encapsulate()
+        }
+
+        /// The size of the public key in bytes.
+        static let bytesCount = Backing.bytesCount
+
+        fileprivate final class Backing {
+            var key: MLKEM1024_public_key
+
+            init(privateKeyBacking: PrivateKey.Backing) {
+                self.key = .init()
+                CCryptoBoringSSL_MLKEM1024_public_from_private(&self.key, &privateKeyBacking.key)
+            }
+
+            /// Initialize a ML-KEM-1024 public key from a raw representation.
+            ///
+            /// - Parameter rawRepresentation: The public key bytes.
+            ///
+            /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
+            init(rawRepresentation: some DataProtocol) throws {
+                guard rawRepresentation.count == MLKEM1024.PublicKey.bytesCount else {
+                    throw CryptoKitError.incorrectKeySize
+                }
+
+                self.key = .init()
+
+                let bytes: ContiguousBytes =
+                    rawRepresentation.regions.count == 1
+                    ? rawRepresentation.regions.first!
+                    : Array(rawRepresentation)
+                try bytes.withUnsafeBytes { rawBuffer in
+                    try rawBuffer.withMemoryRebound(to: UInt8.self) { buffer in
+                        var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
+                        guard CCryptoBoringSSL_MLKEM1024_parse_public_key(&self.key, &cbs) == 1 else {
+                            throw CryptoKitError.internalBoringSSLError()
+                        }
+                    }
+                }
+            }
+
+            /// The raw binary representation of the public key.
+            var rawRepresentation: Data {
+                var cbb = CBB()
+                // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM1024.PublicKey.Backing.bytesCount)
+                defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
+                CCryptoBoringSSL_MLKEM1024_marshal_public_key(&cbb, &self.key)
+                return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
+            }
+
+            /// Encapsulate a shared secret.
+            ///
+            /// - Returns: The shared secret and its encapsulated version.
+            func encapsulate() -> KEM.EncapsulationResult {
+                withUnsafeTemporaryAllocation(
+                    of: UInt8.self,
+                    capacity: MLKEM1024.ciphertextSizeInBytes
+                ) { encapsulatedPtr in
+                    withUnsafeTemporaryAllocation(
+                        of: UInt8.self,
+                        capacity: MLKEM1024.sharedSecretSizeInBytes
+                    ) { secretPtr in
+                        CCryptoBoringSSL_MLKEM1024_encap(
+                            encapsulatedPtr.baseAddress,
+                            secretPtr.baseAddress,
+                            &self.key
+                        )
+
+                        return KEM.EncapsulationResult(
+                            sharedSecret: SymmetricKey(
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM1024.sharedSecretSizeInBytes)
+                            ),
+                            encapsulated: Data(
+                                bytes: encapsulatedPtr.baseAddress!,
+                                count: MLKEM1024.ciphertextSizeInBytes
+                            )
+                        )
+                    }
+                }
+            }
+
+            /// The size of the public key in bytes.
+            static let bytesCount = 1568
+        }
+    }
+}
+
+@available(macOS 14.0, *)
+extension MLKEM1024 {
+    /// The size of the encapsulated shared secret in bytes.
+    private static let ciphertextSizeInBytes = 1568
+}
+
+@available(macOS 14.0, *)
+extension MLKEM1024 {
+    /// The size of the seed in bytes.
+    private static let seedSizeInBytes = 64
+
+    // The size of the shared secret in bytes.
+    private static let sharedSecretSizeInBytes = 32
+}
diff --git a/Tests/_CryptoExtrasTests/MLKEM1024Tests.swift b/Tests/_CryptoExtrasTests/MLKEM1024Tests.swift
new file mode 100644
index 00000000..6ca20372
--- /dev/null
+++ b/Tests/_CryptoExtrasTests/MLKEM1024Tests.swift
@@ -0,0 +1,62 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftCrypto open source project
+//
+// Copyright (c) 2024 Apple Inc. and the SwiftCrypto project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftCrypto project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+import XCTest
+
+@testable import _CryptoExtras
+
+@available(macOS 14.0, *)
+final class MLKEM1024Tests: XCTestCase {
+    func testMLKEM1024() throws {
+        // Generate a key pair
+        let privateKey = MLKEM1024.PrivateKey()
+        let publicKey = privateKey.publicKey
+
+        // Serialize and deserialize the private key
+        let seed = privateKey.seed
+        let privateKey2 = try MLKEM1024.PrivateKey(seed: seed)
+        XCTAssertEqual(privateKey.seed, privateKey2.seed)
+
+        // Serialize and deserialize the public key
+        let publicKeyBytes = publicKey.rawRepresentation
+        var modifiedPublicKeyBytes = publicKeyBytes
+        modifiedPublicKeyBytes[0] = 0xff
+        modifiedPublicKeyBytes[1] = 0xff
+        // Parsing should fail because the first coefficient is >= kPrime;
+        XCTAssertThrowsError(try MLKEM1024.PublicKey(rawRepresentation: modifiedPublicKeyBytes))
+
+        let publicKey2 = try MLKEM1024.PublicKey(rawRepresentation: publicKeyBytes)
+        XCTAssertEqual(publicKeyBytes, publicKey2.rawRepresentation)
+
+        // Ensure public key derived from private key matches the original public key
+        let derivedPublicKey = privateKey.publicKey
+        XCTAssertEqual(publicKeyBytes, derivedPublicKey.rawRepresentation)
+
+        // Serialize and deserialize the private key with modifications
+        var modifiedSeed = privateKey.seed
+        modifiedSeed[0] = 0xff
+        XCTAssertNotEqual(
+            try MLKEM1024.PrivateKey(seed: modifiedSeed).publicKey.rawRepresentation,
+            publicKeyBytes
+        )
+
+        // Encapsulation and decapsulation
+        let encapsulationResult = publicKey.encapsulate()
+        let sharedSecret1 = encapsulationResult.sharedSecret
+        let ciphertext = encapsulationResult.encapsulated
+
+        let sharedSecret2 = try privateKey.decapsulate(ciphertext)
+        XCTAssertEqual(sharedSecret1, sharedSecret2)
+    }
+}

From 66dd1379b67719425aa797b3a560dd24f79e8eb8 Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Thu, 12 Dec 2024 17:29:37 +0100
Subject: [PATCH 08/13] Make PublicKey init public

---
 Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift | 4 +++-
 Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift  | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift
index 1e286208..506aa5d4 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift
@@ -17,6 +17,8 @@ import Crypto
 import Foundation
 
 /// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
+///
+/// > Note: You should prefer ``MLKEM768`` where possible.
 @available(macOS 14.0, *)
 public enum MLKEM1024 {}
 
@@ -179,7 +181,7 @@ extension MLKEM1024 {
         /// - Parameter rawRepresentation: The public key bytes.
         ///
         /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
-        init(rawRepresentation: some DataProtocol) throws {
+        public init(rawRepresentation: some DataProtocol) throws {
             self.backing = try Backing(rawRepresentation: rawRepresentation)
         }
 
diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
index c0673bb9..ffd537a4 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
@@ -179,7 +179,7 @@ extension MLKEM768 {
         /// - Parameter rawRepresentation: The public key bytes.
         ///
         /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
-        init(rawRepresentation: some DataProtocol) throws {
+        public init(rawRepresentation: some DataProtocol) throws {
             self.backing = try Backing(rawRepresentation: rawRepresentation)
         }
 

From f6673f15b41b2a05a3b23d1b4e0b98bb8ec97bac Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Thu, 12 Dec 2024 17:48:22 +0100
Subject: [PATCH 09/13] Update CMake lists

---
 Sources/_CryptoExtras/CMakeLists.txt | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/Sources/_CryptoExtras/CMakeLists.txt b/Sources/_CryptoExtras/CMakeLists.txt
index 9c578708..4f74432a 100644
--- a/Sources/_CryptoExtras/CMakeLists.txt
+++ b/Sources/_CryptoExtras/CMakeLists.txt
@@ -32,6 +32,8 @@ add_library(_CryptoExtras
   "Key Derivation/PBKDF2/PBKDF2.swift"
   "Key Derivation/Scrypt/BoringSSL/Scrypt_boring.swift"
   "Key Derivation/Scrypt/Scrypt.swift"
+  "MLKEM/MLKEM1024_boring.swift"
+  "MLKEM/MLKEM768_boring.swift"
   "OPRFs/OPRF.swift"
   "OPRFs/OPRFClient.swift"
   "OPRFs/OPRFServer.swift"

From e7011aa0fc92dbe2ae60a6758655fdcb3867d1a0 Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Thu, 12 Dec 2024 21:14:07 +0100
Subject: [PATCH 10/13] Unify tests

---
 Tests/_CryptoExtrasTests/MLKEM768Tests.swift  | 62 -------------------
 ...{MLKEM1024Tests.swift => MLKEMTests.swift} | 44 ++++++++++++-
 2 files changed, 43 insertions(+), 63 deletions(-)
 delete mode 100644 Tests/_CryptoExtrasTests/MLKEM768Tests.swift
 rename Tests/_CryptoExtrasTests/{MLKEM1024Tests.swift => MLKEMTests.swift} (56%)

diff --git a/Tests/_CryptoExtrasTests/MLKEM768Tests.swift b/Tests/_CryptoExtrasTests/MLKEM768Tests.swift
deleted file mode 100644
index 875ef08f..00000000
--- a/Tests/_CryptoExtrasTests/MLKEM768Tests.swift
+++ /dev/null
@@ -1,62 +0,0 @@
-//===----------------------------------------------------------------------===//
-//
-// This source file is part of the SwiftCrypto open source project
-//
-// Copyright (c) 2024 Apple Inc. and the SwiftCrypto project authors
-// Licensed under Apache License v2.0
-//
-// See LICENSE.txt for license information
-// See CONTRIBUTORS.txt for the list of SwiftCrypto project authors
-//
-// SPDX-License-Identifier: Apache-2.0
-//
-//===----------------------------------------------------------------------===//
-
-import XCTest
-
-@testable import _CryptoExtras
-
-@available(macOS 14.0, *)
-final class MLKEM768Tests: XCTestCase {
-    func testMLKEM768() throws {
-        // Generate a key pair
-        let privateKey = MLKEM768.PrivateKey()
-        let publicKey = privateKey.publicKey
-
-        // Serialize and deserialize the private key
-        let seed = privateKey.seed
-        let privateKey2 = try MLKEM768.PrivateKey(seed: seed)
-        XCTAssertEqual(privateKey.seed, privateKey2.seed)
-
-        // Serialize and deserialize the public key
-        let publicKeyBytes = publicKey.rawRepresentation
-        var modifiedPublicKeyBytes = publicKeyBytes
-        modifiedPublicKeyBytes[0] = 0xff
-        modifiedPublicKeyBytes[1] = 0xff
-        // Parsing should fail because the first coefficient is >= kPrime;
-        XCTAssertThrowsError(try MLKEM768.PublicKey(rawRepresentation: modifiedPublicKeyBytes))
-
-        let publicKey2 = try MLKEM768.PublicKey(rawRepresentation: publicKeyBytes)
-        XCTAssertEqual(publicKeyBytes, publicKey2.rawRepresentation)
-
-        // Ensure public key derived from private key matches the original public key
-        let derivedPublicKey = privateKey.publicKey
-        XCTAssertEqual(publicKeyBytes, derivedPublicKey.rawRepresentation)
-
-        // Serialize and deserialize the private key with modifications
-        var modifiedSeed = privateKey.seed
-        modifiedSeed[0] = 0xff
-        XCTAssertNotEqual(
-            try MLKEM768.PrivateKey(seed: modifiedSeed).publicKey.rawRepresentation,
-            publicKeyBytes
-        )
-
-        // Encapsulation and decapsulation
-        let encapsulationResult = publicKey.encapsulate()
-        let sharedSecret1 = encapsulationResult.sharedSecret
-        let ciphertext = encapsulationResult.encapsulated
-
-        let sharedSecret2 = try privateKey.decapsulate(ciphertext)
-        XCTAssertEqual(sharedSecret1, sharedSecret2)
-    }
-}
diff --git a/Tests/_CryptoExtrasTests/MLKEM1024Tests.swift b/Tests/_CryptoExtrasTests/MLKEMTests.swift
similarity index 56%
rename from Tests/_CryptoExtrasTests/MLKEM1024Tests.swift
rename to Tests/_CryptoExtrasTests/MLKEMTests.swift
index 6ca20372..e383a115 100644
--- a/Tests/_CryptoExtrasTests/MLKEM1024Tests.swift
+++ b/Tests/_CryptoExtrasTests/MLKEMTests.swift
@@ -17,7 +17,49 @@ import XCTest
 @testable import _CryptoExtras
 
 @available(macOS 14.0, *)
-final class MLKEM1024Tests: XCTestCase {
+final class MLKEMTests: XCTestCase {
+    func testMLKEM768() throws {
+        // Generate a key pair
+        let privateKey = MLKEM768.PrivateKey()
+        let publicKey = privateKey.publicKey
+
+        // Serialize and deserialize the private key
+        let seed = privateKey.seed
+        let privateKey2 = try MLKEM768.PrivateKey(seed: seed)
+        XCTAssertEqual(privateKey.seed, privateKey2.seed)
+
+        // Serialize and deserialize the public key
+        let publicKeyBytes = publicKey.rawRepresentation
+        var modifiedPublicKeyBytes = publicKeyBytes
+        modifiedPublicKeyBytes[0] = 0xff
+        modifiedPublicKeyBytes[1] = 0xff
+        // Parsing should fail because the first coefficient is >= kPrime;
+        XCTAssertThrowsError(try MLKEM768.PublicKey(rawRepresentation: modifiedPublicKeyBytes))
+
+        let publicKey2 = try MLKEM768.PublicKey(rawRepresentation: publicKeyBytes)
+        XCTAssertEqual(publicKeyBytes, publicKey2.rawRepresentation)
+
+        // Ensure public key derived from private key matches the original public key
+        let derivedPublicKey = privateKey.publicKey
+        XCTAssertEqual(publicKeyBytes, derivedPublicKey.rawRepresentation)
+
+        // Serialize and deserialize the private key with modifications
+        var modifiedSeed = privateKey.seed
+        modifiedSeed[0] = 0xff
+        XCTAssertNotEqual(
+            try MLKEM768.PrivateKey(seed: modifiedSeed).publicKey.rawRepresentation,
+            publicKeyBytes
+        )
+
+        // Encapsulation and decapsulation
+        let encapsulationResult = publicKey.encapsulate()
+        let sharedSecret1 = encapsulationResult.sharedSecret
+        let ciphertext = encapsulationResult.encapsulated
+
+        let sharedSecret2 = try privateKey.decapsulate(ciphertext)
+        XCTAssertEqual(sharedSecret1, sharedSecret2)
+    }
+
     func testMLKEM1024() throws {
         // Generate a key pair
         let privateKey = MLKEM1024.PrivateKey()

From 1003e8d74622f6c4b1b12a1076a2c5d8a48b99b6 Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Sat, 14 Dec 2024 12:49:43 +0100
Subject: [PATCH 11/13] Generate all parameter sets with `gyb`

---
 .../_CryptoExtras/MLKEM/MLKEM768_boring.swift | 295 ---------
 .../_CryptoExtras/MLKEM/MLKEM_boring.swift    | 575 ++++++++++++++++++
 ...24_boring.swift => MLKEM_boring.swift.gyb} |  98 +--
 3 files changed, 629 insertions(+), 339 deletions(-)
 delete mode 100644 Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
 create mode 100644 Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
 rename Sources/_CryptoExtras/MLKEM/{MLKEM1024_boring.swift => MLKEM_boring.swift.gyb} (73%)

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
deleted file mode 100644
index ffd537a4..00000000
--- a/Sources/_CryptoExtras/MLKEM/MLKEM768_boring.swift
+++ /dev/null
@@ -1,295 +0,0 @@
-//===----------------------------------------------------------------------===//
-//
-// This source file is part of the SwiftCrypto open source project
-//
-// Copyright (c) 2024 Apple Inc. and the SwiftCrypto project authors
-// Licensed under Apache License v2.0
-//
-// See LICENSE.txt for license information
-// See CONTRIBUTORS.txt for the list of SwiftCrypto project authors
-//
-// SPDX-License-Identifier: Apache-2.0
-//
-//===----------------------------------------------------------------------===//
-
-@_implementationOnly import CCryptoBoringSSL
-import Crypto
-import Foundation
-
-/// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
-@available(macOS 14.0, *)
-public enum MLKEM768 {}
-
-@available(macOS 14.0, *)
-extension MLKEM768 {
-    /// A ML-KEM-768 private key.
-    public struct PrivateKey: Sendable, KEMPrivateKey {
-        private var backing: Backing
-
-        /// Initialize a ML-KEM-768 private key from a random seed.
-        public init() {
-            self.backing = Backing()
-        }
-
-        /// Generate a ML-KEM-768 private key from a random seed.
-        ///
-        /// - Returns: The generated private key.
-        public static func generate() -> MLKEM768.PrivateKey {
-            .init()
-        }
-
-        /// Initialize a ML-KEM-768 private key from a seed.
-        ///
-        /// - Parameter seed: The seed to use to generate the private key.
-        ///
-        /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
-        public init(seed: some DataProtocol) throws {
-            self.backing = try Backing(seed: seed)
-        }
-
-        /// The seed from which this private key was generated.
-        public var seed: Data {
-            self.backing.seed
-        }
-
-        /// The public key associated with this private key.
-        public var publicKey: PublicKey {
-            self.backing.publicKey
-        }
-
-        /// Decapsulate a shared secret and create a symmetric key.
-        ///
-        /// - Parameter encapsulated: The encapsulated shared secret.
-        ///
-        /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
-        ///
-        /// - Returns: The symmetric key.
-        public func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
-            try self.backing.decapsulate(encapsulated)
-        }
-
-        fileprivate final class Backing {
-            var key: MLKEM768_private_key
-            var seed: Data
-
-            /// Initialize a ML-KEM-768 private key from a random seed.
-            init() {
-                self.key = .init()
-                self.seed = Data()
-
-                self.seed = withUnsafeTemporaryAllocation(
-                    of: UInt8.self,
-                    capacity: MLKEM768.seedSizeInBytes
-                ) { seedPtr in
-                    withUnsafeTemporaryAllocation(
-                        of: UInt8.self,
-                        capacity: MLKEM768.PublicKey.bytesCount
-                    ) { publicKeyPtr in
-                        CCryptoBoringSSL_MLKEM768_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
-
-                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM768.seedSizeInBytes)
-                    }
-                }
-            }
-
-            /// Initialize a ML-KEM-768 private key from a seed.
-            ///
-            /// - Parameter seed: The seed to use to generate the private key.
-            ///
-            /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
-            init(seed: some DataProtocol) throws {
-                guard seed.count == MLKEM768.seedSizeInBytes else {
-                    throw CryptoKitError.incorrectKeySize
-                }
-
-                self.key = .init()
-                self.seed = Data(seed)
-
-                guard
-                    self.seed.withUnsafeBytes({ seedPtr in
-                        CCryptoBoringSSL_MLKEM768_private_key_from_seed(
-                            &self.key,
-                            seedPtr.baseAddress,
-                            seedPtr.count
-                        )
-                    }) == 1
-                else {
-                    throw CryptoKitError.internalBoringSSLError()
-                }
-            }
-
-            /// The public key associated with this private key.
-            var publicKey: PublicKey {
-                PublicKey(privateKeyBacking: self)
-            }
-
-            /// Decapsulate a shared secret and create a symmetric key.
-            ///
-            /// - Parameter encapsulated: The encapsulated shared secret.
-            ///
-            /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
-            ///
-            /// - Returns: The symmetric key.
-            func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
-                guard encapsulated.count == MLKEM768.ciphertextSizeInBytes else {
-                    throw CryptoKitError.incorrectParameterSize
-                }
-
-                let output = try [UInt8](
-                    unsafeUninitializedCapacity: MLKEM768.sharedSecretSizeInBytes
-                ) { bufferPtr, length in
-                    let bytes: ContiguousBytes =
-                        encapsulated.regions.count == 1
-                        ? encapsulated.regions.first!
-                        : Array(encapsulated)
-                    let result = bytes.withUnsafeBytes { encapsulatedPtr in
-                        CCryptoBoringSSL_MLKEM768_decap(
-                            bufferPtr.baseAddress,
-                            encapsulatedPtr.baseAddress,
-                            encapsulatedPtr.count,
-                            &self.key
-                        )
-                    }
-
-                    guard result == 1 else {
-                        throw CryptoKitError.internalBoringSSLError()
-                    }
-
-                    length = MLKEM768.sharedSecretSizeInBytes
-                }
-
-                return SymmetricKey(data: Data(output))
-            }
-        }
-    }
-}
-
-@available(macOS 14.0, *)
-extension MLKEM768 {
-    /// A ML-KEM-768 public key.
-    public struct PublicKey: Sendable, KEMPublicKey {
-        private var backing: Backing
-
-        fileprivate init(privateKeyBacking: PrivateKey.Backing) {
-            self.backing = Backing(privateKeyBacking: privateKeyBacking)
-        }
-
-        /// Initialize a ML-KEM-768 public key from a raw representation.
-        ///
-        /// - Parameter rawRepresentation: The public key bytes.
-        ///
-        /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
-        public init(rawRepresentation: some DataProtocol) throws {
-            self.backing = try Backing(rawRepresentation: rawRepresentation)
-        }
-
-        /// The raw binary representation of the public key.
-        public var rawRepresentation: Data {
-            self.backing.rawRepresentation
-        }
-
-        /// Encapsulate a shared secret.
-        ///
-        /// - Returns: The shared secret and its encapsulated version.
-        public func encapsulate() -> KEM.EncapsulationResult {
-            self.backing.encapsulate()
-        }
-
-        /// The size of the public key in bytes.
-        static let bytesCount = Backing.bytesCount
-
-        fileprivate final class Backing {
-            var key: MLKEM768_public_key
-
-            init(privateKeyBacking: PrivateKey.Backing) {
-                self.key = .init()
-                CCryptoBoringSSL_MLKEM768_public_from_private(&self.key, &privateKeyBacking.key)
-            }
-
-            /// Initialize a ML-KEM-768 public key from a raw representation.
-            ///
-            /// - Parameter rawRepresentation: The public key bytes.
-            ///
-            /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
-            init(rawRepresentation: some DataProtocol) throws {
-                guard rawRepresentation.count == MLKEM768.PublicKey.bytesCount else {
-                    throw CryptoKitError.incorrectKeySize
-                }
-
-                self.key = .init()
-
-                let bytes: ContiguousBytes =
-                    rawRepresentation.regions.count == 1
-                    ? rawRepresentation.regions.first!
-                    : Array(rawRepresentation)
-                try bytes.withUnsafeBytes { rawBuffer in
-                    try rawBuffer.withMemoryRebound(to: UInt8.self) { buffer in
-                        var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
-                        guard CCryptoBoringSSL_MLKEM768_parse_public_key(&self.key, &cbs) == 1 else {
-                            throw CryptoKitError.internalBoringSSLError()
-                        }
-                    }
-                }
-            }
-
-            /// The raw binary representation of the public key.
-            var rawRepresentation: Data {
-                var cbb = CBB()
-                // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
-                CCryptoBoringSSL_CBB_init(&cbb, MLKEM768.PublicKey.Backing.bytesCount)
-                defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
-                CCryptoBoringSSL_MLKEM768_marshal_public_key(&cbb, &self.key)
-                return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
-            }
-
-            /// Encapsulate a shared secret.
-            ///
-            /// - Returns: The shared secret and its encapsulated version.
-            func encapsulate() -> KEM.EncapsulationResult {
-                withUnsafeTemporaryAllocation(
-                    of: UInt8.self,
-                    capacity: MLKEM768.ciphertextSizeInBytes
-                ) { encapsulatedPtr in
-                    withUnsafeTemporaryAllocation(
-                        of: UInt8.self,
-                        capacity: MLKEM768.sharedSecretSizeInBytes
-                    ) { secretPtr in
-                        CCryptoBoringSSL_MLKEM768_encap(
-                            encapsulatedPtr.baseAddress,
-                            secretPtr.baseAddress,
-                            &self.key
-                        )
-
-                        return KEM.EncapsulationResult(
-                            sharedSecret: SymmetricKey(
-                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM768.sharedSecretSizeInBytes)
-                            ),
-                            encapsulated: Data(
-                                bytes: encapsulatedPtr.baseAddress!,
-                                count: MLKEM768.ciphertextSizeInBytes
-                            )
-                        )
-                    }
-                }
-            }
-
-            /// The size of the public key in bytes.
-            static let bytesCount = 1184
-        }
-    }
-}
-
-@available(macOS 14.0, *)
-extension MLKEM768 {
-    /// The size of the encapsulated shared secret in bytes.
-    private static let ciphertextSizeInBytes = 1088
-}
-
-@available(macOS 14.0, *)
-extension MLKEM768 {
-    /// The size of the seed in bytes.
-    private static let seedSizeInBytes = 64
-
-    // The size of the shared secret in bytes.
-    private static let sharedSecretSizeInBytes = 32
-}
diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
new file mode 100644
index 00000000..91134458
--- /dev/null
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
@@ -0,0 +1,575 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the SwiftCrypto open source project
+//
+// Copyright (c) 2024 Apple Inc. and the SwiftCrypto project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of SwiftCrypto project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+// MARK: - Generated file, do NOT edit
+// any edits of this file WILL be overwritten and thus discarded
+// see section `gyb` in `README` for details.
+
+@_implementationOnly import CCryptoBoringSSL
+import Crypto
+import Foundation
+
+/// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
+@available(macOS 14.0, *)
+public enum MLKEM768 {}
+
+@available(macOS 14.0, *)
+extension MLKEM768 {
+    /// A ML-KEM-768 private key.
+    public struct PrivateKey: Sendable, KEMPrivateKey {
+        private var backing: Backing
+
+        /// Initialize a ML-KEM-768 private key from a random seed.
+        public init() {
+            self.backing = Backing()
+        }
+
+        /// Generate a ML-KEM-768 private key from a random seed.
+        ///
+        /// - Returns: The generated private key.
+        public static func generate() -> MLKEM768.PrivateKey {
+            .init()
+        }
+
+        /// Initialize a ML-KEM-768 private key from a seed.
+        ///
+        /// - Parameter seed: The seed to use to generate the private key.
+        ///
+        /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
+        public init(seed: some DataProtocol) throws {
+            self.backing = try Backing(seed: seed)
+        }
+
+        /// The seed from which this private key was generated.
+        public var seed: Data {
+            self.backing.seed
+        }
+
+        /// The public key associated with this private key.
+        public var publicKey: PublicKey {
+            self.backing.publicKey
+        }
+
+        /// Decapsulate a shared secret and create a symmetric key.
+        ///
+        /// - Parameter encapsulated: The encapsulated shared secret.
+        ///
+        /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
+        ///
+        /// - Returns: The symmetric key.
+        public func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
+            try self.backing.decapsulate(encapsulated)
+        }
+
+        fileprivate final class Backing {
+            var key: MLKEM768_private_key
+            var seed: Data
+
+            /// Initialize a ML-KEM-768 private key from a random seed.
+            init() {
+                self.key = .init()
+                self.seed = Data()
+
+                self.seed = withUnsafeTemporaryAllocation(
+                    of: UInt8.self,
+                    capacity: MLKEM.seedSizeInBytes
+                ) { seedPtr in
+                    withUnsafeTemporaryAllocation(
+                        of: UInt8.self,
+                        capacity: MLKEM768.PublicKey.bytesCount
+                    ) { publicKeyPtr in
+                        CCryptoBoringSSL_MLKEM768_generate_key(
+                            publicKeyPtr.baseAddress,
+                            seedPtr.baseAddress,
+                            &self.key
+                        )
+
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
+                    }
+                }
+            }
+
+            /// Initialize a ML-KEM-768 private key from a seed.
+            ///
+            /// - Parameter seed: The seed to use to generate the private key.
+            ///
+            /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
+            init(seed: some DataProtocol) throws {
+                guard seed.count == MLKEM.seedSizeInBytes else {
+                    throw CryptoKitError.incorrectKeySize
+                }
+
+                self.key = .init()
+                self.seed = Data(seed)
+
+                guard
+                    self.seed.withUnsafeBytes({ seedPtr in
+                        CCryptoBoringSSL_MLKEM768_private_key_from_seed(
+                            &self.key,
+                            seedPtr.baseAddress,
+                            seedPtr.count
+                        )
+                    }) == 1
+                else {
+                    throw CryptoKitError.internalBoringSSLError()
+                }
+            }
+
+            /// The public key associated with this private key.
+            var publicKey: PublicKey {
+                PublicKey(privateKeyBacking: self)
+            }
+
+            /// Decapsulate a shared secret and create a symmetric key.
+            ///
+            /// - Parameter encapsulated: The encapsulated shared secret.
+            ///
+            /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
+            ///
+            /// - Returns: The symmetric key.
+            func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
+                guard encapsulated.count == MLKEM768.ciphertextSizeInBytes else {
+                    throw CryptoKitError.incorrectParameterSize
+                }
+
+                let output = try [UInt8](
+                    unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes
+                ) { bufferPtr, length in
+                    let bytes: ContiguousBytes =
+                        encapsulated.regions.count == 1
+                        ? encapsulated.regions.first!
+                        : Array(encapsulated)
+                    let result = bytes.withUnsafeBytes { encapsulatedPtr in
+                        CCryptoBoringSSL_MLKEM768_decap(
+                            bufferPtr.baseAddress,
+                            encapsulatedPtr.baseAddress,
+                            encapsulatedPtr.count,
+                            &self.key
+                        )
+                    }
+
+                    guard result == 1 else {
+                        throw CryptoKitError.internalBoringSSLError()
+                    }
+
+                    length = MLKEM.sharedSecretSizeInBytes
+                }
+
+                return SymmetricKey(data: Data(output))
+            }
+        }
+    }
+}
+
+@available(macOS 14.0, *)
+extension MLKEM768 {
+    /// A ML-KEM-768 public key.
+    public struct PublicKey: Sendable, KEMPublicKey {
+        private var backing: Backing
+
+        fileprivate init(privateKeyBacking: PrivateKey.Backing) {
+            self.backing = Backing(privateKeyBacking: privateKeyBacking)
+        }
+
+        /// Initialize a ML-KEM-768 public key from a raw representation.
+        ///
+        /// - Parameter rawRepresentation: The public key bytes.
+        ///
+        /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
+        public init(rawRepresentation: some DataProtocol) throws {
+            self.backing = try Backing(rawRepresentation: rawRepresentation)
+        }
+
+        /// The raw binary representation of the public key.
+        public var rawRepresentation: Data {
+            self.backing.rawRepresentation
+        }
+
+        /// Encapsulate a shared secret.
+        ///
+        /// - Returns: The shared secret and its encapsulated version.
+        public func encapsulate() -> KEM.EncapsulationResult {
+            self.backing.encapsulate()
+        }
+
+        /// The size of the public key in bytes.
+        static let bytesCount = Backing.bytesCount
+
+        fileprivate final class Backing {
+            var key: MLKEM768_public_key
+
+            init(privateKeyBacking: PrivateKey.Backing) {
+                self.key = .init()
+                CCryptoBoringSSL_MLKEM768_public_from_private(&self.key, &privateKeyBacking.key)
+            }
+
+            /// Initialize a ML-KEM-768 public key from a raw representation.
+            ///
+            /// - Parameter rawRepresentation: The public key bytes.
+            ///
+            /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
+            init(rawRepresentation: some DataProtocol) throws {
+                guard rawRepresentation.count == MLKEM768.PublicKey.bytesCount else {
+                    throw CryptoKitError.incorrectKeySize
+                }
+
+                self.key = .init()
+
+                let bytes: ContiguousBytes =
+                    rawRepresentation.regions.count == 1
+                    ? rawRepresentation.regions.first!
+                    : Array(rawRepresentation)
+                try bytes.withUnsafeBytes { rawBuffer in
+                    try rawBuffer.withMemoryRebound(to: UInt8.self) { buffer in
+                        var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
+                        guard CCryptoBoringSSL_MLKEM768_parse_public_key(&self.key, &cbs) == 1 else {
+                            throw CryptoKitError.internalBoringSSLError()
+                        }
+                    }
+                }
+            }
+
+            /// The raw binary representation of the public key.
+            var rawRepresentation: Data {
+                var cbb = CBB()
+                // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM768.PublicKey.Backing.bytesCount)
+                defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
+                CCryptoBoringSSL_MLKEM768_marshal_public_key(&cbb, &self.key)
+                return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
+            }
+
+            /// Encapsulate a shared secret.
+            ///
+            /// - Returns: The shared secret and its encapsulated version.
+            func encapsulate() -> KEM.EncapsulationResult {
+                withUnsafeTemporaryAllocation(
+                    of: UInt8.self,
+                    capacity: MLKEM768.ciphertextSizeInBytes
+                ) { encapsulatedPtr in
+                    withUnsafeTemporaryAllocation(
+                        of: UInt8.self,
+                        capacity: MLKEM.sharedSecretSizeInBytes
+                    ) { secretPtr in
+                        CCryptoBoringSSL_MLKEM768_encap(
+                            encapsulatedPtr.baseAddress,
+                            secretPtr.baseAddress,
+                            &self.key
+                        )
+
+                        return KEM.EncapsulationResult(
+                            sharedSecret: SymmetricKey(
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)
+                            ),
+                            encapsulated: Data(
+                                bytes: encapsulatedPtr.baseAddress!,
+                                count: MLKEM768.ciphertextSizeInBytes
+                            )
+                        )
+                    }
+                }
+            }
+
+            /// The size of the public key in bytes.
+            static let bytesCount = Int(MLKEM768_PUBLIC_KEY_BYTES)
+        }
+    }
+}
+
+@available(macOS 14.0, *)
+extension MLKEM768 {
+    /// The size of the encapsulated shared secret in bytes.
+    private static let ciphertextSizeInBytes = Int(MLKEM768_CIPHERTEXT_BYTES)
+}
+
+/// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
+@available(macOS 14.0, *)
+public enum MLKEM1024 {}
+
+@available(macOS 14.0, *)
+extension MLKEM1024 {
+    /// A ML-KEM-1024 private key.
+    public struct PrivateKey: Sendable, KEMPrivateKey {
+        private var backing: Backing
+
+        /// Initialize a ML-KEM-1024 private key from a random seed.
+        public init() {
+            self.backing = Backing()
+        }
+
+        /// Generate a ML-KEM-1024 private key from a random seed.
+        ///
+        /// - Returns: The generated private key.
+        public static func generate() -> MLKEM1024.PrivateKey {
+            .init()
+        }
+
+        /// Initialize a ML-KEM-1024 private key from a seed.
+        ///
+        /// - Parameter seed: The seed to use to generate the private key.
+        ///
+        /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
+        public init(seed: some DataProtocol) throws {
+            self.backing = try Backing(seed: seed)
+        }
+
+        /// The seed from which this private key was generated.
+        public var seed: Data {
+            self.backing.seed
+        }
+
+        /// The public key associated with this private key.
+        public var publicKey: PublicKey {
+            self.backing.publicKey
+        }
+
+        /// Decapsulate a shared secret and create a symmetric key.
+        ///
+        /// - Parameter encapsulated: The encapsulated shared secret.
+        ///
+        /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
+        ///
+        /// - Returns: The symmetric key.
+        public func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
+            try self.backing.decapsulate(encapsulated)
+        }
+
+        fileprivate final class Backing {
+            var key: MLKEM1024_private_key
+            var seed: Data
+
+            /// Initialize a ML-KEM-1024 private key from a random seed.
+            init() {
+                self.key = .init()
+                self.seed = Data()
+
+                self.seed = withUnsafeTemporaryAllocation(
+                    of: UInt8.self,
+                    capacity: MLKEM.seedSizeInBytes
+                ) { seedPtr in
+                    withUnsafeTemporaryAllocation(
+                        of: UInt8.self,
+                        capacity: MLKEM1024.PublicKey.bytesCount
+                    ) { publicKeyPtr in
+                        CCryptoBoringSSL_MLKEM1024_generate_key(
+                            publicKeyPtr.baseAddress,
+                            seedPtr.baseAddress,
+                            &self.key
+                        )
+
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
+                    }
+                }
+            }
+
+            /// Initialize a ML-KEM-1024 private key from a seed.
+            ///
+            /// - Parameter seed: The seed to use to generate the private key.
+            ///
+            /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
+            init(seed: some DataProtocol) throws {
+                guard seed.count == MLKEM.seedSizeInBytes else {
+                    throw CryptoKitError.incorrectKeySize
+                }
+
+                self.key = .init()
+                self.seed = Data(seed)
+
+                guard
+                    self.seed.withUnsafeBytes({ seedPtr in
+                        CCryptoBoringSSL_MLKEM1024_private_key_from_seed(
+                            &self.key,
+                            seedPtr.baseAddress,
+                            seedPtr.count
+                        )
+                    }) == 1
+                else {
+                    throw CryptoKitError.internalBoringSSLError()
+                }
+            }
+
+            /// The public key associated with this private key.
+            var publicKey: PublicKey {
+                PublicKey(privateKeyBacking: self)
+            }
+
+            /// Decapsulate a shared secret and create a symmetric key.
+            ///
+            /// - Parameter encapsulated: The encapsulated shared secret.
+            ///
+            /// - Throws: `CryptoKitError.incorrectParameterSize` if the encapsulated shared secret is not 1088 bytes long.
+            ///
+            /// - Returns: The symmetric key.
+            func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
+                guard encapsulated.count == MLKEM1024.ciphertextSizeInBytes else {
+                    throw CryptoKitError.incorrectParameterSize
+                }
+
+                let output = try [UInt8](
+                    unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes
+                ) { bufferPtr, length in
+                    let bytes: ContiguousBytes =
+                        encapsulated.regions.count == 1
+                        ? encapsulated.regions.first!
+                        : Array(encapsulated)
+                    let result = bytes.withUnsafeBytes { encapsulatedPtr in
+                        CCryptoBoringSSL_MLKEM1024_decap(
+                            bufferPtr.baseAddress,
+                            encapsulatedPtr.baseAddress,
+                            encapsulatedPtr.count,
+                            &self.key
+                        )
+                    }
+
+                    guard result == 1 else {
+                        throw CryptoKitError.internalBoringSSLError()
+                    }
+
+                    length = MLKEM.sharedSecretSizeInBytes
+                }
+
+                return SymmetricKey(data: Data(output))
+            }
+        }
+    }
+}
+
+@available(macOS 14.0, *)
+extension MLKEM1024 {
+    /// A ML-KEM-1024 public key.
+    public struct PublicKey: Sendable, KEMPublicKey {
+        private var backing: Backing
+
+        fileprivate init(privateKeyBacking: PrivateKey.Backing) {
+            self.backing = Backing(privateKeyBacking: privateKeyBacking)
+        }
+
+        /// Initialize a ML-KEM-1024 public key from a raw representation.
+        ///
+        /// - Parameter rawRepresentation: The public key bytes.
+        ///
+        /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
+        public init(rawRepresentation: some DataProtocol) throws {
+            self.backing = try Backing(rawRepresentation: rawRepresentation)
+        }
+
+        /// The raw binary representation of the public key.
+        public var rawRepresentation: Data {
+            self.backing.rawRepresentation
+        }
+
+        /// Encapsulate a shared secret.
+        ///
+        /// - Returns: The shared secret and its encapsulated version.
+        public func encapsulate() -> KEM.EncapsulationResult {
+            self.backing.encapsulate()
+        }
+
+        /// The size of the public key in bytes.
+        static let bytesCount = Backing.bytesCount
+
+        fileprivate final class Backing {
+            var key: MLKEM1024_public_key
+
+            init(privateKeyBacking: PrivateKey.Backing) {
+                self.key = .init()
+                CCryptoBoringSSL_MLKEM1024_public_from_private(&self.key, &privateKeyBacking.key)
+            }
+
+            /// Initialize a ML-KEM-1024 public key from a raw representation.
+            ///
+            /// - Parameter rawRepresentation: The public key bytes.
+            ///
+            /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
+            init(rawRepresentation: some DataProtocol) throws {
+                guard rawRepresentation.count == MLKEM1024.PublicKey.bytesCount else {
+                    throw CryptoKitError.incorrectKeySize
+                }
+
+                self.key = .init()
+
+                let bytes: ContiguousBytes =
+                    rawRepresentation.regions.count == 1
+                    ? rawRepresentation.regions.first!
+                    : Array(rawRepresentation)
+                try bytes.withUnsafeBytes { rawBuffer in
+                    try rawBuffer.withMemoryRebound(to: UInt8.self) { buffer in
+                        var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
+                        guard CCryptoBoringSSL_MLKEM1024_parse_public_key(&self.key, &cbs) == 1 else {
+                            throw CryptoKitError.internalBoringSSLError()
+                        }
+                    }
+                }
+            }
+
+            /// The raw binary representation of the public key.
+            var rawRepresentation: Data {
+                var cbb = CBB()
+                // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM1024.PublicKey.Backing.bytesCount)
+                defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
+                CCryptoBoringSSL_MLKEM1024_marshal_public_key(&cbb, &self.key)
+                return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
+            }
+
+            /// Encapsulate a shared secret.
+            ///
+            /// - Returns: The shared secret and its encapsulated version.
+            func encapsulate() -> KEM.EncapsulationResult {
+                withUnsafeTemporaryAllocation(
+                    of: UInt8.self,
+                    capacity: MLKEM1024.ciphertextSizeInBytes
+                ) { encapsulatedPtr in
+                    withUnsafeTemporaryAllocation(
+                        of: UInt8.self,
+                        capacity: MLKEM.sharedSecretSizeInBytes
+                    ) { secretPtr in
+                        CCryptoBoringSSL_MLKEM1024_encap(
+                            encapsulatedPtr.baseAddress,
+                            secretPtr.baseAddress,
+                            &self.key
+                        )
+
+                        return KEM.EncapsulationResult(
+                            sharedSecret: SymmetricKey(
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)
+                            ),
+                            encapsulated: Data(
+                                bytes: encapsulatedPtr.baseAddress!,
+                                count: MLKEM1024.ciphertextSizeInBytes
+                            )
+                        )
+                    }
+                }
+            }
+
+            /// The size of the public key in bytes.
+            static let bytesCount = Int(MLKEM1024_PUBLIC_KEY_BYTES)
+        }
+    }
+}
+
+@available(macOS 14.0, *)
+extension MLKEM1024 {
+    /// The size of the encapsulated shared secret in bytes.
+    private static let ciphertextSizeInBytes = Int(MLKEM1024_CIPHERTEXT_BYTES)
+}
+
+private enum MLKEM {
+    /// The size of the seed in bytes.
+    fileprivate static let seedSizeInBytes = 64
+
+    // The size of the shared secret in bytes.
+    fileprivate static let sharedSecretSizeInBytes = 32
+}
diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift.gyb
similarity index 73%
rename from Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift
rename to Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift.gyb
index 506aa5d4..e4becdce 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM1024_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift.gyb
@@ -12,35 +12,41 @@
 //
 //===----------------------------------------------------------------------===//
 
+// MARK: - Generated file, do NOT edit
+// any edits of this file WILL be overwritten and thus discarded
+// see section `gyb` in `README` for details.
+
 @_implementationOnly import CCryptoBoringSSL
 import Crypto
 import Foundation
+%{
+    parameter_sets = ["768", "1024"]
+}%
+% for parameter_set in parameter_sets:
 
 /// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
-///
-/// > Note: You should prefer ``MLKEM768`` where possible.
 @available(macOS 14.0, *)
-public enum MLKEM1024 {}
+public enum MLKEM${parameter_set} {}
 
 @available(macOS 14.0, *)
-extension MLKEM1024 {
-    /// A ML-KEM-1024 private key.
+extension MLKEM${parameter_set} {
+    /// A ML-KEM-${parameter_set} private key.
     public struct PrivateKey: Sendable, KEMPrivateKey {
         private var backing: Backing
 
-        /// Initialize a ML-KEM-1024 private key from a random seed.
+        /// Initialize a ML-KEM-${parameter_set} private key from a random seed.
         public init() {
             self.backing = Backing()
         }
 
-        /// Generate a ML-KEM-1024 private key from a random seed.
+        /// Generate a ML-KEM-${parameter_set} private key from a random seed.
         ///
         /// - Returns: The generated private key.
-        public static func generate() -> MLKEM1024.PrivateKey {
+        public static func generate() -> MLKEM${parameter_set}.PrivateKey {
             .init()
         }
 
-        /// Initialize a ML-KEM-1024 private key from a seed.
+        /// Initialize a ML-KEM-${parameter_set} private key from a seed.
         ///
         /// - Parameter seed: The seed to use to generate the private key.
         ///
@@ -71,36 +77,40 @@ extension MLKEM1024 {
         }
 
         fileprivate final class Backing {
-            var key: MLKEM1024_private_key
+            var key: MLKEM${parameter_set}_private_key
             var seed: Data
 
-            /// Initialize a ML-KEM-1024 private key from a random seed.
+            /// Initialize a ML-KEM-${parameter_set} private key from a random seed.
             init() {
                 self.key = .init()
                 self.seed = Data()
 
                 self.seed = withUnsafeTemporaryAllocation(
                     of: UInt8.self,
-                    capacity: MLKEM1024.seedSizeInBytes
+                    capacity: MLKEM.seedSizeInBytes
                 ) { seedPtr in
                     withUnsafeTemporaryAllocation(
                         of: UInt8.self,
-                        capacity: MLKEM1024.PublicKey.bytesCount
+                        capacity: MLKEM${parameter_set}.PublicKey.bytesCount
                     ) { publicKeyPtr in
-                        CCryptoBoringSSL_MLKEM1024_generate_key(publicKeyPtr.baseAddress, seedPtr.baseAddress, &self.key)
+                        CCryptoBoringSSL_MLKEM${parameter_set}_generate_key(
+                            publicKeyPtr.baseAddress,
+                            seedPtr.baseAddress,
+                            &self.key
+                        )
 
-                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM1024.seedSizeInBytes)
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
                     }
                 }
             }
 
-            /// Initialize a ML-KEM-1024 private key from a seed.
+            /// Initialize a ML-KEM-${parameter_set} private key from a seed.
             ///
             /// - Parameter seed: The seed to use to generate the private key.
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
             init(seed: some DataProtocol) throws {
-                guard seed.count == MLKEM1024.seedSizeInBytes else {
+                guard seed.count == MLKEM.seedSizeInBytes else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -109,7 +119,7 @@ extension MLKEM1024 {
 
                 guard
                     self.seed.withUnsafeBytes({ seedPtr in
-                        CCryptoBoringSSL_MLKEM1024_private_key_from_seed(
+                        CCryptoBoringSSL_MLKEM${parameter_set}_private_key_from_seed(
                             &self.key,
                             seedPtr.baseAddress,
                             seedPtr.count
@@ -133,19 +143,19 @@ extension MLKEM1024 {
             ///
             /// - Returns: The symmetric key.
             func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
-                guard encapsulated.count == MLKEM1024.ciphertextSizeInBytes else {
+                guard encapsulated.count == MLKEM${parameter_set}.ciphertextSizeInBytes else {
                     throw CryptoKitError.incorrectParameterSize
                 }
 
                 let output = try [UInt8](
-                    unsafeUninitializedCapacity: MLKEM1024.sharedSecretSizeInBytes
+                    unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes
                 ) { bufferPtr, length in
                     let bytes: ContiguousBytes =
                         encapsulated.regions.count == 1
                         ? encapsulated.regions.first!
                         : Array(encapsulated)
                     let result = bytes.withUnsafeBytes { encapsulatedPtr in
-                        CCryptoBoringSSL_MLKEM1024_decap(
+                        CCryptoBoringSSL_MLKEM${parameter_set}_decap(
                             bufferPtr.baseAddress,
                             encapsulatedPtr.baseAddress,
                             encapsulatedPtr.count,
@@ -157,7 +167,7 @@ extension MLKEM1024 {
                         throw CryptoKitError.internalBoringSSLError()
                     }
 
-                    length = MLKEM1024.sharedSecretSizeInBytes
+                    length = MLKEM.sharedSecretSizeInBytes
                 }
 
                 return SymmetricKey(data: Data(output))
@@ -167,8 +177,8 @@ extension MLKEM1024 {
 }
 
 @available(macOS 14.0, *)
-extension MLKEM1024 {
-    /// A ML-KEM-1024 public key.
+extension MLKEM${parameter_set} {
+    /// A ML-KEM-${parameter_set} public key.
     public struct PublicKey: Sendable, KEMPublicKey {
         private var backing: Backing
 
@@ -176,7 +186,7 @@ extension MLKEM1024 {
             self.backing = Backing(privateKeyBacking: privateKeyBacking)
         }
 
-        /// Initialize a ML-KEM-1024 public key from a raw representation.
+        /// Initialize a ML-KEM-${parameter_set} public key from a raw representation.
         ///
         /// - Parameter rawRepresentation: The public key bytes.
         ///
@@ -201,20 +211,20 @@ extension MLKEM1024 {
         static let bytesCount = Backing.bytesCount
 
         fileprivate final class Backing {
-            var key: MLKEM1024_public_key
+            var key: MLKEM${parameter_set}_public_key
 
             init(privateKeyBacking: PrivateKey.Backing) {
                 self.key = .init()
-                CCryptoBoringSSL_MLKEM1024_public_from_private(&self.key, &privateKeyBacking.key)
+                CCryptoBoringSSL_MLKEM${parameter_set}_public_from_private(&self.key, &privateKeyBacking.key)
             }
 
-            /// Initialize a ML-KEM-1024 public key from a raw representation.
+            /// Initialize a ML-KEM-${parameter_set} public key from a raw representation.
             ///
             /// - Parameter rawRepresentation: The public key bytes.
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
             init(rawRepresentation: some DataProtocol) throws {
-                guard rawRepresentation.count == MLKEM1024.PublicKey.bytesCount else {
+                guard rawRepresentation.count == MLKEM${parameter_set}.PublicKey.bytesCount else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -227,7 +237,7 @@ extension MLKEM1024 {
                 try bytes.withUnsafeBytes { rawBuffer in
                     try rawBuffer.withMemoryRebound(to: UInt8.self) { buffer in
                         var cbs = CBS(data: buffer.baseAddress, len: buffer.count)
-                        guard CCryptoBoringSSL_MLKEM1024_parse_public_key(&self.key, &cbs) == 1 else {
+                        guard CCryptoBoringSSL_MLKEM${parameter_set}_parse_public_key(&self.key, &cbs) == 1 else {
                             throw CryptoKitError.internalBoringSSLError()
                         }
                     }
@@ -238,9 +248,9 @@ extension MLKEM1024 {
             var rawRepresentation: Data {
                 var cbb = CBB()
                 // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
-                CCryptoBoringSSL_CBB_init(&cbb, MLKEM1024.PublicKey.Backing.bytesCount)
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM${parameter_set}.PublicKey.Backing.bytesCount)
                 defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
-                CCryptoBoringSSL_MLKEM1024_marshal_public_key(&cbb, &self.key)
+                CCryptoBoringSSL_MLKEM${parameter_set}_marshal_public_key(&cbb, &self.key)
                 return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
             }
 
@@ -250,13 +260,13 @@ extension MLKEM1024 {
             func encapsulate() -> KEM.EncapsulationResult {
                 withUnsafeTemporaryAllocation(
                     of: UInt8.self,
-                    capacity: MLKEM1024.ciphertextSizeInBytes
+                    capacity: MLKEM${parameter_set}.ciphertextSizeInBytes
                 ) { encapsulatedPtr in
                     withUnsafeTemporaryAllocation(
                         of: UInt8.self,
-                        capacity: MLKEM1024.sharedSecretSizeInBytes
+                        capacity: MLKEM.sharedSecretSizeInBytes
                     ) { secretPtr in
-                        CCryptoBoringSSL_MLKEM1024_encap(
+                        CCryptoBoringSSL_MLKEM${parameter_set}_encap(
                             encapsulatedPtr.baseAddress,
                             secretPtr.baseAddress,
                             &self.key
@@ -264,11 +274,11 @@ extension MLKEM1024 {
 
                         return KEM.EncapsulationResult(
                             sharedSecret: SymmetricKey(
-                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM1024.sharedSecretSizeInBytes)
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)
                             ),
                             encapsulated: Data(
                                 bytes: encapsulatedPtr.baseAddress!,
-                                count: MLKEM1024.ciphertextSizeInBytes
+                                count: MLKEM${parameter_set}.ciphertextSizeInBytes
                             )
                         )
                     }
@@ -276,22 +286,22 @@ extension MLKEM1024 {
             }
 
             /// The size of the public key in bytes.
-            static let bytesCount = 1568
+            static let bytesCount = Int(MLKEM${parameter_set}_PUBLIC_KEY_BYTES)
         }
     }
 }
 
 @available(macOS 14.0, *)
-extension MLKEM1024 {
+extension MLKEM${parameter_set} {
     /// The size of the encapsulated shared secret in bytes.
-    private static let ciphertextSizeInBytes = 1568
+    private static let ciphertextSizeInBytes = Int(MLKEM${parameter_set}_CIPHERTEXT_BYTES)
 }
+% end
 
-@available(macOS 14.0, *)
-extension MLKEM1024 {
+private enum MLKEM {
     /// The size of the seed in bytes.
-    private static let seedSizeInBytes = 64
+    fileprivate static let seedSizeInBytes = 64
 
     // The size of the shared secret in bytes.
-    private static let sharedSecretSizeInBytes = 32
+    fileprivate static let sharedSecretSizeInBytes = 32
 }

From 3d1731f02c660f8de0fb59802d048cebd7537fd7 Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Sat, 14 Dec 2024 12:54:39 +0100
Subject: [PATCH 12/13] Update CMake lists

---
 Sources/_CryptoExtras/CMakeLists.txt | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Sources/_CryptoExtras/CMakeLists.txt b/Sources/_CryptoExtras/CMakeLists.txt
index 4f74432a..74a31282 100644
--- a/Sources/_CryptoExtras/CMakeLists.txt
+++ b/Sources/_CryptoExtras/CMakeLists.txt
@@ -32,8 +32,7 @@ add_library(_CryptoExtras
   "Key Derivation/PBKDF2/PBKDF2.swift"
   "Key Derivation/Scrypt/BoringSSL/Scrypt_boring.swift"
   "Key Derivation/Scrypt/Scrypt.swift"
-  "MLKEM/MLKEM1024_boring.swift"
-  "MLKEM/MLKEM768_boring.swift"
+  "MLKEM/MLKEM_boring.swift"
   "OPRFs/OPRF.swift"
   "OPRFs/OPRFClient.swift"
   "OPRFs/OPRFServer.swift"

From eacea95a0052f5b74cb37e4dad390972c6aa01fb Mon Sep 17 00:00:00 2001
From: Francesco Paolo Severino <francescopaolo.severino@gmail.com>
Date: Fri, 3 Jan 2025 17:06:38 +0100
Subject: [PATCH 13/13] Rename constants and improve `decapsulate` method

---
 .../_CryptoExtras/MLKEM/MLKEM_boring.swift    | 100 +++++++++---------
 .../MLKEM/MLKEM_boring.swift.gyb              |  52 +++++----
 2 files changed, 73 insertions(+), 79 deletions(-)

diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
index 91134458..b9fe8e89 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
@@ -83,11 +83,11 @@ extension MLKEM768 {
 
                 self.seed = withUnsafeTemporaryAllocation(
                     of: UInt8.self,
-                    capacity: MLKEM.seedSizeInBytes
+                    capacity: MLKEM.seedByteCount
                 ) { seedPtr in
                     withUnsafeTemporaryAllocation(
                         of: UInt8.self,
-                        capacity: MLKEM768.PublicKey.bytesCount
+                        capacity: MLKEM768.PublicKey.byteCount
                     ) { publicKeyPtr in
                         CCryptoBoringSSL_MLKEM768_generate_key(
                             publicKeyPtr.baseAddress,
@@ -95,7 +95,7 @@ extension MLKEM768 {
                             &self.key
                         )
 
-                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedByteCount)
                     }
                 }
             }
@@ -106,7 +106,7 @@ extension MLKEM768 {
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
             init(seed: some DataProtocol) throws {
-                guard seed.count == MLKEM.seedSizeInBytes else {
+                guard seed.count == MLKEM.seedByteCount else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -139,34 +139,32 @@ extension MLKEM768 {
             ///
             /// - Returns: The symmetric key.
             func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
-                guard encapsulated.count == MLKEM768.ciphertextSizeInBytes else {
+                guard encapsulated.count == MLKEM768.ciphertextByteCount else {
                     throw CryptoKitError.incorrectParameterSize
                 }
 
-                let output = try [UInt8](
-                    unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes
-                ) { bufferPtr, length in
+                var symmetricKeyData = Data(repeating: 0, count: MLKEM.sharedSecretByteCount)
+
+                let rc: CInt = symmetricKeyData.withUnsafeMutableBytes { symmetricKeyDataPtr in
                     let bytes: ContiguousBytes =
                         encapsulated.regions.count == 1
                         ? encapsulated.regions.first!
                         : Array(encapsulated)
-                    let result = bytes.withUnsafeBytes { encapsulatedPtr in
+                    return bytes.withUnsafeBytes { encapsulatedPtr in
                         CCryptoBoringSSL_MLKEM768_decap(
-                            bufferPtr.baseAddress,
+                            symmetricKeyDataPtr.baseAddress,
                             encapsulatedPtr.baseAddress,
                             encapsulatedPtr.count,
                             &self.key
                         )
                     }
+                }
 
-                    guard result == 1 else {
-                        throw CryptoKitError.internalBoringSSLError()
-                    }
-
-                    length = MLKEM.sharedSecretSizeInBytes
+                guard rc == 1 else {
+                    throw CryptoKitError.internalBoringSSLError()
                 }
 
-                return SymmetricKey(data: Data(output))
+                return SymmetricKey(data: symmetricKeyData)
             }
         }
     }
@@ -204,7 +202,7 @@ extension MLKEM768 {
         }
 
         /// The size of the public key in bytes.
-        static let bytesCount = Backing.bytesCount
+        static let byteCount = Backing.byteCount
 
         fileprivate final class Backing {
             var key: MLKEM768_public_key
@@ -220,7 +218,7 @@ extension MLKEM768 {
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
             init(rawRepresentation: some DataProtocol) throws {
-                guard rawRepresentation.count == MLKEM768.PublicKey.bytesCount else {
+                guard rawRepresentation.count == MLKEM768.PublicKey.byteCount else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -244,7 +242,7 @@ extension MLKEM768 {
             var rawRepresentation: Data {
                 var cbb = CBB()
                 // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
-                CCryptoBoringSSL_CBB_init(&cbb, MLKEM768.PublicKey.Backing.bytesCount)
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM768.PublicKey.Backing.byteCount)
                 defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
                 CCryptoBoringSSL_MLKEM768_marshal_public_key(&cbb, &self.key)
                 return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
@@ -256,11 +254,11 @@ extension MLKEM768 {
             func encapsulate() -> KEM.EncapsulationResult {
                 withUnsafeTemporaryAllocation(
                     of: UInt8.self,
-                    capacity: MLKEM768.ciphertextSizeInBytes
+                    capacity: MLKEM768.ciphertextByteCount
                 ) { encapsulatedPtr in
                     withUnsafeTemporaryAllocation(
                         of: UInt8.self,
-                        capacity: MLKEM.sharedSecretSizeInBytes
+                        capacity: MLKEM.sharedSecretByteCount
                     ) { secretPtr in
                         CCryptoBoringSSL_MLKEM768_encap(
                             encapsulatedPtr.baseAddress,
@@ -270,11 +268,11 @@ extension MLKEM768 {
 
                         return KEM.EncapsulationResult(
                             sharedSecret: SymmetricKey(
-                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretByteCount)
                             ),
                             encapsulated: Data(
                                 bytes: encapsulatedPtr.baseAddress!,
-                                count: MLKEM768.ciphertextSizeInBytes
+                                count: MLKEM768.ciphertextByteCount
                             )
                         )
                     }
@@ -282,7 +280,7 @@ extension MLKEM768 {
             }
 
             /// The size of the public key in bytes.
-            static let bytesCount = Int(MLKEM768_PUBLIC_KEY_BYTES)
+            static let byteCount = Int(MLKEM768_PUBLIC_KEY_BYTES)
         }
     }
 }
@@ -290,7 +288,7 @@ extension MLKEM768 {
 @available(macOS 14.0, *)
 extension MLKEM768 {
     /// The size of the encapsulated shared secret in bytes.
-    private static let ciphertextSizeInBytes = Int(MLKEM768_CIPHERTEXT_BYTES)
+    private static let ciphertextByteCount = Int(MLKEM768_CIPHERTEXT_BYTES)
 }
 
 /// A module-lattice-based key encapsulation mechanism that provides security against quantum computing attacks.
@@ -356,11 +354,11 @@ extension MLKEM1024 {
 
                 self.seed = withUnsafeTemporaryAllocation(
                     of: UInt8.self,
-                    capacity: MLKEM.seedSizeInBytes
+                    capacity: MLKEM.seedByteCount
                 ) { seedPtr in
                     withUnsafeTemporaryAllocation(
                         of: UInt8.self,
-                        capacity: MLKEM1024.PublicKey.bytesCount
+                        capacity: MLKEM1024.PublicKey.byteCount
                     ) { publicKeyPtr in
                         CCryptoBoringSSL_MLKEM1024_generate_key(
                             publicKeyPtr.baseAddress,
@@ -368,7 +366,7 @@ extension MLKEM1024 {
                             &self.key
                         )
 
-                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedByteCount)
                     }
                 }
             }
@@ -379,7 +377,7 @@ extension MLKEM1024 {
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
             init(seed: some DataProtocol) throws {
-                guard seed.count == MLKEM.seedSizeInBytes else {
+                guard seed.count == MLKEM.seedByteCount else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -412,34 +410,32 @@ extension MLKEM1024 {
             ///
             /// - Returns: The symmetric key.
             func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
-                guard encapsulated.count == MLKEM1024.ciphertextSizeInBytes else {
+                guard encapsulated.count == MLKEM1024.ciphertextByteCount else {
                     throw CryptoKitError.incorrectParameterSize
                 }
 
-                let output = try [UInt8](
-                    unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes
-                ) { bufferPtr, length in
+                var symmetricKeyData = Data(repeating: 0, count: MLKEM.sharedSecretByteCount)
+
+                let rc: CInt = symmetricKeyData.withUnsafeMutableBytes { symmetricKeyDataPtr in
                     let bytes: ContiguousBytes =
                         encapsulated.regions.count == 1
                         ? encapsulated.regions.first!
                         : Array(encapsulated)
-                    let result = bytes.withUnsafeBytes { encapsulatedPtr in
+                    return bytes.withUnsafeBytes { encapsulatedPtr in
                         CCryptoBoringSSL_MLKEM1024_decap(
-                            bufferPtr.baseAddress,
+                            symmetricKeyDataPtr.baseAddress,
                             encapsulatedPtr.baseAddress,
                             encapsulatedPtr.count,
                             &self.key
                         )
                     }
+                }
 
-                    guard result == 1 else {
-                        throw CryptoKitError.internalBoringSSLError()
-                    }
-
-                    length = MLKEM.sharedSecretSizeInBytes
+                guard rc == 1 else {
+                    throw CryptoKitError.internalBoringSSLError()
                 }
 
-                return SymmetricKey(data: Data(output))
+                return SymmetricKey(data: symmetricKeyData)
             }
         }
     }
@@ -477,7 +473,7 @@ extension MLKEM1024 {
         }
 
         /// The size of the public key in bytes.
-        static let bytesCount = Backing.bytesCount
+        static let byteCount = Backing.byteCount
 
         fileprivate final class Backing {
             var key: MLKEM1024_public_key
@@ -493,7 +489,7 @@ extension MLKEM1024 {
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
             init(rawRepresentation: some DataProtocol) throws {
-                guard rawRepresentation.count == MLKEM1024.PublicKey.bytesCount else {
+                guard rawRepresentation.count == MLKEM1024.PublicKey.byteCount else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -517,7 +513,7 @@ extension MLKEM1024 {
             var rawRepresentation: Data {
                 var cbb = CBB()
                 // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
-                CCryptoBoringSSL_CBB_init(&cbb, MLKEM1024.PublicKey.Backing.bytesCount)
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM1024.PublicKey.Backing.byteCount)
                 defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
                 CCryptoBoringSSL_MLKEM1024_marshal_public_key(&cbb, &self.key)
                 return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
@@ -529,11 +525,11 @@ extension MLKEM1024 {
             func encapsulate() -> KEM.EncapsulationResult {
                 withUnsafeTemporaryAllocation(
                     of: UInt8.self,
-                    capacity: MLKEM1024.ciphertextSizeInBytes
+                    capacity: MLKEM1024.ciphertextByteCount
                 ) { encapsulatedPtr in
                     withUnsafeTemporaryAllocation(
                         of: UInt8.self,
-                        capacity: MLKEM.sharedSecretSizeInBytes
+                        capacity: MLKEM.sharedSecretByteCount
                     ) { secretPtr in
                         CCryptoBoringSSL_MLKEM1024_encap(
                             encapsulatedPtr.baseAddress,
@@ -543,11 +539,11 @@ extension MLKEM1024 {
 
                         return KEM.EncapsulationResult(
                             sharedSecret: SymmetricKey(
-                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretByteCount)
                             ),
                             encapsulated: Data(
                                 bytes: encapsulatedPtr.baseAddress!,
-                                count: MLKEM1024.ciphertextSizeInBytes
+                                count: MLKEM1024.ciphertextByteCount
                             )
                         )
                     }
@@ -555,7 +551,7 @@ extension MLKEM1024 {
             }
 
             /// The size of the public key in bytes.
-            static let bytesCount = Int(MLKEM1024_PUBLIC_KEY_BYTES)
+            static let byteCount = Int(MLKEM1024_PUBLIC_KEY_BYTES)
         }
     }
 }
@@ -563,13 +559,13 @@ extension MLKEM1024 {
 @available(macOS 14.0, *)
 extension MLKEM1024 {
     /// The size of the encapsulated shared secret in bytes.
-    private static let ciphertextSizeInBytes = Int(MLKEM1024_CIPHERTEXT_BYTES)
+    private static let ciphertextByteCount = Int(MLKEM1024_CIPHERTEXT_BYTES)
 }
 
 private enum MLKEM {
     /// The size of the seed in bytes.
-    fileprivate static let seedSizeInBytes = 64
+    fileprivate static let seedByteCount = 64
 
     // The size of the shared secret in bytes.
-    fileprivate static let sharedSecretSizeInBytes = 32
+    fileprivate static let sharedSecretByteCount = 32
 }
diff --git a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift.gyb b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift.gyb
index e4becdce..6d440590 100644
--- a/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift.gyb
+++ b/Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift.gyb
@@ -87,11 +87,11 @@ extension MLKEM${parameter_set} {
 
                 self.seed = withUnsafeTemporaryAllocation(
                     of: UInt8.self,
-                    capacity: MLKEM.seedSizeInBytes
+                    capacity: MLKEM.seedByteCount
                 ) { seedPtr in
                     withUnsafeTemporaryAllocation(
                         of: UInt8.self,
-                        capacity: MLKEM${parameter_set}.PublicKey.bytesCount
+                        capacity: MLKEM${parameter_set}.PublicKey.byteCount
                     ) { publicKeyPtr in
                         CCryptoBoringSSL_MLKEM${parameter_set}_generate_key(
                             publicKeyPtr.baseAddress,
@@ -99,7 +99,7 @@ extension MLKEM${parameter_set} {
                             &self.key
                         )
 
-                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedSizeInBytes)
+                        return Data(bytes: seedPtr.baseAddress!, count: MLKEM.seedByteCount)
                     }
                 }
             }
@@ -110,7 +110,7 @@ extension MLKEM${parameter_set} {
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the seed is not 64 bytes long.
             init(seed: some DataProtocol) throws {
-                guard seed.count == MLKEM.seedSizeInBytes else {
+                guard seed.count == MLKEM.seedByteCount else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -143,34 +143,32 @@ extension MLKEM${parameter_set} {
             ///
             /// - Returns: The symmetric key.
             func decapsulate(_ encapsulated: Data) throws -> SymmetricKey {
-                guard encapsulated.count == MLKEM${parameter_set}.ciphertextSizeInBytes else {
+                guard encapsulated.count == MLKEM${parameter_set}.ciphertextByteCount else {
                     throw CryptoKitError.incorrectParameterSize
                 }
 
-                let output = try [UInt8](
-                    unsafeUninitializedCapacity: MLKEM.sharedSecretSizeInBytes
-                ) { bufferPtr, length in
+                var symmetricKeyData = Data(repeating: 0, count: MLKEM.sharedSecretByteCount)
+
+                let rc: CInt = symmetricKeyData.withUnsafeMutableBytes { symmetricKeyDataPtr in
                     let bytes: ContiguousBytes =
                         encapsulated.regions.count == 1
                         ? encapsulated.regions.first!
                         : Array(encapsulated)
-                    let result = bytes.withUnsafeBytes { encapsulatedPtr in
+                    return bytes.withUnsafeBytes { encapsulatedPtr in
                         CCryptoBoringSSL_MLKEM${parameter_set}_decap(
-                            bufferPtr.baseAddress,
+                            symmetricKeyDataPtr.baseAddress,
                             encapsulatedPtr.baseAddress,
                             encapsulatedPtr.count,
                             &self.key
                         )
                     }
+                }
 
-                    guard result == 1 else {
-                        throw CryptoKitError.internalBoringSSLError()
-                    }
-
-                    length = MLKEM.sharedSecretSizeInBytes
+                guard rc == 1 else {
+                    throw CryptoKitError.internalBoringSSLError()
                 }
 
-                return SymmetricKey(data: Data(output))
+                return SymmetricKey(data: symmetricKeyData)
             }
         }
     }
@@ -208,7 +206,7 @@ extension MLKEM${parameter_set} {
         }
 
         /// The size of the public key in bytes.
-        static let bytesCount = Backing.bytesCount
+        static let byteCount = Backing.byteCount
 
         fileprivate final class Backing {
             var key: MLKEM${parameter_set}_public_key
@@ -224,7 +222,7 @@ extension MLKEM${parameter_set} {
             ///
             /// - Throws: `CryptoKitError.incorrectKeySize` if the raw representation is not the correct size.
             init(rawRepresentation: some DataProtocol) throws {
-                guard rawRepresentation.count == MLKEM${parameter_set}.PublicKey.bytesCount else {
+                guard rawRepresentation.count == MLKEM${parameter_set}.PublicKey.byteCount else {
                     throw CryptoKitError.incorrectKeySize
                 }
 
@@ -248,7 +246,7 @@ extension MLKEM${parameter_set} {
             var rawRepresentation: Data {
                 var cbb = CBB()
                 // The following BoringSSL functions can only fail on allocation failure, which we define as impossible.
-                CCryptoBoringSSL_CBB_init(&cbb, MLKEM${parameter_set}.PublicKey.Backing.bytesCount)
+                CCryptoBoringSSL_CBB_init(&cbb, MLKEM${parameter_set}.PublicKey.Backing.byteCount)
                 defer { CCryptoBoringSSL_CBB_cleanup(&cbb) }
                 CCryptoBoringSSL_MLKEM${parameter_set}_marshal_public_key(&cbb, &self.key)
                 return Data(bytes: CCryptoBoringSSL_CBB_data(&cbb), count: CCryptoBoringSSL_CBB_len(&cbb))
@@ -260,11 +258,11 @@ extension MLKEM${parameter_set} {
             func encapsulate() -> KEM.EncapsulationResult {
                 withUnsafeTemporaryAllocation(
                     of: UInt8.self,
-                    capacity: MLKEM${parameter_set}.ciphertextSizeInBytes
+                    capacity: MLKEM${parameter_set}.ciphertextByteCount
                 ) { encapsulatedPtr in
                     withUnsafeTemporaryAllocation(
                         of: UInt8.self,
-                        capacity: MLKEM.sharedSecretSizeInBytes
+                        capacity: MLKEM.sharedSecretByteCount
                     ) { secretPtr in
                         CCryptoBoringSSL_MLKEM${parameter_set}_encap(
                             encapsulatedPtr.baseAddress,
@@ -274,11 +272,11 @@ extension MLKEM${parameter_set} {
 
                         return KEM.EncapsulationResult(
                             sharedSecret: SymmetricKey(
-                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretSizeInBytes)
+                                data: Data(bytes: secretPtr.baseAddress!, count: MLKEM.sharedSecretByteCount)
                             ),
                             encapsulated: Data(
                                 bytes: encapsulatedPtr.baseAddress!,
-                                count: MLKEM${parameter_set}.ciphertextSizeInBytes
+                                count: MLKEM${parameter_set}.ciphertextByteCount
                             )
                         )
                     }
@@ -286,7 +284,7 @@ extension MLKEM${parameter_set} {
             }
 
             /// The size of the public key in bytes.
-            static let bytesCount = Int(MLKEM${parameter_set}_PUBLIC_KEY_BYTES)
+            static let byteCount = Int(MLKEM${parameter_set}_PUBLIC_KEY_BYTES)
         }
     }
 }
@@ -294,14 +292,14 @@ extension MLKEM${parameter_set} {
 @available(macOS 14.0, *)
 extension MLKEM${parameter_set} {
     /// The size of the encapsulated shared secret in bytes.
-    private static let ciphertextSizeInBytes = Int(MLKEM${parameter_set}_CIPHERTEXT_BYTES)
+    private static let ciphertextByteCount = Int(MLKEM${parameter_set}_CIPHERTEXT_BYTES)
 }
 % end
 
 private enum MLKEM {
     /// The size of the seed in bytes.
-    fileprivate static let seedSizeInBytes = 64
+    fileprivate static let seedByteCount = 64
 
     // The size of the shared secret in bytes.
-    fileprivate static let sharedSecretSizeInBytes = 32
+    fileprivate static let sharedSecretByteCount = 32
 }