Partial extension host, some restructuring

I didn't like how the inner objects accessed parent objects, so I
restructured all that.

connection.ts

@@ -1,69 +1,176 @@
-import { ClientConnectionEvent } from "vs/base/parts/ipc/common/ipc";
+import * as cp from "child_process";
-import { ConnectionType } from "vs/platform/remote/common/remoteAgentConnection";
-import { Emitter } from "vs/base/common/event";
-import { PersistentProtocol, ISocket } from "vs/base/parts/ipc/common/ipc.net";
-import { VSBuffer } from "vs/base/common/buffer";
 
-export interface Server {
+import { getPathFromAmdModule } from "vs/base/common/amd";
-	readonly _onDidClientConnect: Emitter<ClientConnectionEvent>;
+import { VSBuffer } from "vs/base/common/buffer";
-	readonly connections: Map<ConnectionType, Map<string, Connection>>;
+import { Emitter } from "vs/base/common/event";
-}
+import { ISocket } from "vs/base/parts/ipc/common/ipc.net";
+import { NodeSocket, WebSocketNodeSocket } from "vs/base/parts/ipc/node/ipc.net";
+import { ILogService } from "vs/platform/log/common/log";
+import { IExtHostReadyMessage, IExtHostSocketMessage } from "vs/workbench/services/extensions/common/extensionHostProtocol";
+
+import { Protocol } from "vs/server/protocol";
 
 export abstract class Connection {
 	private readonly _onClose = new Emitter<void>();
 	public readonly onClose = this._onClose.event;
 
 	private timeout: NodeJS.Timeout | undefined;
-	private readonly wait = 1000 * 60 * 60;
+	private readonly wait = 1000 * 60;
 
-	public constructor(
+	private closed: boolean = false;
-		protected readonly server: Server,
-		private readonly protocol: PersistentProtocol,
-	) {
-		// onClose seems to mean we want to disconnect, so dispose immediately.
-		this.protocol.onClose(() => this.dispose());
 
-		// If the socket closes, we want to wait before disposing so we can
+	public constructor(protected protocol: Protocol) {
-		// reconnect.
+		// onClose seems to mean we want to disconnect, so close immediately.
-		this.protocol.onSocketClose(() => {
+		protocol.onClose(() => this.close());
+
+		// If the socket closes, we want to wait before closing so we can
+		// reconnect in the meantime.
+		protocol.onSocketClose(() => {
 			this.timeout = setTimeout(() => {
-				this.dispose();
+				this.close();
 			}, this.wait);
 		});
 	}
 
 	/**
-	 * Completely close and clean up the connection. Should only do this once we
+	 * Set up the connection on a new socket.
-	 * don't need or want the connection. It cannot be re-used after this.
 	 */
-	public dispose(): void {
+	public reconnect(protocol: Protocol, buffer: VSBuffer): void {
-		this.protocol.sendDisconnect();
+		if (this.closed) {
-		this.protocol.getSocket().end();
+			throw new Error("Cannot reconnect to closed connection");
-		this.protocol.dispose();
+		}
-		this._onClose.fire();
+		clearTimeout(this.timeout as any); // Not sure why the type doesn't work.
+		this.protocol = protocol;
+		this.connect(protocol.getSocket(), buffer);
 	}
 
-	public reconnect(socket: ISocket, buffer: VSBuffer): void {
+	/**
-		clearTimeout(this.timeout as any); // Not sure why the type doesn't work.
+	 * Close and clean up connection. This will also kill the socket the
+	 * connection is on. Probably not safe to reconnect once this has happened.
+	 */
+	protected close(): void {
+		if (!this.closed) {
+			this.closed = true;
+			this.protocol.sendDisconnect();
+			this.dispose();
+			this.protocol.dispose();
+			this._onClose.fire();
+		}
+	}
+
+	/**
+	 * Clean up the connection.
+	 */
+	protected abstract dispose(): void;
+
+	/**
+	 * Connect to a new socket.
+	 */
+	protected abstract connect(socket: ISocket, buffer: VSBuffer): void;
+}
+
+/**
+ * Used for all the IPC channels.
+ */
+export class ManagementConnection extends Connection {
+	protected dispose(): void {
+		// Nothing extra to do here.
+	}
+
+	protected connect(socket: ISocket, buffer: VSBuffer): void {
 		this.protocol.beginAcceptReconnection(socket, buffer);
 		this.protocol.endAcceptReconnection();
 	}
 }
 
 /**
- * The management connection is used for all the IPC channels.
+ * Manage the extension host process.
  */
-export class ManagementConnection extends Connection {
+export class ExtensionHostConnection extends Connection {
-	public constructor(server: Server, protocol: PersistentProtocol) {
+	private process: cp.ChildProcess;
-		super(server, protocol);
+
-		// This will communicate back to the IPCServer that a new client has
+	public constructor(protocol: Protocol, private readonly log: ILogService) {
-		// connected.
+		super(protocol);
-		this.server._onDidClientConnect.fire({
+		const socket = this.protocol.getSocket();
-			protocol,
+		const buffer = this.protocol.readEntireBuffer();
-			onDidClientDisconnect: this.onClose,
+		this.process = this.spawn(socket, buffer);
+	}
+
+	protected dispose(): void {
+		this.process.kill();
+	}
+
+	protected connect(socket: ISocket, buffer: VSBuffer): void {
+		this.sendInitMessage(socket, buffer);
+	}
+
+	private sendInitMessage(nodeSocket: ISocket, buffer: VSBuffer): void {
+		const socket = nodeSocket instanceof NodeSocket
+			? nodeSocket.socket
+			: (nodeSocket as WebSocketNodeSocket).socket.socket;
+
+		socket.pause();
+
+		const initMessage: IExtHostSocketMessage = {
+			type: "VSCODE_EXTHOST_IPC_SOCKET",
+			initialDataChunk: (buffer.buffer as Buffer).toString("base64"),
+			skipWebSocketFrames: nodeSocket instanceof NodeSocket,
+		};
+
+		this.process.send(initMessage, socket);
+	}
+
+	private spawn(socket: ISocket, buffer: VSBuffer): cp.ChildProcess {
+		const proc = cp.fork(
+			getPathFromAmdModule(require, "bootstrap-fork"),
+			[
+				"--type=extensionHost",
+				`--uriTransformerPath=${getPathFromAmdModule(require, "vs/server/transformer")}`
+			],
+			{
+				env: {
+					...process.env,
+					AMD_ENTRYPOINT: "vs/workbench/services/extensions/node/extensionHostProcess",
+					PIPE_LOGGING: "true",
+					VERBOSE_LOGGING: "true",
+					VSCODE_EXTHOST_WILL_SEND_SOCKET: "true",
+					VSCODE_HANDLES_UNCAUGHT_ERRORS: "true",
+					VSCODE_LOG_STACK: "false",
+				},
+				silent: true,
+			},
+		);
+
+		proc.on("error", (error) => {
+			console.error(error);
+			this.close();
 		});
+
+		proc.on("exit", (code, signal) => {
+			console.error("Extension host exited", { code, signal });
+			this.close();
+		});
+
+		proc.stdout.setEncoding("utf8");
+		proc.stderr.setEncoding("utf8");
+		proc.stdout.on("data", (data) => this.log.info("Extension host stdout", data));
+		proc.stderr.on("data", (data) => this.log.error("Extension host stderr", data));
+		proc.on("message", (event) => {
+			if (event && event.type === "__$console") {
+				const severity = this.log[event.severity] ? event.severity : "info";
+				this.log[severity]("Extension host", event.arguments);
+			}
+		});
+
+		const listen = (message: IExtHostReadyMessage) => {
+			if (message.type === "VSCODE_EXTHOST_IPC_READY") {
+				proc.removeListener("message", listen);
+				this.sendInitMessage(socket, buffer);
+			}
+		};
+
+		proc.on("message", listen);
+
+		return proc;
 	}
 }
-
-export class ExtensionHostConnection extends Connection {
-}

protocol.ts

@@ -0,0 +1,98 @@
+import * as crypto from "crypto";
+import * as net from "net";
+
+import { VSBuffer } from "vs/base/common/buffer";
+import { NodeSocket, WebSocketNodeSocket } from "vs/base/parts/ipc/node/ipc.net";
+import { PersistentProtocol } from "vs/base/parts/ipc/common/ipc.net";
+import { AuthRequest, ConnectionTypeRequest, HandshakeMessage } from "vs/platform/remote/common/remoteAgentConnection";
+
+export interface SocketOptions {
+	readonly reconnectionToken: string;
+	readonly reconnection: boolean;
+	readonly skipWebSocketFrames: boolean;
+}
+
+export class Protocol extends PersistentProtocol {
+	public constructor(
+		secWebsocketKey: string,
+		socket: net.Socket,
+		public readonly options: SocketOptions,
+	) {
+		super(
+			options.skipWebSocketFrames
+				? new NodeSocket(socket)
+				: new WebSocketNodeSocket(new NodeSocket(socket)),
+		);
+		socket.on("error", () => this.dispose());
+		socket.on("end", () => this.dispose());
+
+		// This magic value is specified by the websocket spec.
+		const magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
+		const reply = crypto.createHash("sha1")
+			.update(secWebsocketKey + magic)
+			.digest("base64");
+
+		socket.write([
+			"HTTP/1.1 101 Switching Protocols",
+			"Upgrade: websocket",
+			"Connection: Upgrade",
+			`Sec-WebSocket-Accept: ${reply}`,
+		].join("\r\n") + "\r\n\r\n");
+	}
+
+	public dispose(error?: Error): void {
+		if (error) {
+			this.sendMessage({ type: "error", reason: error.message });
+		}
+		super.dispose();
+		this.getSocket().dispose();
+	}
+
+	/**
+	 * Perform a handshake to get a connection request.
+	 */
+	public handshake(): Promise<ConnectionTypeRequest> {
+		return new Promise((resolve, reject) => {
+			const handler = this.onControlMessage((rawMessage) => {
+				try {
+					const message = JSON.parse(rawMessage.toString());
+					switch (message.type) {
+						case "auth": return this.authenticate(message);
+						case "connectionType":
+							handler.dispose();
+							return resolve(message);
+						default: throw new Error("Unrecognized message type");
+					}
+				} catch (error) {
+					handler.dispose();
+					reject(error);
+				}
+			});
+		});
+	}
+
+	/**
+	 * TODO: This ignores the authentication process entirely for now.
+	 */
+	private authenticate(_message: AuthRequest): void {
+		this.sendMessage({
+			type: "sign",
+			data: "",
+		});
+	}
+
+	/**
+	 * TODO: implement.
+	 */
+	public tunnel(): void {
+		throw new Error("Tunnel is not implemented yet");
+	}
+
+	/**
+	 * Send a handshake message. In the case of the extension host, it just sends
+	 * back a debug port.
+	 */
+	public sendMessage(message: HandshakeMessage | { debugPort?: number } ): void {
+		this.sendControl(VSBuffer.fromString(JSON.stringify(message)));
+	}
+}

server.ts

@@ -16,18 +16,18 @@ import { parseMainProcessArgv } from "vs/platform/environment/node/argvHelper";
 import { ParsedArgs } from "vs/platform/environment/common/environment";
 import { EnvironmentService } from "vs/platform/environment/node/environmentService";
 import { InstantiationService } from "vs/platform/instantiation/common/instantiationService";
-import { getLogLevel } from "vs/platform/log/common/log";
+import { getLogLevel, ILogService } from "vs/platform/log/common/log";
 import { LogLevelSetterChannel } from "vs/platform/log/common/logIpc";
 import { SpdLogService } from "vs/platform/log/node/spdlogService";
 import { IProductConfiguration } from "vs/platform/product/common/product";
-import { ConnectionType } from "vs/platform/remote/common/remoteAgentConnection";
+import { ConnectionType, ConnectionTypeRequest } from "vs/platform/remote/common/remoteAgentConnection";
 import { REMOTE_FILE_SYSTEM_CHANNEL_NAME } from "vs/platform/remote/common/remoteAgentFileSystemChannel";
 import { RemoteExtensionLogFileName } from "vs/workbench/services/remote/common/remoteAgentService";
 import { IWorkbenchConstructionOptions } from "vs/workbench/workbench.web.api";
 
-import { Connection, Server as IServer } from "vs/server/connection";
+import { Connection, ManagementConnection, ExtensionHostConnection } from "vs/server/connection";
 import { ExtensionEnvironmentChannel, FileProviderChannel } from "vs/server/channel";
-import { Socket } from "vs/server/socket";
+import { Protocol } from "vs/server/protocol";
 
 export enum HttpCode {
 	Ok = 200,
@@ -51,9 +51,8 @@ export class HttpError extends Error {
 	}
 }
 
-export class Server implements IServer {
+export class Server {
-	// When a new client connects, it will fire this event which is used in the
+	// Used to notify the IPC server that there is a new client.
-	// IPC server which manages channels.
 	public readonly _onDidClientConnect = new Emitter<ClientConnectionEvent>();
 	public readonly onDidClientConnect = this._onDidClientConnect.event;
 
@@ -67,11 +66,10 @@ export class Server implements IServer {
 	private readonly server: http.Server;
 
 	private readonly environmentService: EnvironmentService;
+	private readonly logService: ILogService;
 
-	// Persistent connections. These can reconnect within a timeout. Individual
+	// Persistent connections. These can reconnect within a timeout.
-	// sockets will add connections made through them to this map and remove them
+	private readonly connections = new Map<ConnectionType, Map<string, Connection>>();
-	// when they close.
-	public readonly connections = new Map<ConnectionType, Map<string, Connection>>();
 
 	public constructor() {
 		this.server = http.createServer(async (request, response): Promise<void> => {
@@ -89,12 +87,12 @@ export class Server implements IServer {
 			}
 		});
 
-		this.server.on("upgrade", (request, socket) => {
+		this.server.on("upgrade", async (request, socket) => {
+			const protocol = this.createProtocol(request, socket);
 			try {
-				const nodeSocket = this.handleUpgrade(request, socket);
+				await this.connect(await protocol.handshake(), protocol);
-				nodeSocket.handshake(this);
 			} catch (error) {
-				socket.end(error.message);
+				protocol.dispose(error);
 			}
 		});
 
@@ -114,18 +112,18 @@ export class Server implements IServer {
 
 		this.environmentService = new EnvironmentService(args, process.execPath);
 
-		const logService = new SpdLogService(
+		this.logService = new SpdLogService(
 			RemoteExtensionLogFileName,
 			this.environmentService.logsPath,
 			getLogLevel(this.environmentService),
 		);
-		this.ipc.registerChannel("loglevel", new LogLevelSetterChannel(logService));
+		this.ipc.registerChannel("loglevel", new LogLevelSetterChannel(this.logService));
 
 		const instantiationService = new InstantiationService();
 		instantiationService.invokeFunction(() => {
 			this.ipc.registerChannel(
 				REMOTE_FILE_SYSTEM_CHANNEL_NAME,
-				new FileProviderChannel(logService),
+				new FileProviderChannel(this.logService),
 			);
 			this.ipc.registerChannel(
 				"remoteextensionsenvironment",
@@ -191,7 +189,7 @@ export class Server implements IServer {
 		}
 	}
 
-	private handleUpgrade(request: http.IncomingMessage, socket: net.Socket): Socket {
+	private createProtocol(request: http.IncomingMessage, socket: net.Socket): Protocol {
 		if (request.headers.upgrade !== "websocket") {
 			throw new Error("HTTP/1.1 400 Bad Request");
 		}
@@ -215,10 +213,11 @@ export class Server implements IServer {
 			}
 		}
 
-		const nodeSocket = new Socket(socket, options);
+		return new Protocol(
-		nodeSocket.upgrade(request.headers["sec-websocket-key"] as string);
+			request.headers["sec-websocket-key"] as string,
-
+			socket,
-		return nodeSocket;
+			options,
+		);
 	}
 
 	public listen(port: number = 8443): void {
@@ -231,4 +230,63 @@ export class Server implements IServer {
 			console.log(`Serving ${this.rootPath}`);
 		});
 	}
+
+	private async connect(message: ConnectionTypeRequest, protocol: Protocol): Promise<void> {
+		switch (message.desiredConnectionType) {
+			case ConnectionType.ExtensionHost:
+			case ConnectionType.Management:
+				const debugPort = await this.getDebugPort();
+				const ok = message.desiredConnectionType === ConnectionType.ExtensionHost
+					? (debugPort ? { debugPort } : {})
+					: { type: "ok" };
+
+				if (!this.connections.has(message.desiredConnectionType)) {
+					this.connections.set(message.desiredConnectionType, new Map());
+				}
+
+				const connections = this.connections.get(message.desiredConnectionType)!;
+				const token = protocol.options.reconnectionToken;
+
+				if (protocol.options.reconnection && connections.has(token)) {
+					protocol.sendMessage(ok);
+					const buffer = protocol.readEntireBuffer();
+					protocol.dispose();
+					return connections.get(token)!.reconnect(protocol, buffer);
+				}
+
+				if (protocol.options.reconnection || connections.has(token)) {
+					throw new Error(protocol.options.reconnection
+						? "Unrecognized reconnection token"
+						: "Duplicate reconnection token"
+					);
+				}
+
+				protocol.sendMessage(ok);
+
+				let connection: Connection;
+				if (message.desiredConnectionType === ConnectionType.Management) {
+					connection = new ManagementConnection(protocol);
+					this._onDidClientConnect.fire({
+						protocol,
+						onDidClientDisconnect: connection.onClose,
+					});
+				} else {
+					connection = new ExtensionHostConnection(protocol, this.logService);
+				}
+				connections.set(protocol.options.reconnectionToken, connection);
+				connection.onClose(() => {
+					connections.delete(protocol.options.reconnectionToken);
+				});
+				break;
+			case ConnectionType.Tunnel: return protocol.tunnel();
+			default: throw new Error("Unrecognized connection type");
+		}
+	}
+
+	/**
+	 * TODO: implement.
+	 */
+	private async getDebugPort(): Promise<number | undefined> {
+		return undefined;
+	}
 }

socket.ts

@@ -1,159 +0,0 @@
-import * as crypto from "crypto";
-import * as net from "net";
-
-import { VSBuffer } from "vs/base/common/buffer";
-import { NodeSocket, WebSocketNodeSocket } from "vs/base/parts/ipc/node/ipc.net";
-import { PersistentProtocol, ISocket } from "vs/base/parts/ipc/common/ipc.net";
-import { AuthRequest, ConnectionType, ConnectionTypeRequest, HandshakeMessage } from "vs/platform/remote/common/remoteAgentConnection";
-
-import { ExtensionHostConnection, ManagementConnection, Server } from "vs/server/connection";
-
-export interface SocketOptions {
-	readonly reconnectionToken: string;
-	readonly reconnection: boolean;
-	readonly skipWebSocketFrames: boolean;
-}
-
-export class Socket {
-	private nodeSocket: ISocket;
-	public protocol: PersistentProtocol;
-
-	public constructor(private readonly socket: net.Socket, private readonly options: SocketOptions) {
-		socket.on("error", () => this.dispose());
-		this.nodeSocket = new NodeSocket(socket);
-		if (!this.options.skipWebSocketFrames) {
-			this.nodeSocket = new WebSocketNodeSocket(this.nodeSocket as NodeSocket);
-		}
-		this.protocol = new PersistentProtocol(this.nodeSocket);
-	}
-
-	/**
-	 * Upgrade the connection into a web socket.
-	 */
-	public upgrade(secWebsocketKey: string): void {
-		// This magic value is specified by the websocket spec.
-		const magic = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
-		const reply = crypto.createHash("sha1")
-			.update(secWebsocketKey + magic)
-			.digest("base64");
-
-		this.socket.write([
-			"HTTP/1.1 101 Switching Protocols",
-			"Upgrade: websocket",
-			"Connection: Upgrade",
-			`Sec-WebSocket-Accept: ${reply}`,
-		].join("\r\n") + "\r\n\r\n");
-	}
-
-	public dispose(): void {
-		this.nodeSocket.dispose();
-		this.protocol.dispose();
-		this.nodeSocket = undefined!;
-		this.protocol = undefined!;
-	}
-
-	public handshake(server: Server): void {
-		const handler = this.protocol.onControlMessage((rawMessage) => {
-			const message = JSON.parse(rawMessage.toString());
-			switch (message.type) {
-				case "auth": return this.authenticate(message);
-				case "connectionType":
-					handler.dispose();
-					return this.connect(message, server);
-				case "default":
-					return this.dispose();
-			}
-		});
-	}
-
-	/**
-	 * TODO: This ignores the authentication process entirely for now.
-	 */
-	private authenticate(_message: AuthRequest): void {
-		this.sendControl({
-			type: "sign",
-			data: "",
-		});
-	}
-
-	private connect(message: ConnectionTypeRequest, server: Server): void {
-		switch (message.desiredConnectionType) {
-			case ConnectionType.ExtensionHost:
-			case ConnectionType.Management:
-				const debugPort = this.getDebugPort();
-				const ok = message.desiredConnectionType === ConnectionType.ExtensionHost
-					? (debugPort ? { debugPort } : {})
-					: { type: "ok" };
-
-				if (!server.connections.has(message.desiredConnectionType)) {
-					server.connections.set(message.desiredConnectionType, new Map());
-				}
-
-				const connections = server.connections.get(message.desiredConnectionType)!;
-
-				if (this.options.reconnection && connections.has(this.options.reconnectionToken)) {
-					this.sendControl(ok);
-					const buffer = this.protocol.readEntireBuffer();
-					this.protocol.dispose();
-					return connections.get(this.options.reconnectionToken)!
-						.reconnect(this.nodeSocket, buffer);
-				}
-
-				if (this.options.reconnection || connections.has(this.options.reconnectionToken)) {
-					this.sendControl({
-						type: "error",
-						reason: this.options.reconnection
-							? "Unrecognized reconnection token"
-							: "Duplicate reconnection token",
-					});
-					return this.dispose();
-				}
-
-				this.sendControl(ok);
-
-				const connection = message.desiredConnectionType === ConnectionType.Management
-					? new ManagementConnection(server, this.protocol)
-					: new ExtensionHostConnection(server, this.protocol);
-
-				connections.set(this.options.reconnectionToken, connection);
-				connection.onClose(() => {
-					connections.delete(this.options.reconnectionToken);
-				});
-				break;
-			case ConnectionType.Tunnel:
-				return this.tunnel();
-			default:
-				this.sendControl({
-					type: "error",
-					reason: "Unrecognized connection type",
-				});
-				return this.dispose();
-		}
-	}
-
-	/**
-	 * TODO: implement.
-	 */
-	private tunnel(): void {
-		this.sendControl({
-			type: "error",
-			reason: "Tunnel is not implemented yet",
-		});
-		this.dispose();
-	}
-
-	/**
-	 * TODO: implement.
-	 */
-	private getDebugPort(): number | undefined {
-		return undefined;
-	}
-
-	/**
-	 * Send a handshake message. In the case of the extension host, it just sends
-	 * back a debug port.
-	 */
-	private sendControl(message: HandshakeMessage | { debugPort?: number } ): void {
-		this.protocol.sendControl(VSBuffer.fromString(JSON.stringify(message)));
-	}
-}

transformer.js

@@ -0,0 +1,27 @@
+// This file is included via a regular Node require. I'm not sure how (or if)
+// we can write this in Typescript and have it compile to non-AMD syntax.
+module.exports = (remoteAuthority) => {
+	return {
+		transformIncoming: (uri) => {
+			switch (uri.scheme) {
+				case "vscode-remote": return { scheme: "file", path: uri.path };
+				case "file ": return { scheme: "vscode-local", path: uri.path };
+				default: return uri;
+			}
+		},
+		transformOutgoing: (uri) => {
+			switch (uri.scheme) {
+				case "vscode-local": return { scheme: "file", path: uri.path };
+				case "file ": return { scheme: "vscode-remote", authority: remoteAuthority, path: uri.path };
+				default: return uri;
+			}
+		},
+		transformOutgoingScheme: (scheme) => {
+			switch (scheme) {
+				case "vscode-local": return "file";
+				case "file": return "vscode-remote";
+				default: return scheme;
+			}
+		},
+	};
+};