Implement SOCKS proxy functionality

Package.swift

@@ -21,22 +21,22 @@ let package = Package(
         .library(name: "AsyncHTTPClient", targets: ["AsyncHTTPClient"]),
     ],
     dependencies: [
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.27.0"),
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.29.0"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.13.0"),
-        .package(url: "https://github.com/apple/swift-nio-extras.git", from: "1.3.0"),
+        .package(url: "https://github.com/apple/swift-nio-extras.git", from: "1.9.0"),
         .package(url: "https://github.com/apple/swift-nio-transport-services.git", from: "1.5.1"),
         .package(url: "https://github.com/apple/swift-log.git", from: "1.4.0"),
     ],
     targets: [
         .target(
             name: "AsyncHTTPClient",
             dependencies: ["NIO", "NIOHTTP1", "NIOSSL", "NIOConcurrencyHelpers", "NIOHTTPCompression",
-                           "NIOFoundationCompat", "NIOTransportServices", "Logging"]
+                           "NIOFoundationCompat", "NIOTransportServices", "Logging", "NIOSOCKS"]
         ),
         .testTarget(
             name: "AsyncHTTPClientTests",
             dependencies: ["NIO", "NIOConcurrencyHelpers", "NIOSSL", "AsyncHTTPClient", "NIOFoundationCompat",
-                           "NIOTestUtils", "Logging"]
+                           "NIOTestUtils", "Logging", "NIOSOCKS"]
         ),
     ]
 )

Sources/AsyncHTTPClient/HTTPClient.swift

@@ -21,6 +21,7 @@ import NIOHTTPCompression
 import NIOSSL
 import NIOTLS
 import NIOTransportServices
+import NIOSOCKS
 
 extension Logger {
     private func requestInfo(_ request: HTTPClient.Request) -> Logger.Metadata.Value {
@@ -899,6 +900,13 @@ extension ChannelPipeline {
         try sync.addHandler(decoder)
         try sync.addHandler(handler)
     }
+
+    func syncAddSOCKSProxyHandler(host: String, port: Int) throws {
+        let address = try SocketAddress(ipAddress: host, port: port)
+        let handler = SOCKSClientHandler(targetAddress: .address(address))
+        let sync = self.syncOperations
+        try sync.addHandler(handler)
+    }
 
     func syncAddLateSSLHandlerIfNeeded(for key: ConnectionPool.Key,
                                        sslContext: NIOSSLContext,

Sources/AsyncHTTPClient/HTTPClientProxyHandler.swift

@@ -14,6 +14,7 @@
 
 import NIO
 import NIOHTTP1
+import NIOSOCKS
 
 public extension HTTPClient.Configuration {
     /// Proxy server configuration
@@ -27,20 +28,28 @@ public extension HTTPClient.Configuration {
     /// TLS will be established _after_ successful proxy, between your client
     /// and the destination server.
     struct Proxy {
+
+        enum ProxyType: Hashable {
+            case http
+            case socks
+        }
+
         /// Specifies Proxy server host.
         public var host: String
         /// Specifies Proxy server port.
         public var port: Int
         /// Specifies Proxy server authorization.
         public var authorization: HTTPClient.Authorization?
 
+        var type: ProxyType
+
         /// Create proxy.
         ///
         /// - parameters:
         ///     - host: proxy server host.
         ///     - port: proxy server port.
         public static func server(host: String, port: Int) -> Proxy {
-            return .init(host: host, port: port, authorization: nil)
+            return .init(host: host, port: port, authorization: nil, type: .http)
         }
 
         /// Create proxy.
@@ -50,7 +59,11 @@ public extension HTTPClient.Configuration {
         ///     - port: proxy server port.
         ///     - authorization: proxy server authorization.
         public static func server(host: String, port: Int, authorization: HTTPClient.Authorization? = nil) -> Proxy {
-            return .init(host: host, port: port, authorization: authorization)
+            return .init(host: host, port: port, authorization: authorization, type: .http)
+        }
+
+        public static func socksServer(host: String, port: Int = 1080) -> Proxy {
+            return .init(host: host, port: port, authorization: nil, type: .socks)
         }
     }
 }

Sources/AsyncHTTPClient/Utils.swift

@@ -148,9 +148,14 @@ extension NIOClientTCPBootstrap {
                 return bootstrap.channelInitializer { channel in
                     do {
                         if let proxy = configuration.proxy {
-                            try channel.pipeline.syncAddProxyHandler(host: host,
-                                                                     port: port,
-                                                                     authorization: proxy.authorization)
+                            switch proxy.type {
+                            case .http:
+                                try channel.pipeline.syncAddProxyHandler(host: host,
+                                                                         port: port,
+                                                                         authorization: proxy.authorization)
+                            case .socks:
+                                try channel.pipeline.syncAddSOCKSProxyHandler(host: host, port: port)
+                            }
                         } else if requiresTLS {
                             // We only add the handshake verifier if we need TLS and we're not going through a proxy.
                             // If we're going through a proxy we add it later (outside of this method).

Tests/AsyncHTTPClientTests/HTTPClientTests.swift

@@ -26,6 +26,7 @@ import NIOSSL
 import NIOTestUtils
 import NIOTransportServices
 import XCTest
+import NIOSOCKS
 
 class HTTPClientTests: XCTestCase {
     typealias Request = HTTPClient.Request
@@ -708,6 +709,23 @@ class HTTPClientTests: XCTestCase {
             }
         }
     }
+
+    func testProxySOCKS() throws {
+        let socksBin = try MockSocksServer(expectedURL: "/socks/test", expectedResponse: "it works!")
+        let localClient = HTTPClient(eventLoopGroupProvider: .shared(self.clientGroup),
+                                     configuration: .init(proxy: .socksServer(host: "127.0.0.1")))
+
+        do  {
+            let response = try localClient.get(url: "http://127.0.0.1/socks/test").wait()
+            XCTAssertEqual(.ok, response.status)
+            XCTAssertEqual(ByteBuffer(string: "it works!"), response.body)
+        } catch {
+            XCTFail("\(error)")
+        }
+
+        XCTAssertNoThrow(try localClient.syncShutdown())
+        XCTAssertNoThrow(try socksBin.shutdown())
+    }
 
     func testUploadStreaming() throws {
         let body: HTTPClient.Body = .stream(length: 8) { writer in

Tests/AsyncHTTPClientTests/SOCKSTestUtils.swift

@@ -0,0 +1,111 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2020 Apple Inc. and the AsyncHTTPClient project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of AsyncHTTPClient project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+@testable import AsyncHTTPClient
+import NIO
+import NIOHTTP1
+import NIOSOCKS
+import XCTest
+
+class MockSocksServer {
+
+    let channel: Channel
+
+    public init(expectedURL: String, expectedResponse: String, file: String = (#file), line: UInt = #line) throws {
+        let elg = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        let bootstrap = ServerBootstrap.init(group: elg)
+            .serverChannelOption(ChannelOptions.socket(SocketOptionLevel(SOL_SOCKET), SO_REUSEADDR), value: 1)
+            .childChannelInitializer { channel in
+            let handshakeHandler = SOCKSServerHandshakeHandler()
+            return channel.pipeline.addHandlers([
+                handshakeHandler,
+                SOCKSTestHandler(handshakeHandler: handshakeHandler),
+                SOCKSTestHTTPClient(expectedURL: expectedURL, expectedResponse: expectedResponse, file: file, line: line)
+            ])
+        }
+        self.channel = try bootstrap.bind(host: "127.0.0.1", port: 1080).wait()
+    }
+
+    func shutdown() throws {
+        try self.channel.close().wait()
+    }
+
+}
+
+class SOCKSTestHTTPClient: ChannelInboundHandler {
+
+    typealias InboundIn = HTTPServerRequestPart
+    typealias OutboundOut = HTTPServerResponsePart
+
+    let expectedURL: String
+    let expectedResponse: String
+    let file: String
+    let line: UInt
+
+    init(expectedURL: String, expectedResponse: String, file: String, line: UInt) {
+        self.expectedURL = expectedURL
+        self.expectedResponse = expectedResponse
+        self.file = file
+        self.line = line
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let message = self.unwrapInboundIn(data)
+        switch message {
+        case .head(let head):
+            XCTAssertEqual(head.uri, self.expectedURL)
+        case .body:
+            break
+        case .end:
+            context.write(self.wrapOutboundOut(.head(.init(version: .http1_1, status: .ok))), promise: nil)
+            context.write(self.wrapOutboundOut(.body(.byteBuffer(context.channel.allocator.buffer(string: self.expectedResponse)))), promise: nil)
+            context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
+        }
+    }
+}
+
+class SOCKSTestHandler: ChannelInboundHandler, RemovableChannelHandler {
+
+    typealias InboundIn = ClientMessage
+
+    let handshakeHandler: SOCKSServerHandshakeHandler
+
+    init(handshakeHandler: SOCKSServerHandshakeHandler) {
+        self.handshakeHandler = handshakeHandler
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        let message = self.unwrapInboundIn(data)
+        switch message {
+        case .greeting:
+            context.writeAndFlush(.init(
+                ServerMessage.selectedAuthenticationMethod(.init(method: .noneRequired))), promise: nil)
+            context.writeAndFlush(.init(
+                ServerMessage.authenticationData(context.channel.allocator.buffer(capacity: 0), complete: true)), promise: nil)
+        case .authenticationData:
+            break
+        case .request(let request):
+            context.writeAndFlush(.init(
+                ServerMessage.response(.init(reply: .succeeded, boundAddress: request.addressType))), promise: nil)
+            context.channel.pipeline.addHandlers([
+                ByteToMessageHandler(HTTPRequestDecoder()),
+                HTTPResponseEncoder(),
+            ], position: .after(self)).whenSuccess {
+                context.channel.pipeline.removeHandler(self, promise: nil)
+                context.channel.pipeline.removeHandler(self.handshakeHandler, promise: nil)
+            }
+        }
+    }
+
+}