From aac208defb1796311c278fa158195d03b63e45bf Mon Sep 17 00:00:00 2001 From: Paul Toffoloni Date: Mon, 9 Feb 2026 12:37:31 +0100 Subject: [PATCH 1/4] Add key removal functions to key collection --- Sources/JWTKit/JWTKeyCollection.swift | 44 +++++++++++- Tests/JWTKitTests/JWTKitTests.swift | 96 +++++++++++++++++++++++++++ 2 files changed, 139 insertions(+), 1 deletion(-) diff --git a/Sources/JWTKit/JWTKeyCollection.swift b/Sources/JWTKit/JWTKeyCollection.swift index 3dd611f4..ca27c256 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,40 @@ 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. + /// - Returns: True if the default signer was found and removed, false otherwise. + @discardableResult + public func removeDefault() -> 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 after removal. + /// - Parameter withDefault: Whether to remove the default signer too. Defaults to true. + public func removeAll(except kids: [JWKIdentifier] = [], withDefault: Bool = true) { + let kidsToKeep = Set(kids) + for kid in storage.keys where !kidsToKeep.contains(kid) { + self.remove(kid: kid) + } + } + /// 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..98a921ee 100644 --- a/Tests/JWTKitTests/JWTKitTests.swift +++ b/Tests/JWTKitTests/JWTKitTests.swift @@ -852,6 +852,102 @@ 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.removeDefault() + #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.self) { + _ = try await keys.verify(token1, as: TestPayload.self) + } + await #expect(throws: JWTError.noKeyProvided.self) { + _ = try await keys.verify(token3, as: TestPayload.self) + } + } } let microsoftJWKS = """ From e4bbf3eb111caed3ede281898c0d378f2d5bd223 Mon Sep 17 00:00:00 2001 From: Paul Toffoloni Date: Mon, 9 Feb 2026 14:45:38 +0100 Subject: [PATCH 2/4] Add removeAll includingDefault test --- Sources/JWTKit/JWTKeyCollection.swift | 11 +++++--- Tests/JWTKitTests/JWTKitTests.swift | 36 +++++++++++++++++++++++++-- 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/Sources/JWTKit/JWTKeyCollection.swift b/Sources/JWTKit/JWTKeyCollection.swift index ca27c256..59035093 100644 --- a/Sources/JWTKit/JWTKeyCollection.swift +++ b/Sources/JWTKit/JWTKeyCollection.swift @@ -103,11 +103,14 @@ public actor JWTKeyCollection: Sendable { /// Removes all keys from the collection except the ones defined in the `kids` parameter. /// - Parameter kids: The KIDs that should be kept after removal. - /// - Parameter withDefault: Whether to remove the default signer too. Defaults to true. - public func removeAll(except kids: [JWKIdentifier] = [], withDefault: Bool = true) { + /// - Parameter includingDefault: Whether to remove the default signer too. Defaults to true. + /// - Note: If `includingDefault` is set to `true` but the KID of the default key is in the exception list, + /// the key will be cleared from the default but kept in the collection. + public func removeAll(except kids: [JWKIdentifier] = [], includingDefault: Bool = true) { let kidsToKeep = Set(kids) - for kid in storage.keys where !kidsToKeep.contains(kid) { - self.remove(kid: kid) + self.storage = self.storage.filter { kidsToKeep.contains($0.key) } + if includingDefault { // N.B.: We assume that includingDefault == true means remove it even if it's in the exception list. + self.default = nil } } diff --git a/Tests/JWTKitTests/JWTKitTests.swift b/Tests/JWTKitTests/JWTKitTests.swift index 98a921ee..4186dbfb 100644 --- a/Tests/JWTKitTests/JWTKitTests.swift +++ b/Tests/JWTKitTests/JWTKitTests.swift @@ -941,13 +941,45 @@ struct JWTKitTests { _ = try await keys.verify(token4, as: TestPayload.self) // v1 and v3 no longer work - await #expect(throws: JWTError.noKeyProvided.self) { + await #expect(throws: JWTError.noKeyProvided) { _ = try await keys.verify(token1, as: TestPayload.self) } - await #expect(throws: JWTError.noKeyProvided.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"], includingDefault: 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 = """ From c3ab44c450c53e7020e4dd81184a259cccc5582e Mon Sep 17 00:00:00 2001 From: Paul Toffoloni Date: Tue, 10 Feb 2026 11:22:24 +0100 Subject: [PATCH 3/4] Adjust API --- Sources/JWTKit/JWTKeyCollection.swift | 18 +++++++++++------- Tests/JWTKitTests/JWTKitTests.swift | 4 ++-- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/Sources/JWTKit/JWTKeyCollection.swift b/Sources/JWTKit/JWTKeyCollection.swift index 59035093..0f8e99ce 100644 --- a/Sources/JWTKit/JWTKeyCollection.swift +++ b/Sources/JWTKit/JWTKeyCollection.swift @@ -93,7 +93,7 @@ public actor JWTKeyCollection: Sendable { /// Removes the default signer. /// - Returns: True if the default signer was found and removed, false otherwise. @discardableResult - public func removeDefault() -> Bool { + public func clearDefault() -> Bool { if self.default != nil { self.default = nil return true @@ -102,16 +102,20 @@ public actor JWTKeyCollection: Sendable { } /// Removes all keys from the collection except the ones defined in the `kids` parameter. - /// - Parameter kids: The KIDs that should be kept after removal. - /// - Parameter includingDefault: Whether to remove the default signer too. Defaults to true. - /// - Note: If `includingDefault` is set to `true` but the KID of the default key is in the exception list, - /// the key will be cleared from the default but kept in the collection. - public func removeAll(except kids: [JWKIdentifier] = [], includingDefault: Bool = true) { + /// - 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. 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) } - if includingDefault { // N.B.: We assume that includingDefault == true means remove it even if it's in the exception list. + 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. diff --git a/Tests/JWTKitTests/JWTKitTests.swift b/Tests/JWTKitTests/JWTKitTests.swift index 4186dbfb..793630ce 100644 --- a/Tests/JWTKitTests/JWTKitTests.swift +++ b/Tests/JWTKitTests/JWTKitTests.swift @@ -906,7 +906,7 @@ struct JWTKitTests { let token = try await keys.sign(payload) _ = try await keys.verify(token, as: TestPayload.self) - let removed = await keys.removeDefault() + let removed = await keys.clearDefault() #expect(removed == true) await #expect(throws: JWTError.noKeyProvided) { @@ -967,7 +967,7 @@ struct JWTKitTests { let token1 = try await keys.sign(payload, kid: "v1") let token2 = try await keys.sign(payload, kid: "v2") - await keys.removeAll(except: ["v2"], includingDefault: true) + await keys.removeAll(except: ["v2"], clearingDefault: true) // v2 still works _ = try await keys.verify(token2, as: TestPayload.self) From f00a7ae2cb9d3154b8d6a26070ff75e0f2c5a877 Mon Sep 17 00:00:00 2001 From: Paul Toffoloni Date: Wed, 25 Feb 2026 13:20:53 +0100 Subject: [PATCH 4/4] Clarify API docs --- Sources/JWTKit/JWTKeyCollection.swift | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/Sources/JWTKit/JWTKeyCollection.swift b/Sources/JWTKit/JWTKeyCollection.swift index 0f8e99ce..74c27c6e 100644 --- a/Sources/JWTKit/JWTKeyCollection.swift +++ b/Sources/JWTKit/JWTKeyCollection.swift @@ -91,6 +91,7 @@ public actor JWTKeyCollection: Sendable { } /// 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 { @@ -104,7 +105,8 @@ public actor JWTKeyCollection: Sendable { /// 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. If false, preserves the default signer. + /// 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 {