Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add ML-KEM post-quantum key agreement to _CryptoExtras #314

Open
wants to merge 18 commits into
base: main
Choose a base branch
from
Open
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Swift Format
fpseverino committed Nov 30, 2024
commit dee35107835fe5d4733c35bba27bf758e7754711
69 changes: 43 additions & 26 deletions Sources/_CryptoExtras/MLKEM/MLKEM_boring.swift
Original file line number Diff line number Diff line change
@@ -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,21 +237,27 @@ 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,
&self.key
)

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
}
}
16 changes: 8 additions & 8 deletions Tests/_CryptoExtrasTests/MLKEMTests.swift
Original file line number Diff line number Diff line change
@@ -22,41 +22,41 @@ 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
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 sharedSecret1 = encapsulationResult.sharedSecret
let ciphertext = encapsulationResult.encapsulated

let sharedSecret2 = try privateKey.decapsulate(ciphertext)
XCTAssertEqual(sharedSecret1, sharedSecret2)
}
}
}