From cd68188eb56add625a3bda4b72edf0ee8f17ba16 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Sat, 19 Jul 2025 10:32:54 +0800
Subject: [PATCH] feat: ecsign support sha256 digest

---
 swift-secure-enclave-tool-v2.swift | 29 ++++++++++++++++++++++++-----
 1 file changed, 24 insertions(+), 5 deletions(-)

diff --git a/swift-secure-enclave-tool-v2.swift b/swift-secure-enclave-tool-v2.swift
index 2c0c7a0..3e60a0a 100644
--- a/swift-secure-enclave-tool-v2.swift
+++ b/swift-secure-enclave-tool-v2.swift
@@ -40,11 +40,13 @@ func parseGenerateSecureEnclaveP256KeyPairRequest() -> GenerateSecureEnclaveP256
 struct ComputeP256EcSignRequest {
     var dataRepresentationBase64: String
     var messageBase64: String
+    var messageType: String
 }
 
 func parseComputeP256EcSignRequest() -> ComputeP256EcSignRequest? {
     var dataRepresentationBase64Opt: String?
     var messageBase64Opt: String?
+    var messageTypeOpt: String?
     let len = CommandLine.arguments.count;
     if CommandLine.arguments.count > 2 {
         var i = 2
@@ -56,6 +58,9 @@ func parseComputeP256EcSignRequest() -> ComputeP256EcSignRequest? {
             } else if (k == "--message-base64") {
                 messageBase64Opt = CommandLine.arguments[i + 1]
                 i += 2
+            } else if (k == "--message-type") {
+                messageTypeOpt = CommandLine.arguments[i + 1]
+                i += 2;
             } else {
                 i += 1
             }
@@ -71,7 +76,8 @@ func parseComputeP256EcSignRequest() -> ComputeP256EcSignRequest? {
     }
     return ComputeP256EcSignRequest(
         dataRepresentationBase64: dataRepresentationBase64,
-        messageBase64: messageBase64
+        messageBase64: messageBase64,
+        messageType: messageTypeOpt ?? "raw"
     )
 }
 
@@ -79,6 +85,7 @@ func parseExternalSignRequest() -> ComputeP256EcSignRequest? {
     var algOpt: String?
     var dataRepresentationBase64Opt: String?
     var messageBase64Opt: String?
+    var messageTypeOpt: String?
     let len = CommandLine.arguments.count;
     if CommandLine.arguments.count > 2 {
         var i = 2
@@ -93,6 +100,9 @@ func parseExternalSignRequest() -> ComputeP256EcSignRequest? {
             } else if (k == "--alg") {
                 algOpt = CommandLine.arguments[i + 1]
                 i += 2
+            } else if (k == "--message-type") {
+                messageTypeOpt = CommandLine.arguments[i + 1]
+                i += 2;
             } else {
                 i += 1
             }
@@ -116,7 +126,8 @@ func parseExternalSignRequest() -> ComputeP256EcSignRequest? {
     }
     return ComputeP256EcSignRequest(
         dataRepresentationBase64: dataRepresentationBase64,
-        messageBase64: messageBase64
+        messageBase64: messageBase64,
+        messageType: messageTypeOpt ?? "raw"
     )
 }
 
@@ -400,8 +411,16 @@ func computeSecureEnclaveP256Ecsign(request: ComputeP256EcSignRequest) -> Comput
             authenticationContext: context
         )
 
-        let digest = SHA256.hash(data: contentData)
-        let signature = try p.signature(for: digest)
+        let signature: P256.Signing.ECDSASignature
+        if (request.messageType == "raw") {
+            let digest = SHA256.hash(data: contentData)
+            signature = try p.signature(for: digest)
+        } else if (request.messageType == "sha256") {
+            signature = try p.signature(for: contentData)
+        } else {
+            exitError("not supported message type: \(request.messageType)")
+            return nil
+        }
 
         return ComputeSecureEnclaveP256EcsignResponse(
             success: true,
@@ -539,7 +558,7 @@ if (command == "help" || command == "-h" || command == "--help") {
     print("generate_p256_ecdh_keypair --control-flag <>                   - generate Secure Enclave P256 EC DH key pair")
     print("recover_p256_ecsign_public_key --private-key <>                - recover Secure Enclave P256 EC sign key pair")
     print("recover_p256_ecdh_public_key --private-key <>                  - recover Secure Enclave P256 EC DH key pair")
-    print("compute_p256_ecsign --private-key <> --message-base64 <>       - compure Secure Enclave P256 EC sign")
+    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("external_spec                                                  - external specification")
     print("external_public_key --parameter <>                             - external public key")