From 0bbe3359d25d754b932b1f290893f22d500189a2 Mon Sep 17 00:00:00 2001
From: mceIdo <58237675+mceIdo@users.noreply.github.com>
Date: Tue, 9 May 2023 11:56:47 +0300
Subject: [PATCH] [proxy-agent] Support for `getProxyForUrl` option (#160)

---
 .changeset/odd-mirrors-nail.md    |  5 +++++
 packages/proxy-agent/src/index.ts | 14 ++++++++++++--
 packages/proxy-agent/test/test.ts | 26 ++++++++++++++++++++++++++
 3 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 .changeset/odd-mirrors-nail.md

diff --git a/.changeset/odd-mirrors-nail.md b/.changeset/odd-mirrors-nail.md
new file mode 100644
index 00000000..58eb38d3
--- /dev/null
+++ b/.changeset/odd-mirrors-nail.md
@@ -0,0 +1,5 @@
+---
+'proxy-agent': minor
+---
+
+Support for `getProxyForUrl` option, to provide proxy address dynamically per different URLs
diff --git a/packages/proxy-agent/src/index.ts b/packages/proxy-agent/src/index.ts
index 29ade631..6477cc71 100644
--- a/packages/proxy-agent/src/index.ts
+++ b/packages/proxy-agent/src/index.ts
@@ -3,7 +3,7 @@ import * as https from 'https';
 import LRUCache from 'lru-cache';
 import { Agent, AgentConnectOpts } from 'agent-base';
 import createDebug from 'debug';
-import { getProxyForUrl } from 'proxy-from-env';
+import { getProxyForUrl as envGetProxyForUrl } from 'proxy-from-env';
 import { PacProxyAgent, PacProxyAgentOptions } from 'pac-proxy-agent';
 import { HttpProxyAgent, HttpProxyAgentOptions } from 'http-proxy-agent';
 import { HttpsProxyAgent, HttpsProxyAgentOptions } from 'https-proxy-agent';
@@ -21,6 +21,8 @@ type ValidProtocol = (typeof PROTOCOLS)[number];
 
 type AgentConstructor = new (...args: never[]) => Agent;
 
+type GetProxyForUrlCallback = (url: string) => string;
+
 /**
  * Supported proxy types.
  */
@@ -61,6 +63,12 @@ export type ProxyAgentOptions = HttpProxyAgentOptions<''> &
 		 * instance with the proxy agent options passed in.
 		 */
 		httpsAgent?: http.Agent;
+		/**
+		 * A callback for dynamic provision of proxy for url.
+		 * Defaults to standard proxy environment variables, 
+		 * see https://www.npmjs.com/package/proxy-from-env for details
+		 */
+		getProxyForUrl?: GetProxyForUrlCallback;
 	};
 
 /**
@@ -79,6 +87,7 @@ export class ProxyAgent extends Agent {
 	connectOpts?: ProxyAgentOptions;
 	httpAgent: http.Agent;
 	httpsAgent: http.Agent;
+	getProxyForUrl: GetProxyForUrlCallback;
 
 	constructor(opts?: ProxyAgentOptions) {
 		super(opts);
@@ -87,6 +96,7 @@ export class ProxyAgent extends Agent {
 		this.httpAgent = opts?.httpAgent || new http.Agent(opts);
 		this.httpsAgent =
 			opts?.httpsAgent || new https.Agent(opts as https.AgentOptions);
+		this.getProxyForUrl = opts?.getProxyForUrl || envGetProxyForUrl;
 	}
 
 	async connect(
@@ -97,7 +107,7 @@ export class ProxyAgent extends Agent {
 		const protocol = secureEndpoint ? 'https:' : 'http:';
 		const host = req.getHeader('host');
 		const url = new URL(req.path, `${protocol}//${host}`).href;
-		const proxy = getProxyForUrl(url);
+		const proxy = this.getProxyForUrl(url);
 
 		if (!proxy) {
 			debug('Proxy not enabled for URL: %o', url);
diff --git a/packages/proxy-agent/test/test.ts b/packages/proxy-agent/test/test.ts
index 8d565397..9c52fb8e 100644
--- a/packages/proxy-agent/test/test.ts
+++ b/packages/proxy-agent/test/test.ts
@@ -251,5 +251,31 @@ describe('ProxyAgent', () => {
 				HttpsProxyAgent
 			);
 		});
+
+		it('should call provided function with getProxyForUrl option', async () => {
+			let gotCall = false;
+			let urlParameter = "";
+			httpsServer.once('request', function (req, res) {
+				res.end(JSON.stringify(req.headers));
+			});
+
+			const agent = new ProxyAgent({
+				rejectUnauthorized: false,
+				getProxyForUrl: (u) => {
+					gotCall = true;
+					urlParameter = u;
+					return httpsProxyServerUrl.href;
+				}
+			});
+			const requestUrl = new URL('/test', httpsServerUrl);
+			const res = await req(requestUrl, {
+				agent,
+				rejectUnauthorized: false,
+			});
+			const body = await json(res);
+			assert(httpsServerUrl.host === body.host);
+			assert(gotCall);
+			assert(requestUrl.href === urlParameter);
+		});
 	});
 });