feat: ML-KEM-768, ML-KEM-1024 ecdh

swift-secure-enclave-tool-v2.swift

@@ -9,6 +9,12 @@ import CryptoKit
 import Foundation
 import LocalAuthentication
 
+extension Data {
+    var hexEncodedString: String {
+        return map { String(format: "%02x", $0) }.joined()
+    }
+}
+
 struct GenerateSecureEnclaveKeyPairRequest {
     var controlFlag: String
 }
@@ -131,12 +137,12 @@ func parseExternalSignRequest() -> ComputeP256EcSignRequest? {
     )
 }
 
-struct ComputeP256EcdhRequest {
+struct ComputeEcdhRequest {
     var dataRepresentationBase64: String
     var ephemeraPublicKeyBase64: String
 }
 
-func parseComputeP256EcdhRequest() -> ComputeP256EcdhRequest? {
+func parseComputeEcdhRequest() -> ComputeEcdhRequest? {
     var dataRepresentationBase64Opt: String?
     var ephemeraPublicKeyBase64Opt: String?
     let len = CommandLine.arguments.count;
@@ -163,7 +169,7 @@ func parseComputeP256EcdhRequest() -> ComputeP256EcdhRequest? {
         exitError("parameter --ephemera-public-key-base64 required.")
         return nil
     }
-    return ComputeP256EcdhRequest(
+    return ComputeEcdhRequest(
         dataRepresentationBase64: dataRepresentationBase64,
         ephemeraPublicKeyBase64: ephemeraPublicKeyBase64
     )
@@ -259,7 +265,7 @@ struct ComputeSecureEnclaveP256EcsignResponse: Codable {
     var signature_base64: String
 }
 
-struct ComputeSecureEnclaveP256EcdhResponse: Codable {
+struct ComputeSecureEnclaveEcdhResponse: Codable {
     var success: Bool
     var shared_secret_hex: String
 }
@@ -580,7 +586,7 @@ func computeSecureEnclaveP256Ecsign(request: ComputeP256EcSignRequest) -> Comput
     }
 }
 
-func computeSecureEnclaveP256Ecdh(request: ComputeP256EcdhRequest) -> ComputeSecureEnclaveP256EcdhResponse? {
+func computeSecureEnclaveP256Ecdh(request: ComputeEcdhRequest) -> ComputeSecureEnclaveEcdhResponse? {
     guard let privateKeyDataRepresentation = Data(
         base64Encoded: request.dataRepresentationBase64
     ) else {
@@ -614,7 +620,7 @@ func computeSecureEnclaveP256Ecdh(request: ComputeP256EcdhRequest) -> ComputeSec
             shared_secret_hex = sharedSecret.description
         }
 
-        return ComputeSecureEnclaveP256EcdhResponse(
+        return ComputeSecureEnclaveEcdhResponse(
             success: true,
             shared_secret_hex: shared_secret_hex
         )
@@ -624,6 +630,53 @@ func computeSecureEnclaveP256Ecdh(request: ComputeP256EcdhRequest) -> ComputeSec
     }
 }
 
+func computeSecureEnclaveMlKemEcdh(request: ComputeEcdhRequest, keyLen: Int) -> ComputeSecureEnclaveEcdhResponse? {
+    guard let privateKeyDataRepresentation = Data(
+        base64Encoded: request.dataRepresentationBase64
+    ) else {
+        exitError("private key base64 decode failed")
+        return nil
+    }
+    guard let ephemeralPublicKeyRepresentation = Data(
+        base64Encoded: request.ephemeraPublicKeyBase64
+    ) else {
+        exitError("ephemeral public key base64 decode failed")
+        return nil
+    }
+    do {
+        let context = LAContext()
+        if keyLen == 768 {
+            let p =  try SecureEnclave.MLKEM768.PrivateKey(
+                dataRepresentation: privateKeyDataRepresentation,
+                authenticationContext: context
+            )
+            let sharedSecret = try p.decapsulate(ephemeralPublicKeyRepresentation)
+            let sharedSecretData = sharedSecret.withUnsafeBytes { Data($0) }
+            return ComputeSecureEnclaveEcdhResponse(
+                success: true,
+                shared_secret_hex: sharedSecretData.hexEncodedString
+            )
+        } else if (keyLen == 1024) {
+            let p =  try SecureEnclave.MLKEM1024.PrivateKey(
+                dataRepresentation: privateKeyDataRepresentation,
+                authenticationContext: context
+            )
+            let sharedSecret = try p.decapsulate(ephemeralPublicKeyRepresentation)
+            let sharedSecretData = sharedSecret.withUnsafeBytes { Data($0) }
+            return ComputeSecureEnclaveEcdhResponse(
+                success: true,
+                shared_secret_hex: sharedSecretData.hexEncodedString
+            )
+        } else {
+            exitError("Invalid algorithm: ML-KEM-\(keyLen)")
+            return nil
+        }
+    } catch {
+        exitError("\(error)")
+        return nil
+    }
+}
+
 func externalSpec() -> ExternalSpecResponse {
     return ExternalSpecResponse(
         success: true,
@@ -694,11 +747,23 @@ if (command == "compute_p256_ecsign") {
 }
 
 if (command == "compute_p256_ecdh") {
-    let request = parseComputeP256EcdhRequest()!;
+    let request = parseComputeEcdhRequest()!;
     let response =  computeSecureEnclaveP256Ecdh(request: request)
     exitOkWithJson(response)
 }
 
+if (command == "compute_mlkem768_ecdh") {
+    let request = parseComputeEcdhRequest()!;
+    let response =  computeSecureEnclaveMlKemEcdh(request: request, keyLen: 768)
+    exitOkWithJson(response)
+}
+
+if (command == "compute_mlkem1024_ecdh") {
+    let request = parseComputeEcdhRequest()!;
+    let response =  computeSecureEnclaveMlKemEcdh(request: request, keyLen: 1024)
+    exitOkWithJson(response)
+}
+
 if (command == "external_spec") {
     exitOkWithJson(externalSpec())
 }
@@ -734,6 +799,8 @@ if (command == "help" || command == "-h" || command == "--help") {
     print("recover_mlkem1024_public_key --private-key <>                   - recover Secure Enclave ML-KEM-1024 key pair")
     print("compute_p256_ecsign --private-key <> --message-base64 <> [--message-type <>]  - compure Secure Enclave P256 EC sign")
     print("compute_p256_ecdh --private-key <> --ephemera-public-key <>      - compure Secure Enclave P256 EC DH")
+    print("compute_mlkem768_ecdh --private-key <> --ephemera-public-key <>  - compure Secure Enclave ML-KEM-768")
+    print("compute_mlkem1024_ecdh --private-key <> --ephemera-public-key <> - compure Secure Enclave ML-KEM-1024")
     print("external_spec                                                    - external specification")
     print("external_public_key --parameter <>                               - external public key")
     print("external_sign --parameter <> --alg ES256  --message-base64 <>    - external sign")