From f4e128f8e34a75fe9cd2136aa83e29f879eecdc0 Mon Sep 17 00:00:00 2001
From: Kevin Wooten <kevin@outfoxx.io>
Date: Sun, 25 Jun 2023 22:10:36 -0700
Subject: [PATCH] Use CryptoKit to derive the EC public key during PKCS8 import

This ensures the PKCS8 import works when the encoded without the optional public key point, since the full public key point and private key are required for `SecKeyCreateWithData`.
---
 .../ShieldSecurity/AlgorithmIdentifier.swift  |  3 +-
 Sources/ShieldSecurity/SecKeyPair.swift       | 23 ++++++++---
 Tests/SecKeyPairTests.swift                   | 41 +++++++++++++++++--
 3 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/Sources/ShieldSecurity/AlgorithmIdentifier.swift b/Sources/ShieldSecurity/AlgorithmIdentifier.swift
index d8fd743ae..ce935c58a 100644
--- a/Sources/ShieldSecurity/AlgorithmIdentifier.swift
+++ b/Sources/ShieldSecurity/AlgorithmIdentifier.swift
@@ -19,6 +19,7 @@ public extension AlgorithmIdentifier {
 
   enum Error: Swift.Error {
     case unsupportedAlgorithm
+    @available(*, deprecated, message: "No longer used")
     case unsupportedECKeySize
   }
 
@@ -81,7 +82,7 @@ public extension AlgorithmIdentifier {
         // P-521, secp521r1
         curve = iso.org.certicom.curve.ansip521r1.oid
       default:
-        throw Error.unsupportedECKeySize
+        throw Error.unsupportedAlgorithm
       }
 
       self.init(
diff --git a/Sources/ShieldSecurity/SecKeyPair.swift b/Sources/ShieldSecurity/SecKeyPair.swift
index 1ce54b864..cbf2a4ea9 100644
--- a/Sources/ShieldSecurity/SecKeyPair.swift
+++ b/Sources/ShieldSecurity/SecKeyPair.swift
@@ -9,6 +9,7 @@
 //
 
 import Algorithms
+import CryptoKit
 import Foundation
 import PotentASN1
 import Security
@@ -518,13 +519,25 @@ private extension SecKey {
       importKeyData = privateKeyInfo.privateKey
 
     case iso.memberBody.us.ansix962.keyType.ecPublicKey.oid:
+      keyType = .ec
+
       let ecPrivateKey = try ASN1.Decoder.decode(ECPrivateKey.self, from: privateKeyInfo.privateKey)
-      guard let publicKey = ecPrivateKey.publicKey else {
-        throw SecKeyPair.Error.invalidEncodedPrivateKey
+      guard
+        let curveOID = privateKeyInfo.privateKeyAlgorithm.parameters?.objectIdentifierValue
+      else {
+        throw SecKey.Error.importFailed
       }
 
-      keyType = .ec
-      importKeyData = publicKey.bytes + ecPrivateKey.privateKey
+      switch curveOID {
+      case iso.memberBody.us.ansix962.curves.prime.prime256v1.oid:
+        importKeyData = try P256.Signing.PrivateKey(rawRepresentation: ecPrivateKey.privateKey).x963Representation
+      case iso.org.certicom.curve.ansip384r1.oid:
+        importKeyData = try P384.Signing.PrivateKey(rawRepresentation: ecPrivateKey.privateKey).x963Representation
+      case iso.org.certicom.curve.ansip521r1.oid:
+        importKeyData = try P521.Signing.PrivateKey(rawRepresentation: ecPrivateKey.privateKey).x963Representation
+      default:
+        throw AlgorithmIdentifier.Error.unsupportedAlgorithm
+      }
 
     default:
       throw AlgorithmIdentifier.Error.unsupportedAlgorithm
@@ -594,7 +607,7 @@ private extension SecKey {
       // P-521, secp521r1
       return (iso.org.certicom.curve.ansip521r1.oid, 66)
     default:
-      throw AlgorithmIdentifier.Error.unsupportedECKeySize
+      throw AlgorithmIdentifier.Error.unsupportedAlgorithm
     }
   }
 
diff --git a/Tests/SecKeyPairTests.swift b/Tests/SecKeyPairTests.swift
index b7df44fa7..44c98f331 100644
--- a/Tests/SecKeyPairTests.swift
+++ b/Tests/SecKeyPairTests.swift
@@ -213,11 +213,46 @@ class SecKeyPairTests: XCTestCase {
     XCTAssertThrowsError(try SecKeyPair.import(fromData: exportedKeyData, withPassword: "456"))
   }
 
-  func testImportExportEC() throws {
+  func testImportExportEC192() throws {
 
-    let exportedKeyData = try ecKeyPair.export()
+    let ecKeyPair =
+      try SecKeyPair.Builder(type: .ec, keySize: 192)
+        .generate(label: "Test 192 EC Key")
+    defer { try? ecKeyPair.delete() }
+
+    XCTAssertThrowsError(try SecKeyPair.import(fromData: ecKeyPair.export())) { error in
+      XCTAssertTrue(error is AlgorithmIdentifier.Error)
+    }
+  }
+
+  func testImportExportEC256() throws {
+
+    let ecKeyPair =
+      try SecKeyPair.Builder(type: .ec, keySize: 256)
+        .generate(label: "Test 256 EC Key")
+    defer { try? ecKeyPair.delete() }
+
+    _ = try SecKeyPair.import(fromData: ecKeyPair.export())
+  }
+
+  func testImportExportEC384() throws {
+
+    let ecKeyPair =
+      try SecKeyPair.Builder(type: .ec, keySize: 384)
+        .generate(label: "Test 384 EC Key")
+    defer { try? ecKeyPair.delete() }
+
+    _ = try SecKeyPair.import(fromData: ecKeyPair.export())
+  }
+
+  func testImportExportEC521() throws {
+
+    let ecKeyPair =
+      try SecKeyPair.Builder(type: .ec, keySize: 521)
+        .generate(label: "Test 521 EC Key")
+    defer { try? ecKeyPair.delete() }
 
-    _ = try SecKeyPair.import(fromData: exportedKeyData)
+    _ = try SecKeyPair.import(fromData: ecKeyPair.export())
   }
 
   func testCodable() throws {