add proxy support

Author: tanner0101

.gitignore

@@ -1,3 +1,4 @@
 .build
 Package.resolved
 *.xcodeproj
+DerivedData

Sources/NIOHTTPClient/HTTPClientProxyHandler.swift

@@ -0,0 +1,114 @@
+import Foundation
+import NIO
+import NIOHTTP1
+
+/// Specifies the remote address of an HTTP proxy.
+///
+/// Adding an `HTTPClientProxy` to your client's `HTTPClientConfiguration`
+/// will cause requests to be passed through the specified proxy using the
+/// HTTP `CONNECT` method.
+///
+/// If a `TLSConfiguration` is used in conjunction with `HTTPClientProxy`,
+/// TLS will be established _after_ successful proxy, between your client
+/// and the destination server.
+public struct HTTPClientProxy {
+    internal let host: String
+    internal let port: Int
+
+    public static func server(host: String, port: Int) -> HTTPClientProxy {
+        return .init(host: host, port: port)
+    }
+}
+
+internal final class HTTPClientProxyHandler: ChannelDuplexHandler, RemovableChannelHandler {
+    typealias InboundIn = HTTPClientResponsePart
+    typealias OutboundIn = HTTPClientRequestPart
+    typealias OutboundOut = HTTPClientRequestPart
+
+    enum BufferItem {
+        case write(NIOAny, EventLoopPromise<Void>?)
+        case flush
+    }
+
+    private let host: String
+    private let port: Int
+    private var onConnect: (Channel) -> EventLoopFuture<Void>
+    private var buffer: [BufferItem]
+
+    init(host: String, port: Int, onConnect: @escaping (Channel) -> EventLoopFuture<Void>) {
+        self.host = host
+        self.port = port
+        self.onConnect = onConnect
+        self.buffer = []
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let res = self.unwrapInboundIn(data)
+        switch res {
+        case .head(let head):
+            switch head.status.code {
+            case 200..<300:
+                // Any 2xx (Successful) response indicates that the sender (and all
+                // inbound proxies) will switch to tunnel mode immediately after the
+                // blank line that concludes the successful response's header section
+                break
+            default:
+                // Any response other than a successful response
+                // indicates that the tunnel has not yet been formed and that the
+                // connection remains governed by HTTP.
+                context.fireErrorCaught(HTTPClientErrors.InvalidProxyResponseError())
+            }
+        case .end:
+            _ = self.handleConnect(context: context)
+        case .body:
+            break
+        }
+    }
+
+    func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
+        self.buffer.append(.write(data, promise))
+    }
+
+    func flush(context: ChannelHandlerContext) {
+        self.buffer.append(.flush)
+    }
+
+    func channelActive(context: ChannelHandlerContext) {
+        self.sendConnect(context: context)
+        context.fireChannelActive()
+    }
+
+    // MARK: Private
+
+    private func handleConnect(context: ChannelHandlerContext) -> EventLoopFuture<Void> {
+        return self.onConnect(context.channel).flatMap {
+            while self.buffer.count > 0 {
+                // make a copy of the current buffer and clear it in case any
+                // calls to context.write cause more requests to be buffered
+                let buffer = self.buffer
+                self.buffer = []
+                buffer.forEach { item in
+                    switch item {
+                    case .flush:
+                        context.flush()
+                    case .write(let data, let promise):
+                        context.write(data, promise: promise)
+                    }
+                }
+            }
+            return context.pipeline.removeHandler(self)
+        }
+    }
+
+    private func sendConnect(context: ChannelHandlerContext) {
+        var head = HTTPRequestHead(
+            version: .init(major: 1, minor: 1),
+            method: .CONNECT,
+            uri: "\(self.host):\(self.port)"
+        )
+        head.headers.add(name: "proxy-connection", value: "keep-alive")
+        context.write(self.wrapOutboundOut(.head(head)), promise: nil)
+        context.write(self.wrapOutboundOut(.end(nil)), promise: nil)
+        context.flush()
+    }
+}

Sources/NIOHTTPClient/HTTPHandler.swift

@@ -46,6 +46,9 @@ public struct HTTPClientErrors {
 
     public struct CancelledError : HTTPClientError {
     }
+
+    public struct InvalidProxyResponseError : HTTPClientError {
+    }
 }
 
 public enum HTTPBody : Equatable {

Sources/NIOHTTPClient/SwiftNIOHTTP.swift

@@ -37,17 +37,20 @@ public struct HTTPClientConfiguration {
     public var tlsConfiguration: TLSConfiguration?
     public var followRedirects: Bool
     public var timeout: Timeout
+    public var proxy: HTTPClientProxy?
 
-    public init(tlsConfiguration: TLSConfiguration? = nil, followRedirects: Bool = false, timeout: Timeout = Timeout()) {
+    public init(tlsConfiguration: TLSConfiguration? = nil, followRedirects: Bool = false, timeout: Timeout = Timeout(), proxy: HTTPClientProxy? = nil) {
         self.tlsConfiguration = tlsConfiguration
         self.followRedirects = followRedirects
         self.timeout = timeout
+        self.proxy = proxy
     }
 
-    public init(certificateVerification: CertificateVerification, followRedirects: Bool = false, timeout: Timeout = Timeout()) {
+    public init(certificateVerification: CertificateVerification, followRedirects: Bool = false, timeout: Timeout = Timeout(), proxy: HTTPClientProxy? = nil) {
         self.tlsConfiguration = TLSConfiguration.forClient(certificateVerification: certificateVerification)
         self.followRedirects = followRedirects
         self.timeout = timeout
+        self.proxy = proxy
     }
 }
 
@@ -151,16 +154,24 @@ public class HTTPClient {
         var bootstrap = ClientBootstrap(group: group)
             .channelOption(ChannelOptions.socket(SocketOptionLevel(IPPROTO_TCP), TCP_NODELAY), value: 1)
             .channelInitializer { channel in
-                channel.pipeline.addHTTPClientHandlers().flatMap {
-                    self.configureSSL(channel: channel, useTLS: request.useTLS, hostname: request.host)
+                let encoder = HTTPRequestEncoder()
+                let decoder = ByteToMessageHandler(HTTPResponseDecoder(leftOverBytesStrategy: .forwardBytes))
+                return channel.pipeline.addHandlers([encoder, decoder], position: .first).flatMap {
+                    switch self.configuration.proxy {
+                    case .none:
+                        return channel.pipeline.addSSLHandlerIfNeeded(for: request, tlsConfiguration: self.configuration.tlsConfiguration)
+                    case .some:
+                        return channel.pipeline.addProxyHandler(for: request, decoder: decoder, encoder: encoder, tlsConfiguration: self.configuration.tlsConfiguration)
+                    }
                 }.flatMap {
                     if let readTimeout = timeout.read {
                         return channel.pipeline.addHandler(IdleStateHandler(readTimeout: readTimeout))
                     } else {
                         return channel.eventLoop.makeSucceededFuture(())
                     }
                 }.flatMap {
-                    channel.pipeline.addHandler(HTTPTaskHandler(delegate: delegate, promise: promise, redirectHandler: redirectHandler))
+                    let taskHandler = HTTPTaskHandler(delegate: delegate, promise: promise, redirectHandler: redirectHandler)
+                    return channel.pipeline.addHandler(taskHandler)
                 }
             }
 
@@ -170,7 +181,19 @@ public class HTTPClient {
 
         let task = HTTPTask(future: promise.futureResult)
 
-        bootstrap.connect(host: request.host, port: request.port)
+        let host: String
+        let port: Int
+
+        switch self.configuration.proxy {
+        case .none:
+            host = request.host
+            port = request.port
+        case .some(let proxy):
+            host = proxy.host
+            port = proxy.port
+        }
+
+        bootstrap.connect(host: host, port: port)
             .map { channel in
                 task.setChannel(channel)
             }
@@ -183,19 +206,35 @@ public class HTTPClient {
 
         return task
     }
+}
 
-    private func configureSSL(channel: Channel, useTLS: Bool, hostname: String) -> EventLoopFuture<Void> {
-        if useTLS {
-            do {
-                let tlsConfiguration = self.configuration.tlsConfiguration ?? TLSConfiguration.forClient()
-                let context = try NIOSSLContext(configuration: tlsConfiguration)
-                return channel.pipeline.addHandler(try NIOSSLClientHandler(context: context, serverHostname: hostname),
-                                                   position: .first)
-            } catch {
-                return channel.eventLoop.makeFailedFuture(error)
+private extension ChannelPipeline {
+    func addProxyHandler(for request: HTTPRequest, decoder: ByteToMessageHandler<HTTPResponseDecoder>, encoder: HTTPRequestEncoder, tlsConfiguration: TLSConfiguration?) -> EventLoopFuture<Void> {
+        let handler = HTTPClientProxyHandler(host: request.host, port: request.port, onConnect: { channel in
+            return channel.pipeline.removeHandler(decoder).flatMap {
+                return channel.pipeline.addHandler(
+                    ByteToMessageHandler(HTTPResponseDecoder(leftOverBytesStrategy: .forwardBytes)),
+                    position: .after(encoder)
+                )
+            }.flatMap {
+                return channel.pipeline.addSSLHandlerIfNeeded(for: request, tlsConfiguration: tlsConfiguration)
             }
-        } else {
-            return channel.eventLoop.makeSucceededFuture(())
+        })
+        return self.addHandler(handler)
+    }
+
+    func addSSLHandlerIfNeeded(for request: HTTPRequest, tlsConfiguration: TLSConfiguration?) -> EventLoopFuture<Void> {
+        guard request.useTLS else {
+            return self.eventLoop.makeSucceededFuture(())
+        }
+
+        do {
+            let tlsConfiguration = tlsConfiguration ?? TLSConfiguration.forClient()
+            let context = try NIOSSLContext(configuration: tlsConfiguration)
+            return self.addHandler(try NIOSSLClientHandler(context: context, serverHostname: request.host),
+                                   position: .first)
+        } catch {
+            return self.eventLoop.makeFailedFuture(error)
         }
     }
 

Tests/NIOHTTPClientTests/HTTPClientTestUtils.swift

@@ -87,17 +87,27 @@ internal class HttpBin {
         return self.serverChannel.localAddress!
     }
 
-    init(ssl: Bool = false) {
+    static func configureTLS(channel: Channel) -> EventLoopFuture<Void> {
+        let configuration = TLSConfiguration.forServer(certificateChain: [.certificate(try! NIOSSLCertificate(buffer: cert.utf8.map(Int8.init), format: .pem))],
+                                                       privateKey: .privateKey(try! NIOSSLPrivateKey(buffer: key.utf8.map(Int8.init), format: .pem)))
+        let context = try! NIOSSLContext(configuration: configuration)
+        return channel.pipeline.addHandler(try! NIOSSLServerHandler(context: context), position: .first)
+    }
+
+    init(ssl: Bool = false, simulateProxy: HTTPProxySimulator.Option? = nil) {
         self.serverChannel = try! ServerBootstrap(group: group)
             .serverChannelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEADDR), value: 1)
             .childChannelOption(ChannelOptions.socket(IPPROTO_TCP, TCP_NODELAY), value: 1)
             .childChannelInitializer { channel in
-                channel.pipeline.configureHTTPServerPipeline(withPipeliningAssistance: true, withErrorHandling: true).flatMap {
+                return channel.pipeline.configureHTTPServerPipeline(withPipeliningAssistance: true, withErrorHandling: true).flatMap {
+                    if let simulateProxy = simulateProxy {
+                        return channel.pipeline.addHandler(HTTPProxySimulator(option: simulateProxy), position: .first)
+                    } else {
+                        return channel.eventLoop.makeSucceededFuture(())
+                    }
+                }.flatMap {
                     if ssl {
-                        let configuration = TLSConfiguration.forServer(certificateChain: [.certificate(try! NIOSSLCertificate(buffer: cert.utf8.map(Int8.init), format: .pem))],
-                                                                       privateKey: .privateKey(try! NIOSSLPrivateKey(buffer: key.utf8.map(Int8.init), format: .pem)))
-                        let context = try! NIOSSLContext(configuration: configuration)
-                        return channel.pipeline.addHandler(try! NIOSSLServerHandler(context: context), position: .first).flatMap {
+                        return HttpBin.configureTLS(channel: channel).flatMap {
                             channel.pipeline.addHandler(HttpBinHandler())
                         }
                     } else {
@@ -113,6 +123,48 @@ internal class HttpBin {
 
 }
 
+final class HTTPProxySimulator: ChannelInboundHandler, RemovableChannelHandler {
+    typealias InboundIn = ByteBuffer
+    typealias InboundOut = ByteBuffer
+    typealias OutboundOut = ByteBuffer
+
+    enum Option {
+        case plaintext
+        case tls
+    }
+
+    let option: Option
+
+    init(option: Option) {
+        self.option = option
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let response = """
+        HTTP/1.1 200 OK\r\n\
+        Content-Length: 0\r\n\
+        Connection: close\r\n\
+        \r\n
+        """
+        var buffer = self.unwrapInboundIn(data)
+        let request = buffer.readString(length: buffer.readableBytes)!
+        if request.hasPrefix("CONNECT") {
+            var buffer = context.channel.allocator.buffer(capacity: 0)
+            buffer.writeString(response)
+            context.write(self.wrapInboundOut(buffer), promise: nil)
+            context.flush()
+            context.channel.pipeline.removeHandler(self, promise: nil)
+            switch self.option {
+            case .tls:
+                _ = HttpBin.configureTLS(channel: context.channel)
+            case .plaintext: break
+            }
+        } else {
+            fatalError("Expected a CONNECT request")
+        }
+    }
+}
+
 internal struct HTTPResponseBuilder {
     let head: HTTPResponseHead
     var body: ByteBuffer?

Tests/NIOHTTPClientTests/SwiftNIOHTTPTests+XCTest.swift

@@ -39,6 +39,8 @@ extension SwiftHTTPTests {
                 ("testRemoteClose", testRemoteClose),
                 ("testReadTimeout", testReadTimeout),
                 ("testCancel", testCancel),
+                ("testProxyPlaintext", testProxyPlaintext),
+                ("testProxyTLS", testProxyTLS),
            ]
    }
 }

Tests/NIOHTTPClientTests/SwiftNIOHTTPTests.swift

@@ -267,4 +267,35 @@ class SwiftHTTPTests: XCTestCase {
             XCTFail("Unexpected error: \(error)")
         }
     }
+
+    func testProxyPlaintext() throws {
+        let httpBin = HttpBin(simulateProxy: .plaintext)
+        let httpClient = HTTPClient(
+            eventLoopGroupProvider: .createNew,
+            configuration: .init(proxy: .server(host: "localhost", port: httpBin.port))
+        )
+        defer {
+            try! httpClient.syncShutdown()
+            httpBin.shutdown()
+        }
+        let res = try httpClient.get(url: "http://test/ok").wait()
+        XCTAssertEqual(res.status, .ok)
+    }
+
+    func testProxyTLS() throws {
+        let httpBin = HttpBin(simulateProxy: .tls)
+        let httpClient = HTTPClient(
+            eventLoopGroupProvider: .createNew,
+            configuration: .init(
+                certificateVerification: .none,
+                proxy: .server(host: "localhost", port: httpBin.port)
+            )
+        )
+        defer {
+            try! httpClient.syncShutdown()
+            httpBin.shutdown()
+        }
+        let res = try httpClient.get(url: "https://test/ok").wait()
+        XCTAssertEqual(res.status, .ok)
+    }
 }