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

@@ -18,6 +18,7 @@ import NIO
 import NIOConcurrencyHelpers
 import NIOHTTP1
 import NIOHTTPCompression
+import NIOSOCKS
 import NIOSSL
 import NIOTLS
 import NIOTransportServices
@@ -900,6 +901,13 @@ extension ChannelPipeline {
         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,
                                        handshakePromise: EventLoopPromise<Void>) {

Sources/AsyncHTTPClient/HTTPClientProxyHandler.swift

@@ -14,8 +14,9 @@
 
 import NIO
 import NIOHTTP1
+import NIOSOCKS
 
-public extension HTTPClient.Configuration {
+extension HTTPClient.Configuration {
     /// Proxy server configuration
     /// Specifies the remote address of an HTTP proxy.
     ///
@@ -26,21 +27,42 @@ public extension HTTPClient.Configuration {
     /// If a `TLSConfiguration` is used in conjunction with `HTTPClient.Configuration.Proxy`,
     /// TLS will be established _after_ successful proxy, between your client
     /// and the destination server.
-    struct Proxy {
+    public struct Proxy {
+        enum ProxyType: Hashable {
+            case http(HTTPClient.Authorization?)
+            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?
+        public var authorization: HTTPClient.Authorization? {
+            set {
+                precondition(self.type == .http(self.authorization), "Only HTTP proxies can have authorization set.")
+                self.type = .http(newValue)
+            }
+
+            get {
+                switch self.type {
+                case .http(let authorization):
+                    return authorization
+                case .socks:
+                    return nil
+                }
+            }
+        }
+
+        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, type: .http(nil))
         }
 
         /// Create proxy.
@@ -50,7 +72,15 @@ 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, type: .http(authorization))
+        }
+
+        /// Create a SOCKSv5 proxy.
+        /// - parameter host: The SOCKSv5 proxy address.
+        /// - parameter port: The SOCKSv5 proxy port, defaults to 1080.
+        /// - returns: A new instance of `Proxy` configured to connect to a `SOCKSv5` server.
+        public static func socksServer(host: String, port: Int = 1080) -> Proxy {
+            return .init(host: host, port: port, type: .socks)
         }
     }
 }

Sources/AsyncHTTPClient/HTTPHandler.swift

@@ -342,8 +342,8 @@ extension HTTPClient {
     }
 
     /// HTTP authentication
-    public struct Authorization {
-        private enum Scheme {
+    public struct Authorization: Hashable {
+        private enum Scheme: Hashable {
             case Basic(String)
             case Bearer(String)
         }

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+XCTest.swift

@@ -58,6 +58,10 @@ extension HTTPClientTests {
             ("testProxyTLS", testProxyTLS),
             ("testProxyPlaintextWithCorrectlyAuthorization", testProxyPlaintextWithCorrectlyAuthorization),
             ("testProxyPlaintextWithIncorrectlyAuthorization", testProxyPlaintextWithIncorrectlyAuthorization),
+            ("testProxySOCKS", testProxySOCKS),
+            ("testProxySOCKSFailureNoServer", testProxySOCKSFailureNoServer),
+            ("testProxySOCKSFailureInvalidServer", testProxySOCKSFailureInvalidServer),
+            ("testProxySOCKSMisbehavingServer", testProxySOCKSMisbehavingServer),
             ("testUploadStreaming", testUploadStreaming),
             ("testNoContentLengthForSSLUncleanShutdown", testNoContentLengthForSSLUncleanShutdown),
             ("testNoContentLengthWithIgnoreErrorForSSLUncleanShutdown", testNoContentLengthWithIgnoreErrorForSSLUncleanShutdown),

Tests/AsyncHTTPClientTests/HTTPClientTests.swift

@@ -22,6 +22,7 @@ import NIOConcurrencyHelpers
 import NIOFoundationCompat
 import NIOHTTP1
 import NIOHTTPCompression
+import NIOSOCKS
 import NIOSSL
 import NIOTestUtils
 import NIOTransportServices
@@ -709,6 +710,59 @@ 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")))
+
+        defer {
+            XCTAssertNoThrow(try localClient.syncShutdown())
+            XCTAssertNoThrow(try socksBin.shutdown())
+        }
+
+        var response: HTTPClient.Response?
+        XCTAssertNoThrow(response = try localClient.get(url: "http://127.0.0.1/socks/test").wait())
+        XCTAssertEqual(.ok, response?.status)
+        XCTAssertEqual(ByteBuffer(string: "it works!"), response?.body)
+    }
+
+    // there is no socks server, so we should fail
+    func testProxySOCKSFailureNoServer() throws {
+        let localHTTPBin = HTTPBin()
+        let localClient = HTTPClient(eventLoopGroupProvider: .shared(self.clientGroup),
+                                     configuration: .init(proxy: .socksServer(host: "127.0.0.1", port: localHTTPBin.port)))
+        defer {
+            XCTAssertNoThrow(try localClient.syncShutdown())
+            XCTAssertNoThrow(try localHTTPBin.shutdown())
+        }
+        XCTAssertThrowsError(try localClient.get(url: "http://127.0.0.1/socks/test").wait())
+    }
+
+    // speak to a server that doesn't speak SOCKS
+    func testProxySOCKSFailureInvalidServer() throws {
+        let localClient = HTTPClient(eventLoopGroupProvider: .shared(self.clientGroup),
+                                     configuration: .init(proxy: .socksServer(host: "127.0.0.1")))
+        defer {
+            XCTAssertNoThrow(try localClient.syncShutdown())
+        }
+        XCTAssertThrowsError(try localClient.get(url: "http://127.0.0.1/socks/test").wait())
+    }
+
+    // test a handshake failure with a misbehaving server
+    func testProxySOCKSMisbehavingServer() throws {
+        let socksBin = try MockSOCKSServer(expectedURL: "/socks/test", expectedResponse: "it works!", misbehave: true)
+        let localClient = HTTPClient(eventLoopGroupProvider: .shared(self.clientGroup),
+                                     configuration: .init(proxy: .socksServer(host: "127.0.0.1")))
+
+        defer {
+            XCTAssertNoThrow(try localClient.syncShutdown())
+            XCTAssertNoThrow(try socksBin.shutdown())
+        }
+
+        // the server will send a bogus message in response to the clients request
+        XCTAssertThrowsError(try localClient.get(url: "http://127.0.0.1/socks/test").wait())
+    }
+
     func testUploadStreaming() throws {
         let body: HTTPClient.Body = .stream(length: 8) { writer in
             let buffer = ByteBuffer(string: "1234")

Tests/AsyncHTTPClientTests/SOCKSTestUtils.swift

@@ -0,0 +1,130 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the AsyncHTTPClient open source project
+//
+// Copyright (c) 2021 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
+//
+//===----------------------------------------------------------------------===//
+
+import AsyncHTTPClient
+import NIO
+import NIOHTTP1
+import NIOSOCKS
+import XCTest
+
+struct MockSOCKSError: Error, Hashable {
+    var description: String
+}
+
+class MockSOCKSServer {
+    let channel: Channel
+
+    public init(expectedURL: String, expectedResponse: String, misbehave: Bool = false, file: String = #file, line: UInt = #line) throws {
+        let elg = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        let bootstrap = ServerBootstrap(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, misbehave: misbehave),
+                    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
+    var requestCount = 0
+
+    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):
+            guard self.requestCount == 0 else {
+                return
+            }
+            XCTAssertEqual(head.uri, self.expectedURL)
+            self.requestCount += 1
+        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)
+        }
+    }
+
+    func errorCaught(context: ChannelHandlerContext, error: Error) {
+        context.fireErrorCaught(error)
+        context.close(promise: nil)
+    }
+}
+
+class SOCKSTestHandler: ChannelInboundHandler, RemovableChannelHandler {
+    typealias InboundIn = ClientMessage
+
+    let handshakeHandler: SOCKSServerHandshakeHandler
+    let misbehave: Bool
+
+    init(handshakeHandler: SOCKSServerHandshakeHandler, misbehave: Bool) {
+        self.handshakeHandler = handshakeHandler
+        self.misbehave = misbehave
+    }
+
+    func channelRead(context: ChannelHandlerContext, data: NIOAny) {
+        guard context.channel.isActive else {
+            return
+        }
+
+        let message = self.unwrapInboundIn(data)
+        switch message {
+        case .greeting:
+            context.writeAndFlush(.init(
+                ServerMessage.selectedAuthenticationMethod(.init(method: .noneRequired))), promise: nil)
+        case .authenticationData:
+            context.fireErrorCaught(MockSOCKSError(description: "Received authentication data but didn't receive any."))
+        case .request(let request):
+            guard !self.misbehave else {
+                context.writeAndFlush(
+                    .init(ServerMessage.authenticationData(context.channel.allocator.buffer(string: "bad server!"), complete: true)), promise: nil
+                )
+                return
+            }
+            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)
+            }
+        }
+    }
+}