From 76b52a8b711c9b8c216f20f5e58fdc0ff30b630d Mon Sep 17 00:00:00 2001
From: Max Isom <hi@maxisom.me>
Date: Fri, 17 Nov 2023 16:41:57 -0800
Subject: [PATCH] Allow SOCKS proxies to be chained

---
 packages/socks-proxy-agent/src/index.ts | 83 ++++++++++++++++++-------
 packages/socks-proxy-agent/test/test.ts | 43 +++++++++++--
 2 files changed, 100 insertions(+), 26 deletions(-)

diff --git a/packages/socks-proxy-agent/src/index.ts b/packages/socks-proxy-agent/src/index.ts
index 07359872..5dcd3b85 100644
--- a/packages/socks-proxy-agent/src/index.ts
+++ b/packages/socks-proxy-agent/src/index.ts
@@ -1,4 +1,9 @@
-import { SocksClient, SocksProxy, SocksClientOptions } from 'socks';
+import {
+	SocksClient,
+	SocksProxy,
+	SocksClientOptions,
+	SocksClientChainOptions,
+} from 'socks';
 import { Agent, AgentConnectOpts } from 'agent-base';
 import createDebug from 'debug';
 import * as dns from 'dns';
@@ -87,18 +92,31 @@ export class SocksProxyAgent extends Agent {
 		'socks5h',
 	] as const;
 
-	readonly shouldLookup: boolean;
-	readonly proxy: SocksProxy;
+	readonly shouldLookup!: boolean;
+	readonly proxies: SocksProxy[];
 	timeout: number | null;
 
-	constructor(uri: string | URL, opts?: SocksProxyAgentOptions) {
+	constructor(
+		uri: string | URL | string[] | URL[],
+		opts?: SocksProxyAgentOptions
+	) {
 		super(opts);
 
-		const url = typeof uri === 'string' ? new URL(uri) : uri;
-		const { proxy, lookup } = parseSocksURL(url);
+		const uriList = Array.isArray(uri) ? uri : [uri];
+
+		if (uriList.length === 0) {
+			throw new Error('At least one proxy server URI must be specified.');
+		}
+
+		this.proxies = [];
+		for (const [i, uri] of uriList.entries()) {
+			const { proxy, lookup } = parseSocksURL(new URL(uri.toString()));
+			this.proxies.push(proxy);
+			if (i === 0) {
+				this.shouldLookup = lookup;
+			}
+		}
 
-		this.shouldLookup = lookup;
-		this.proxy = proxy;
 		this.timeout = opts?.timeout ?? null;
 	}
 
@@ -110,7 +128,7 @@ export class SocksProxyAgent extends Agent {
 		req: http.ClientRequest,
 		opts: AgentConnectOpts
 	): Promise<net.Socket> {
-		const { shouldLookup, proxy, timeout } = this;
+		const { shouldLookup, proxies, timeout } = this;
 
 		if (!opts.host) {
 			throw new Error('No `host` defined!');
@@ -133,25 +151,46 @@ export class SocksProxyAgent extends Agent {
 			});
 		}
 
-		const socksOpts: SocksClientOptions = {
-			proxy,
-			destination: {
-				host,
-				port: typeof port === 'number' ? port : parseInt(port, 10),
-			},
-			command: 'connect',
-			timeout: timeout ?? undefined,
-		};
-
+		let socket: net.Socket;
 		const cleanup = (tlsSocket?: tls.TLSSocket) => {
 			req.destroy();
 			socket.destroy();
 			if (tlsSocket) tlsSocket.destroy();
 		};
 
-		debug('Creating socks proxy connection: %o', socksOpts);
-		const { socket } = await SocksClient.createConnection(socksOpts);
-		debug('Successfully created socks proxy connection');
+		if (proxies.length === 1) {
+			const socksOpts: SocksClientOptions = {
+				proxy: proxies[0],
+				destination: {
+					host,
+					port: typeof port === 'number' ? port : parseInt(port, 10),
+				},
+				command: 'connect',
+				timeout: timeout ?? undefined,
+			};
+
+			debug('Creating socks proxy connection: %o', socksOpts);
+			const connection = await SocksClient.createConnection(socksOpts);
+			socket = connection.socket;
+			debug('Successfully created socks proxy connection');
+		} else {
+			const socksOpts: SocksClientChainOptions = {
+				proxies: proxies,
+				destination: {
+					host,
+					port: typeof port === 'number' ? port : parseInt(port, 10),
+				},
+				command: 'connect',
+				timeout: timeout ?? undefined,
+			};
+
+			debug('Creating chained socks proxy connection: %o', socksOpts);
+			const connection = await SocksClient.createConnectionChain(
+				socksOpts
+			);
+			socket = connection.socket;
+			debug('Successfully created chained socks proxy connection');
+		}
 
 		if (timeout !== null) {
 			socket.setTimeout(timeout);
diff --git a/packages/socks-proxy-agent/test/test.ts b/packages/socks-proxy-agent/test/test.ts
index 737e8834..15155b1d 100644
--- a/packages/socks-proxy-agent/test/test.ts
+++ b/packages/socks-proxy-agent/test/test.ts
@@ -68,13 +68,13 @@ describe('SocksProxyAgent', () => {
 	describe('constructor', () => {
 		it('should accept a "string" proxy argument', () => {
 			const agent = new SocksProxyAgent(socksServerUrl.href);
-			assert.equal(socksServerUrl.hostname, agent.proxy.host);
-			assert.equal(+socksServerUrl.port, agent.proxy.port);
+			assert.equal(socksServerUrl.hostname, agent.proxies[0].host);
+			assert.equal(+socksServerUrl.port, agent.proxies[0].port);
 		});
 		it('should accept a `new URL()` result object argument', () => {
 			const agent = new SocksProxyAgent(socksServerUrl);
-			assert.equal(socksServerUrl.hostname, agent.proxy.host);
-			assert.equal(+socksServerUrl.port, agent.proxy.port);
+			assert.equal(socksServerUrl.hostname, agent.proxies[0].host);
+			assert.equal(+socksServerUrl.port, agent.proxies[0].port);
 		});
 		it('should respect `timeout` option during connection to socks server', async () => {
 			const agent = new SocksProxyAgent(socksServerUrl, { timeout: 1 });
@@ -107,6 +107,41 @@ describe('SocksProxyAgent', () => {
 			const body = await json(res);
 			assert.equal('bar', body.foo);
 		});
+
+		it('should work against an HTTP endpoint with multiple SOCKS proxies', async () => {
+			const secondSocksServer = socks.createServer(function (
+				// @ts-expect-error no types for `socksv5`
+				_info,
+				// @ts-expect-error no types for `socksv5`
+				accept
+			) {
+				accept();
+			});
+			await listen(secondSocksServer);
+			const port = secondSocksServer.address().port;
+			const secondSocksServerUrl = new URL(`socks://127.0.0.1:${port}`);
+			secondSocksServer.useAuth(socks.auth.None());
+
+			httpServer.once('request', function (req, res) {
+				assert.equal('/foo', req.url);
+				res.statusCode = 404;
+				res.end(JSON.stringify(req.headers));
+			});
+
+			const res = await req(new URL('/foo', httpServerUrl), {
+				agent: new SocksProxyAgent([
+					socksServerUrl,
+					secondSocksServerUrl,
+				]),
+				headers: { foo: 'bar' },
+			});
+			assert.equal(404, res.statusCode);
+
+			const body = await json(res);
+			assert.equal('bar', body.foo);
+
+			secondSocksServer.close();
+		});
 	});
 
 	describe('"https" module', () => {