From ae6089f852f3d92de778cb8f81e948a2fd4ca48b Mon Sep 17 00:00:00 2001 From: Asher Date: Tue, 20 Apr 2021 11:52:21 -0500 Subject: [PATCH] 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. --- .../src/vs/base/parts/ipc/common/ipc.net.ts | 5 ++ lib/vscode/src/vs/server/node/connection.ts | 30 ++++----- lib/vscode/src/vs/server/node/protocol.ts | 64 +++++++++++++++---- lib/vscode/src/vs/server/node/server.ts | 5 +- 4 files changed, 72 insertions(+), 32 deletions(-) diff --git a/lib/vscode/src/vs/base/parts/ipc/common/ipc.net.ts b/lib/vscode/src/vs/base/parts/ipc/common/ipc.net.ts index 777d4379..dcce7c0e 100644 --- a/lib/vscode/src/vs/base/parts/ipc/common/ipc.net.ts +++ b/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; } diff --git a/lib/vscode/src/vs/server/node/connection.ts b/lib/vscode/src/vs/server/node/connection.ts index 3b4d54ab..9af7736a 100644 --- a/lib/vscode/src/vs/server/node/connection.ts +++ b/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"); + } - const wrapperSocket = this.protocol.getSocket(); + this.protocol.getUnderlyingSocket().pause(); - this.logger.trace('Sending socket'); - this.process!.send({ // Process must be set at this point. + this.logger.debug('Sending socket'); + + 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 { diff --git a/lib/vscode/src/vs/server/node/protocol.ts b/lib/vscode/src/vs/server/node/protocol.ts index 15fa414e..ec32ad3a 100644 --- a/lib/vscode/src/vs/server/node/protocol.ts +++ b/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 { - 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; + } } diff --git a/lib/vscode/src/vs/server/node/server.ts b/lib/vscode/src/vs/server/node/server.ts index a44cbd0b..a5a24494 100644 --- a/lib/vscode/src/vs/server/node/server.ts +++ b/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; }