Fix leaked TLS handshake promise (#180) (#180)

motivation:
TLS handshake promise was leaked in some cases of failure (see #179)

changes:
- Avoid leaking promise
- Clearer completion flow for related futures
- Add testAvoidLeakingTLSHandshakeCompletionPromise test

Author: Trevör

Sources/AsyncHTTPClient/ConnectionPool.swift

@@ -383,32 +383,33 @@ final class ConnectionPool {
             }
 
             return channel.flatMap { channel -> EventLoopFuture<ConnectionPool.Connection> in
-                channel.pipeline.addSSLHandlerIfNeeded(for: self.key, tlsConfiguration: self.configuration.tlsConfiguration, handshakePromise: handshakePromise).flatMap {
+                channel.pipeline.addSSLHandlerIfNeeded(for: self.key, tlsConfiguration: self.configuration.tlsConfiguration, handshakePromise: handshakePromise)
+                return handshakePromise.futureResult.flatMap {
                     channel.pipeline.addHTTPClientHandlers(leftOverBytesStrategy: .forwardBytes)
                 }.map {
                     let connection = Connection(key: self.key, channel: channel, parentPool: self.parentPool)
                     connection.isLeased = true
                     return connection
                 }
-            }.flatMap { connection in
-                handshakePromise.futureResult.map {
-                    self.configureCloseCallback(of: connection)
-                    return connection
-                }.flatMapError { error in
-                    connection.closePromise.succeed(())
-                    let action = self.parentPool.connectionProvidersLock.withLock {
-                        self.stateLock.withLock {
-                            self.state.failedConnectionAction()
-                        }
-                    }
-                    switch action {
-                    case .makeConnectionAndComplete(let el, let promise):
-                        self.makeConnection(on: el).cascade(to: promise)
-                    case .none:
-                        break
+            }.map { connection in
+                self.configureCloseCallback(of: connection)
+                return connection
+            }.flatMapError { error in
+                // This promise may not have been completed if we reach this
+                // so we fail it to avoid any leak
+                handshakePromise.fail(error)
+                let action = self.parentPool.connectionProvidersLock.withLock {
+                    self.stateLock.withLock {
+                        self.state.failedConnectionAction()
                     }
-                    return self.eventLoop.makeFailedFuture(error)
                 }
+                switch action {
+                case .makeConnectionAndComplete(let el, let promise):
+                    self.makeConnection(on: el).cascade(to: promise)
+                case .none:
+                    break
+                }
+                return self.eventLoop.makeFailedFuture(error)
             }
         }
 

Sources/AsyncHTTPClient/HTTPClient.swift

@@ -625,10 +625,10 @@ extension ChannelPipeline {
         return addHandlers([encoder, decoder, handler])
     }
 
-    func addSSLHandlerIfNeeded(for key: ConnectionPool.Key, tlsConfiguration: TLSConfiguration?, handshakePromise: EventLoopPromise<Void>) -> EventLoopFuture<Void> {
+    func addSSLHandlerIfNeeded(for key: ConnectionPool.Key, tlsConfiguration: TLSConfiguration?, handshakePromise: EventLoopPromise<Void>) {
         guard key.scheme == .https else {
             handshakePromise.succeed(())
-            return self.eventLoop.makeSucceededFuture(())
+            return
         }
 
         do {
@@ -638,10 +638,9 @@ extension ChannelPipeline {
                 try NIOSSLClientHandler(context: context, serverHostname: key.host.isIPAddress ? nil : key.host),
                 TLSEventsHandler(completionPromise: handshakePromise),
             ]
-
-            return self.addHandlers(handlers)
+            self.addHandlers(handlers).cascadeFailure(to: handshakePromise)
         } catch {
-            return self.eventLoop.makeFailedFuture(error)
+            handshakePromise.fail(error)
         }
     }
 }

Tests/AsyncHTTPClientTests/HTTPClientTests+XCTest.swift

@@ -95,6 +95,7 @@ extension HTTPClientTests {
             ("testWeRecoverFromServerThatClosesTheConnectionOnUs", testWeRecoverFromServerThatClosesTheConnectionOnUs),
             ("testPoolClosesIdleConnections", testPoolClosesIdleConnections),
             ("testRacePoolIdleConnectionsAndGet", testRacePoolIdleConnectionsAndGet),
+            ("testAvoidLeakingTLSHandshakeCompletionPromise", testAvoidLeakingTLSHandshakeCompletionPromise),
         ]
     }
 }

Tests/AsyncHTTPClientTests/HTTPClientTests.swift

@@ -1656,4 +1656,21 @@ class HTTPClientTests: XCTestCase {
             Thread.sleep(forTimeInterval: 0.01 + .random(in: -0.05...0.05))
         }
     }
+
+    func testAvoidLeakingTLSHandshakeCompletionPromise() {
+        let httpClient = HTTPClient(eventLoopGroupProvider: .createNew)
+        let httpBin = HTTPBin()
+        let port = httpBin.port
+        XCTAssertNoThrow(try httpBin.shutdown())
+        defer {
+            XCTAssertNoThrow(try httpClient.syncShutdown(requiresCleanClose: true))
+        }
+
+        XCTAssertThrowsError(try httpClient.get(url: "http://localhost:\(port)").wait()) { error in
+            guard error is NIOConnectionError else {
+                XCTFail("Unexpected error: \(error)")
+                return
+            }
+        }
+    }
 }