Add support for upstream proxies with websockets too

Author: pimterry

src/rules/requests/request-handlers.ts

@@ -9,7 +9,6 @@ import tls = require('tls');
 import http = require('http');
 import http2 = require('http2');
 import https = require('https');
-import ProxyAgent = require('proxy-agent');
 import * as h2Client from 'http2-wrapper';
 import CacheableLookup from 'cacheable-lookup';
 import { encode as encodeBase64, decode as decodeBase64 } from 'base64-arraybuffer';
@@ -18,6 +17,18 @@ import { stripIndent, oneLine } from 'common-tags';
 import { TypedError } from 'typed-error';
 import { encodeBuffer, SUPPORTED_ENCODING } from 'http-encoding';
 
+import {
+    Headers,
+    OngoingRequest,
+    CompletedRequest,
+    OngoingResponse,
+    CompletedBody,
+    Explainable
+} from "../../types";
+
+import { byteLength } from '../../util/util';
+import { MaybePromise, Replace } from '../../util/type-utils';
+import { readFile } from '../../util/fs';
 import {
     waitForCompletedRequest,
     setHeaders,
@@ -29,8 +40,7 @@ import {
     h2HeadersToH1,
     isAbsoluteUrl,
     cleanUpHeaders,
-    isMockttpBody,
-    matchesNoProxy
+    isMockttpBody
 } from '../../util/request-utils';
 import { streamToBuffer, asBuffer } from '../../util/buffer-utils';
 import { isLocalPortActive, isSocketLoop } from '../../util/socket-util';
@@ -45,18 +55,7 @@ import {
     serializeBuffer,
     deserializeBuffer
 } from "../../util/serialization";
-import { MaybePromise, Replace } from '../../util/type-utils';
-import { readFile } from '../../util/fs';
-
-import {
-    Headers,
-    OngoingRequest,
-    CompletedRequest,
-    OngoingResponse,
-    CompletedBody,
-    Explainable
-} from "../../types";
-import { byteLength, isNode } from '../../util/util';
+import { getAgent, ProxyConfig } from '../../util/http-agents';
 
 // An error that indicates that the handler is aborting the request.
 // This could be intentional, or an upstream server aborting the request.
@@ -416,36 +415,6 @@ export interface ForwardingOptions {
     updateHostHeader?: true | false | string // Change automatically/ignore/change to custom value
 }
 
-export interface ProxyConfig {
-    /**
-     * The URL for the proxy to pass through.
-     *
-     * This can be any URL supported by https://www.npmjs.com/package/proxy-agent.
-     * For example: http://..., socks5://..., pac+http://...
-     */
-    proxyUrl: string;
-
-    /**
-     * A list of no-proxy values, matching URLs which should not be proxied.
-     *
-     * This is a common proxy feature, but unfortunately isn't standardized. See
-     * https://about.gitlab.com/blog/2021/01/27/we-need-to-talk-no-proxy/ for some
-     * background. This implementation is intended to match Curl's behaviour, and
-     * any differences are a bug.
-     *
-     * The currently supported formats are:
-     * - example.com (matches domain and all subdomains)
-     * - example.com:443 (matches domain and all subdomains, but only on that port)
-     * - 10.0.0.1 (matches IP, but only when used directly - does not resolve domains)
-     *
-     * Some other formats (e.g. leading dots or *.) will work, but the leading
-     * characters are ignored. More formats may be added in future, e.g. CIDR ranges.
-     * To maximize compatibility with values used elsewhere, unrecognized formats
-     * will generally be ignored, but may match in unexpected ways.
-     */
-    noProxy?: string[];
-}
-
 export interface PassThroughLookupOptions {
     /**
      * The maximum time to cache a DNS response. Up to this limit,
@@ -883,50 +852,6 @@ function validateCustomHeaders(
 const OMIT_SYMBOL = Symbol('omit-value');
 const SERIALIZED_OMIT = "__mockttp__transform__omit__";
 
-const KeepAliveAgents = isNode
-    ? { // These are only used (and only available) on the node server side
-        'http:': new http.Agent({
-            keepAlive: true
-        }),
-        'https:': new https.Agent({
-            keepAlive: true
-        })
-    } : {};
-
-function getAgent({
-    protocol, hostname, port, tryHttp2, keepAlive, proxyConfig
-}: {
-    protocol: 'http:' | 'https:' | undefined,
-    hostname: string,
-    port: number,
-    tryHttp2: boolean,
-    keepAlive: boolean
-    proxyConfig: ProxyConfig | undefined,
-}): {} | undefined {
-    if (proxyConfig && proxyConfig.proxyUrl) {
-        // If there's a (non-empty) proxy configured, use it. We require non-empty because empty strings
-        // will fall back to detecting from the environment, which is likely to behave unexpectedly.
-
-        if (!matchesNoProxy(hostname, port, proxyConfig.noProxy)) {
-            // We notably ignore HTTP/2 upstream in this case: it's complicated to mix that up with proxying
-            // so for now we ignore it entirely.
-            return new ProxyAgent(proxyConfig.proxyUrl);
-        }
-    }
-
-    if (tryHttp2 && protocol === 'https:') {
-        // H2 wrapper takes multiple agents, uses the appropriate one for the detected protocol.
-        // We notably never use H2 upstream for plaintext, it's rare and we can't use ALPN to detect it.
-        return { https: KeepAliveAgents['https:'], http2: undefined };
-    } else if (keepAlive) {
-        // HTTP/1.1 or HTTP/1 with explicit keep-alive
-        return KeepAliveAgents[protocol || 'http:']
-    } else {
-        // HTTP/1 without KA - just send the request with no agent
-        return undefined;
-    }
-}
-
 export class PassThroughHandler extends Serializable implements RequestHandler {
     readonly type = 'passthrough';
 
@@ -1265,7 +1190,7 @@ export class PassThroughHandler extends Serializable implements RequestHandler {
             tryHttp2: shouldTryH2Upstream,
             keepAlive: shouldKeepAlive(clientReq),
             proxyConfig: this.proxyConfig
-        }) as http.Agent;
+        });
 
         if (isH2Downstream && shouldTryH2Upstream) {
             // We drop all incoming pseudoheaders, and regenerate them (except legally modified ones)

src/rules/websockets/websocket-handlers.ts

@@ -27,6 +27,7 @@ import {
 } from '../requests/request-handlers';
 import { isHttp2 } from '../../util/request-utils';
 import { streamToBuffer } from '../../util/buffer-utils';
+import { getAgent, ProxyConfig } from '../../util/http-agents';
 
 export interface WebSocketHandler extends Explainable, Serializable {
     type: keyof typeof WsHandlerLookup;
@@ -162,6 +163,11 @@ export interface PassThroughWebSocketHandlerOptions {
      */
     ignoreHostCertificateErrors?: string[];
 
+    /**
+     * Upstream proxy configuration: pass through requests via this proxy
+     */
+    proxyConfig?: ProxyConfig;
+
     /**
      * Custom DNS options, to allow configuration of the resolver used
      * when forwarding requests upstream. Passing any option switches
@@ -189,6 +195,8 @@ export class PassThroughWebSocketHandler extends Serializable implements WebSock
     // Same lookup configuration as normal request PassThroughHandler:
     public readonly lookupOptions: PassThroughLookupOptions | undefined;
 
+    public readonly proxyConfig?: ProxyConfig;
+
     private _cacheableLookupInstance: CacheableLookup | undefined;
     private lookup() {
         if (!this.lookupOptions) return undefined;
@@ -235,6 +243,7 @@ export class PassThroughWebSocketHandler extends Serializable implements WebSock
         this.forwarding = options.forwarding;
 
         this.lookupOptions = options.lookupOptions;
+        this.proxyConfig = options.proxyConfig;
     }
 
     explain() {
@@ -330,9 +339,23 @@ export class PassThroughWebSocketHandler extends Serializable implements WebSock
         const checkServerCertificate = !_.includes(this.ignoreHostHttpsErrors, parsedUrl.hostname) &&
             !_.includes(this.ignoreHostHttpsErrors, parsedUrl.host);
 
+        const effectivePort = !!parsedUrl.port
+            ? parseInt(parsedUrl.port, 10)
+            : parsedUrl.protocol == 'wss:' ? 443 : 80;
+
+        const agent = getAgent({
+            protocol: parsedUrl.protocol as 'ws:' | 'wss:',
+            hostname: parsedUrl.hostname!,
+            port: effectivePort,
+            proxyConfig: this.proxyConfig,
+            tryHttp2: false, // We don't support websockets over H2 yet
+            keepAlive: false // Not a thing for websockets: they take over the whole connection
+        });
+
         const upstreamWebSocket = new WebSocket(wsUrl, {
             rejectUnauthorized: checkServerCertificate,
             maxPayload: 0,
+            agent,
             lookup: this.lookup(),
             headers: _.omitBy(headers, (_v, headerName) =>
                 headerName.toLowerCase().startsWith('sec-websocket') ||

src/util/http-agents.ts

@@ -0,0 +1,109 @@
+import * as http from 'http';
+import * as https from 'https';
+import ProxyAgent = require('proxy-agent');
+
+import { isNode } from "./util";
+
+const KeepAliveAgents = isNode
+    ? { // These are only used (and only available) on the node server side
+        'http:': new http.Agent({
+            keepAlive: true
+        }),
+        'https:': new https.Agent({
+            keepAlive: true
+        })
+    } : {};
+
+export interface ProxyConfig {
+    /**
+     * The URL for the proxy to forward traffic through.
+     *
+     * This can be any URL supported by https://www.npmjs.com/package/proxy-agent.
+     * For example: http://..., socks5://..., pac+http://...
+     */
+    proxyUrl: string;
+
+    /**
+     * A list of no-proxy values, matching hosts' traffic should *not* be proxied.
+     *
+     * This is a common proxy feature, but unfortunately isn't standardized. See
+     * https://about.gitlab.com/blog/2021/01/27/we-need-to-talk-no-proxy/ for some
+     * background. This implementation is intended to match Curl's behaviour, and
+     * any differences are a bug.
+     *
+     * The currently supported formats are:
+     * - example.com (matches domain and all subdomains)
+     * - example.com:443 (matches domain and all subdomains, but only on that port)
+     * - 10.0.0.1 (matches IP, but only when used directly - does not resolve domains)
+     *
+     * Some other formats (e.g. leading dots or *.) will work, but the leading
+     * characters are ignored. More formats may be added in future, e.g. CIDR ranges.
+     * To maximize compatibility with values used elsewhere, unrecognized formats
+     * will generally be ignored, but may match in unexpected ways.
+     */
+    noProxy?: string[];
+}
+
+export function getAgent({
+    protocol, hostname, port, tryHttp2, keepAlive, proxyConfig
+}: {
+    protocol: 'http:' | 'https:' | 'ws:' | 'wss:' | undefined,
+    hostname: string,
+    port: number,
+    tryHttp2: boolean,
+    keepAlive: boolean
+    proxyConfig: ProxyConfig | undefined,
+}): http.Agent | undefined { // <-- We force this cast for convenience in various different uses later
+    if (proxyConfig && proxyConfig.proxyUrl) {
+        // If there's a (non-empty) proxy configured, use it. We require non-empty because empty strings
+        // will fall back to detecting from the environment, which is likely to behave unexpectedly.
+
+        if (!matchesNoProxy(hostname, port, proxyConfig.noProxy)) {
+            // We notably ignore HTTP/2 upstream in this case: it's complicated to mix that up with proxying
+            // so for now we ignore it entirely.
+            return new ProxyAgent(proxyConfig.proxyUrl) as http.Agent;
+        }
+    }
+
+    if (tryHttp2 && (protocol === 'https:' || protocol === 'wss:')) {
+        // H2 wrapper takes multiple agents, uses the appropriate one for the detected protocol.
+        // We notably never use H2 upstream for plaintext, it's rare and we can't use ALPN to detect it.
+        return { https: KeepAliveAgents['https:'], http2: undefined } as any as http.Agent;
+    } else if (keepAlive && protocol !== 'wss:' && protocol !== 'ws:') {
+        // HTTP/1.1 or HTTP/1 with explicit keep-alive
+        return KeepAliveAgents[protocol || 'http:']
+    } else {
+        // HTTP/1 without KA - just send the request with no agent
+        return undefined;
+    }
+}
+
+export const matchesNoProxy = (hostname: string, portNum: number, noProxyValues: string[] | undefined) => {
+    if (!noProxyValues || noProxyValues.length === 0) return false; // Skip everything in the common case.
+
+    const port = portNum.toString();
+    const hostParts = hostname.split('.').reverse();
+
+    return noProxyValues.some((noProxy) => {
+        const [noProxyHost, noProxyPort] = noProxy.split(':') as [string, string | undefined];
+
+        let noProxyParts = noProxyHost.split('.').reverse();
+        const lastPart = noProxyParts[noProxyParts.length - 1];
+        if (lastPart === '' || lastPart === '*') {
+            noProxyParts = noProxyParts.slice(0, -1);
+        }
+
+        if (noProxyPort && port !== noProxyPort) return false;
+
+        for (let i = 0; i < noProxyParts.length; i++) {
+            let noProxyPart = noProxyParts[i];
+            let hostPart = hostParts[i];
+
+            if (hostPart === undefined) return false; // No-proxy is longer than hostname
+            if (noProxyPart !== hostPart) return false; // Mismatch
+        }
+
+        // If we run out of no-proxy parts with no mismatch then we've matched
+        return true;
+    });
+}

src/util/request-utils.ts

@@ -65,36 +65,6 @@ export const shouldKeepAlive = (req: OngoingRequest): boolean =>
     req.headers['connection'] !== 'close' &&
     req.headers['proxy-connection'] !== 'close';
 
-export const matchesNoProxy = (hostname: string, portNum: number, noProxyValues: string[] | undefined) => {
-    if (!noProxyValues || noProxyValues.length === 0) return false; // Skip everything in the common case.
-
-    const port = portNum.toString();
-    const hostParts = hostname.split('.').reverse();
-
-    return noProxyValues.some((noProxy) => {
-        const [noProxyHost, noProxyPort] = noProxy.split(':') as [string, string | undefined];
-
-        let noProxyParts = noProxyHost.split('.').reverse();
-        const lastPart = noProxyParts[noProxyParts.length - 1];
-        if (lastPart === '' || lastPart === '*') {
-            noProxyParts = noProxyParts.slice(0, -1);
-        }
-
-        if (noProxyPort && port !== noProxyPort) return false;
-
-        for (let i = 0; i < noProxyParts.length; i++) {
-            let noProxyPart = noProxyParts[i];
-            let hostPart = hostParts[i];
-
-            if (hostPart === undefined) return false; // No-proxy is longer than hostname
-            if (noProxyPart !== hostPart) return false; // Mismatch
-        }
-
-        // If we run out of no-proxy parts with no mismatch then we've matched
-        return true;
-    });
-}
-
 export const setHeaders = (response: http.ServerResponse, headers: Headers) => {
     Object.keys(headers).forEach((header) => {
         let value = headers[header];