Verify ECDSA Digital Signature in Swift using a PEM-encoded Key

ECPublicKey.swift

/**
 *  Copyright (c) 2017 Håvard Fossli.
 *
 *  Licensed under the MIT license, as follows:
 *
 *  Permission is hereby granted, free of charge, to any person obtaining a copy
 *  of this software and associated documentation files (the "Software"), to deal
 *  in the Software without restriction, including without limitation the rights
 *  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 *  copies of the Software, and to permit persons to whom the Software is
 *  furnished to do so, subject to the following conditions:
 *
 *  The above copyright notice and this permission notice shall be included in all
 *  copies or substantial portions of the Software.
 *
 *  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 *  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 *  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 *  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 *  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 *  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 *  SOFTWARE.
 */

import Foundation
import Security

struct KeyConfig {
    
    enum KeyType {
        case secp256r1
    }
    
}

class ECPublicKey {
    
    let key: SecKey
    
    enum Error: Swift.Error {
        case osstatus(OSStatus)
        case unknown
    }
    
    convenience init(data: ECPublicKeyData) throws {
        let keyData = data.raw
        let addQuery = Queries.add(data: keyData) as CFDictionary
        let deleteQuery = Queries.delete(data: keyData) as CFDictionary
        var result: CFTypeRef? = nil
        var status = SecItemCopyMatching(addQuery, &result)
        if status != errSecSuccess {
            status = SecItemAdd(addQuery, &result)
            if status == errSecDuplicateItem {
                let _ = SecItemDelete(deleteQuery)
                status = SecItemAdd(addQuery, &result)
            }
        }
        guard status == errSecSuccess else {
            throw Error.osstatus(status)
        }
        guard let resultUnwrapped = result else {
            throw Error.unknown
        }
        let key = resultUnwrapped as! SecKey
        status = SecItemDelete(deleteQuery)
        guard status == errSecSuccess || status == errSecItemNotFound else {
            throw Error.osstatus(status)
        }
        self.init(key: key)
    }
    
    private init(key: SecKey) {
        self.key = key
    }
    
    // https://crypto.stackexchange.com/a/1797
    public struct EcdsaAsn1Signature {
        
        static func smallestBigEndian(_ bytes: Data) -> Data {
            var smallest = bytes
            while smallest.first == 0x00 {
                smallest.removeFirst()
            }
            if let firstByte = smallest.first, firstByte > 0x7f {
                return Data(bytes: [0x00]) + smallest
            } else {
                return smallest
            }
        }
        
        static func isProperlyAsn1Encoded(_ bytes: Data) -> Bool {
            var parser = bytes
            guard bytes.count > 2 else { return false }
            guard parser.popFirst() == 0x30 else { return false }
            guard let sequenceLength = parser.popFirst() else { return false }
            guard parser.count == sequenceLength else { return false }
            guard parser.popFirst() == 0x02 else { return false }
            guard let rLength = parser.popFirst() else { return false }
            guard rLength < parser.count else { return false }
            parser.removeFirst(Int(rLength))
            guard parser.count > 2 else { return false }
            guard parser.popFirst() == 0x02 else { return false }
            guard let sLength = parser.popFirst() else { return false }
            guard sLength == parser.count else { return false }
            return true
        }
        
        static func encode(_ bytes: Data) -> Data {
            guard !isProperlyAsn1Encoded(bytes) else { return bytes }
            let r = smallestBigEndian(bytes.prefix(bytes.count / 2))
            let s = smallestBigEndian(bytes.suffix(bytes.count / 2))
            guard r.count <= 64 && s.count <= 64 else { return bytes }
            var asn1 = Data()
            asn1.append(0x30)
            asn1.append(UInt8(2 + r.count + 2 + s.count))
            asn1.append(0x02)
            asn1.append(UInt8(r.count))
            asn1.append(r)
            asn1.append(0x02)
            asn1.append(UInt8(s.count))
            asn1.append(s)
            return asn1
        }
    }
    
    func verify(signature: Data, messageData: Data) throws {
        let asn1Signature = EcdsaAsn1Signature.encode(signature)
        let sha = messageData.sha256()
        var shaBytes = [UInt8](repeating: 0, count: sha.count)
        sha.copyBytes(to: &shaBytes, count: sha.count)
        
        var signatureBytes = [UInt8](repeating: 0, count: asn1Signature.count)
        asn1Signature.copyBytes(to: &signatureBytes, count: asn1Signature.count)
        
        let status = SecKeyRawVerify(key, .PKCS1, &shaBytes, shaBytes.count, &signatureBytes, signatureBytes.count)
        guard status == errSecSuccess else {
            throw Error.osstatus(status)
        }
    }
    
    private struct Queries {
        static func add(data: Data) -> [String:Any] {
            return [
                kSecClass as String: kSecClassKey,
                kSecAttrApplicationTag as String: "imported_by_value_" + data.sha256().base64EncodedString(),
                kSecAttrKeyClass as String: kSecAttrKeyClassPublic,
                kSecAttrKeyType as String: kSecAttrKeyTypeEC,
                kSecValueData as String: data,
                kSecReturnRef as String: true,
                kSecAttrAccessible as String: kSecAttrAccessibleWhenUnlocked,
                kSecAttrIsPermanent as String: true
            ]
        }
        static func delete(data: Data) -> [String:Any] {
            return [
                kSecClass as String: kSecClassKey,
                kSecAttrApplicationTag as String: "imported_by_value_" + data.sha256().base64EncodedString(),
                kSecAttrKeyClass as String: kSecAttrKeyClassPublic,
                kSecAttrKeyType as String: kSecAttrKeyTypeEC,
                kSecValueData as String: data
            ]
        }
    }
}

class ECPublicKeyData {
    
    enum Error: Swift.Error {
        case unknownOrBadFormat
    }
    
    let raw: Data
    lazy var rawWithHeader: Data = {
       return Data(Constants.x9_62ECHeader) + raw
    }()
    
    struct Constants {
        static let x9_62ECHeader = [UInt8]([
            /* sequence          */ 0x30, 0x59,
            /* |-> sequence      */ 0x30, 0x13,
            /* |---> ecPublicKey */ 0x06, 0x07, 0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x02, 0x01, // http://oid-info.com/get/1.2.840.10045.2.1 (ANSI X9.62 public key type)
            /* |---> prime256v1  */ 0x06, 0x08, 0x2A, 0x86, 0x48, 0xCE, 0x3D, 0x03, 0x01, // http://oid-info.com/get/1.2.840.10045.3.1.7 (ANSI X9.62 named elliptic curve)
            /* |-> bit headers   */ 0x07, 0x03, 0x42, 0x00
            ])
    }
    
    convenience init(PEM: String) throws {
        guard let armoringBegin = PEM.range(of: "-----BEGIN PUBLIC KEY-----"),
            let armoringEnd = PEM.range(of: "-----END PUBLIC KEY-----") else {
                throw Error.unknownOrBadFormat
        }
        guard armoringBegin.upperBound < armoringEnd.lowerBound else {
            throw Error.unknownOrBadFormat
        }
        let contentsRange = armoringBegin.upperBound..<armoringEnd.lowerBound
        let contents = PEM[contentsRange]
            .trimmingCharacters(in: .whitespacesAndNewlines)
            .replacingOccurrences(of: "\n", with: "")
        guard let DER = Data(base64Encoded: contents, options: []) else {
            throw Error.unknownOrBadFormat
        }
        try self.init(DER: DER)
    }
    
    convenience init(DER: Data) throws {
        let x9_62ECHeaderAsData = Data(Constants.x9_62ECHeader)
        if let headerRange = DER.range(of: x9_62ECHeaderAsData) {
            let strippedForHeader = DER.suffix(from: headerRange.upperBound)
            try self.init(DER: strippedForHeader)
        } else if DER.count == 65 {
            try self.init(data: DER)
        } else {
            throw Error.unknownOrBadFormat
        }
    }
    
    private init(data: Data) throws {
        raw = data
    }
}