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.

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,17 +53,17 @@ 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':
@@ -58,6 +71,7 @@ export class Protocol extends PersistentProtocol {
 						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;
+	}
 }