Refactor code to use `await`

src/vs/platform/remote/common/remoteAgentConnection.ts

@@ -8,7 +8,7 @@ import { generateUuid } from 'vs/base/common/uuid';
 import { RemoteAgentConnectionContext } from 'vs/platform/remote/common/remoteAgentEnvironment';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { VSBuffer } from 'vs/base/common/buffer';
-import { Emitter } from 'vs/base/common/event';
+import { Emitter, Event } from 'vs/base/common/event';
 import { RemoteAuthorityResolverError } from 'vs/platform/remote/common/remoteAuthorityResolver';
 import { isPromiseCanceledError, onUnexpectedError } from 'vs/base/common/errors';
 import { ISignService } from 'vs/platform/sign/common/sign';
@@ -86,6 +86,38 @@ export interface ISocketFactory {
 	connect(host: string, port: number, query: string, callback: IConnectCallback): void;
 }
 
+async function readOneControlMessage<T>(protocol: PersistentProtocol): Promise<T> {
+	const raw = await Event.toPromise(protocol.onControlMessage);
+	const msg = JSON.parse(raw.toString());
+	const error = getErrorFromMessage(msg);
+	if (error) {
+		throw error;
+	}
+	return msg;
+}
+
+function waitWithTimeout<T>(promise: Promise<T>, timeout: number): Promise<T> {
+	return new Promise<T>((resolve, reject) => {
+		const timeoutToken = setTimeout(() => {
+			const error: any = new Error('Timeout');
+			error.code = 'ETIMEDOUT';
+			error.syscall = 'connect';
+			reject(error);
+		}, timeout);
+
+		promise.then(
+			(result) => {
+				clearTimeout(timeoutToken);
+				resolve(result);
+			},
+			(error) => {
+				clearTimeout(timeoutToken);
+				reject(error);
+			}
+		);
+	});
+}
+
 async function connectToRemoteExtensionHostAgent(options: ISimpleConnectionOptions, connectionType: ConnectionType, args: any | undefined): Promise<{ protocol: PersistentProtocol; ownsProtocol: boolean; }> {
 	const logPrefix = connectLogPrefix(options, connectionType);
 	const { protocol, ownsProtocol } = await new Promise<{ protocol: PersistentProtocol; ownsProtocol: boolean; }>((c, e) => {
@@ -113,69 +145,54 @@ async function connectToRemoteExtensionHostAgent(options: ISimpleConnectionOptio
 		);
 	});
 
-	return new Promise<{ protocol: PersistentProtocol; ownsProtocol: boolean; }>((c, e) => {
-
-		const errorTimeoutToken = setTimeout(() => {
-			const error: any = new Error('handshake timeout');
-			error.code = 'ETIMEDOUT';
-			error.syscall = 'connect';
-			options.logService.error(`${logPrefix} the handshake took longer than 10 seconds. Error:`);
-			options.logService.error(error);
-			if (ownsProtocol) {
-				safeDisposeProtocolAndSocket(protocol);
-			}
-			e(error);
-		}, 10000);
+	options.logService.trace(`${logPrefix} 3/6. sending AuthRequest control message.`);
+	const authRequest: AuthRequest = {
+		type: 'auth',
+		auth: options.connectionToken || '00000000000000000000'
+	};
+	protocol.sendControl(VSBuffer.fromString(JSON.stringify(authRequest)));
 
-		const messageRegistration = protocol.onControlMessage(async raw => {
-			const msg = <HandshakeMessage>JSON.parse(raw.toString());
-			// Stop listening for further events
-			messageRegistration.dispose();
+	try {
+		const msg = await waitWithTimeout(readOneControlMessage<HandshakeMessage>(protocol), 10000);
 
-			const error = getErrorFromMessage(msg);
-			if (error) {
-				options.logService.error(`${logPrefix} received error control message when negotiating connection. Error:`);
-				options.logService.error(error);
-				if (ownsProtocol) {
-					safeDisposeProtocolAndSocket(protocol);
-				}
-				return e(error);
-			}
+		if (msg.type !== 'sign' || typeof msg.data !== 'string') {
+			const error: any = new Error('Unexpected handshake message');
+			error.code = 'VSCODE_CONNECTION_ERROR';
+			throw error;
+		}
 
-			if (msg.type === 'sign') {
-				options.logService.trace(`${logPrefix} 4/6. received SignRequest control message.`);
-				const signed = await options.signService.sign(msg.data);
-				const connTypeRequest: ConnectionTypeRequest = {
-					type: 'connectionType',
-					commit: options.commit,
-					signedData: signed,
-					desiredConnectionType: connectionType
-				};
-				if (args) {
-					connTypeRequest.args = args;
-				}
-				options.logService.trace(`${logPrefix} 5/6. sending ConnectionTypeRequest control message.`);
-				protocol.sendControl(VSBuffer.fromString(JSON.stringify(connTypeRequest)));
-				clearTimeout(errorTimeoutToken);
-				c({ protocol, ownsProtocol });
-			} else {
-				const error = new Error('handshake error');
-				options.logService.error(`${logPrefix} received unexpected control message. Error:`);
-				options.logService.error(error);
-				if (ownsProtocol) {
-					safeDisposeProtocolAndSocket(protocol);
-				}
-				e(error);
-			}
-		});
+		options.logService.trace(`${logPrefix} 4/6. received SignRequest control message.`);
 
-		options.logService.trace(`${logPrefix} 3/6. sending AuthRequest control message.`);
-		const authRequest: AuthRequest = {
-			type: 'auth',
-			auth: options.connectionToken || '00000000000000000000'
+		const signed = await options.signService.sign(msg.data);
+		const connTypeRequest: ConnectionTypeRequest = {
+			type: 'connectionType',
+			commit: options.commit,
+			signedData: signed,
+			desiredConnectionType: connectionType
 		};
-		protocol.sendControl(VSBuffer.fromString(JSON.stringify(authRequest)));
-	});
+		if (args) {
+			connTypeRequest.args = args;
+		}
+
+		options.logService.trace(`${logPrefix} 5/6. sending ConnectionTypeRequest control message.`);
+		protocol.sendControl(VSBuffer.fromString(JSON.stringify(connTypeRequest)));
+
+		return { protocol, ownsProtocol };
+
+	} catch (error) {
+		if (error && error.code === 'ETIMEDOUT') {
+			options.logService.error(`${logPrefix} the handshake took longer than 10 seconds. Error:`);
+			options.logService.error(error);
+		}
+		if (error && error.code === 'VSCODE_CONNECTION_ERROR') {
+			options.logService.error(`${logPrefix} received error control message when negotiating connection. Error:`);
+			options.logService.error(error);
+		}
+		if (ownsProtocol) {
+			safeDisposeProtocolAndSocket(protocol);
+		}
+		throw error;
+	}
 }
 
 interface IManagementConnectionResult {