SSL pinning validator with implementation for the Subject public key info (SPKI), based on the one at https://github.com/datatheorem/TrustKit

SSLPinningValidator.swift

//
//  SSLPinningValidator.swift
//
//  Created by Angel Garcia on 17/08/16.
//

import Foundation
import Security
import CommonCrypto

// swiftlint:disable force_unwrapping

protocol SSLPinningValidator {
    func canHandleChallenge(challenge: URLAuthenticationChallenge) -> Bool
    func isChallengeValid(challenge: URLAuthenticationChallenge) -> Bool
}

extension SSLPinningValidator {
    func canHandleChallenge(challenge: URLAuthenticationChallenge) -> Bool {
        return challenge.protectionSpace.authenticationMethod == NSURLAuthenticationMethodServerTrust
    }
}

private let lockQueue = DispatchQueue(label: "com.sslpinningvalidator")
private let certificateCache = NSCache<NSData, NSData>()

//Code for SPKI validation based on the one at https://github.com/datatheorem/TrustKit
class SSLPinningSPKIValidator: NSObject, SSLPinningValidator {
    
    enum PublicKeyAlgorithm: String {
        case rsa2048
        case rsa4096
        case ecDsaSecp256r1
        
        var asn1HeaderBytes: [UInt8] {
            switch self {
            case .rsa2048:
                return [ 0x30, 0x82, 0x01, 0x22, 0x30, 0x0d, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86,
                         0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x03, 0x82, 0x01, 0x0f, 0x00 ]
            case .rsa4096:
                return [ 0x30, 0x82, 0x02, 0x22, 0x30, 0x0d, 0x06, 0x09, 0x2a, 0x86, 0x48, 0x86,
                0xf7, 0x0d, 0x01, 0x01, 0x01, 0x05, 0x00, 0x03, 0x82, 0x02, 0x0f, 0x00 ]
            case .ecDsaSecp256r1:
                return [ 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 ]
            }
        }
    }
    
    typealias SPKI = (hostname: String, algorithm: PublicKeyAlgorithm, sha256: NSData)
    
    let validSPKIs: [SPKI]
    
    init(validSPKIs: [SPKI]) {
        self.validSPKIs = validSPKIs
    }
    
    func isChallengeValid(challenge: URLAuthenticationChallenge) -> Bool {
        let hostname = challenge.protectionSpace.host
        let spkis = self.validSPKIs.filter({ hostname.contains($0.hostname) })
        let serverTrust = challenge.protectionSpace.serverTrust!
        
        //Domain not pinned, then is valid
        guard spkis.count > 0 else {
            return true
        }
        
        // First re-check the certificate chain using the default SSL validation in case it was disabled
        // This gives us revocation (only for EV certs I think?) and also ensures the certificate chain is sane
        // And also gives us the exact path that successfully validated the chain
        let sslPolicy = SecPolicyCreateSSL(true, hostname as NSString?)
        SecTrustSetPolicies(serverTrust, sslPolicy)
        var trustResult: SecTrustResultType = .unspecified
        guard SecTrustEvaluate(serverTrust, &trustResult) == errSecSuccess else {
            return false
        }
        guard trustResult == .unspecified || trustResult == .proceed else { return false }
        
        // Check each certificate in the server's certificate chain (the trust object); start with the CA all the way down to the leaf
        let certificateChainLen = SecTrustGetCertificateCount(serverTrust)
        for i in (0..<certificateChainLen).reversed() {
            
            // Extract the certificate
            if  let certificate: SecCertificate = SecTrustGetCertificateAtIndex(serverTrust, i) {
            
                // For each spki key configuration, generate the subject public key info hash
                for spki in spkis {
                    if  let subjectPublicKeyInfoHash = sha256SubjectPublicKeyInfoFromCertificate(certificate: certificate, algorithm: spki.algorithm),
                        subjectPublicKeyInfoHash == spki.sha256 {
                        return true
                    }
                }
            }
        }
        return false
    }
    
    func sha256SubjectPublicKeyInfoFromCertificate(certificate: SecCertificate, algorithm: PublicKeyAlgorithm) -> NSData? {
        
        //Check the cache for already processed sha256
        let cacheKey = cacheKeyForCertificate(certificate: certificate, algorithm: algorithm)
        if let data = certificateCache.object(forKey: cacheKey) {
            return data
        }
        
        // First extract the public key bytes
        guard let publicKeyData = extractPublicKeyDataFromCertificate(certificate: certificate) else { return nil }
        
        // Generate a hash of the subject public key info
        let subjectPublicKeyInfoHash = NSMutableData(length: Int(CC_SHA256_DIGEST_LENGTH))!
        var shaCtx = CC_SHA256_CTX()
        CC_SHA256_Init(&shaCtx)
        
        // Add the missing ASN1 header for public keys to re-create the subject public key info
        let header = algorithm.asn1HeaderBytes
        CC_SHA256_Update(&shaCtx, header, CC_LONG(header.count))
        
        // Add the public key
        CC_SHA256_Update(&shaCtx, publicKeyData.bytes, CC_LONG(publicKeyData.length))
        
        let subjectPublicKeyInfoHashBytes = UnsafeMutableRawPointer(subjectPublicKeyInfoHash.mutableBytes).assumingMemoryBound(to: UInt8.self)
        CC_SHA256_Final(subjectPublicKeyInfoHashBytes, &shaCtx)
        
        //Save in the cache for later usage
        certificateCache.setObject(subjectPublicKeyInfoHash, forKey: cacheKey)
        
        return subjectPublicKeyInfoHash
    }
    
    func extractPublicKeyDataFromCertificate(certificate: SecCertificate) -> NSData? {
        var tempTrust: SecTrust? = nil
        let policy = SecPolicyCreateBasicX509()
        
        // Get a public key reference from the certificate
        SecTrustCreateWithCertificates(certificate, policy, &tempTrust)
        SecTrustEvaluate(tempTrust!, nil)
        let publicKey = SecTrustCopyPublicKey(tempTrust!)!
        
        // Extract the actual bytes from the key reference using the Keychain
        // Prepare the dictionary to add the key
        let peerPublicKeyAdd: [NSString: Any] = [
            kSecClass: kSecClassKey,
            kSecAttrApplicationTag: "SSLPinningSPKIValidator",
            kSecValueRef: publicKey,
            kSecAttrAccessible: kSecAttrAccessibleAfterFirstUnlockThisDeviceOnly,
            kSecReturnData: kCFBooleanTrue
        ]
        
        // Prepare the dictionary to retrieve and delete the key
        let publicKeyGet: [NSString: Any] = [
            kSecClass: kSecClassKey,
            kSecAttrApplicationTag: "SSLPinningSPKIValidator",
            kSecReturnData: kCFBooleanTrue
        ]
        
        var publicKeyData: NSData? = nil
        lockQueue.sync {
            var data: AnyObject? = nil
            SecItemAdd(peerPublicKeyAdd as CFDictionary, &data)
            SecItemDelete(publicKeyGet as CFDictionary)
            publicKeyData = data as? NSData
        }
        return publicKeyData
    }
    
    func cacheKeyForCertificate(certificate: SecCertificate, algorithm: PublicKeyAlgorithm) -> NSData {
        var data = SecCertificateCopyData(certificate) as Data
        data.append(algorithm.rawValue.data(using: .utf8)!)
        return data as NSData
    }
}

Usage In URLSessionDelegate implementation:

var sslPinningValidator: SSLPinningValidator? {
    // Pins are base64 SHA-256 hashes as in HTTP Public Key Pinning (HPKP)
    let validSPKIs: [SSLPinningSPKIValidator.SPKI] =
        pins.map { SSLPinningSPKIValidator.SPKI(hostname: self.hostname,
                                                algorithm: SSLPinningSPKIValidator.PublicKeyAlgorithm.rsa2048,
                                                sha256: $0) }
    return SSLPinningSPKIValidator(validSPKIs: validSPKIs)
}

func urlSession(_ session: URLSession, didReceive challenge: URLAuthenticationChallenge, completionHandler: @escaping (URLSession.AuthChallengeDisposition, URLCredential?) -> Swift.Void) {
    
    guard challenge.protectionSpace.host.contains(hostname),
            let sslPinningValidator = sslPinningValidator,
            sslPinningValidator.canHandleChallenge(challenge: challenge)
    else {
        completionHandler(URLSession.AuthChallengeDisposition.performDefaultHandling, nil)
        return
    }
    
    if sslPinningValidator.isChallengeValid(challenge: challenge),
        let serverTrust = challenge.protectionSpace.serverTrust {
        let credential = URLCredential(trust: serverTrust)
        completionHandler(.useCredential, credential)
    } else {
        // cancel now or do something about it and call completionHandler eventually.
        completionHandler(URLSession.AuthChallengeDisposition.cancelAuthenticationChallenge, nil)
    }
}

To generate pins from cert, look at get_pin_from_certificate.py in TrustKit repo.