Fix incorrect callback caching (#541)

Motivation:

The certificate override callback incorrectly cached results on the
SSLContext, instead of the SSLConnection. The result is that it would
only fire once for a given SSLContext, instead of on every connection,
incorrectly caching the result indefinitely.

Modifications:

Move the callback state to SSLConnection
Write tests to validate that the callback is called the right
    number of times.

Result:

Better behaviour.

Author: Lukasa

Sources/NIOSSL/NIOSSLHandler.swift

@@ -411,7 +411,7 @@ public class NIOSSLHandler: ChannelInboundHandler, ChannelOutboundHandler, Remov
             }
 
             // If there's a failed custom context operation, we fire both errors.
-            if let customContextError = self.connection.parentContext.customContextManager?.loadContextError {
+            if let customContextError = self.connection.customContextManager?.loadContextError {
                 context.fireErrorCaught(customContextError)
             }
 

Sources/NIOSSL/SSLCallbacks.swift

@@ -342,7 +342,7 @@ extension CustomContextManager {
                 // Ensure we execute any completion on the next event loop tick
                 // This ensures that we suspend before calling resume
                 eventLoop.assumeIsolated().execute {
-                    connection.parentContext.customContextManager?.state = .complete(result)
+                    connection.customContextManager?.state = .complete(result)
                     connection.parentHandler?.resumeHandshake()
                 }
             }

Sources/NIOSSL/SSLConnection.swift

@@ -58,6 +58,7 @@ internal final class SSLConnection {
     private var verificationCallback: NIOSSLVerificationCallback?
     internal var customVerificationManager: CustomVerifyManager?
     internal var customPrivateKeyResult: Result<ByteBuffer, Error>?
+    internal var customContextManager: CustomContextManager?
 
     /// Whether certificate hostnames should be validated.
     var validateHostnames: Bool {
@@ -71,6 +72,10 @@ internal final class SSLConnection {
         self.ssl = ownedSSL
         self.parentContext = parentContext
 
+        if let customContextCallback = parentContext.configuration.sslContextCallback {
+            self.customContextManager = CustomContextManager(callback: customContextCallback)
+        }
+
         // We pass the SSL object an unowned reference to this object.
         let pointerToSelf = Unmanaged.passUnretained(self).toOpaque()
         CNIOBoringSSL_SSL_set_ex_data(self.ssl, sslConnectionExDataIndex, pointerToSelf)

Sources/NIOSSL/SSLContext.swift

@@ -250,7 +250,7 @@ private func sslContextCallback(ssl: OpaquePointer?, arg: UnsafeMutableRawPointe
         )
     }
 
-    let parentSwiftContext = NIOSSLContext.lookupFromRawContext(ssl: ssl)
+    let parentSwiftContext = SSLConnection.loadConnectionFromSSL(ssl)
 
     // This is a safe force unwrap as this callback is only register directly after setting the manager
     var contextManager = parentSwiftContext.customContextManager!
@@ -295,7 +295,6 @@ public final class NIOSSLContext {
     fileprivate let sslContext: OpaquePointer
     private let callbackManager: CallbackManagerProtocol?
     private var keyLogManager: KeyLogCallbackManager?
-    internal var customContextManager: CustomContextManager?
     internal var pskClientConfigurationCallback: _NIOPSKClientIdentityProvider?
     internal var pskServerConfigurationCallback: _NIOPSKServerIdentityProvider?
     internal let configuration: TLSConfiguration
@@ -372,8 +371,8 @@ public final class NIOSSLContext {
         }
 
         // Set the SSL Context Configuration callback.
-        if let sslContextConfigurationCallback = configuration.sslContextCallback {
-            self.customContextManager = CustomContextManager(callback: sslContextConfigurationCallback)
+        // The state is managed on the connection.
+        if configuration.sslContextCallback != nil {
             CNIOBoringSSL_SSL_CTX_set_cert_cb(context, sslContextCallback, nil)
         }
 

Tests/NIOSSLTests/TLSConfigurationTest.swift

@@ -242,7 +242,24 @@ class TLSConfigurationTest: XCTestCase {
     ) throws {
         let clientContext = try assertNoThrowWithValue(NIOSSLContext(configuration: clientConfig))
         let serverContext = try assertNoThrowWithValue(NIOSSLContext(configuration: serverConfig))
+        try self.assertHandshakeSucceededInMemory(
+            withClientContext: clientContext,
+            andServerContext: serverContext,
+            file: file,
+            line: line
+        )
+    }
 
+    /// Performs a connection in memory and validates that the handshake was successful.
+    ///
+    /// - NOTE: This function should only be used when you know that there is no custom verification
+    /// callback in use, otherwise it will not be thread-safe.
+    func assertHandshakeSucceededInMemory(
+        withClientContext clientContext: NIOSSLContext,
+        andServerContext serverContext: NIOSSLContext,
+        file: StaticString = #filePath,
+        line: UInt = #line
+    ) throws {
         let serverChannel = EmbeddedChannel()
         let clientChannel = EmbeddedChannel()
 
@@ -294,7 +311,23 @@ class TLSConfigurationTest: XCTestCase {
             file: file,
             line: line
         )
+        try self.assertHandshakeSucceededEventLoop(
+            withClientContext: clientContext,
+            andServerContext: serverContext,
+            file: file,
+            line: line
+        )
+    }
 
+    /// Performs a connection using a real event loop and validates that the handshake was successful.
+    ///
+    /// This function is thread-safe in the presence of custom verification callbacks.
+    func assertHandshakeSucceededEventLoop(
+        withClientContext clientContext: NIOSSLContext,
+        andServerContext serverContext: NIOSSLContext,
+        file: StaticString = #filePath,
+        line: UInt = #line
+    ) throws {
         let group = MultiThreadedEventLoopGroup(numberOfThreads: 1)
         defer {
             XCTAssertNoThrow(try group.syncShutdownGracefully())
@@ -358,6 +391,31 @@ class TLSConfigurationTest: XCTestCase {
         #endif
     }
 
+    func assertHandshakeSucceeded(
+        withClientContext clientContext: NIOSSLContext,
+        andServerContext serverContext: NIOSSLContext,
+        file: StaticString = #filePath,
+        line: UInt = #line
+    ) throws {
+        // The only use of a custom callback is on Darwin...
+        #if os(Linux)
+        return try self.assertHandshakeSucceededInMemory(
+            withClientContext: clientContext,
+            andServerContext: serverContext,
+            file: file,
+            line: line
+        )
+
+        #else
+        return try self.assertHandshakeSucceededEventLoop(
+            withClientContext: clientContext,
+            andServerContext: serverContext,
+            file: file,
+            line: line
+        )
+        #endif
+    }
+
     func setupTLSLeafandClientIdentitiesFromCustomCARoot() throws -> (
         leafCert: NIOSSLCertificate, leafKey: NIOSSLPrivateKey,
         clientCert: NIOSSLCertificate, clientKey: NIOSSLPrivateKey
@@ -2062,6 +2120,73 @@ class TLSConfigurationTest: XCTestCase {
             errorTextContains: "TLSV1_ALERT_INTERNAL_ERROR"
         )
     }
+
+    func testClientSideCertSelection_eachConnectionSelectsAgain() throws {
+        let callbackCount = NIOLockedValueBox(0)
+        var clientConfig = TLSConfiguration.makeClientConfiguration()
+        clientConfig.certificateVerification = .noHostnameVerification
+        clientConfig.trustRoots = .certificates([TLSConfigurationTest.cert1])
+        clientConfig.sslContextCallback = { _, promise in
+            callbackCount.withLockedValue { $0 += 1 }
+
+            var `override` = NIOSSLContextConfigurationOverride()
+            override.certificateChain = [.certificate(TLSConfigurationTest.cert2)]
+            override.privateKey = .privateKey(TLSConfigurationTest.key2)
+            promise.succeed(override)
+        }
+
+        var serverConfig = TLSConfiguration.makeServerConfiguration(
+            certificateChain: [.certificate(TLSConfigurationTest.cert1)],
+            privateKey: .privateKey(TLSConfigurationTest.key1)
+        )
+        serverConfig.certificateVerification = .noHostnameVerification
+        serverConfig.trustRoots = .certificates([TLSConfigurationTest.cert2])
+
+        let clientContext = try assertNoThrowWithValue(
+            NIOSSLContext(configuration: clientConfig)
+        )
+        let serverContext = try assertNoThrowWithValue(
+            NIOSSLContext(configuration: serverConfig)
+        )
+
+        for _ in 0..<5 {
+            try assertHandshakeSucceeded(withClientContext: clientContext, andServerContext: serverContext)
+        }
+
+        XCTAssertEqual(callbackCount.withLockedValue { $0 }, 5)
+    }
+
+    func testServerSideCertSelection_eachConnectionSelectsAgain() throws {
+        let callbackCount = NIOLockedValueBox(0)
+        var clientConfig = TLSConfiguration.makeClientConfiguration()
+        clientConfig.certificateVerification = .noHostnameVerification
+        clientConfig.trustRoots = .certificates([TLSConfigurationTest.cert1])
+
+        var serverConfig = TLSConfiguration.makeServerConfiguration(
+            certificateChain: [.certificate(TLSConfigurationTest.cert2)],
+            privateKey: .privateKey(TLSConfigurationTest.key2)
+        )
+        serverConfig.sslContextCallback = { _, promise in
+            var `override` = NIOSSLContextConfigurationOverride()
+            override.certificateChain = [.certificate(TLSConfigurationTest.cert1)]
+            override.privateKey = .privateKey(TLSConfigurationTest.key1)
+            callbackCount.withLockedValue { $0 += 1 }
+            promise.succeed(override)
+        }
+
+        let clientContext = try assertNoThrowWithValue(
+            NIOSSLContext(configuration: clientConfig)
+        )
+        let serverContext = try assertNoThrowWithValue(
+            NIOSSLContext(configuration: serverConfig)
+        )
+
+        for _ in 0..<5 {
+            try assertHandshakeSucceeded(withClientContext: clientContext, andServerContext: serverContext)
+        }
+
+        XCTAssertEqual(callbackCount.withLockedValue { $0 }, 5)
+    }
 }
 
 extension EmbeddedChannel {