From e1681a3a89ce6cc5c62dbba0a11001f156eb08d4 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Thu, 15 Sep 2022 09:53:55 -0700
Subject: [PATCH] cherry-pick(#17349): fix(socks proxy): destroy sockets on
 close to avoid hanging (#17366)

---
 packages/playwright-core/src/common/socksProxy.ts | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/packages/playwright-core/src/common/socksProxy.ts b/packages/playwright-core/src/common/socksProxy.ts
index 2c8c5a82708ef..a46d20af55edf 100644
--- a/packages/playwright-core/src/common/socksProxy.ts
+++ b/packages/playwright-core/src/common/socksProxy.ts
@@ -294,6 +294,8 @@ export class SocksProxy extends EventEmitter implements SocksConnectionClient {
 
   private _server: net.Server;
   private _connections = new Map<string, SocksConnection>();
+  private _sockets = new Set<net.Socket>();
+  private _closed = false;
 
   constructor() {
     super();
@@ -302,6 +304,14 @@ export class SocksProxy extends EventEmitter implements SocksConnectionClient {
       const connection = new SocksConnection(uid, socket, this);
       this._connections.set(uid, connection);
     });
+    this._server.on('connection', socket => {
+      if (this._closed) {
+        socket.destroy();
+        return;
+      }
+      this._sockets.add(socket);
+      socket.once('close', () => this._sockets.delete(socket));
+    });
   }
 
   async listen(port: number): Promise<number> {
@@ -315,6 +325,10 @@ export class SocksProxy extends EventEmitter implements SocksConnectionClient {
   }
 
   async close() {
+    this._closed = true;
+    for (const socket of this._sockets)
+      socket.destroy();
+    this._sockets.clear();
     await new Promise(f => this._server.close(f));
   }