From 7385b31f13343f40a3d05710513df0fa9225674d Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Wed, 16 Dec 2020 14:21:59 -0800
Subject: [PATCH] fix(driver): stop sending protocol messages after disconnect
 (#4688)

When the client only closes the input pipe, we are still
sending protocol messages over the output pipe. This could
probably lead to some errors, e.g. write buffer being full.
---
 src/driver.ts             | 6 ++++--
 src/protocol/transport.ts | 6 +++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/driver.ts b/src/driver.ts
index 69d97edeb0..ad190a133e 100644
--- a/src/driver.ts
+++ b/src/driver.ts
@@ -50,15 +50,17 @@ export function runServer() {
 
   const dispatcherConnection = new DispatcherConnection();
   const transport = new Transport(process.stdout, process.stdin);
+  transport.onmessage = message => dispatcherConnection.dispatch(JSON.parse(message));
+  dispatcherConnection.onmessage = message => transport.send(JSON.stringify(message));
   transport.onclose = async () => {
+    // Drop any messages during shutdown on the floor.
+    dispatcherConnection.onmessage = () => {};
     // Force exit after 30 seconds.
     setTimeout(() => process.exit(0), 30000);
     // Meanwhile, try to gracefully close all browsers.
     await gracefullyCloseAll();
     process.exit(0);
   };
-  transport.onmessage = message => dispatcherConnection.dispatch(JSON.parse(message));
-  dispatcherConnection.onmessage = message => transport.send(JSON.stringify(message));
 
   const playwright = new Playwright(__dirname, require('../browsers.json')['browsers']);
   new PlaywrightDispatcher(dispatcherConnection.rootDispatcher(), playwright);
diff --git a/src/protocol/transport.ts b/src/protocol/transport.ts
index c23329a082..9b4be17edf 100644
--- a/src/protocol/transport.ts
+++ b/src/protocol/transport.ts
@@ -47,7 +47,11 @@ export class Transport {
     this._endian = endian;
     this._closeableStream = closeable;
     pipeRead.on('data', buffer => this._dispatch(buffer));
-    pipeRead.on('close', () => this.onclose && this.onclose());
+    pipeRead.on('close', () => {
+      this._closed = true;
+      if (this.onclose)
+        this.onclose();
+    });
     this.onmessage = undefined;
     this.onclose = undefined;
   }