From 59d591bce6feebe60642782d890854e04131ddc9 Mon Sep 17 00:00:00 2001
From: Max Schmitt <max@schmitt.mx>
Date: Wed, 26 May 2021 16:51:38 +0200
Subject: [PATCH] chore(port-forwarding): validate forwarded ports on the
 client side (#6756)

---
 src/browserServerImpl.ts  | 2 +-
 src/client/browserType.ts | 1 +
 src/client/playwright.ts  | 1 +
 src/client/socksSocket.ts | 9 +++++++--
 4 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/browserServerImpl.ts b/src/browserServerImpl.ts
index 7793bdb010..a5a9bd8e14 100644
--- a/src/browserServerImpl.ts
+++ b/src/browserServerImpl.ts
@@ -97,7 +97,7 @@ export class BrowserServerLauncherImpl implements BrowserServerLauncher {
       portForwardingServer.enablePortForwarding(ports);
     });
     const incomingSocksSocketHandler = (socket: SocksInterceptedSocketHandler) => {
-      playwrightDispatcher._dispatchEvent('incomingSocksSocket', { socket: new SocksSocketDispatcher(playwrightDispatcher._scope, socket) });
+      playwrightDispatcher._dispatchEvent('incomingSocksSocket', { socket: new SocksSocketDispatcher(playwrightDispatcher, socket) });
     };
     portForwardingServer.on('incomingSocksSocket', incomingSocksSocketHandler);
 
diff --git a/src/client/browserType.ts b/src/client/browserType.ts
index 4d5854cc52..10122efd08 100644
--- a/src/client/browserType.ts
+++ b/src/client/browserType.ts
@@ -189,6 +189,7 @@ export class BrowserType extends ChannelOwner<channels.BrowserTypeChannel, chann
               await playwright._channel.enablePortForwarding({
                 ports: params._forwardPorts,
               });
+              playwright._forwardPorts = params._forwardPorts;
             } catch (err) {
               reject(err);
               return;
diff --git a/src/client/playwright.ts b/src/client/playwright.ts
index 6802d3a740..09faae15cb 100644
--- a/src/client/playwright.ts
+++ b/src/client/playwright.ts
@@ -44,6 +44,7 @@ export class Playwright extends ChannelOwner<channels.PlaywrightChannel, channel
   readonly selectors: Selectors;
   readonly errors: { TimeoutError: typeof TimeoutError };
   private _selectorsOwner: SelectorsOwner;
+  _forwardPorts: number[] = [];
 
   constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.PlaywrightInitializer) {
     super(parent, type, guid, initializer);
diff --git a/src/client/socksSocket.ts b/src/client/socksSocket.ts
index b10a638958..8af0ce499e 100644
--- a/src/client/socksSocket.ts
+++ b/src/client/socksSocket.ts
@@ -17,6 +17,7 @@
 import net from 'net';
 
 import * as channels from '../protocol/channels';
+import { Playwright } from './playwright';
 import { assert, isLocalIpAddress, isUnderTest } from '../utils/utils';
 import { ChannelOwner } from './channelOwner';
 
@@ -28,11 +29,13 @@ export class SocksSocket extends ChannelOwner<channels.SocksSocketChannel, chann
 
   constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.SocksSocketInitializer) {
     super(parent, type, guid, initializer);
-    this._connection.on('disconnect', () => this._socket.end());
+    assert(parent instanceof Playwright);
+
+    assert(parent._forwardPorts.includes(this._initializer.dstPort));
+    assert(isLocalIpAddress(this._initializer.dstAddr));
 
     if (isUnderTest() && process.env.PW_TEST_PROXY_TARGET)
       this._initializer.dstPort = Number(process.env.PW_TEST_PROXY_TARGET);
-    assert(isLocalIpAddress(this._initializer.dstAddr));
 
     this._socket = net.createConnection(this._initializer.dstPort, this._initializer.dstAddr);
     this._socket.on('error', (err: Error) => this._channel.error({error: String(err)}));
@@ -50,6 +53,8 @@ export class SocksSocket extends ChannelOwner<channels.SocksSocketChannel, chann
       this._socket.write(Buffer.from(data, 'base64'));
     });
     this._channel.on('close', () => this._socket.end());
+
+    this._connection.on('disconnect', () => this._socket.end());
   }
 
   async write(data: Buffer): Promise<void> {