Catch socket close during protocol handshake

zlib errors can cause the protocol handshake to not reject (until the
timeout).

Author: code-asher

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);
 

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);
 				}
 			});