Improve protocol class

- Move destroy logic into the class itself
- Improve logging a bit
- Remove the record option; we should always do this when using
  permessage-deflate.
- Let debug port be null (it can be null in the message args).
- Add setSocket so we don't have to initiate a connection to set it.
- Move inflate bytes logic into the class itself.

Author: code-asher

lib/vscode/src/vs/base/parts/ipc/common/ipc.net.ts

@@ -743,6 +743,11 @@ export class PersistentProtocol implements IMessagePassingProtocol {
 		}, Math.max(ProtocolConstants.KeepAliveTimeoutTime - timeSinceLastIncomingMsg, 0) + 5);
 	}
 
+	// NOTE@coder: Set the socket without initiating a reconnect.
+	public setSocket(socket: ISocket): void {
+		this._socket = socket;
+	}
+
 	public getSocket(): ISocket {
 		return this._socket;
 	}

lib/vscode/src/vs/server/node/connection.ts

@@ -4,7 +4,6 @@ import { VSBuffer } from 'vs/base/common/buffer';
 import { Emitter } from 'vs/base/common/event';
 import { FileAccess } from 'vs/base/common/network';
 import { ISocket } from 'vs/base/parts/ipc/common/ipc.net';
-import { WebSocketNodeSocket } from 'vs/base/parts/ipc/node/ipc.net';
 import { INativeEnvironmentService } from 'vs/platform/environment/common/environment';
 import { getNlsConfiguration } from 'vs/server/node/nls';
 import { Protocol } from 'vs/server/node/protocol';
@@ -59,9 +58,7 @@ export class ManagementConnection extends Connection {
 	}
 
 	protected doDispose(): void {
-		this.protocol.sendDisconnect();
-		this.protocol.dispose();
-		this.protocol.getUnderlyingSocket().destroy();
+		this.protocol.destroy();
 	}
 
 	protected doReconnect(socket: ISocket, buffer: VSBuffer): void {
@@ -99,35 +96,34 @@ export class ExtensionHostConnection extends Connection {
 	}
 
 	protected doDispose(): void {
+		this.protocol.destroy();
 		if (this.process) {
 			this.process.kill();
 		}
-		this.protocol.getUnderlyingSocket().destroy();
 	}
 
 	protected doReconnect(socket: ISocket, buffer: VSBuffer): void {
-		// This is just to set the new socket.
-		this.protocol.beginAcceptReconnection(socket, null);
+		this.protocol.setSocket(socket);
 		this.protocol.dispose();
 		this.sendInitMessage(buffer);
 	}
 
 	private sendInitMessage(buffer: VSBuffer): void {
-		const socket = this.protocol.getUnderlyingSocket();
-		socket.pause();
+		if (!this.process) {
+			throw new Error("Tried to initialize VS Code before spawning");
+		}
+
+		this.protocol.getUnderlyingSocket().pause();
 
-		const wrapperSocket = this.protocol.getSocket();
+		this.logger.debug('Sending socket');
 
-		this.logger.trace('Sending socket');
-		this.process!.send({ // Process must be set at this point.
+		this.process.send({
 			type: 'VSCODE_EXTHOST_IPC_SOCKET',
 			initialDataChunk: Buffer.from(buffer.buffer).toString('base64'),
-			skipWebSocketFrames: !(wrapperSocket instanceof WebSocketNodeSocket),
+			skipWebSocketFrames: this.protocol.options.skipWebSocketFrames,
 			permessageDeflate: this.protocol.options.permessageDeflate,
-			inflateBytes: wrapperSocket instanceof WebSocketNodeSocket
-				? Buffer.from(wrapperSocket.recordedInflateBytes.buffer).toString('base64')
-				: undefined,
-		}, socket);
+			inflateBytes: this.protocol.inflateBytes,
+		}, this.protocol.getUnderlyingSocket());
 	}
 
 	private async spawn(locale: string, buffer: VSBuffer): Promise<cp.ChildProcess> {

lib/vscode/src/vs/server/node/protocol.ts

@@ -1,21 +1,31 @@
-import { field } from '@coder/logger';
+import { field, logger, Logger } from '@coder/logger';
 import * as net from 'net';
 import { VSBuffer } from 'vs/base/common/buffer';
 import { PersistentProtocol } from 'vs/base/parts/ipc/common/ipc.net';
 import { NodeSocket, WebSocketNodeSocket } from 'vs/base/parts/ipc/node/ipc.net';
 import { AuthRequest, ConnectionTypeRequest, HandshakeMessage } from 'vs/platform/remote/common/remoteAgentConnection';
-import { logger } from 'vs/server/node/logger';
 
 export interface SocketOptions {
+	/** The token is how we identify and connect to existing sessions. */
 	readonly reconnectionToken: string;
+	/** Specifies that the client is trying to reconnect. */
 	readonly reconnection: boolean;
+	/** If true assume this is not a web socket (always false for code-server). */
 	readonly skipWebSocketFrames: boolean;
+	/** Whether to support compression (web socket only). */
 	readonly permessageDeflate?: boolean;
+	/**
+	 * Seed zlib with these bytes (web socket only). If parts of inflating was
+	 * done in a different zlib instance we need to pass all those bytes into zlib
+	 * otherwise the inflate might hit an inflated portion referencing a distance
+	 * too far back.
+	 */
 	readonly inflateBytes?: VSBuffer;
-	readonly recordInflateBytes?: boolean;
 }
 
 export class Protocol extends PersistentProtocol {
+	private readonly logger: Logger;
+
 	public constructor(socket: net.Socket, public readonly options: SocketOptions) {
 		super(
 			options.skipWebSocketFrames
@@ -24,9 +34,12 @@ export class Protocol extends PersistentProtocol {
 					new NodeSocket(socket),
 					options.permessageDeflate || false,
 					options.inflateBytes || null,
-					options.recordInflateBytes || false,
+					// Always record inflate bytes if using permessage-deflate.
+					options.permessageDeflate || false,
 				),
 		);
+
+		this.logger = logger.named('protocol', field("token", this.options.reconnectionToken));
 	}
 
 	public getUnderlyingSocket(): net.Socket {
@@ -40,24 +53,25 @@ export class Protocol extends PersistentProtocol {
 	 * Perform a handshake to get a connection request.
 	 */
 	public handshake(): Promise<ConnectionTypeRequest> {
-		logger.trace('Protocol handshake', field('token', this.options.reconnectionToken));
+		this.logger.debug('Initiating handshake...');
 		return new Promise((resolve, reject) => {
 			const timeout = setTimeout(() => {
-				logger.error('Handshake timed out', field('token', this.options.reconnectionToken));
-				reject(new Error('timed out'));
+				this.logger.debug('Handshake timed out');
+				reject(new Error('protocol handshake timed out'));
 			}, 10000); // Matches the client timeout.
 
 			const handler = this.onControlMessage((rawMessage) => {
 				try {
 					const raw = rawMessage.toString();
-					logger.trace('Protocol message', field('token', this.options.reconnectionToken), field('message', raw));
+					this.logger.trace('Got message', field('message', raw));
 					const message = JSON.parse(raw);
 					switch (message.type) {
 						case 'auth':
 							return this.authenticate(message);
 						case 'connectionType':
 							handler.dispose();
 							clearTimeout(timeout);
+							this.logger.debug('Handshake completed');
 							return resolve(message);
 						default:
 							throw new Error('Unrecognized message type');
@@ -90,10 +104,38 @@ export class Protocol extends PersistentProtocol {
 	}
 
 	/**
-	 * Send a handshake message. In the case of the extension host, it just sends
-	 * back a debug port.
+	 * Send a handshake message. In the case of the extension host it should just
+	 * send a debug port.
 	 */
-	public sendMessage(message: HandshakeMessage | { debugPort?: number } ): void {
+	public sendMessage(message: HandshakeMessage | { debugPort?: number | null } ): void {
 		this.sendControl(VSBuffer.fromString(JSON.stringify(message)));
 	}
+
+	/**
+	 * Disconnect and dispose everything including the underlying socket.
+	 */
+	public destroy(reason?: string): void {
+		try {
+			if (reason) {
+				this.sendMessage({ type: 'error', reason });
+			}
+			// If still connected try notifying the client.
+			this.sendDisconnect();
+		} catch (error) {
+			// I think the write might fail if already disconnected.
+			this.logger.warn(error.message || error);
+		}
+		this.dispose(); // This disposes timers and socket event handlers.
+		this.getSocket().dispose(); // This will destroy() the socket.
+	}
+
+	/**
+	 * Get inflateBytes in base64 format from the current socket.
+	 */
+	public get inflateBytes(): string | undefined {
+		const socket = this.getSocket();
+		return socket instanceof WebSocketNodeSocket
+			? Buffer.from(socket.recordedInflateBytes.buffer).toString('base64')
+			: undefined;
+	}
 }

lib/vscode/src/vs/server/node/server.ts

@@ -123,14 +123,11 @@ export class Vscode {
 			reconnection: query.reconnection === 'true',
 			skipWebSocketFrames: query.skipWebSocketFrames === 'true',
 			permessageDeflate,
-			recordInflateBytes: permessageDeflate,
 		});
 		try {
 			await this.connect(await protocol.handshake(), protocol);
 		} catch (error) {
-			protocol.sendMessage({ type: 'error', reason: error.message });
-			protocol.dispose();
-			protocol.getSocket().dispose();
+			protocol.destroy(error.message);
 		}
 		return true;
 	}