Partial extension host, some restructuring

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

Author: code-asher

connection.ts

@@ -1,69 +1,176 @@
-import { ClientConnectionEvent } from "vs/base/parts/ipc/common/ipc";
-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 * as cp from "child_process";
+
+import { getPathFromAmdModule } from "vs/base/common/amd";
 import { VSBuffer } from "vs/base/common/buffer";
+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";
 
-export interface Server {
-	readonly _onDidClientConnect: Emitter<ClientConnectionEvent>;
-	readonly connections: Map<ConnectionType, Map<string, Connection>>;
-}
+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;
-
-	public constructor(
-		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
-		// reconnect.
-		this.protocol.onSocketClose(() => {
+	private readonly wait = 1000 * 60;
+
+	private closed: boolean = false;
+
+	public constructor(protected protocol: Protocol) {
+		// onClose seems to mean we want to disconnect, so close immediately.
+		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
-	 * don't need or want the connection. It cannot be re-used after this.
+	 * Set up the connection on a new socket.
 	 */
-	public dispose(): void {
-		this.protocol.sendDisconnect();
-		this.protocol.getSocket().end();
-		this.protocol.dispose();
-		this._onClose.fire();
+	public reconnect(protocol: Protocol, buffer: VSBuffer): void {
+		if (this.closed) {
+			throw new Error("Cannot reconnect to closed connection");
+		}
+		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.
-		this.protocol.beginAcceptReconnection(socket, buffer);
-		this.protocol.endAcceptReconnection();
+	/**
+	 * 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;
 }
 
 /**
- * The management connection is used for all the IPC channels.
+ * Used for all the IPC channels.
  */
 export class ManagementConnection extends Connection {
-	public constructor(server: Server, protocol: PersistentProtocol) {
-		super(server, protocol);
-		// This will communicate back to the IPCServer that a new client has
-		// connected.
-		this.server._onDidClientConnect.fire({
-			protocol,
-			onDidClientDisconnect: this.onClose,
-		});
+	protected dispose(): void {
+		// Nothing extra to do here.
+	}
+
+	protected connect(socket: ISocket, buffer: VSBuffer): void {
+		this.protocol.beginAcceptReconnection(socket, buffer);
+		this.protocol.endAcceptReconnection();
 	}
 }
 
+/**
+ * Manage the extension host process.
+ */
 export class ExtensionHostConnection extends Connection {
+	private process: cp.ChildProcess;
+
+	public constructor(protocol: Protocol, private readonly log: ILogService) {
+		super(protocol);
+		const socket = this.protocol.getSocket();
+		const buffer = this.protocol.readEntireBuffer();
+		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;
+	}
 }

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