Add error simulation to websocket passthrough & fix WS rejection crash

Previously it was possible to crash the proxy if the remote server
rejected a websocket connection under some circumstances. That's now
resolved, and mirroring of connection errors in general is improved in
proxied websockets when the (HTTP-equivalent) simulateConnectionErrors
option is provided.

Author: pimterry

src/rules/passthrough-handling-definitions.ts

@@ -108,4 +108,29 @@ export interface PassThroughHandlerConnectionOptions {
      * cacheable-lookup module, which will cache responses.
      */
     lookupOptions?: PassThroughLookupOptions;
+
+    /**
+     * Whether to simulate connection errors back to the client.
+     *
+     * By default (in most cases - see below) when an upstream request fails
+     * outright a 502 "Bad Gateway" response is sent to the downstream client,
+     * explicitly indicating the failure and containing the error that caused
+     * the issue in the response body.
+     *
+     * Only in the case of upstream HTTP connection reset errors is a connection
+     * reset normally sent back downstream to existing clients (this behaviour
+     * exists for backward compatibility, and will change to match other error
+     * behaviour in a future version).
+     *
+     * When this option is set to `true`, low-level connection failures will
+     * always trigger a downstream connection close/reset, rather than a 502
+     * response.
+     *
+     * This includes DNS failures, TLS connection errors, TCP connection resets,
+     * etc (but not HTTP non-200 responses, which are still proxied as normal).
+     * This is less convenient for debugging in a testing environment or when
+     * using a proxy intentionally, but can be more accurate when trying to
+     * transparently proxy network traffic, errors and all.
+     */
+    simulateConnectionErrors?: boolean;
 }

src/rules/requests/request-handler-definitions.ts

@@ -479,31 +479,6 @@ export interface PassThroughResponse {
 }
 
 export interface PassThroughHandlerOptions extends PassThroughHandlerConnectionOptions {
-    /**
-     * Whether to simulate connection errors back to the client.
-     *
-     * By default (in most cases - see below) when an upstream request fails
-     * outright a 502 "Bad Gateway" response is sent to the downstream client,
-     * explicitly indicating the failure and containing the error that caused
-     * the issue in the response body.
-     *
-     * Only in the case of upstream connection reset errors is a connection reset
-     * normally sent back downstream to existing clients (this behaviour exists
-     * for backward compatibility, and will change to match other error behaviour
-     * in a future version).
-     *
-     * When this option is set to `true`, low-level connection failures will
-     * always trigger a downstream connection close/reset, rather than a 502
-     * response.
-     *
-     * This includes DNS failures, TLS connection errors, TCP connection resets,
-     * etc (but not HTTP non-200 responses, which are still proxied as normal).
-     * This is less convenient for debugging in a testing environment or when
-     * using a proxy intentionally, but can be more accurate when trying to
-     * transparently proxy network traffic, errors and all.
-     */
-    simulateConnectionErrors?: boolean;
-
     /**
      * A set of data to automatically transform a request. This includes properties
      * to support many transformation common use cases.

src/rules/websockets/websocket-handler-definitions.ts

@@ -52,6 +52,7 @@ export interface SerializedPassThroughWebSocketData {
     forwarding?: ForwardingOptions;
     lookupOptions?: PassThroughLookupOptions;
     proxyConfig?: SerializedProxyConfig;
+    simulateConnectionErrors?: boolean;
     ignoreHostCertificateErrors?: string[] | boolean; // Doesn't match option name, backward compat
     extraCACertificates?: Array<{ cert: string } | { certPath: string }>;
     clientCertificateHostMap?: { [host: string]: { pfx: string, passphrase?: string } };
@@ -63,6 +64,7 @@ export class PassThroughWebSocketHandlerDefinition extends Serializable implemen
     // Same lookup configuration as normal request PassThroughHandler:
     public readonly lookupOptions: PassThroughLookupOptions | undefined;
     public readonly proxyConfig?: ProxyConfig;
+    public readonly simulateConnectionErrors: boolean;
 
     public readonly forwarding?: ForwardingOptions;
     public readonly ignoreHostHttpsErrors: string[] | boolean = [];
@@ -98,6 +100,7 @@ export class PassThroughWebSocketHandlerDefinition extends Serializable implemen
 
         this.lookupOptions = options.lookupOptions;
         this.proxyConfig = options.proxyConfig;
+        this.simulateConnectionErrors = !!options.simulateConnectionErrors;
 
         this.extraCACertificates =
             options.additionalTrustedCAs ||
@@ -121,6 +124,7 @@ export class PassThroughWebSocketHandlerDefinition extends Serializable implemen
             forwarding: this.forwarding,
             lookupOptions: this.lookupOptions,
             proxyConfig: serializeProxyConfig(this.proxyConfig, channel),
+            simulateConnectionErrors: this.simulateConnectionErrors,
             ignoreHostCertificateErrors: this.ignoreHostHttpsErrors,
             extraCACertificates: this.extraCACertificates.map((certObject) => {
                 // We use toString to make sure that buffers always end up as

src/rules/websockets/websocket-handlers.ts

@@ -51,6 +51,7 @@ import {
     WebSocketHandlerDefinition,
     WsHandlerDefinitionLookup,
 } from './websocket-handler-definitions';
+import { resetOrDestroy } from '../../util/socket-util';
 
 export interface WebSocketHandler extends WebSocketHandlerDefinition {
     handle(
@@ -157,19 +158,38 @@ function pipeWebSocket(inSocket: WebSocket, outSocket: WebSocket) {
     });
 }
 
-async function mirrorRejection(socket: net.Socket, rejectionResponse: http.IncomingMessage) {
-    if (socket.writable) {
-        const { statusCode, statusMessage, rawHeaders } = rejectionResponse;
-
-        socket.write(
-            rawResponse(statusCode || 500, statusMessage || 'Unknown error', pairFlatRawHeaders(rawHeaders))
-        );
-
-        const body = await streamToBuffer(rejectionResponse);
-        if (socket.writable) socket.write(body);
-    }
+function mirrorRejection(
+    downstreamSocket: net.Socket,
+    upstreamRejectionResponse: http.IncomingMessage,
+    simulateConnectionErrors: boolean
+) {
+    return new Promise<void>((resolve) => {
+        if (downstreamSocket.writable) {
+            const { statusCode, statusMessage, rawHeaders } = upstreamRejectionResponse;
+
+            downstreamSocket.write(
+                rawResponse(statusCode || 500, statusMessage || 'Unknown error', pairFlatRawHeaders(rawHeaders))
+            );
+
+            upstreamRejectionResponse.pipe(downstreamSocket);
+            upstreamRejectionResponse.on('end', resolve);
+            upstreamRejectionResponse.on('error', (error) => {
+                console.warn('Error receiving WebSocket upstream rejection response:', error);
+                if (simulateConnectionErrors) {
+                    resetOrDestroy(downstreamSocket);
+                } else {
+                    downstreamSocket.destroy();
+                }
+                resolve();
+            });
 
-    socket.destroy();
+            // The socket is being optimistically written to and then killed - we don't care
+            // about any more errors occuring here.
+            downstreamSocket.on('error', () => {
+                resolve();
+            });
+        }
+    }).catch(() => {});
 }
 
 const rawResponse = (
@@ -402,25 +422,29 @@ export class PassThroughWebSocketHandler extends PassThroughWebSocketHandlerDefi
             console.log(`Unexpected websocket response from ${wsUrl}: ${res.statusCode}`);
 
             // Clean up the downstream connection
-            mirrorRejection(incomingSocket, res);
-
-            // Clean up the upstream connection (WS would do this automatically, but doesn't if you listen to this event)
-            // See https://github.com/websockets/ws/blob/45e17acea791d865df6b255a55182e9c42e5877a/lib/websocket.js#L1050
-            // We don't match that perfectly, but this should be effectively equivalent:
-            req.destroy();
-            if (req.socket && !req.socket.destroyed) {
-                res.socket.destroy();
-            }
-            unexpectedResponse = true; // So that we ignore this in the error handler
-            upstreamWebSocket.terminate();
+            mirrorRejection(incomingSocket, res, this.simulateConnectionErrors).then(() => {
+                // Clean up the upstream connection (WS would do this automatically, but doesn't if you listen to this event)
+                // See https://github.com/websockets/ws/blob/45e17acea791d865df6b255a55182e9c42e5877a/lib/websocket.js#L1050
+                // We don't match that perfectly, but this should be effectively equivalent:
+                req.destroy();
+                if (res.socket?.destroyed === false) {
+                    res.socket.destroy();
+                }
+                unexpectedResponse = true; // So that we ignore this in the error handler
+                upstreamWebSocket.terminate();
+            });
         });
 
         // If there's some other error, we just kill the socket:
         upstreamWebSocket.on('error', (e) => {
             if (unexpectedResponse) return; // Handled separately above
 
             console.warn(e);
-            incomingSocket.end();
+            if (this.simulateConnectionErrors) {
+                resetOrDestroy(incomingSocket);
+            } else {
+                incomingSocket.end();
+            }
         });
 
         incomingSocket.on('error', () => upstreamWebSocket.close(1011)); // Internal error
@@ -438,6 +462,7 @@ export class PassThroughWebSocketHandler extends PassThroughWebSocketHandlerDefi
         return _.create(this.prototype, {
             ...data,
             proxyConfig: deserializeProxyConfig(data.proxyConfig, channel, ruleParams),
+            simulateConnectionErrors: data.simulateConnectionErrors ?? false,
             extraCACertificates: data.extraCACertificates || [],
             ignoreHostHttpsErrors: data.ignoreHostCertificateErrors,
             clientCertificateHostMap: _.mapValues(data.clientCertificateHostMap,