Invalidate pending signatures by per-key policy generation

Sources/ClavisCore/Protocol/AgentControlDispatcher.swift

@@ -3,6 +3,7 @@ import Darwin
 
 internal struct AgentControlDispatcher {
     let applicationRegistry: ApplicationRegistryStoring
+    let keyPolicyGenerations: KeyPolicyGenerations
     let keyManager: KeychainManager
     let auditRecorder: AuditRecording
     let agentSessions: AgentSessionRegistry
@@ -21,9 +22,11 @@ internal struct AgentControlDispatcher {
         promptGate: SigningPromptGate,
         peerProcessValidator: @escaping (pid_t, String, UInt64) -> Bool,
         gitSigningGraceManager: GitSigningGraceManager = .shared,
-        applicationRegistry: ApplicationRegistryStoring = KeychainApplicationRegistry.shared
+        applicationRegistry: ApplicationRegistryStoring = KeychainApplicationRegistry.shared,
+        keyPolicyGenerations: KeyPolicyGenerations = .shared
     ) {
         self.applicationRegistry = applicationRegistry
+        self.keyPolicyGenerations = keyPolicyGenerations
         self.keyManager = keyManager
         self.auditRecorder = auditRecorder
         self.agentSessions = agentSessions
@@ -110,6 +113,8 @@ internal struct AgentControlDispatcher {
         guard let label = reader.readWireString(), !label.isEmpty, reader.isEOF else {
             return Data([5])
         }
+        let fingerprint = (try? keyManager.fetchKeyInfo(label: label))?.fingerprint
+        keyPolicyGenerations.invalidate(label: label, fingerprint: fingerprint)
         gitSigningGraceManager.invalidate(keyLabel: label)
         agentSessions.endAll(keyLabel: label, reason: .keyChanged)
         return Data([6]) // SSH_AGENT_SUCCESS

Sources/ClavisCore/Protocol/ApplicationRequestAuthorization.swift

@@ -5,8 +5,12 @@ public struct ApplicationRequestLease: Equatable, Sendable {
     public let eligibility: ApplicationEligibility
     public let revision: UUID
     public let caller: VerifiedApplicationCaller?
-    public init(eligibility: ApplicationEligibility, revision: UUID, caller: VerifiedApplicationCaller?) {
+    public let keyFingerprint: String?
+    public let keyGeneration: UUID?
+    public init(eligibility: ApplicationEligibility, revision: UUID, caller: VerifiedApplicationCaller?,
+                keyFingerprint: String? = nil, keyGeneration: UUID? = nil) {
         self.eligibility = eligibility; self.revision = revision; self.caller = caller
+        self.keyFingerprint = keyFingerprint; self.keyGeneration = keyGeneration
     }
     public var approvalIdentity: String? {
         caller.map { "\($0.identity.kind.rawValue):\($0.identity.value)" }
@@ -21,8 +25,10 @@ public protocol ApplicationRequestAuthorizing {
 public final class ApplicationRequestAuthorization: ApplicationRequestAuthorizing {
     private let store: ApplicationRegistryStoring
     private let resolver: ApplicationCallerResolver
-    public init(store: ApplicationRegistryStoring = KeychainApplicationRegistry.shared, resolver: ApplicationCallerResolver = .init()) {
-        self.store = store; self.resolver = resolver
+    private let generations: KeyPolicyGenerations
+    public init(store: ApplicationRegistryStoring = KeychainApplicationRegistry.shared,
+                resolver: ApplicationCallerResolver = .init(), generations: KeyPolicyGenerations = .shared) {
+        self.store = store; self.resolver = resolver; self.generations = generations
     }
 
     public func authorize(key: SSHKeyInfo, operation: KeyOperation, role: AgentSocketRole, pid: pid_t, path: String, start: UInt64) throws -> ApplicationRequestLease {
@@ -35,12 +41,15 @@ public final class ApplicationRequestAuthorization: ApplicationRequestAuthorizin
             profile: key.effectiveAccessProfile, operation: operation,
             classification: caller?.application.classification, restrictedBotEndpoint: role == .agent
         )
-        return ApplicationRequestLease(eligibility: eligibility, revision: snapshot.revision, caller: caller)
+        return ApplicationRequestLease(eligibility: eligibility, revision: snapshot.revision, caller: caller,
+                                       keyFingerprint: key.fingerprint, keyGeneration: generations.current(for: key))
     }
 
     public func validate(_ lease: ApplicationRequestLease, pid: pid_t, path: String, start: UInt64) throws -> Bool {
         let snapshot = try store.load()
         guard !snapshot.isPaused, snapshot.revision == lease.revision else { return false }
+        guard let fingerprint = lease.keyFingerprint, let generation = lease.keyGeneration,
+              generations.matches(fingerprint: fingerprint, generation: generation) else { return false }
         let caller = resolver.resolve(peerPID: pid, peerPath: path, peerStart: start, registry: snapshot)
         return caller == lease.caller
     }

Sources/ClavisCore/Protocol/KeyPolicyGenerations.swift

@@ -0,0 +1,38 @@
+import Foundation
+
+/// Process-local generations protect pending requests from policy ABA changes.
+/// Restarting the daemon discards every request and grant, so no epoch persists.
+/// Both listeners must use the same instance; routing is not an identity boundary.
+public final class KeyPolicyGenerations: @unchecked Sendable {
+    public static let shared = KeyPolicyGenerations()
+    private let lock = NSLock()
+    private var generations: [String: UUID] = [:]
+    private var fingerprintsByLabel: [String: String] = [:]
+    public init() {}
+
+    public func current(for key: SSHKeyInfo) -> UUID {
+        lock.lock()
+        defer { lock.unlock() }
+        fingerprintsByLabel[key.label] = key.fingerprint
+        if let generation = generations[key.fingerprint] { return generation }
+        let generation = UUID()
+        generations[key.fingerprint] = generation
+        return generation
+    }
+
+    public func matches(fingerprint: String, generation: UUID) -> Bool {
+        lock.lock()
+        defer { lock.unlock() }
+        return generations[fingerprint] == generation
+    }
+
+    public func invalidate(label: String, fingerprint: String? = nil) {
+        lock.lock()
+        defer { lock.unlock() }
+        guard let fingerprint = fingerprint ?? fingerprintsByLabel[label] else { return }
+        generations[fingerprint] = UUID()
+        // Keep the old-label mapping so rename revocation also invalidates a
+        // request that was authorized before the public index was renamed.
+        fingerprintsByLabel[label] = fingerprint
+    }
+}

Sources/ClavisCore/Protocol/SSHAgentServer.swift

@@ -94,6 +94,7 @@ public class SSHAgentServer {
     public init(
         applicationAccess: ApplicationRequestAuthorizing? = nil,
         applicationRegistry: ApplicationRegistryStoring = KeychainApplicationRegistry.shared,
+        keyPolicyGenerations: KeyPolicyGenerations = .shared,
         socketPath: String = SSHAgentServer.defaultSocketPath,
         role: AgentSocketRole = .personal,
         maxConcurrentClients: Int = 32,
@@ -134,7 +135,7 @@ public class SSHAgentServer {
 
         let resolvedPeerValidator = peerProcessValidator ?? { SSHAgentServer.peerProcessUnchanged(pid: $0, path: $1, startTime: $2) }
 
-        let applicationAccess = applicationAccess ?? ApplicationRequestAuthorization(store: applicationRegistry)
+        let applicationAccess = applicationAccess ?? ApplicationRequestAuthorization(store: applicationRegistry, generations: keyPolicyGenerations)
         self.controlDispatcher = AgentControlDispatcher(
             keyManager: keyManager,
             auditRecorder: auditRecorder,
@@ -144,7 +145,8 @@ public class SSHAgentServer {
             promptGate: promptGate,
             peerProcessValidator: resolvedPeerValidator,
             gitSigningGraceManager: gitSigningGraceManager,
-            applicationRegistry: applicationRegistry
+            applicationRegistry: applicationRegistry,
+            keyPolicyGenerations: keyPolicyGenerations
         )
 
         self.applicationAccess = applicationAccess

Tests/ClavisCoreTests/KeyPolicyGenerationTests.swift

@@ -0,0 +1,213 @@
+import Foundation
+import Testing
+import LocalAuthentication
+@testable import ClavisCore
+@testable import ClavisTestSupport
+
+@Suite struct KeyPolicyGenerationTests {
+    private func key(label: String = "key", fingerprint: String = "fp") -> SSHKeyInfo {
+        SSHKeyInfo(label: label, publicKeyOpenSSH: "", publicKeyBlob: Data(), fingerprint: fingerprint)
+    }
+
+    private func access(_ generations: KeyPolicyGenerations) -> ApplicationRequestAuthorization {
+        ApplicationRequestAuthorization(store: GenerationRegistry(), resolver: .init(
+            inspect: { _ in nil }, verify: { _, _ in false }, transparentHelper: { _ in false }), generations: generations)
+    }
+
+    @Test func copiedGenerationDoesNotValidateAnotherFingerprint() {
+        let generations = KeyPolicyGenerations()
+        let token = generations.current(for: key())
+        #expect(!generations.matches(fingerprint: "different-key", generation: token))
+    }
+
+    @Test func unchangedKeyKeepsGeneration() {
+        let generations = KeyPolicyGenerations()
+        let key = key()
+        let initial = generations.current(for: key)
+        #expect(generations.current(for: key) == initial)
+        #expect(generations.matches(fingerprint: key.fingerprint, generation: initial))
+    }
+
+    @Test func invalidatingOneKeyDoesNotChangeAnotherKey() {
+        let generations = KeyPolicyGenerations()
+        let first = key(label: "first", fingerprint: "one")
+        let second = key(label: "second", fingerprint: "two")
+        let original = generations.current(for: first)
+        let unaffected = generations.current(for: second)
+        generations.invalidate(label: first.label)
+        #expect(!generations.matches(fingerprint: first.fingerprint, generation: original))
+        #expect(generations.matches(fingerprint: second.fingerprint, generation: unaffected))
+    }
+
+    @Test func renameCannotRestoreOldGeneration() {
+        let generations = KeyPolicyGenerations()
+        let old = key(label: "old")
+        let initial = generations.current(for: old)
+        _ = generations.current(for: key(label: "new"))
+        generations.invalidate(label: "old")
+        #expect(generations.current(for: key(label: "new")) != initial)
+        #expect(!generations.matches(fingerprint: old.fingerprint, generation: initial))
+    }
+
+    @Test func changingAwayAndBackNeverRevalidatesPendingLease() throws {
+        let generations = KeyPolicyGenerations()
+        let authorizer = access(generations)
+        let key = key()
+        let lease = try authorizer.authorize(key: key, operation: .ssh, role: .personal, pid: 2, path: "/peer", start: 7)
+        #expect(try authorizer.validate(lease, pid: 2, path: "/peer", start: 7))
+        generations.invalidate(label: key.label)
+        generations.invalidate(label: key.label)
+        #expect(try !authorizer.validate(lease, pid: 2, path: "/peer", start: 7))
+        let fresh = try authorizer.authorize(key: key, operation: .ssh, role: .personal, pid: 2, path: "/peer", start: 7)
+        #expect(try authorizer.validate(fresh, pid: 2, path: "/peer", start: 7))
+    }
+
+    @Test func generationsAreSharedAcrossEndpointAuthorizers() throws {
+        let generations = KeyPolicyGenerations()
+        let primary = access(generations)
+        let compatibility = access(generations)
+        let bot = SSHKeyInfo(label: "bot", publicKeyOpenSSH: "", publicKeyBlob: Data(), fingerprint: "fp", accessProfile: .init(type: .bot, capabilities: .init(ssh: true, signature: true)))
+        let first = try primary.authorize(key: bot, operation: .ssh, role: .personal, pid: 2, path: "/peer", start: 7)
+        let second = try compatibility.authorize(key: bot, operation: .signature, role: .agent, pid: 2, path: "/peer", start: 7)
+        generations.invalidate(label: bot.label)
+        #expect(try !primary.validate(first, pid: 2, path: "/peer", start: 7))
+        #expect(try !compatibility.validate(second, pid: 2, path: "/peer", start: 7))
+    }
+
+    @Test func leaseWithoutGenerationCannotPassProductionValidation() throws {
+        let authorizer = access(KeyPolicyGenerations())
+        let lease = ApplicationRequestLease(eligibility: .authenticatedOnceOnly, revision: GenerationRegistry.revision, caller: nil)
+        #expect(try !authorizer.validate(lease, pid: 2, path: "/peer", start: 7))
+    }
+
+    @Test func trustedKeyRevocationAdvancesGenerationWithoutChangingRegistry() async throws {
+        let fixture = try ClavisTestFixture()
+        defer { fixture.cleanup() }
+        try await fixture.run {
+            let generations = KeyPolicyGenerations()
+            let manager = fixture.makeKeyManager()
+            let key = try manager.generateKey(label: "revoked-key")
+            let registry = GenerationRegistry()
+            let server = SSHAgentServer(applicationAccess: TestApplicationAccess(), applicationRegistry: registry,
+                keyPolicyGenerations: generations, keyManager: manager, auditRecorder: InMemoryAuditRecorder())
+            let old = generations.current(for: key)
+            var payload = Data([SSHAgentServer.invalidateKeyRequest])
+            payload.appendWireString(key.label)
+            #expect(server.processAgentRequest(payload: payload, isTrustedControlPeer: { true }) == Data([6]))
+            #expect(!generations.matches(fingerprint: key.fingerprint, generation: old))
+            #expect(try registry.load().revision == GenerationRegistry.revision)
+        }
+    }
+
+    @Test func untrustedRevocationCannotAdvanceGeneration() async throws {
+        let fixture = try ClavisTestFixture()
+        defer { fixture.cleanup() }
+        try await fixture.run {
+            let generations = KeyPolicyGenerations()
+            let manager = fixture.makeKeyManager()
+            let key = try manager.generateKey(label: "untrusted-revoke")
+            let server = SSHAgentServer(applicationAccess: TestApplicationAccess(), keyPolicyGenerations: generations,
+                keyManager: manager, auditRecorder: InMemoryAuditRecorder())
+            let old = generations.current(for: key)
+            var payload = Data([SSHAgentServer.invalidateKeyRequest])
+            payload.appendWireString(key.label)
+            #expect(server.processAgentRequest(payload: payload, isTrustedControlPeer: { false }) == Data([5]))
+            #expect(generations.matches(fingerprint: key.fingerprint, generation: old))
+        }
+    }
+    @Test(arguments: [false, true])
+    func actualPolicyABADuringPromptDropsOldRequestAndAllowsFreshRetry(capabilities: Bool) async throws {
+        let fixture = try ClavisTestFixture()
+        defer { fixture.cleanup() }
+        try await fixture.run {
+            let generations = KeyPolicyGenerations()
+            let manager = fixture.makeKeyManager(agentGrantRevoker: { generations.invalidate(label: $0) })
+            let key = try manager.generateKey(label: "policy-aba")
+            var changed = false
+            let auth = GenerationMutationAuthenticator {
+                guard !changed else { return }
+                changed = true
+                if capabilities {
+                    try manager.changeAccessProfile(label: key.label, capabilities: .init(ssh: false, signature: true))
+                    try manager.changeAccessProfile(label: key.label, capabilities: key.effectiveAccessProfile.capabilities)
+                } else {
+                    try manager.changeAccessProfile(label: key.label, type: .bot)
+                    try manager.changeAccessProfile(label: key.label, type: .personal)
+                }
+            }
+            let recorder = InMemoryAuditRecorder()
+            let evaluator = SigningPolicyEvaluator(applicationAccess: access(generations), keyManager: manager,
+                auditRecorder: recorder, authenticator: auth, peerProcessValidator: { _, _, _ in true })
+            let old = try evaluator.evaluateSigningRequest(keyBlob: key.publicKeyBlob, dataToSign: Data("ssh".utf8), flags: 0,
+                role: .personal, clientPid: getpid(), clientExecutablePath: "/peer", clientStartTime: 7)
+            #expect(old == .dropConnection)
+            #expect(try manager.fetchKeyInfo(label: key.label)?.effectiveAccessProfile == key.effectiveAccessProfile)
+            let retry = try evaluator.evaluateSigningRequest(keyBlob: key.publicKeyBlob, dataToSign: Data("ssh".utf8), flags: 0,
+                role: .personal, clientPid: getpid(), clientExecutablePath: "/peer", clientStartTime: 7)
+            guard case .approved = retry else { Issue.record("Fresh request should use current generation"); return }
+            #expect(recorder.events.filter { $0.type == .signature }.map(\.result) == [.denied, .allowed])
+        }
+    }
+
+    @Test func daemonRestartCannotAcceptPriorProcessLease() throws {
+        let oldAuthorizer = access(KeyPolicyGenerations())
+        let lease = try oldAuthorizer.authorize(key: key(), operation: .ssh, role: .personal, pid: 2, path: "/peer", start: 7)
+        let restarted = access(KeyPolicyGenerations())
+        _ = try restarted.authorize(key: key(), operation: .ssh, role: .personal, pid: 2, path: "/peer", start: 7)
+        #expect(try !restarted.validate(lease, pid: 2, path: "/peer", start: 7))
+    }
+
+    @Test func concurrentInvalidationAlwaysRejectsOriginalGeneration() async {
+        let generations = KeyPolicyGenerations()
+        let key = key()
+        let original = generations.current(for: key)
+        await withTaskGroup(of: Void.self) { group in
+            for _ in 0..<32 {
+                group.addTask {
+                    generations.invalidate(label: key.label)
+                    _ = generations.current(for: key)
+                }
+            }
+        }
+        #expect(!generations.matches(fingerprint: key.fingerprint, generation: original))
+        let final = generations.current(for: key)
+        #expect(generations.matches(fingerprint: key.fingerprint, generation: final))
+    }
+
+    @Test func renameRevocationWorksAfterOldPublicIndexDisappears() async throws {
+        let fixture = try ClavisTestFixture()
+        defer { fixture.cleanup() }
+        try await fixture.run {
+            let generations = KeyPolicyGenerations()
+            let old = key(label: "old-label")
+            let epoch = generations.current(for: old)
+            let server = SSHAgentServer(applicationAccess: TestApplicationAccess(),
+                keyPolicyGenerations: generations, keyManager: fixture.makeKeyManager(),
+                auditRecorder: InMemoryAuditRecorder())
+            var payload = Data([SSHAgentServer.invalidateKeyRequest])
+            payload.appendWireString(old.label)
+            #expect(server.processAgentRequest(payload: payload, isTrustedControlPeer: { true }) == Data([6]))
+            #expect(!generations.matches(fingerprint: old.fingerprint, generation: epoch))
+            #expect(generations.current(for: key(label: "new-label")) != epoch)
+        }
+    }
+
+}
+
+private struct GenerationRegistry: ApplicationRegistryStoring {
+    static let revision = UUID()
+    func load() throws -> ApplicationRegistrySnapshot { .init(revision: Self.revision) }
+    func save(_ snapshot: ApplicationRegistrySnapshot, replacing revision: UUID) throws { throw ApplicationRegistryError.staleRevision }
+}
+
+private struct GenerationMutationAuthenticator: UserAuthenticating {
+    let action: () throws -> Void
+    func authenticate(reason: String, policy: LAPolicy) throws -> LAContext {
+        try action()
+        return LAContext()
+    }
+    func authenticate(reason: String, policy: LAPolicy) async throws -> LAContext {
+        try action()
+        return LAContext()
+    }
+}