Add support for tagging request metadata via SOCKS auth

Author: pimterry

src/main.ts

@@ -2,6 +2,7 @@ export type { MaybePromise } from '@httptoolkit/util';
 
 import { Mockttp, MockttpOptions, MockttpHttpsOptions, SubscribableEvent, PortRange } from "./mockttp";
 import { MockttpServer } from "./server/mockttp-server";
+import { SocksServerOptions } from "./server/socks-server";
 import {
     MockttpClient,
     MockttpClientOptions
@@ -19,7 +20,8 @@ export type {
     MockttpClientOptions,
     MockttpAdminServerOptions,
     SubscribableEvent,
-    PortRange
+    PortRange,
+    SocksServerOptions
 };
 
 // Export now-renamed types with the old aliases to provide backward compat and

src/mockttp.ts

@@ -26,6 +26,7 @@ import {
 } from "./types";
 import type { RequestRuleData } from "./rules/requests/request-rule";
 import type { WebSocketRuleData } from "./rules/websockets/websocket-rule";
+import type { SocksServerOptions } from "./server/socks-server";
 
 export type PortRange = { startPort: number, endPort: number };
 
@@ -811,12 +812,15 @@ export interface MockttpOptions {
 
     /**
      * Should the server accept incoming SOCKS connections? Defaults to false.
-     * If set to true, the server will listen for incoming SOCKS connections
-     * on the same port as the HTTP server, unwrap received connections, and
-     * handle them like any other incoming TCP connection (intercepting HTTP(S)
-     * from within the SOCKS connection as normal).
+     *
+     * If set to true or if detailed options are provided, the server will listen
+     * for incoming SOCKS connections on the same port as the HTTP server, unwrap
+     * received connections, and handle them like any other incoming TCP connection
+     * (intercepting HTTP(S) from within the SOCKS connection as normal).
+     *
+     * The only supported option for now is `authMethods`.
      */
-    socks?: boolean;
+    socks?: boolean | SocksServerOptions;
 
     /**
      * An array of rules for traffic that should be passed through the proxy

src/rules/requests/request-handlers.ts

@@ -1121,7 +1121,10 @@ export class PassThroughHandler extends PassThroughHandlerDefinition {
 
                 options.emitEventCallback('passthrough-abort', {
                     downstreamAborted: !!(serverReq?.aborted),
-                    tags: buildUpstreamErrorTags(e),
+                    tags: [
+                        ...clientReq.tags,
+                        buildUpstreamErrorTags(e)
+                    ],
                     error: {
                         name: e.name,
                         code: e.code,

src/server/http-combo-server.ts

@@ -33,10 +33,11 @@ import {
     TlsMetadata,
     TlsSetupCompleted,
     getAddressAndPort,
-    resetOrDestroy
+    resetOrDestroy,
+    SocketMetadata
 } from '../util/socket-util';
 import { MockttpHttpsOptions } from '../mockttp';
-import { buildSocksServer, SocksTcpAddress } from './socks-server';
+import { buildSocksServer, SocksServerOptions, SocksTcpAddress } from './socks-server';
 
 // Hardcore monkey-patching: force TLSSocket to link servername & remoteAddress to
 // sockets as soon as they're available, without waiting for the handshake to fully
@@ -145,7 +146,7 @@ export interface ComboServerOptions {
     debug: boolean;
     https: MockttpHttpsOptions | undefined;
     http2: boolean | 'fallback';
-    socks: boolean;
+    socks: boolean | SocksServerOptions;
     passthroughUnknownProtocols: boolean;
 
     requestListener: (req: http.IncomingMessage, res: http.ServerResponse) => void;
@@ -225,7 +226,7 @@ export async function createComboServer(options: ComboServerOptions): Promise<De
     }
 
     if (options.socks) {
-        socksServer = buildSocksServer();
+        socksServer = buildSocksServer(options.socks === true ? {} : options.socks);
         socksServer.on('socks-tcp-connect', (socket: net.Socket, address: SocksTcpAddress) => {
             const addressString =
                 address.type === 'ipv4'
@@ -291,8 +292,7 @@ export async function createComboServer(options: ComboServerOptions): Promise<De
         if (parentSocket) {
             // Sometimes wrapper TLS sockets created by the HTTP/2 server don't include the
             // underlying socket details, so it's better to make sure we copy them up.
-            copyAddressDetails(parentSocket, socket);
-            copyTimingDetails(parentSocket, socket);
+            inheritSocketDetails(parentSocket, socket);
             // With TLS metadata, we only propagate directly from parent sockets, not through
             // CONNECT etc - we only want it if the final hop is TLS, previous values don't matter.
             socket[TlsMetadata] ??= parentSocket[TlsMetadata];
@@ -371,8 +371,7 @@ export async function createComboServer(options: ComboServerOptions): Promise<De
 
         // Send a 200 OK response, and start the tunnel:
         res.writeHead(200, {});
-        copyAddressDetails(res.socket, res.stream);
-        copyTimingDetails(res.socket, res.stream);
+        inheritSocketDetails(res.socket, res.stream);
         res.stream[LastTunnelAddress] = connectUrl;
 
         // When layering HTTP/2 on JS streams, we have to make sure the JS stream won't autoclose
@@ -390,39 +389,37 @@ export async function createComboServer(options: ComboServerOptions): Promise<De
 }
 
 
-const SOCKET_ADDRESS_METADATA_FIELDS = [
+const SOCKET_METADATA = [
     'localAddress',
     'localPort',
     'remoteAddress',
     'remotePort',
+    SocketTimingInfo,
+    SocketMetadata,
     LastTunnelAddress
 ] as const;
 
-// Update the target socket(-ish) with the address details from the source socket,
-// iff the target has no details of its own.
-function copyAddressDetails(
-    source: SocketIsh<typeof SOCKET_ADDRESS_METADATA_FIELDS[number]>,
-    target: SocketIsh<typeof SOCKET_ADDRESS_METADATA_FIELDS[number]>
+function inheritSocketDetails(
+    source: SocketIsh<typeof SOCKET_METADATA[number]>,
+    target: SocketIsh<typeof SOCKET_METADATA[number]>
 ) {
+    // Update the target socket(-ish) with the assorted metadata from the source socket,
+    // iff the target has no details of its own.
+
+    // Make sure all properties are writable - HTTP/2 streams notably try to block this.
     Object.defineProperties(target, _.zipObject(
-        SOCKET_ADDRESS_METADATA_FIELDS,
-        _.range(SOCKET_ADDRESS_METADATA_FIELDS.length).map(() => ({ writable: true }))
+        SOCKET_METADATA,
+        _.range(SOCKET_METADATA.length).map(() => ({ writable: true }))
     ) as PropertyDescriptorMap);
 
-    SOCKET_ADDRESS_METADATA_FIELDS.forEach((fieldName) => {
+    for (let fieldName of SOCKET_METADATA) {
         if (target[fieldName] === undefined) {
-            (target as any)[fieldName] = source[fieldName];
+            if (typeof source[fieldName] === 'object') {
+                (target as any)[fieldName] = _.cloneDeep(source[fieldName]);
+            } else {
+                (target as any)[fieldName] = source[fieldName];
+            }
         }
-    });
-}
-
-function copyTimingDetails<T extends SocketIsh<typeof SocketTimingInfo>>(
-    source: SocketIsh<typeof SocketTimingInfo>,
-    target: T
-): asserts target is T & { [SocketTimingInfo]: Required<net.Socket>[typeof SocketTimingInfo] } {
-    if (!target[SocketTimingInfo]) {
-        // Clone timing info, don't copy it - child sockets get their own independent timing stats
-        target[SocketTimingInfo] = Object.assign({}, source[SocketTimingInfo]);
     }
 }
 

src/server/mockttp-server.ts

@@ -56,7 +56,9 @@ import {
     LastTunnelAddress,
     TlsSetupCompleted,
     isSocketLoop,
-    resetOrDestroy
+    resetOrDestroy,
+    SocketMetadata,
+    getSocketMetadataTags
 } from "../util/socket-util";
 import {
     parseRequestBody,
@@ -77,6 +79,7 @@ import {
 import { AbortError } from "../rules/requests/request-handlers";
 import { WebSocketRuleData, WebSocketRule } from "../rules/websockets/websocket-rule";
 import { RejectWebSocketHandler, WebSocketHandler } from "../rules/websockets/websocket-handlers";
+import { SocksServerOptions } from "./socks-server";
 
 type ExtendedRawRequest = (http.IncomingMessage | http2.Http2ServerRequest) & {
     protocol?: string;
@@ -99,7 +102,7 @@ export class MockttpServer extends AbstractMockttp implements Mockttp {
 
     private httpsOptions: MockttpHttpsOptions | undefined;
     private isHttp2Enabled: boolean | 'fallback';
-    private socksEnabled: boolean;
+    private socksOptions: boolean | SocksServerOptions;
     private passthroughUnknownProtocols: boolean;
     private maxBodySize: number;
 
@@ -119,7 +122,7 @@ export class MockttpServer extends AbstractMockttp implements Mockttp {
 
         this.httpsOptions = options.https;
         this.isHttp2Enabled = options.http2 ?? 'fallback';
-        this.socksEnabled = options.socks ?? false;
+        this.socksOptions = options.socks ?? false;
         this.passthroughUnknownProtocols = options.passthrough?.includes('unknown-protocol') ?? false;
         this.maxBodySize = options.maxBodySize ?? Infinity;
         this.eventEmitter = new EventEmitter();
@@ -146,7 +149,7 @@ export class MockttpServer extends AbstractMockttp implements Mockttp {
             debug: this.debug,
             https: this.httpsOptions,
             http2: this.isHttp2Enabled,
-            socks: this.socksEnabled,
+            socks: this.socksOptions,
             passthroughUnknownProtocols: this.passthroughUnknownProtocols,
 
             requestListener: this.app,
@@ -629,7 +632,7 @@ export class MockttpServer extends AbstractMockttp implements Mockttp {
         }
 
         const id = uuid();
-        const tags: string[] = [];
+        const tags: string[] = getSocketMetadataTags(req.socket);
 
         const timingEvents: TimingEvents = {
             startTime: Date.now(),
@@ -986,7 +989,10 @@ ${await this.suggestRule(request)}`
 
             const commonParams = {
                 id: uuid(),
-                tags: [`client-error:${error.code || 'UNKNOWN'}`],
+                tags: [
+                    `client-error:${error.code || 'UNKNOWN'}`,
+                    ...getSocketMetadataTags(socket)
+                ],
                 timingEvents: { startTime: Date.now(), startTimestamp: now() } as TimingEvents
             };
 
@@ -1076,7 +1082,8 @@ ${await this.suggestRule(request)}`
                 id: uuid(),
                 tags: [
                     `client-error:${error.code || 'UNKNOWN'}`,
-                    ...(isBadPreface ? ['client-error:bad-preface'] : [])
+                    ...(isBadPreface ? ['client-error:bad-preface'] : []),
+                    ...getSocketMetadataTags(socket)
                 ],
                 httpVersion: '2',