web-terminal/ios/Packages/ClientTLS/Sources/ClientTLS/KeychainClientIdentityStore.swift

import Foundation
import os
import Security

/// Persists the device identity so it survives relaunch. The raw `.p12` bytes
/// **and** its passphrase are stored together (the passphrase is required to
/// re-import via `SecPKCS12Import` at every launch), then re-imported on load.
public protocol ClientIdentityStore: Sendable {
    /// Validate (`SecPKCS12Import`) then persist the `.p12` + passphrase.
    /// Throws `PKCS12ImportError` on a bad passphrase / corrupt file (nothing is
    /// persisted in that case) and `ClientIdentityStoreError` on a storage fault.
    func save(p12Data: Data, passphrase: String) throws
    /// Re-import and return the stored identity; `nil` if none is installed.
    func loadIdentity() throws -> ClientIdentity?
    /// Display summary of the stored certificate; `nil` if none is installed.
    func loadSummary() throws -> ClientCertificateSummary?
    /// Delete the stored identity (rotation / removal). Idempotent.
    func remove() throws
    /// Cheap existence check for gating (does NOT re-import).
    func hasInstalledIdentity() -> Bool
}

public enum ClientIdentityStoreError: Error, Equatable, Sendable {
    /// A Keychain `SecItem*` call failed with this `OSStatus`.
    case keychain(OSStatus)
    /// The stored blob was present but could not be decoded.
    case corruptStoredBlob
}

public extension ClientIdentityStore {
    /// Convenience for composition roots: load the identity, logging and
    /// swallowing errors into `nil`. A missing cert is the normal pre-install
    /// state; a genuine fault must not crash launch, but is logged (never
    /// silently dropped).
    func loadedIdentityOrNil() -> ClientIdentity? {
        do {
            return try loadIdentity()
        } catch {
            ClientTLSLog.identity.error(
                "loadIdentity failed: \(String(describing: error), privacy: .public)"
            )
            return nil
        }
    }
}

/// The stored payload — `.p12` bytes plus the passphrase needed to re-import.
private struct StoredP12Blob: Codable {
    let p12: Data
    let passphrase: String
}

/// C-iOS · The enrolled-identity record stored beside the Secure-Enclave leaf:
/// the issuer chain (presented on the handshake), the `deviceId` (drives renew),
/// and the rotation timing. The private key itself never lives here — it stays
/// non-exportable in the Secure Enclave.
private struct StoredEnrollment: Codable {
    let deviceId: String
    let deviceName: String
    let caChain: [Data]
    let notAfter: Date?
    let renewAfter: Date?
}

/// Keychain-backed store: one `kSecClassGenericPassword` item holding the
/// JSON-encoded `StoredP12Blob` in `kSecValueData`, protected with
/// `kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly` (available after the first
/// unlock post-boot, never migrates off this device).
public struct KeychainClientIdentityStore: ClientIdentityStore {
    public static let defaultService = "com.yaojia.webterm.clienttls"
    public static let defaultAccount = "device-identity"

    private let service: String
    private let account: String

    public init(
        service: String = defaultService, account: String = defaultAccount
    ) {
        self.service = service
        self.account = account
    }

    public func save(p12Data: Data, passphrase: String) throws {
        // Validate BEFORE persisting — a wrong passphrase / corrupt file must
        // surface to the install UI and leave any prior identity untouched.
        _ = try PKCS12Importer.importIdentity(data: p12Data, passphrase: passphrase)
        let blob = try encode(StoredP12Blob(p12: p12Data, passphrase: passphrase))
        try writeItem(blob)
    }

    public func loadIdentity() throws -> ClientIdentity? {
        // C-iOS · Prefer the Secure-Enclave-backed enrolled identity (the .p12-free
        // path). During the dual-trust migration window a device may still carry a
        // legacy imported `.p12`; fall back to it so existing installs keep working.
        if let enrolled = try loadDeviceIdentity() { return enrolled }
        guard let blob = try readBlob() else { return nil }
        return try PKCS12Importer.importIdentity(
            data: blob.p12, passphrase: blob.passphrase
        )
    }

    public func loadSummary() throws -> ClientCertificateSummary? {
        try loadIdentity()?.summary()
    }

    public func remove() throws {
        // Clear BOTH paths so removal is unconditional: the SE key + enrolled leaf
        // and the legacy `.p12` blob.
        try removeEnrolled()
        let status = SecItemDelete(baseQuery() as CFDictionary)
        guard status == errSecSuccess || status == errSecItemNotFound else {
            throw ClientIdentityStoreError.keychain(status)
        }
    }

    public func hasInstalledIdentity() -> Bool {
        if hasEnrolledLeaf() { return true }
        var query = baseQuery()
        query[kSecReturnData as String] = false
        query[kSecMatchLimit as String] = kSecMatchLimitOne
        return SecItemCopyMatching(query as CFDictionary, nil) == errSecSuccess
    }

    // MARK: - C-iOS · Secure-Enclave enrollment (the .p12-free path)

    /// Keychain tag of the device's Secure-Enclave private key. The enrolled leaf
    /// binds to this key so `SecItemCopyMatching(kSecClassIdentity)` can assemble
    /// the `SecIdentity` presented on the existing mTLS path (unchanged).
    private var deviceKeyTag: Data { Data("\(service).device-key".utf8) }
    /// Label under which the enrolled leaf certificate is stored.
    private var leafLabel: String { "\(service).device-leaf" }
    /// Account of the generic-password item holding the enrollment record
    /// (deviceId + issuer chain + rotation timing) alongside the leaf.
    private var enrollmentAccount: String { "\(account).enrollment" }

    /// One-time enrollment: generate a NON-EXPORTABLE Secure-Enclave P-256 key,
    /// self-sign a CSR with it, POST it, and store the returned leaf against that
    /// key. Returns the installed cert's summary. `keyProvider` is injectable so
    /// non-device builds can supply a software key; production defaults to the SE.
    public func enroll(
        using client: DeviceEnrollmentClient,
        subdomain: String,
        deviceName: String,
        keyProvider: (@Sendable () throws -> any P256HardwareKey)? = nil
    ) async throws -> ClientCertificateSummary? {
        let key = try keyProvider?()
            ?? SecureEnclaveKeyFactory.generateSecureEnclave(tag: deviceKeyTag)
        let csr = try CertificateSigningRequest.der(subjectCommonName: deviceName, signer: key)
        let result = try await client.enroll(
            csrDER: csr, subdomain: subdomain, deviceName: deviceName
        )
        try storeEnrolledLeaf(result, deviceName: deviceName)
        return try loadSummary()
    }

    /// Rotation: re-CSR from the SAME Secure-Enclave key and replace the leaf.
    /// The caller (rotation scheduler) tears down live connections afterward so
    /// the new cert is presented on the next handshake (plan §3.3).
    public func renew(
        using client: DeviceEnrollmentClient
    ) async throws -> ClientCertificateSummary? {
        guard let record = try readEnrollment(),
              let key = try SecureEnclaveKeyFactory.load(tag: deviceKeyTag)
        else {
            throw ClientIdentityStoreError.corruptStoredBlob // nothing to renew
        }
        let csr = try CertificateSigningRequest.der(
            subjectCommonName: record.deviceName, signer: key
        )
        let result = try await client.renew(deviceId: record.deviceId, csrDER: csr)
        try storeEnrolledLeaf(result, deviceName: record.deviceName)
        return try loadSummary()
    }

    /// The Secure-Enclave-backed identity: the assembled `SecIdentity` (leaf bound
    /// to the SE key) plus the stored issuer chain. `nil` when not enrolled.
    func loadDeviceIdentity() throws -> ClientIdentity? {
        let query: [String: Any] = [
            kSecClass as String: kSecClassIdentity,
            kSecAttrApplicationTag as String: deviceKeyTag,
            kSecReturnRef as String: true,
            kSecMatchLimit as String: kSecMatchLimitOne,
        ]
        var result: CFTypeRef?
        let status = SecItemCopyMatching(query as CFDictionary, &result)
        if status == errSecItemNotFound { return nil }
        guard status == errSecSuccess, let value = result else {
            throw ClientIdentityStoreError.keychain(status)
        }
        // Safe: the query pins kSecClassIdentity, so a match is a SecIdentity.
        let identity = value as! SecIdentity
        let issuers = ((try? readEnrollment())?.caChain ?? []).compactMap {
            SecCertificateCreateWithData(nil, $0 as CFData)
        }
        return ClientIdentity(secIdentity: identity, issuerCertificates: issuers)
    }

    /// Cheap check: is a Secure-Enclave identity installed? (No cert re-parse.)
    private func hasEnrolledLeaf() -> Bool {
        let query: [String: Any] = [
            kSecClass as String: kSecClassIdentity,
            kSecAttrApplicationTag as String: deviceKeyTag,
            kSecMatchLimit as String: kSecMatchLimitOne,
        ]
        return SecItemCopyMatching(query as CFDictionary, nil) == errSecSuccess
    }

    /// Persist the enrolled leaf (delete-then-add so rotation replaces the prior
    /// one) plus the enrollment record. The leaf binds to the permanent SE key,
    /// letting the keychain form the identity on load.
    private func storeEnrolledLeaf(_ result: EnrollmentResult, deviceName: String) throws {
        guard let certificate = SecCertificateCreateWithData(nil, result.certificate as CFData)
        else {
            throw ClientIdentityStoreError.corruptStoredBlob // not a valid DER cert
        }
        let deleteLeaf: [String: Any] = [
            kSecClass as String: kSecClassCertificate,
            kSecAttrLabel as String: leafLabel,
        ]
        let deleteStatus = SecItemDelete(deleteLeaf as CFDictionary)
        guard deleteStatus == errSecSuccess || deleteStatus == errSecItemNotFound else {
            throw ClientIdentityStoreError.keychain(deleteStatus)
        }
        let addLeaf: [String: Any] = [
            kSecClass as String: kSecClassCertificate,
            kSecValueRef as String: certificate,
            kSecAttrLabel as String: leafLabel,
        ]
        let addStatus = SecItemAdd(addLeaf as CFDictionary, nil)
        guard addStatus == errSecSuccess else {
            throw ClientIdentityStoreError.keychain(addStatus)
        }
        try writeEnrollment(
            StoredEnrollment(
                deviceId: result.deviceId,
                deviceName: deviceName,
                caChain: result.caChain,
                notAfter: result.notAfter,
                renewAfter: result.renewAfter
            )
        )
    }

    /// Delete the SE key, the enrolled leaf, and the enrollment record. Idempotent.
    private func removeEnrolled() throws {
        let deleteLeaf: [String: Any] = [
            kSecClass as String: kSecClassCertificate,
            kSecAttrLabel as String: leafLabel,
        ]
        let leafStatus = SecItemDelete(deleteLeaf as CFDictionary)
        guard leafStatus == errSecSuccess || leafStatus == errSecItemNotFound else {
            throw ClientIdentityStoreError.keychain(leafStatus)
        }
        try SecureEnclaveKeyFactory.delete(tag: deviceKeyTag)
        let recordStatus = SecItemDelete(enrollmentQuery() as CFDictionary)
        guard recordStatus == errSecSuccess || recordStatus == errSecItemNotFound else {
            throw ClientIdentityStoreError.keychain(recordStatus)
        }
    }

    private func enrollmentQuery() -> [String: Any] {
        [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: service,
            kSecAttrAccount as String: enrollmentAccount,
        ]
    }

    private func writeEnrollment(_ record: StoredEnrollment) throws {
        let data: Data
        do {
            data = try JSONEncoder().encode(record)
        } catch {
            throw ClientIdentityStoreError.corruptStoredBlob
        }
        let deleteStatus = SecItemDelete(enrollmentQuery() as CFDictionary)
        guard deleteStatus == errSecSuccess || deleteStatus == errSecItemNotFound else {
            throw ClientIdentityStoreError.keychain(deleteStatus)
        }
        var attributes = enrollmentQuery()
        attributes[kSecValueData as String] = data
        attributes[kSecAttrAccessible as String] =
            kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly
        let addStatus = SecItemAdd(attributes as CFDictionary, nil)
        guard addStatus == errSecSuccess else {
            throw ClientIdentityStoreError.keychain(addStatus)
        }
    }

    private func readEnrollment() throws -> StoredEnrollment? {
        var query = enrollmentQuery()
        query[kSecReturnData as String] = true
        query[kSecMatchLimit as String] = kSecMatchLimitOne
        var result: CFTypeRef?
        let status = SecItemCopyMatching(query as CFDictionary, &result)
        if status == errSecItemNotFound { return nil }
        guard status == errSecSuccess, let data = result as? Data else {
            throw ClientIdentityStoreError.keychain(status)
        }
        do {
            return try JSONDecoder().decode(StoredEnrollment.self, from: data)
        } catch {
            throw ClientIdentityStoreError.corruptStoredBlob
        }
    }

    // MARK: - Keychain plumbing

    private func baseQuery() -> [String: Any] {
        [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: service,
            kSecAttrAccount as String: account,
        ]
    }

    private func writeItem(_ data: Data) throws {
        // Delete-then-add keeps the item's protection class deterministic
        // (SecItemUpdate can't change kSecAttrAccessible in place).
        let deleteStatus = SecItemDelete(baseQuery() as CFDictionary)
        guard deleteStatus == errSecSuccess || deleteStatus == errSecItemNotFound else {
            throw ClientIdentityStoreError.keychain(deleteStatus)
        }
        var attributes = baseQuery()
        attributes[kSecValueData as String] = data
        attributes[kSecAttrAccessible as String] =
            kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly
        let addStatus = SecItemAdd(attributes as CFDictionary, nil)
        guard addStatus == errSecSuccess else {
            throw ClientIdentityStoreError.keychain(addStatus)
        }
    }

    private func readBlob() throws -> StoredP12Blob? {
        var query = baseQuery()
        query[kSecReturnData as String] = true
        query[kSecMatchLimit as String] = kSecMatchLimitOne
        var result: CFTypeRef?
        let status = SecItemCopyMatching(query as CFDictionary, &result)
        if status == errSecItemNotFound { return nil }
        guard status == errSecSuccess, let data = result as? Data else {
            throw ClientIdentityStoreError.keychain(status)
        }
        do {
            return try JSONDecoder().decode(StoredP12Blob.self, from: data)
        } catch {
            throw ClientIdentityStoreError.corruptStoredBlob
        }
    }

    private func encode(_ blob: StoredP12Blob) throws -> Data {
        do {
            return try JSONEncoder().encode(blob)
        } catch {
            throw ClientIdentityStoreError.corruptStoredBlob
        }
    }
}

/// In-memory store for previews and unit tests: same import/summary code path as
/// the keychain store (so the roundtrip is exercised) without any keychain
/// entitlement. `@unchecked Sendable` — mutable blob guarded by a lock.
public final class InMemoryClientIdentityStore: ClientIdentityStore, @unchecked Sendable {
    private let lock = NSLock()
    private var blob: StoredP12BlobBox?

    /// Boxed so the private `StoredP12Blob` type stays file-private above; this
    /// mirror keeps the two bytes+passphrase without exposing the Codable type.
    private struct StoredP12BlobBox {
        let p12: Data
        let passphrase: String
    }

    public init() {}

    public func save(p12Data: Data, passphrase: String) throws {
        _ = try PKCS12Importer.importIdentity(data: p12Data, passphrase: passphrase)
        lock.withLock { blob = StoredP12BlobBox(p12: p12Data, passphrase: passphrase) }
    }

    public func loadIdentity() throws -> ClientIdentity? {
        guard let stored = lock.withLock({ blob }) else { return nil }
        return try PKCS12Importer.importIdentity(
            data: stored.p12, passphrase: stored.passphrase
        )
    }

    public func loadSummary() throws -> ClientCertificateSummary? {
        try loadIdentity()?.summary()
    }

    public func remove() throws {
        lock.withLock { blob = nil }
    }

    public func hasInstalledIdentity() -> Bool {
        lock.withLock { blob != nil }
    }
}

enum ClientTLSLog {
    static let identity = Logger(subsystem: "com.yaojia.webterm", category: "client-tls")
}