diff --git a/Sources/JWTKit/JWTKeyCollection.swift b/Sources/JWTKit/JWTKeyCollection.swift index 3dd611f4..74c27c6e 100644 --- a/Sources/JWTKit/JWTKeyCollection.swift +++ b/Sources/JWTKit/JWTKeyCollection.swift @@ -11,9 +11,17 @@ import Foundation /// This actor provides methods to manage multiple keys. It can be used to verify and decode JWTs, as well as to sign and encode JWTs. /// It also facilitates the encoding and decoding of JWTs using custom or default JSON encoders and decoders. public actor JWTKeyCollection: Sendable { - private enum Signer { + private enum Signer: Equatable { case jwt(JWTSigner) case jwk(JWKSigner) + + static func == (lhs: JWTKeyCollection.Signer, rhs: JWTKeyCollection.Signer) -> Bool { + switch (lhs, rhs) { + case (.jwt(let lhsSigner), .jwt(let rhsSigner)): lhsSigner === rhsSigner + case (.jwk(let lhsSigner), .jwk(let rhsSigner)): lhsSigner === rhsSigner + default: false + } + } } private var storage: [JWKIdentifier: Signer] @@ -69,6 +77,49 @@ public actor JWTKeyCollection: Sendable { return self } + /// Removes the key with the selected KID from the collection. + /// If the default matches, that one is also removed. + /// - Parameter kid: The KID identifying the signer to the collection. + /// - Returns: True if the key was found and removed, false otherwise. + @discardableResult + public func remove(kid: JWKIdentifier) -> Bool { + let value = self.storage.removeValue(forKey: kid) + if value == self.default { + self.default = nil + } + return value != nil + } + + /// Removes the default signer. + /// The signer might still exist in the collection as non-default one. + /// - Returns: True if the default signer was found and removed, false otherwise. + @discardableResult + public func clearDefault() -> Bool { + if self.default != nil { + self.default = nil + return true + } + return false + } + + /// Removes all keys from the collection except the ones defined in the `kids` parameter. + /// - Parameter kids: The KIDs that should be kept during removal. + /// - Parameter clearingDefault: If true, clears the default signer regardless of whether + /// its KID is in the exception list, however keeping it in the collection if it's + /// KID is defined in the exception list. If false, preserves the default signer. + /// - Returns: The number of keys removed from storage. + @discardableResult + public func removeAll(except kids: [JWKIdentifier] = [], clearingDefault: Bool = true) -> Int { + let kidsToKeep = Set(kids) + self.storage = self.storage.filter { kidsToKeep.contains($0.key) } + let originalCount = self.storage.count + if clearingDefault { + self.default = nil + } + + return originalCount - self.storage.count + } + /// Adds a `JWKS` (JSON Web Key Set) to the collection by decoding a JSON string. /// /// - Parameter json: A JSON string representing a JWKS. diff --git a/Tests/JWTKitTests/JWTKitTests.swift b/Tests/JWTKitTests/JWTKitTests.swift index 260af760..793630ce 100644 --- a/Tests/JWTKitTests/JWTKitTests.swift +++ b/Tests/JWTKitTests/JWTKitTests.swift @@ -852,6 +852,134 @@ struct JWTKitTests { #expect(header.field1 == nil) #expect(header.field2 == .string("value2")) } + + @Test("Remove single key by kid") + func testRemoveSingleKey() async throws { + let keys = await JWTKeyCollection() + .add(hmac: "key-1", digestAlgorithm: .sha256, kid: "v1") + .add(hmac: "key-2", digestAlgorithm: .sha256, kid: "v2") + + let payload = TestPayload( + sub: "vapor", + name: "Test", + admin: false, + exp: .init(value: Date().addingTimeInterval(3600)) + ) + + let token1 = try await keys.sign(payload, kid: "v1") + let token2 = try await keys.sign(payload, kid: "v2") + + _ = try await keys.verify(token1, as: TestPayload.self) + _ = try await keys.verify(token2, as: TestPayload.self) + + let removed = await keys.remove(kid: "v1") + #expect(removed == true) + + _ = try await keys.verify(token2, as: TestPayload.self) + + await #expect(throws: JWTError.self) { + _ = try await keys.verify(token1, as: TestPayload.self) + } + } + + @Test("Remove non-existent key returns false") + func testRemoveNonExistentKey() async throws { + let keys = await JWTKeyCollection() + .add(hmac: "key-1", digestAlgorithm: .sha256, kid: "v1") + + let removed = await keys.remove(kid: "non-existent") + #expect(removed == false) + } + + @Test("Remove default key clears default") + func testRemoveDefaultKey() async throws { + let keys = await JWTKeyCollection() + .add(hmac: "key-1", digestAlgorithm: .sha256) // No kid = default + + let payload = TestPayload( + sub: "vapor", + name: "Test", + admin: false, + exp: .init(value: Date().addingTimeInterval(3600)) + ) + + let token = try await keys.sign(payload) + _ = try await keys.verify(token, as: TestPayload.self) + + let removed = await keys.clearDefault() + #expect(removed == true) + + await #expect(throws: JWTError.noKeyProvided) { + _ = try await keys.sign(payload) + } + } + + @Test("RemoveAll except specified keys") + func testRemoveAllExcept() async throws { + let keys = await JWTKeyCollection() + .add(hmac: "key-1", digestAlgorithm: .sha256, kid: "v1") + .add(hmac: "key-2", digestAlgorithm: .sha256, kid: "v2") + .add(hmac: "key-3", digestAlgorithm: .sha256, kid: "v3") + .add(hmac: "key-4", digestAlgorithm: .sha256, kid: "v4") + + let payload = TestPayload( + sub: "vapor", + name: "Test", + admin: false, + exp: .init(value: Date().addingTimeInterval(3600)) + ) + + let token1 = try await keys.sign(payload, kid: "v1") + let token2 = try await keys.sign(payload, kid: "v2") + let token3 = try await keys.sign(payload, kid: "v3") + let token4 = try await keys.sign(payload, kid: "v4") + + await keys.removeAll(except: ["v2", "v4"]) + + // v2 and v4 still work + _ = try await keys.verify(token2, as: TestPayload.self) + _ = try await keys.verify(token4, as: TestPayload.self) + + // v1 and v3 no longer work + await #expect(throws: JWTError.noKeyProvided) { + _ = try await keys.verify(token1, as: TestPayload.self) + } + await #expect(throws: JWTError.noKeyProvided) { + _ = try await keys.verify(token3, as: TestPayload.self) + } + } + + @Test("RemoveAll including default") + func testRemoveAllWithDefault() async throws { + let keys = await JWTKeyCollection() + .add(hmac: "key-0", digestAlgorithm: .sha256) + .add(hmac: "key-1", digestAlgorithm: .sha256, kid: "v1") + .add(hmac: "key-2", digestAlgorithm: .sha256, kid: "v2") + + let payload = TestPayload( + sub: "vapor", + name: "Test", + admin: false, + exp: .init(value: Date().addingTimeInterval(3600)) + ) + + let token0 = try await keys.sign(payload) + let token1 = try await keys.sign(payload, kid: "v1") + let token2 = try await keys.sign(payload, kid: "v2") + + await keys.removeAll(except: ["v2"], clearingDefault: true) + + // v2 still works + _ = try await keys.verify(token2, as: TestPayload.self) + + // v0 and v1 no longer work + await #expect(throws: JWTError.noKeyProvided) { + _ = try await keys.verify(token0, as: TestPayload.self) + } + await #expect(throws: JWTError.noKeyProvided) { + _ = try await keys.verify(token1, as: TestPayload.self) + } + } } let microsoftJWKS = """