Skip to content

Commit fe629cf

Browse files
pimterrypimterry
committed
Use the latest 'ws' to support perfect WebSocket mirroring
We now correctly mirror raw headers, WebSocket keys, extensions (even unparseable stuff, though we can't accept it if the upstream actually does) and subprotocols (including accepting no-response cases).
1 parent b6983e1 commit fe629cf

4 files changed

Lines changed: 446 additions & 183 deletions

File tree

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,6 @@
200200
"socks-proxy-agent": "^7.0.0",
201201
"typed-error": "^3.0.2",
202202
"urlpattern-polyfill": "^8.0.0",
203-
"ws": "^8.8.0"
203+
"ws": "^8.20.0"
204204
}
205205
}

‎src/rules/websockets/websocket-step-impls.ts‎

Lines changed: 177 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,28 @@ import { Buffer } from 'buffer';
22
import * as net from 'net';
33
import * as url from 'url';
44
import * as http from 'http';
5+
import * as https from 'https';
6+
import { Duplex } from 'stream';
57

68
import * as _ from 'lodash';
79
import * as WebSocket from 'ws';
810

11+
// These were internal ws modules before 8.20.0, now officially exported.
12+
// @types/ws doesn't include types for these yet, so we type them manually:
13+
const { PerMessageDeflate, extension: wsExtension } = WebSocket as any as {
14+
PerMessageDeflate: {
15+
extensionName: string;
16+
new(options?: object, isServer?: boolean, maxPayload?: number): {
17+
accept(offers: object[]): void;
18+
params: object | null;
19+
};
20+
};
21+
extension: {
22+
parse(header: string): Record<string, object[]>;
23+
format(extensions: Record<string, object[]>): string;
24+
};
25+
};
26+
927
import {
1028
ClientServerChannel,
1129
deserializeBuffer,
@@ -30,9 +48,9 @@ import { resetOrDestroy } from '../../util/socket-util';
3048
import { isHttp2 } from '../../util/request-utils';
3149
import {
3250
findRawHeaders,
51+
flattenPairedRawHeaders,
3352
objectHeadersToRaw,
34-
pairFlatRawHeaders,
35-
rawHeadersToObjectPreservingCase
53+
pairFlatRawHeaders
3654
} from '../../util/header-utils';
3755
import { MaybePromise } from '@httptoolkit/util';
3856

@@ -77,10 +95,6 @@ export interface WebSocketStepImpl extends WebSocketStepDefinition {
7795
>;
7896
}
7997

80-
interface InterceptedWebSocketRequest extends http.IncomingMessage {
81-
upstreamWebSocketProtocol?: string | false;
82-
}
83-
8498
interface InterceptedWebSocket extends WebSocket {
8599
upstreamWebSocket: WebSocket;
86100
}
@@ -214,6 +228,38 @@ const rawResponse = (
214228
).join('\r\n') +
215229
'\r\n\r\n';
216230

231+
/**
232+
* Create a client-mode WebSocket on an existing stream, bypassing the normal
233+
* HTTP handshake. This is used when we've already performed the upgrade
234+
* handshake ourselves via http.request. We do this with custom APIs so that
235+
* we can fully control the handshake and mirror exact configurations.
236+
*/
237+
function createWebSocketFromStream(
238+
socket: Duplex,
239+
head: Buffer,
240+
options: {
241+
maxPayload?: number;
242+
extensions?: Record<string, object>;
243+
} = {}
244+
): WebSocket {
245+
const maxPayload = options.maxPayload ?? 0;
246+
247+
const ws = new (WebSocket as any)(null, undefined, { maxPayload });
248+
ws._isServer = false; // Client mode: mask frames per RFC 6455
249+
250+
if (options.extensions) {
251+
ws._extensions = options.extensions;
252+
}
253+
254+
ws.setSocket(socket, head, {
255+
allowSynchronousEvents: true,
256+
maxPayload,
257+
skipUTF8Validation: true // Preserve even invalid weird stuff
258+
});
259+
260+
return ws;
261+
}
262+
217263
export { PassThroughWebSocketStepOptions };
218264

219265
export class PassThroughWebSocketStepImpl extends PassThroughWebSocketStep {
@@ -225,14 +271,8 @@ export class PassThroughWebSocketStepImpl extends PassThroughWebSocketStep {
225271

226272
this.wsServer = new WebSocket.Server({
227273
noServer: true,
228-
// Mirror subprotocols back to the client:
229-
handleProtocols(protocols, request: InterceptedWebSocketRequest) {
230-
return request.upstreamWebSocketProtocol
231-
// If there's no upstream socket, default to mirroring the first protocol. This matches
232-
// WS's default behaviour - we could be stricter, but it'd be a breaking change.
233-
?? protocols.values().next().value
234-
?? false; // If there were no protocols specific and this is called for some reason
235-
},
274+
perMessageDeflate: true,
275+
skipUTF8Validation: true // Preserve even invalid weird stuff
236276
});
237277
this.wsServer.on('connection', (ws: InterceptedWebSocket) => {
238278
pipeWebSocket(ws, ws.upstreamWebSocket);
@@ -331,126 +371,145 @@ export class PassThroughWebSocketStepImpl extends PassThroughWebSocketStep {
331371
keepAlive: false // Not a thing for websockets: they take over the whole connection
332372
});
333373

334-
// We have to flatten the headers, as WS doesn't support raw headers - it builds its own
335-
// header object internally.
336-
const headers = rawHeadersToObjectPreservingCase(rawHeaders);
337-
338-
// Subprotocols have to be handled explicitly. WS takes control of the headers itself,
339-
// and checks the response, so we need to parse the client headers and use them manually:
340-
const originalSubprotocols = findRawHeaders(rawHeaders, 'sec-websocket-protocol')
341-
.flatMap(([_k, value]) => value.split(',').map(p => p.trim()));
342-
343-
// Drop empty subprotocols, to better handle mildly badly behaved clients
344-
const filteredSubprotocols = originalSubprotocols.filter(p => !!p);
345-
346-
// If the subprotocols are invalid (there are some empty strings, or an entirely empty value) then
347-
// WS will reject the upgrade. With this, we reset the header to the 'equivalent' valid version, to
348-
// avoid unnecessarily rejecting clients who send mildly wrong headers (empty protocol values).
349-
if (originalSubprotocols.length !== filteredSubprotocols.length) {
350-
if (filteredSubprotocols.length) {
351-
// Note that req.headers is auto-lowercased by Node, so we can ignore case
352-
req.headers['sec-websocket-protocol'] = filteredSubprotocols.join(',')
353-
} else {
354-
delete req.headers['sec-websocket-protocol'];
374+
// Strip any extension offers we can't handle (i.e. anything other than
375+
// permessage-deflate) to prevent the upstream from accepting them and causing trouble:
376+
const extensionHeaderValues = findRawHeaders(rawHeaders, 'sec-websocket-extensions');
377+
if (extensionHeaderValues.length > 0) {
378+
try {
379+
const parsed = wsExtension.parse(
380+
extensionHeaderValues.map(([_k, v]) => v).join(', ')
381+
);
382+
383+
// This is very unlikely - approximately zero other extensions exist in any form.
384+
const hasUnsupported = Object.keys(parsed)
385+
.some(name => name !== PerMessageDeflate.extensionName);
386+
if (hasUnsupported) {
387+
rawHeaders = rawHeaders.filter(([key]) =>
388+
key.toLowerCase() !== 'sec-websocket-extensions'
389+
);
390+
if (parsed[PerMessageDeflate.extensionName]) {
391+
rawHeaders.push(['Sec-WebSocket-Extensions', wsExtension.format({
392+
[PerMessageDeflate.extensionName]: parsed[PerMessageDeflate.extensionName]
393+
})]);
394+
}
395+
}
396+
} catch {
397+
// If we can't parse the client's offer, forward it as-is and let
398+
// the upstream handle/reject it:
355399
}
356400
}
357401

358-
const upstreamWebSocket = new WebSocket(wsUrl, filteredSubprotocols, {
359-
host: destination.hostname,
360-
port: destination.port,
402+
// Build the upstream request manually, mirroring the input as closely as possible:
403+
const isSecure = parsedUrl.protocol === 'wss:';
404+
const httpModule = isSecure ? https : http;
361405

362-
maxPayload: 0,
406+
const upstreamReqOptions: http.RequestOptions & https.RequestOptions = {
407+
hostname: destination.hostname,
408+
port: destination.port,
409+
path: parsedUrl.path,
410+
headers: flattenPairedRawHeaders(rawHeaders),
411+
setDefaultHeaders: false, // No auto-headers - we exactly mirror the client
412+
method: req.method!,
363413
agent,
364-
lookup: getDnsLookupFunction(this.lookupOptions),
365-
headers: _.omitBy(headers, (_v, headerName) =>
366-
headerName.toLowerCase().startsWith('sec-websocket') ||
367-
headerName.toLowerCase() === 'connection' ||
368-
headerName.toLowerCase() === 'upgrade'
369-
) as { [key: string]: string }, // Simplify to string - doesn't matter though, only used by http module anyway
370-
371-
// TLS options:
372-
...getUpstreamTlsOptions({
414+
lookup: getDnsLookupFunction(this.lookupOptions) as any,
415+
...(isSecure ? getUpstreamTlsOptions({
373416
hostname: effectiveHostname,
374417
port: effectivePort,
375418
ignoreHostHttpsErrors: this.ignoreHostHttpsErrors,
376419
clientCertificateHostMap: this.clientCertificateHostMap,
377420
trustedCAs,
378-
})
379-
} as WebSocket.ClientOptions & { lookup: any, maxPayload: number });
421+
}) : {})
422+
};
380423

381-
const upstreamReq = (upstreamWebSocket as any as { _req: http.ClientRequest })._req;
424+
const upstreamReq = httpModule.request(upstreamReqOptions);
382425

383-
if (options.emitEventCallback) {
384-
// This is slower than req.getHeaders(), but gives us (roughly) the correct casing
385-
// of the headers as sent. Still not perfect (loses dupe ordering) but at least it
386-
// generally matches what's actually sent on the wire.
387-
const rawHeaders = upstreamReq.getRawHeaderNames().map((headerName) => {
388-
const value = upstreamReq.getHeader(headerName);
389-
if (!value) return [];
390-
if (Array.isArray(value)) {
391-
return value.map(v => [headerName, v]);
392-
} else {
393-
return [[headerName, value.toString()]];
394-
}
395-
}).flat() as RawHeaders;
426+
// Track the upstream WebSocket so the incomingSocket error handler can close it:
427+
let upstreamWebSocket: WebSocket | undefined;
396428

429+
if (options.emitEventCallback) {
397430
// This effectively matches the URL preprocessing logic in MockttpServer.preprocessRequest,
398431
// so that the resulting event matches the req.url property elsewhere.
399-
const urlHost = getEffectiveHostname(upstreamReq.host, req.socket, rawHeaders);
432+
const urlHost = getEffectiveHostname(effectiveHostname, req.socket, rawHeaders);
433+
434+
const wsProtocol = parsedUrl.protocol!.replace(/^http/, 'ws').replace(/:$/, '');
435+
436+
const subprotocols = findRawHeaders(rawHeaders, 'sec-websocket-protocol')
437+
.flatMap(([_k, v]) => v.split(',').map(s => s.trim()).filter(s => !!s));
400438

401439
options.emitEventCallback('passthrough-websocket-connect', {
402-
method: upstreamReq.method,
403-
protocol: upstreamReq.protocol
404-
.replace(/:$/, '')
405-
.replace(/^http/, 'ws'),
440+
method: req.method!,
441+
protocol: wsProtocol,
406442
hostname: urlHost,
407443
port: effectivePort.toString(),
408-
path: upstreamReq.path,
409-
rawHeaders: rawHeaders,
410-
subprotocols: filteredSubprotocols
444+
path: parsedUrl.path || '/',
445+
rawHeaders,
446+
subprotocols
411447
});
412448
}
413449

414450
if (options.keyLogStream) {
415451
upstreamReq.on('socket', (socket) => {
416-
socket.on('keylog', (line) => options.keyLogStream!.write(line));
452+
socket.on('keylog', (line: Buffer) => options.keyLogStream!.write(line));
417453
});
418454
}
419455

420-
upstreamWebSocket.once('open', () => {
421-
// Used in the subprotocol selection handler during the upgrade:
422-
(req as InterceptedWebSocketRequest).upstreamWebSocketProtocol = upstreamWebSocket.protocol || false;
456+
upstreamReq.on('upgrade', (upstreamRes, upstreamSocket, upgradeHead) => {
457+
// Handle permessage-deflate extension negotiation. If the upstream server
458+
// committed to extensions we can't set up, we must kill the connection rather
459+
// than silently mishandling compressed frames:
460+
const responseExtensionHeader = upstreamRes.headers['sec-websocket-extensions'];
461+
462+
let extensions: Record<string, object> | undefined;
463+
try {
464+
if (responseExtensionHeader) {
465+
const parsed = wsExtension.parse(responseExtensionHeader);
466+
if (parsed[PerMessageDeflate.extensionName]) {
467+
const pmd = new PerMessageDeflate({}, false); // false = client mode
468+
pmd.accept(parsed[PerMessageDeflate.extensionName]);
469+
extensions = { [PerMessageDeflate.extensionName]: pmd };
470+
}
471+
}
472+
} catch (e) {
473+
console.warn('Failed to negotiate WebSocket extensions:', e);
474+
upstreamSocket.destroy();
475+
incomingSocket.end();
476+
return;
477+
}
478+
479+
upstreamWebSocket = createWebSocketFromStream(upstreamSocket, upgradeHead, {
480+
maxPayload: 0,
481+
extensions
482+
});
483+
484+
// Set req.headers to match exactly what the upstream confirmed, so ws's
485+
// handleUpgrade negotiates the same values downstream without any issues
486+
// from malformed original headers:
487+
if (!extensions) {
488+
delete req.headers['sec-websocket-extensions'];
489+
}
490+
491+
// For WS's sake, we simplify the subprotocol header to only the upstream-selected value so
492+
// that it can just accept as is, and ignore any other badly behaved client's headers.
493+
const serverProtocol = upstreamRes.headers['sec-websocket-protocol'];
494+
if (serverProtocol?.trim()) {
495+
req.headers['sec-websocket-protocol'] = serverProtocol;
496+
} else {
497+
delete req.headers['sec-websocket-protocol'];
498+
}
423499

424500
this.wsServer!.handleUpgrade(req, incomingSocket, head, (ws) => {
425-
(<InterceptedWebSocket> ws).upstreamWebSocket = upstreamWebSocket;
501+
(<InterceptedWebSocket> ws).upstreamWebSocket = upstreamWebSocket!;
426502
incomingSocket.emit('ws-upgrade', ws);
427-
this.wsServer!.emit('connection', ws); // This pipes the connections together
503+
this.wsServer!.emit('connection', ws);
428504
});
429505
});
430506

431-
// If the upstream says no, we say no too.
432-
let unexpectedResponse = false;
433-
upstreamWebSocket.on('unexpected-response', (req, res) => {
434-
console.log(`Unexpected websocket response from ${wsUrl}: ${res.statusCode}`);
435-
436-
// Clean up the downstream connection
437-
mirrorRejection(incomingSocket, res, this.simulateConnectionErrors).then(() => {
438-
// Clean up the upstream connection (WS would do this automatically, but doesn't if you listen to this event)
439-
// See https://github.com/websockets/ws/blob/45e17acea791d865df6b255a55182e9c42e5877a/lib/websocket.js#L1050
440-
// We don't match that perfectly, but this should be effectively equivalent:
441-
req.destroy();
442-
if (res.socket?.destroyed === false) {
443-
res.socket.destroy();
444-
}
445-
unexpectedResponse = true; // So that we ignore this in the error handler
446-
upstreamWebSocket.terminate();
447-
});
507+
upstreamReq.on('response', (upstreamRes) => {
508+
console.log(`Unexpected websocket response from ${wsUrl}: ${upstreamRes.statusCode}`);
509+
mirrorRejection(incomingSocket, upstreamRes, this.simulateConnectionErrors);
448510
});
449511

450-
// If there's some other error, we just kill the socket:
451-
upstreamWebSocket.on('error', (e) => {
452-
if (unexpectedResponse) return; // Handled separately above
453-
512+
upstreamReq.on('error', (e) => {
454513
console.warn(e);
455514
if (this.simulateConnectionErrors) {
456515
resetOrDestroy(incomingSocket);
@@ -459,7 +518,15 @@ export class PassThroughWebSocketStepImpl extends PassThroughWebSocketStep {
459518
}
460519
});
461520

462-
incomingSocket.on('error', () => upstreamWebSocket.close(1011)); // Internal error
521+
incomingSocket.on('error', () => {
522+
if (upstreamWebSocket) {
523+
upstreamWebSocket.close(1011);
524+
} else {
525+
upstreamReq.destroy();
526+
}
527+
});
528+
529+
upstreamReq.end();
463530
}
464531

465532
/**
@@ -516,7 +583,11 @@ export class EchoWebSocketStepImpl extends EchoWebSocketStep {
516583
private initializeWsServer() {
517584
if (this.wsServer) return;
518585

519-
this.wsServer = new WebSocket.Server({ noServer: true });
586+
this.wsServer = new WebSocket.Server({
587+
noServer: true,
588+
perMessageDeflate: true,
589+
skipUTF8Validation: true // Preserve even invalid weird stuff
590+
});
520591
this.wsServer.on('connection', (ws: WebSocket) => {
521592
pipeWebSocket(ws, ws);
522593
});
@@ -539,7 +610,11 @@ export class ListenWebSocketStepImpl extends ListenWebSocketStep {
539610
private initializeWsServer() {
540611
if (this.wsServer) return;
541612

542-
this.wsServer = new WebSocket.Server({ noServer: true });
613+
this.wsServer = new WebSocket.Server({
614+
noServer: true,
615+
perMessageDeflate: true,
616+
skipUTF8Validation: true // Accept even invalid weird stuff
617+
});
543618
this.wsServer.on('connection', (ws: WebSocket) => {
544619
// Accept but ignore the incoming websocket data
545620
ws.resume();

0 commit comments

Comments
 (0)