[proxy-agent] Use HttpsProxyAgent for WebSocket requests (#192)

Fixes #176.

Author: TooTallNate

.changeset/fluffy-bears-admire.md

@@ -0,0 +1,5 @@
+---
+'proxy-agent': patch
+---
+
+Fix WebSocket connections over "http"/"https" proxies

packages/proxy-agent/package.json

@@ -47,12 +47,14 @@
     "@types/jest": "^29.5.1",
     "@types/node": "^14.18.45",
     "@types/proxy-from-env": "^1.0.1",
+    "@types/ws": "^8.5.4",
     "async-listen": "^3.0.0",
     "jest": "^29.5.0",
     "proxy": "workspace:*",
     "socksv5": "github:TooTallNate/socksv5#fix/dstSock-close-event",
     "ts-jest": "^29.1.0",
     "tsconfig": "workspace:*",
-    "typescript": "^5.0.4"
+    "typescript": "^5.0.4",
+    "ws": "^8.13.0"
   }
 }

packages/proxy-agent/src/index.ts

@@ -104,7 +104,14 @@ export class ProxyAgent extends Agent {
 		opts: AgentConnectOpts
 	): Promise<http.Agent> {
 		const { secureEndpoint } = opts;
-		const protocol = secureEndpoint ? 'https:' : 'http:';
+		const isWebSocket = req.getHeader('upgrade') === 'websocket';
+		const protocol = secureEndpoint
+			? isWebSocket
+				? 'wss:'
+				: 'https:'
+			: isWebSocket
+			? 'ws:'
+			: 'http:';
 		const host = req.getHeader('host');
 		const url = new URL(req.path, `${protocol}//${host}`).href;
 		const proxy = this.getProxyForUrl(url);
@@ -126,7 +133,8 @@ export class ProxyAgent extends Agent {
 			if (!isValidProtocol(proxyProto)) {
 				throw new Error(`Unsupported protocol for proxy URL: ${proxy}`);
 			}
-			const ctor = proxies[proxyProto][secureEndpoint ? 1 : 0];
+			const ctor =
+				proxies[proxyProto][secureEndpoint || isWebSocket ? 1 : 0];
 			// @ts-expect-error meh…
 			agent = new ctor(proxy, this.connectOpts);
 			this.cache.set(cacheKey, agent);

packages/proxy-agent/test/test.ts

@@ -3,6 +3,7 @@ import * as http from 'http';
 import * as https from 'https';
 import { once } from 'events';
 import assert from 'assert';
+import WebSocket, { WebSocketServer } from 'ws';
 import { json, req } from 'agent-base';
 import { ProxyServer, createProxy } from 'proxy';
 // @ts-expect-error no types
@@ -20,8 +21,10 @@ const sslOptions = {
 describe('ProxyAgent', () => {
 	// target servers
 	let httpServer: http.Server;
+	let httpWebSocketServer: WebSocketServer;
 	let httpServerUrl: URL;
 	let httpsServer: https.Server;
+	let httpsWebSocketServer: WebSocketServer;
 	let httpsServerUrl: URL;
 
 	// proxy servers
@@ -36,12 +39,14 @@ describe('ProxyAgent', () => {
 	beforeAll(async () => {
 		// setup target HTTP server
 		httpServer = http.createServer();
+		httpWebSocketServer = new WebSocketServer({ server: httpServer });
 		httpServerUrl = await listen(httpServer);
 	});
 
 	beforeAll(async () => {
 		// setup target SSL HTTPS server
 		httpsServer = https.createServer(sslOptions);
+		httpsWebSocketServer = new WebSocketServer({ server: httpsServer });
 		httpsServerUrl = await listen(httpsServer);
 	});
 
@@ -79,9 +84,13 @@ describe('ProxyAgent', () => {
 	beforeEach(() => {
 		delete process.env.HTTP_PROXY;
 		delete process.env.HTTPS_PROXY;
+		delete process.env.WS_PROXY;
+		delete process.env.WSS_PROXY;
 		delete process.env.NO_PROXY;
 		httpServer.removeAllListeners('request');
 		httpsServer.removeAllListeners('request');
+		httpWebSocketServer.removeAllListeners('connection');
+		httpsWebSocketServer.removeAllListeners('connection');
 	});
 
 	describe('"http" module', () => {
@@ -278,4 +287,57 @@ describe('ProxyAgent', () => {
 			assert(requestUrl.href === urlParameter);
 		});
 	});
+
+	describe('"ws" module', () => {
+		it('should work over "http" proxy to `ws:` URL', async () => {
+			let requestCount = 0;
+			let connectionCount = 0;
+			httpServer.once('request', function (req, res) {
+				requestCount++;
+				res.end();
+			});
+			httpWebSocketServer.on('connection', (ws) => {
+				connectionCount++;
+				ws.send('OK');
+			});
+
+			process.env.WS_PROXY = httpProxyServerUrl.href;
+			const agent = new ProxyAgent();
+
+			const ws = new WebSocket(httpServerUrl.href.replace('http', 'ws'), {
+				agent,
+			});
+			const [message] = await once(ws, 'message');
+			expect(connectionCount).toEqual(1);
+			expect(requestCount).toEqual(0);
+			expect(message.toString()).toEqual('OK');
+			ws.close();
+		});
+
+		it('should work over "http" proxy to `wss:` URL', async () => {
+			let requestCount = 0;
+			let connectionCount = 0;
+			httpsServer.once('request', function (req, res) {
+				requestCount++;
+				res.end();
+			});
+			httpsWebSocketServer.on('connection', (ws) => {
+				connectionCount++;
+				ws.send('OK');
+			});
+
+			process.env.WSS_PROXY = httpProxyServerUrl.href;
+			const agent = new ProxyAgent();
+
+			const ws = new WebSocket(httpsServerUrl.href.replace('https', 'wss'), {
+				agent,
+				rejectUnauthorized: false
+			});
+			const [message] = await once(ws, 'message');
+			expect(connectionCount).toEqual(1);
+			expect(requestCount).toEqual(0);
+			expect(message.toString()).toEqual('OK');
+			ws.close();
+		});
+	});
 });