From 16fc315afa7c7090af09f2ef27c5d714a1ace5e0 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Tue, 20 Apr 2021 15:48:24 -0500
Subject: [PATCH] Catch socket close during protocol handshake

zlib errors can cause the protocol handshake to not reject (until the
timeout).
---
 lib/vscode/src/vs/server/entry.ts         |  4 +++-
 lib/vscode/src/vs/server/node/protocol.ts | 22 +++++++++++++++++-----
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/lib/vscode/src/vs/server/entry.ts b/lib/vscode/src/vs/server/entry.ts
index 901bcda6..2f458f40 100644
--- a/lib/vscode/src/vs/server/entry.ts
+++ b/lib/vscode/src/vs/server/entry.ts
@@ -6,7 +6,9 @@ import { logger } from 'vs/server/node/logger';
 import { enableCustomMarketplace } from 'vs/server/node/marketplace';
 import { Vscode } from 'vs/server/node/server';
 
-setUnexpectedErrorHandler((error) => logger.warn(error instanceof Error ? error.message : error));
+setUnexpectedErrorHandler((error) => {
+	logger.warn('Uncaught error', field('error', error instanceof Error ? error.message : error));
+});
 enableCustomMarketplace();
 proxyAgent.monkeyPatch(true);
 
diff --git a/lib/vscode/src/vs/server/node/protocol.ts b/lib/vscode/src/vs/server/node/protocol.ts
index e0ab987c..fd61d4c9 100644
--- a/lib/vscode/src/vs/server/node/protocol.ts
+++ b/lib/vscode/src/vs/server/node/protocol.ts
@@ -54,10 +54,24 @@ export class Protocol extends PersistentProtocol {
 	 */
 	public handshake(): Promise<ConnectionTypeRequest> {
 		this.logger.debug('Initiating handshake...');
+
 		return new Promise((resolve, reject) => {
+			const cleanup = () => {
+				handler.dispose();
+				onClose.dispose();
+				clearTimeout(timeout);
+			};
+
+			const onClose = this.onSocketClose(() => {
+				cleanup();
+				this.logger.debug('Handshake failed');
+				reject(new Error('Protocol socket closed unexpectedly'));
+			});
+
 			const timeout = setTimeout(() => {
+				cleanup();
 				this.logger.debug('Handshake timed out');
-				reject(new Error('protocol handshake timed out'));
+				reject(new Error('Protocol handshake timed out'));
 			}, 10000); // Matches the client timeout.
 
 			const handler = this.onControlMessage((rawMessage) => {
@@ -69,16 +83,14 @@ export class Protocol extends PersistentProtocol {
 						case 'auth':
 							return this.authenticate(message);
 						case 'connectionType':
-							handler.dispose();
-							clearTimeout(timeout);
+							cleanup();
 							this.logger.debug('Handshake completed');
 							return resolve(message);
 						default:
 							throw new Error('Unrecognized message type');
 					}
 				} catch (error) {
-					handler.dispose();
-					clearTimeout(timeout);
+					cleanup();
 					reject(error);
 				}
 			});