PONGo_ios/ServerTrustEvaluator.swift

import Foundation
import Security
import CryptoKit

/// How to evaluate the server's TLS certificate.
///
/// MCMS appliances usually present a self-signed or internal-CA certificate.
/// Pick the *narrowest* policy that works for your deployment. Never ship
/// `.allowSelfSignedForHost` against a host you don't control.
enum ServerTrustPolicy: Equatable, Codable {
    /// Default OS trust evaluation (use this if the box has a real cert or a
    /// CA your devices already trust via MDM/profile — strongly preferred).
    case system

    /// Accept the server's certificate ONLY for this exact host, after a normal
    /// trust evaluation that ignores the "untrusted root" failure. Scoped escape
    /// hatch for self-signed boxes on a controlled mgmt network/VPN.
    case allowSelfSignedForHost(String)

    /// Certificate/public-key pinning: trust only if the leaf's SHA-256
    /// public-key hash (base64) matches one of these. The safest option for a
    /// fixed appliance — capture the hash once and pin it.
    case pinPublicKeySHA256([String], host: String)
}

/// `URLSessionDelegate` that applies a `ServerTrustPolicy`. Retained by the
/// `URLSession` it's attached to.
final class ServerTrustEvaluator: NSObject, URLSessionDelegate {
    private let policy: ServerTrustPolicy

    init(policy: ServerTrustPolicy) {
        self.policy = policy
    }

    func urlSession(
        _ session: URLSession,
        didReceive challenge: URLAuthenticationChallenge,
        completionHandler: @escaping (URLSession.AuthChallengeDisposition, URLCredential?) -> Void
    ) {
        guard challenge.protectionSpace.authenticationMethod == NSURLAuthenticationMethodServerTrust,
              let trust = challenge.protectionSpace.serverTrust else {
            completionHandler(.performDefaultHandling, nil)
            return
        }

        let host = challenge.protectionSpace.host

        switch policy {
        case .system:
            completionHandler(.performDefaultHandling, nil)

        case .allowSelfSignedForHost(let allowedHost):
            // Scoped escape hatch for self-signed appliances on a controlled
            // network. Because a self-signed root fails normal evaluation, this
            // accepts ANY certificate the allowed host presents — including
            // expired or name-mismatched ones. The only gate is the hostname, so
            // use this only on hosts you control; prefer `.pinPublicKeySHA256`
            // once you can capture the appliance's key hash.
            guard host == allowedHost else {
                completionHandler(.cancelAuthenticationChallenge, nil)
                return
            }
            completionHandler(.useCredential, URLCredential(trust: trust))

        case .pinPublicKeySHA256(let pins, let pinnedHost):
            guard host == pinnedHost,
                  let leaf = Self.leafCertificate(from: trust),
                  let hash = Self.publicKeySHA256Base64(for: leaf),
                  pins.contains(hash) else {
                completionHandler(.cancelAuthenticationChallenge, nil)
                return
            }
            completionHandler(.useCredential, URLCredential(trust: trust))
        }
    }

    // MARK: - Pinning helpers

    private static func leafCertificate(from trust: SecTrust) -> SecCertificate? {
        if #available(iOS 15.0, *) {
            return (SecTrustCopyCertificateChain(trust) as? [SecCertificate])?.first
        } else {
            return SecTrustGetCertificateAtIndex(trust, 0)
        }
    }

    /// Base64 of SHA-256 over the DER-encoded **SubjectPublicKeyInfo**, i.e. the
    /// standard public-key pin. `SecKeyCopyExternalRepresentation` returns only the
    /// raw key (PKCS#1 `RSAPublicKey` for RSA, the `04||X||Y` point for EC) — NOT
    /// the SPKI — so we prepend the algorithm-specific ASN.1 SPKI header before
    /// hashing. The result matches the canonical capture:
    ///
    ///   openssl s_client -connect host:443 </dev/null 2>/dev/null \
    ///     | openssl x509 -pubkey -noout \
    ///     | openssl pkey -pubin -outform der \
    ///     | openssl dgst -sha256 -binary | base64
    ///
    /// Returns nil for key types/sizes without a known header (caller fails closed).
    static func publicKeySHA256Base64(for certificate: SecCertificate) -> String? {
        guard let key = SecCertificateCopyKey(certificate),
              let raw = SecKeyCopyExternalRepresentation(key, nil) as Data?,
              let attrs = SecKeyCopyAttributes(key) as? [String: Any],
              let keyType = attrs[kSecAttrKeyType as String] as? String,
              let keySize = attrs[kSecAttrKeySizeInBits as String] as? Int,
              let header = spkiHeader(keyType: keyType, keySizeInBits: keySize)
        else { return nil }
        var spki = Data(header)
        spki.append(raw)
        return CryptoSHA256.base64(of: spki)
    }

    /// ASN.1 SubjectPublicKeyInfo prefix for the raw key returned by
    /// `SecKeyCopyExternalRepresentation`, keyed by (algorithm, key size). These
    /// are the fixed, well-known DER headers; the raw key bits follow.
    private static func spkiHeader(keyType: String, keySizeInBits: Int) -> [UInt8]? {
        if keyType == (kSecAttrKeyTypeRSA as String) {
            switch keySizeInBits {
            case 2048: return rsa2048SPKIHeader
            case 3072: return rsa3072SPKIHeader
            case 4096: return rsa4096SPKIHeader
            default: return nil
            }
        }
        if keyType == (kSecAttrKeyTypeECSECPrimeRandom as String) {
            switch keySizeInBits {
            case 256: return ecP256SPKIHeader
            case 384: return ecP384SPKIHeader
            default: return nil
            }
        }
        return nil
    }

    private static let rsa2048SPKIHeader: [UInt8] = [
        0x30, 0x82, 0x01, 0x22, 0x30, 0x0d, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86,
        0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x03, 0x82, 0x01, 0x0f, 0x00]
    private static let rsa3072SPKIHeader: [UInt8] = [
        0x30, 0x82, 0x01, 0xa2, 0x30, 0x0d, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86,
        0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x03, 0x82, 0x01, 0x8f, 0x00]
    private static let rsa4096SPKIHeader: [UInt8] = [
        0x30, 0x82, 0x02, 0x22, 0x30, 0x0d, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86,
        0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x03, 0x82, 0x02, 0x0f, 0x00]
    private static let ecP256SPKIHeader: [UInt8] = [
        0x30, 0x59, 0x30, 0x13, 0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02,
        0x01, 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, 0x03,
        0x42, 0x00]
    private static let ecP384SPKIHeader: [UInt8] = [
        0x30, 0x76, 0x30, 0x10, 0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x02,
        0x01, 0x06, 0x05, 0x2b, 0x81, 0x04, 0x00, 0x22, 0x03, 0x62, 0x00]
}

/// Minimal SHA-256 wrapper using CryptoKit.
enum CryptoSHA256 {
    static func base64(of data: Data) -> String {
        let digest = SHA256.hash(data: data)
        return Data(digest).base64EncodedString()
    }
}