feat: Custom parser for added key classes in AnyJSONWebKey

!chore: Rename JSONWebKeySymmetricPortable to JSONWebKeySymmetric
amosavian · Mar 8, 2024 · 141b8cd · 141b8cd
1 parent c820129
commit 141b8cd
Show file tree

Hide file tree

Showing 6 changed files with 100 additions and 37 deletions.
diff --git a/Sources/JWSETKit/Cryptography/Algorithms/KeyEncryption.swift b/Sources/JWSETKit/Cryptography/Algorithms/KeyEncryption.swift
@@ -304,12 +304,28 @@ extension JSONWebKeyEncryptionAlgorithm {
             ephemeralKey = try JSONWebECPrivateKey(curve: privateKey.curve ?? .empty).publicKey
             header.ephemeralPublicKey = .init(ephemeralKey)
         }
+
+        let apu: Data
+        let apv: Data
+        if let headerAPU = header.agreementPartyUInfo {
+            apu = headerAPU
+        } else {
+            apu = ephemeralKey.xCoordinate ?? .init()
+            header.agreementPartyUInfo = apu
+        }
+        if let headerAPV = header.agreementPartyVInfo {
+            apv = headerAPV
+        } else {
+            apv = (keyEncryptionKey?.keyId?.utf8).map { Data($0) } ?? .init()
+            header.agreementPartyVInfo = apv
+        }
+
         let secret = try privateKey.sharedSecretFromKeyAgreement(with: ephemeralKey)
         let symmetricKey = try secret.concatDerivedSymmetricKey(
             algorithm: keyEncryptingAlgorithm,
             contentEncryptionAlgorithm: header.encryptionAlgorithm,
-            apu: header.agreementPartyUInfo ?? .init(),
-            apv: header.agreementPartyVInfo ?? .init(),
+            apu: apu,
+            apv: apv,
             hashFunction: hashFunction
         )
         if keyEncryptingAlgorithm == .ecdhEphemeralStatic {

diff --git a/Sources/JWSETKit/Cryptography/KeyExporter.swift b/Sources/JWSETKit/Cryptography/KeyExporter.swift
@@ -81,17 +81,18 @@ extension JSONWebKeyExportable {
     }
 }
 
-public protocol JSONWebKeySymmetricPortable: JSONWebKeyImportable, JSONWebKeyExportable {
+public protocol JSONWebKeySymmetric: JSONWebKeyImportable, JSONWebKeyExportable {
     init(_ key: SymmetricKey) throws
 }
 
-extension JSONWebKeySymmetricPortable {
+extension JSONWebKeySymmetric {
     public init(importing key: Data, format: JSONWebKeyFormat) throws {
         switch format {
         case .raw:
             try self.init(.init(data: key))
         case .jwk:
             self = try JSONDecoder().decode(Self.self, from: key)
+            try validate()
         default:
             throw JSONWebKeyError.invalidKeyFormat
         }

diff --git a/Sources/JWSETKit/Cryptography/KeyParser.swift b/Sources/JWSETKit/Cryptography/KeyParser.swift
@@ -37,21 +37,32 @@ extension AnyJSONWebKey {
 
     /// Deserializes JSON and converts to the most appropriate key.
     ///
-    /// - Parameter jsonWebKey: JWK in JSON string.
+    /// - Parameters:
+    ///   - data: The key data to deserialize.
+    ///   - format: The format of the key data.
     /// - Returns: Related specific key object.
-    public static func deserialize(_ data: Data) throws -> any JSONWebKey {
+    public static func deserialize(_ data: Data, format: JSONWebKeyFormat) throws -> any JSONWebKey {
+
         let webKey = try JSONDecoder().decode(AnyJSONWebKey.self, from: data)
         return webKey.specialized()
     }
 }
 
 /// A specializer that can convert a `AnyJSONWebKey` to a specific `JSONWebKey` type.
 public protocol JSONWebKeySpecializer {
-    /// Specializes a `AnyJSONWebKey` to a specific `JSONWebKey` type, returns nil if key is appropiate.
+    /// Specializes a `AnyJSONWebKey` to a specific `JSONWebKey` type, returns `nil` if key is not appropiate.
     ///
     /// - Parameter key: The key to specialize.
-    /// - Returns: A specific `JSONWebKey` type, or nil if the key is not appropiate.
+    /// - Returns: A specific `JSONWebKey` type, or `nil` if the key is not appropiate.
     static func specialize(_ key: AnyJSONWebKey) throws -> (any JSONWebKey)?
+
+    /// Deserializes a key from a data, returns `nil` if key is not appropiate.
+    ///
+    /// - Parameters:
+    ///   - key: The key data to deserialize.
+    ///   - format: The format of the key data.
+    ///   - Returns: A specific `JSONWebKey` type, or `nil` if the key is not appropiate.
+    static func deserialize(key: Data, format: JSONWebKeyFormat) throws -> (any JSONWebKey)?
 }
 
 enum JSONWebKeyRSASpecializer: JSONWebKeySpecializer {
@@ -63,6 +74,19 @@ enum JSONWebKeyRSASpecializer: JSONWebKeySpecializer {
             return try JSONWebRSAPublicKey.create(storage: key.storage)
         }
     }
+
+    static func deserialize(key: Data, format: JSONWebKeyFormat) throws -> (any JSONWebKey)? {
+        switch format {
+        case .pkcs8:
+            guard try PKCS8PrivateKey(derEncoded: key).keyType == .rsa else { return nil }
+            return try JSONWebRSAPrivateKey(importing: key, format: format)
+        case .spki:
+            guard try SubjectPublicKeyInfo(derEncoded: key).keyType == .rsa else { return nil }
+            return try JSONWebRSAPublicKey(importing: key, format: format)
+        default:
+            return nil
+        }
+    }
 }
 
 enum JSONWebKeyEllipticCurveSpecializer: JSONWebKeySpecializer {
@@ -80,6 +104,23 @@ enum JSONWebKeyEllipticCurveSpecializer: JSONWebKeySpecializer {
             return nil
         }
     }
+
+    static func deserialize(key: Data, format: JSONWebKeyFormat) throws -> (any JSONWebKey)? {
+        switch format {
+        case .pkcs8:
+            let pkcs8 = try PKCS8PrivateKey(derEncoded: key)
+            guard try pkcs8.keyType == .ellipticCurve else { return nil }
+            guard pkcs8.keyCurve != nil else { return nil }
+            return try JSONWebECPrivateKey(importing: key, format: format)
+        case .spki:
+            let spki = try SubjectPublicKeyInfo(derEncoded: key)
+            guard try spki.keyType == .ellipticCurve else { return nil }
+            guard spki.keyCurve != nil else { return nil }
+            return try JSONWebECPublicKey(importing: key, format: format)
+        default:
+            return nil
+        }
+    }
 }
 
 enum JSONWebKeyCurve25519Specializer: JSONWebKeySpecializer {
@@ -97,6 +138,10 @@ enum JSONWebKeyCurve25519Specializer: JSONWebKeySpecializer {
             return nil
         }
     }
+
+    static func deserialize(key: Data, format: JSONWebKeyFormat) throws -> (any JSONWebKey)? {
+        return nil
+    }
 }
 
 enum JSONWebKeySymmetricSpecializer: JSONWebKeySpecializer {
@@ -125,6 +170,15 @@ enum JSONWebKeySymmetricSpecializer: JSONWebKeySpecializer {
             return try SymmetricKey.create(storage: key.storage)
         }
     }
+
+    static func deserialize(key: Data, format: JSONWebKeyFormat) throws -> (any JSONWebKey)? {
+        switch format {
+        case .raw:
+            return try AnyJSONWebKey(storage: SymmetricKey(importing: key, format: .raw).storage).specialized()
+        default:
+            return nil
+        }
+    }
 }
 
 enum JSONWebKeyCertificateChainSpecializer: JSONWebKeySpecializer {
@@ -134,11 +188,15 @@ enum JSONWebKeyCertificateChainSpecializer: JSONWebKeySpecializer {
         }
         return nil
     }
+
+    static func deserialize(key: Data, format: JSONWebKeyFormat) throws -> (any JSONWebKey)? {
+        return nil
+    }
 }
 
 extension AnyJSONWebKey {
     @ReadWriteLocked
-    fileprivate static var specializers: [any JSONWebKeySpecializer.Type] = [
+    static var specializers: [any JSONWebKeySpecializer.Type] = [
         JSONWebKeyRSASpecializer.self,
         JSONWebKeyEllipticCurveSpecializer.self,
         JSONWebKeyCurve25519Specializer.self,

diff --git a/Sources/JWSETKit/Cryptography/Keys.swift b/Sources/JWSETKit/Cryptography/Keys.swift
@@ -165,7 +165,7 @@ extension JSONWebDecryptingKey {
 }
 
 /// A JSON Web Key (JWK) able to decrypt cipher-texts using a symmetric key.
-public protocol JSONWebSymmetricDecryptingKey: JSONWebDecryptingKey, JSONWebKeySymmetricPortable where PublicKey == Self {
+public protocol JSONWebSymmetricDecryptingKey: JSONWebDecryptingKey, JSONWebKeySymmetric where PublicKey == Self {
     init(_ key: SymmetricKey) throws
 }
 
@@ -297,7 +297,7 @@ extension JSONWebSigningKey {
 }
 
 /// A JSON Web Key (JWK) able to generate a signature using a symmetric key.
-public protocol JSONWebSymmetricSigningKey: JSONWebSigningKey, JSONWebKeySymmetricPortable {
+public protocol JSONWebSymmetricSigningKey: JSONWebSigningKey, JSONWebKeySymmetric {
     init(_ key: SymmetricKey) throws
 }
 
@@ -347,27 +347,17 @@ extension AnyJSONWebKey: JSONWebKeyImportable, JSONWebKeyExportable {
     }
 
     public init(importing key: Data, format: JSONWebKeyFormat) throws {
-        switch format {
-        case .raw:
-            // Here is a little tricky. The key can be either a X9.63 formatted EC
-            // or a symmetric key for HMAC/AES.
-            //
-            // First we try to decode it as EC. If fails, we then fall back to SymmetricKey.
-            if let key = try? JSONWebECPrivateKey(importing: key, format: .raw) {
-                self.init(storage: key.storage)
-            } else if let key = try? JSONWebECPublicKey(importing: key, format: .raw) {
-                self.init(storage: key.storage)
-            } else {
-                try self.init(storage: SymmetricKey(importing: key, format: .raw).storage)
+        if format == .jwk {
+            let key = try JSONDecoder().decode(AnyJSONWebKey.self, from: key).specialized()
+            try key.validate()
+            self.storage = key.storage
+        }
+        for specializer in AnyJSONWebKey.specializers {
+            if let specialized = try specializer.deserialize(key: key, format: format) {
+                self.storage = specialized.storage
             }
-        case .pkcs8:
-            try self.init(importing: key, format: format, keyType: PKCS8PrivateKey(derEncoded: key).keyType)
-        case .spki:
-            try self.init(importing: key, format: format, keyType: SubjectPublicKeyInfo(derEncoded: key).keyType)
-        case .jwk:
-            self = try JSONDecoder().decode(Self.self, from: key)
-            try validate()
         }
+        throw JSONWebKeyError.unknownKeyType
     }
 
     public func exportKey(format: JSONWebKeyFormat) throws -> Data {

diff --git a/Sources/JWSETKit/Cryptography/Symmetric/ConcatKDF.swift b/Sources/JWSETKit/Cryptography/Symmetric/ConcatKDF.swift
@@ -47,17 +47,15 @@ extension SharedSecret {
     ) throws -> SymmetricKey where H: HashFunction, APU: DataProtocol, APV: DataProtocol {
         let algorithmID: String
         let keySize: Int
-        if algorithm == .ecdhEphemeralStatic {
+        if let length = algorithm.keyLength {
+            algorithmID = algorithm.rawValue
+            keySize = length
+        } else {
             guard let cek = contentEncryptionAlgorithm, let contentKeySize = cek.keyLength?.bitCount else {
                 throw CryptoKitError.incorrectKeySize
             }
             algorithmID = cek.rawValue
             keySize = contentKeySize
-        } else if let length = algorithm.keyLength {
-            algorithmID = algorithm.rawValue
-            keySize = length
-        } else {
-            throw CryptoKitError.incorrectKeySize
         }
 
         return try SymmetricKey.concatDerivedSymmetricKey(

diff --git a/Sources/JWSETKit/Resources/en.lproj/Localizable.stringsdict b/Sources/JWSETKit/Resources/en.lproj/Localizable.stringsdict
@@ -40,7 +40,7 @@
 	<key>notBeforeToken</key>
 	<dict>
 		<key>NSStringLocalizedFormatKey</key>
-		<string></string>
+		<string>Token is invalid before %@</string>
 	</dict>
 	<key>errorInvalidAudience</key>
 	<dict>