Reduce the number of heap allocations when computing hash digests (apple#238)

* Remove array allocation when computing digest

Instead of creating a temporary array, the digest bytes are stored in a temporary stack-allocated buffer.

This reduces the number of heap allocations needed to hash a message from 7 to 6.

* Store digest context inline

This reduces the number of heap allocations needed to hash a message from 6 to 2.

Before:
  1. the DigestContext object is created
  2. the EVP_MD_CTX struct is allocated by EVP_MD_CTX_new
  3. md_data is allocated when initializing the context

After:
  1. the DigestContext object is created

In both cases, the number is doubled because COW happens during finalize().

* Use temporary variable to finalize context

This reduces the number of heap allocations needed to hash a message from 2 to 1.

---------

Co-authored-by: Cory Benfield <[email protected]>

Author: baarde

Sources/Crypto/Digests/BoringSSL/Digest_boring.swift

@@ -19,44 +19,128 @@
 protocol HashFunctionImplementationDetails: HashFunction where Digest: DigestPrivate {}
 
 protocol BoringSSLBackedHashFunction: HashFunctionImplementationDetails {
-    static var digestType: DigestContext.DigestType { get }
+    associatedtype Context
+    static var digestSize: Int { get }
+    static func initialize() -> Context?
+    static func update(_ context: inout Context, data: UnsafeRawBufferPointer) -> Bool
+    static func finalize(_ context: inout Context, digest: UnsafeMutableRawBufferPointer) -> Bool
 }
 
 extension Insecure.MD5: BoringSSLBackedHashFunction {
-    static var digestType: DigestContext.DigestType {
-        .md5
+    static var digestSize: Int {
+        Int(MD5_DIGEST_LENGTH)
+    }
+
+    static func initialize() -> MD5_CTX? {
+        var context = MD5_CTX()
+        guard CCryptoBoringSSL_MD5_Init(&context) == 1 else {
+            return nil
+        }
+        return context
+    }
+
+    static func update(_ context: inout MD5_CTX, data: UnsafeRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_MD5_Update(&context, data.baseAddress, data.count) == 1
+    }
+
+    static func finalize(_ context: inout MD5_CTX, digest: UnsafeMutableRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_MD5_Final(digest.baseAddress, &context) == 1
     }
 }
 
 extension Insecure.SHA1: BoringSSLBackedHashFunction {
-    static var digestType: DigestContext.DigestType {
-        .sha1
+    static var digestSize: Int {
+        Int(SHA_DIGEST_LENGTH)
+    }
+
+    static func initialize() -> SHA_CTX? {
+        var context = SHA_CTX()
+        guard CCryptoBoringSSL_SHA1_Init(&context) == 1 else {
+            return nil
+        }
+        return context
+    }
+
+    static func update(_ context: inout SHA_CTX, data: UnsafeRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_SHA1_Update(&context, data.baseAddress, data.count) == 1
+    }
+
+    static func finalize(_ context: inout SHA_CTX, digest: UnsafeMutableRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_SHA1_Final(digest.baseAddress, &context) == 1
     }
 }
 
 extension SHA256: BoringSSLBackedHashFunction {
-    static var digestType: DigestContext.DigestType {
-        .sha256
+    static var digestSize: Int {
+        Int(SHA256_DIGEST_LENGTH)
+    }
+
+    static func initialize() -> SHA256_CTX? {
+        var context = SHA256_CTX()
+        guard CCryptoBoringSSL_SHA256_Init(&context) == 1 else {
+            return nil
+        }
+        return context
+    }
+
+    static func update(_ context: inout SHA256_CTX, data: UnsafeRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_SHA256_Update(&context, data.baseAddress, data.count) == 1
+    }
+
+    static func finalize(_ context: inout SHA256_CTX, digest: UnsafeMutableRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_SHA256_Final(digest.baseAddress, &context) == 1
     }
 }
 
 extension SHA384: BoringSSLBackedHashFunction {
-    static var digestType: DigestContext.DigestType {
-        .sha384
+    static var digestSize: Int {
+        Int(SHA384_DIGEST_LENGTH)
+    }
+
+    static func initialize() -> SHA512_CTX? {
+        var context = SHA512_CTX()
+        guard CCryptoBoringSSL_SHA384_Init(&context) == 1 else {
+            return nil
+        }
+        return context
+    }
+
+    static func update(_ context: inout SHA512_CTX, data: UnsafeRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_SHA384_Update(&context, data.baseAddress, data.count) == 1
+    }
+
+    static func finalize(_ context: inout SHA512_CTX, digest: UnsafeMutableRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_SHA384_Final(digest.baseAddress, &context) == 1
     }
 }
 
 extension SHA512: BoringSSLBackedHashFunction {
-    static var digestType: DigestContext.DigestType {
-        .sha512
+    static var digestSize: Int {
+        Int(SHA512_DIGEST_LENGTH)
+    }
+
+    static func initialize() -> SHA512_CTX? {
+        var context = SHA512_CTX()
+        guard CCryptoBoringSSL_SHA512_Init(&context) == 1 else {
+            return nil
+        }
+        return context
+    }
+
+    static func update(_ context: inout SHA512_CTX, data: UnsafeRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_SHA512_Update(&context, data.baseAddress, data.count) == 1
+    }
+
+    static func finalize(_ context: inout SHA512_CTX, digest: UnsafeMutableRawBufferPointer) -> Bool {
+        CCryptoBoringSSL_SHA512_Final(digest.baseAddress, &context) == 1
     }
 }
 
 struct OpenSSLDigestImpl<H: BoringSSLBackedHashFunction> {
-    private var context: DigestContext
+    private var context: DigestContext<H>
 
     init() {
-        self.context = DigestContext(digest: H.digestType)
+        self.context = DigestContext()
     }
 
     internal mutating func update(data: UnsafeRawBufferPointer) {
@@ -67,81 +151,49 @@ struct OpenSSLDigestImpl<H: BoringSSLBackedHashFunction> {
     }
 
     internal func finalize() -> H.Digest {
-        // To have a non-destructive finalize operation we must allocate.
-        let copyContext = DigestContext(copying: self.context)
-        let digestBytes = copyContext.finalize()
-        return digestBytes.withUnsafeBytes {
-            // We force unwrap here because if the digest size is wrong it's an internal error.
-            H.Digest(bufferPointer: $0)!
-        }
+        self.context.finalize()
     }
 }
 
-class DigestContext {
-    private var contextPointer: UnsafeMutablePointer<EVP_MD_CTX>
+fileprivate final class DigestContext<H: BoringSSLBackedHashFunction> {
+    private var context: H.Context
 
-    init(digest: DigestType) {
-        // We force unwrap because we cannot recover from allocation failure.
-        self.contextPointer = CCryptoBoringSSL_EVP_MD_CTX_new()!
-        guard CCryptoBoringSSL_EVP_DigestInit(self.contextPointer, digest.dispatchTable) != 0 else {
-            // We can't do much but crash here.
-            fatalError("Unable to initialize digest state: \(CCryptoBoringSSL_ERR_get_error())")
+    init() {
+        guard let contex = H.initialize() else {
+            preconditionFailure("Unable to initialize digest state")
         }
+        self.context = contex
     }
 
     init(copying original: DigestContext) {
-        // We force unwrap because we cannot recover from allocation failure.
-        self.contextPointer = CCryptoBoringSSL_EVP_MD_CTX_new()!
-        guard CCryptoBoringSSL_EVP_MD_CTX_copy(self.contextPointer, original.contextPointer) != 0 else {
-            // We can't do much but crash here.
-            fatalError("Unable to copy digest state: \(CCryptoBoringSSL_ERR_get_error())")
-        }
+        self.context = original.context
     }
 
     func update(data: UnsafeRawBufferPointer) {
-        guard let baseAddress = data.baseAddress else {
-            return
+        guard H.update(&self.context, data: data) else {
+            preconditionFailure("Unable to update digest state")
         }
-
-        CCryptoBoringSSL_EVP_DigestUpdate(self.contextPointer, baseAddress, data.count)
     }
 
-    // This finalize function is _destructive_: do not call it if you want to reuse the object!
-    func finalize() -> [UInt8] {
-        let digestSize = CCryptoBoringSSL_EVP_MD_size(self.contextPointer.pointee.digest)
-        var digestBytes = Array(repeating: UInt8(0), count: digestSize)
-        var count = UInt32(digestSize)
-
-        digestBytes.withUnsafeMutableBufferPointer { digestPointer in
-            assert(digestPointer.count == count)
-            CCryptoBoringSSL_EVP_DigestFinal(self.contextPointer, digestPointer.baseAddress, &count)
+    func finalize() -> H.Digest {
+        var copyContext = self.context
+        defer {
+            withUnsafeMutablePointer(to: &copyContext) { $0.zeroize() }
+        }
+        return withUnsafeTemporaryAllocation(byteCount: H.digestSize, alignment: 1) { digestPointer in
+            defer {
+                digestPointer.zeroize()
+            }
+            guard H.finalize(&copyContext, digest: digestPointer) else {
+                preconditionFailure("Unable to finalize digest state")
+            }
+            // We force unwrap here because if the digest size is wrong it's an internal error.
+            return H.Digest(bufferPointer: UnsafeRawBufferPointer(digestPointer))!
         }
-
-        return digestBytes
     }
 
     deinit {
-        CCryptoBoringSSL_EVP_MD_CTX_free(self.contextPointer)
-    }
-}
-
-extension DigestContext {
-    struct DigestType {
-        var dispatchTable: OpaquePointer
-
-        private init(_ dispatchTable: OpaquePointer) {
-            self.dispatchTable = dispatchTable
-        }
-
-        static let md5 = DigestType(CCryptoBoringSSL_EVP_md5())
-
-        static let sha1 = DigestType(CCryptoBoringSSL_EVP_sha1())
-
-        static let sha256 = DigestType(CCryptoBoringSSL_EVP_sha256())
-
-        static let sha384 = DigestType(CCryptoBoringSSL_EVP_sha384())
-
-        static let sha512 = DigestType(CCryptoBoringSSL_EVP_sha512())
+        withUnsafeMutablePointer(to: &self.context) { $0.zeroize() }
     }
 }
 #endif // CRYPTO_IN_SWIFTPM && !CRYPTO_IN_SWIFTPM_FORCE_BUILD_API