diff --git a/src/kernels/common/usedPorts.ts b/src/kernels/common/usedPorts.ts index b098691b3ea..5375b1eab74 100644 --- a/src/kernels/common/usedPorts.ts +++ b/src/kernels/common/usedPorts.ts @@ -1,5 +1,45 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +import { + PortAttributes, + type CancellationToken, + type PortAttributesProvider, + PortAutoForwardAction, + type PortAttributesSelector, + workspace +} from 'vscode'; +import { DisposableStore } from '../../platform/common/utils/lifecycle'; +import { traceError } from '../../platform/logging'; + // Keeps track of all ports used by Kernels and other processes spawned by Kernels and related code export const UsedPorts = new Set(); + +export function ignorePortForwarding(...ports: number[]) { + const disposableStore = new DisposableStore(); + try { + const provider = new (class implements PortAttributesProvider { + async providePortAttributes( + attributes: { port: number; pid?: number; commandLine?: string }, + _token: CancellationToken + ) { + if (ports.includes(attributes.port)) { + return new PortAttributes(PortAutoForwardAction.Ignore); + } + return undefined; + } + })(); + + for (const port of ports) { + const portSelector: PortAttributesSelector = { + portRange: port + }; + disposableStore.add(workspace.registerPortAttributesProvider(portSelector, provider)); + } + } catch (ex) { + // In case proposed API changes. + traceError('Failure in registering port attributes', ex); + } + + return disposableStore; +} diff --git a/src/kernels/jupyter/launcher/jupyterServerStarter.node.ts b/src/kernels/jupyter/launcher/jupyterServerStarter.node.ts index 6af9a4fcf06..59cccc8fbc2 100644 --- a/src/kernels/jupyter/launcher/jupyterServerStarter.node.ts +++ b/src/kernels/jupyter/launcher/jupyterServerStarter.node.ts @@ -28,7 +28,7 @@ import { getFilePath } from '../../../platform/common/platform/fs-paths'; import { JupyterNotebookNotInstalled } from '../../../platform/errors/jupyterNotebookNotInstalled'; import { JupyterCannotBeLaunchedWithRootError } from '../../../platform/errors/jupyterCannotBeLaunchedWithRootError'; import { noop } from '../../../platform/common/utils/misc'; -import { UsedPorts } from '../../common/usedPorts'; +import { UsedPorts, ignorePortForwarding } from '../../common/usedPorts'; /** * Responsible for starting a notebook. @@ -132,8 +132,12 @@ export class JupyterServerStarter implements IJupyterServerStarter { try { const port = parseInt(new URL(connection.baseUrl).port || '0', 10); if (port && !isNaN(port)) { + const disposable = ignorePortForwarding(port); if (launchResult.proc) { - launchResult.proc.on('exit', () => UsedPorts.delete(port)); + launchResult.proc.on('exit', () => { + UsedPorts.delete(port); + disposable.dispose(); + }); } UsedPorts.add(port); } diff --git a/src/kernels/portAttributeProvider.node.ts b/src/kernels/portAttributeProvider.node.ts deleted file mode 100644 index bba3d90eb9d..00000000000 --- a/src/kernels/portAttributeProvider.node.ts +++ /dev/null @@ -1,39 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -import { inject, injectable } from 'inversify'; -import { workspace } from 'vscode'; -import { CancellationToken, PortAttributes, PortAttributesProvider, PortAutoForwardAction } from 'vscode'; -import { IExtensionSyncActivationService } from '../platform/activation/types'; -import { traceError } from '../platform/logging'; -import { IDisposableRegistry } from '../platform/common/types'; -import { UsedPorts } from './common/usedPorts'; - -/** - * Used to determine how ports can be used when creating a raw kernel. - */ -@injectable() -export class PortAttributesProviders implements PortAttributesProvider, IExtensionSyncActivationService { - constructor(@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry) {} - activate(): void { - try { - this.disposables.push(workspace.registerPortAttributesProvider({}, this)); - } catch (ex) { - // In case proposed API changes. - traceError('Failure in registering port attributes', ex); - } - } - providePortAttributes( - attributes: { port: number; pid?: number; commandLine?: string }, - _token: CancellationToken - ): PortAttributes | undefined { - try { - if (UsedPorts.has(attributes.port)) { - return new PortAttributes(PortAutoForwardAction.Ignore); - } - } catch (ex) { - // In case proposed API changes. - traceError('Failure in returning port attributes', ex); - } - } -} diff --git a/src/kernels/raw/launcher/kernelLauncher.node.ts b/src/kernels/raw/launcher/kernelLauncher.node.ts index 0ebb852e2c9..728479cd98a 100644 --- a/src/kernels/raw/launcher/kernelLauncher.node.ts +++ b/src/kernels/raw/launcher/kernelLauncher.node.ts @@ -38,8 +38,9 @@ import { StopWatch } from '../../../platform/common/utils/stopWatch'; import { getResourceType } from '../../../platform/common/utils'; import { format, splitLines } from '../../../platform/common/helpers'; import { IPythonExecutionFactory } from '../../../platform/interpreter/types.node'; -import { UsedPorts } from '../../common/usedPorts'; +import { UsedPorts, ignorePortForwarding } from '../../common/usedPorts'; import { isPythonKernelConnection } from '../../helpers'; +import { once } from '../../../platform/common/utils/events'; const PortFormatString = `kernelLauncherPortStart_{0}.tmp`; // Launches and returns a kernel process given a resource or python interpreter. @@ -206,7 +207,13 @@ export class KernelLauncher implements IKernelLauncher { ? window.createOutputChannel(DataScience.kernelConsoleOutputChannel(baseName), 'log') : undefined; outputChannel?.clear(); - + const portAttributeProvider = ignorePortForwarding( + connection.control_port, + connection.hb_port, + connection.iopub_port, + connection.shell_port, + connection.stdin_port + ); // Create the process const kernelProcess = new KernelProcess( this.processExecutionFactory, @@ -223,8 +230,8 @@ export class KernelLauncher implements IKernelLauncher { this.pythonKernelInterruptDaemon, this.platformService ); - - kernelProcess.exited(() => outputChannel?.dispose(), this, this.disposables); + once(kernelProcess.onDidDispose)(() => portAttributeProvider.dispose(), this, this.disposables); + once(kernelProcess.exited)(() => outputChannel?.dispose(), this, this.disposables); try { await raceCancellationError(cancelToken, kernelProcess.launch(workingDirectory, timeout, cancelToken)); } catch (ex) { @@ -233,7 +240,7 @@ export class KernelLauncher implements IKernelLauncher { throw ex; } - const disposable = kernelProcess.exited( + const disposable = once(kernelProcess.exited)( ({ exitCode, reason }) => { sendKernelTelemetryEvent( resource, diff --git a/src/kernels/raw/launcher/kernelLauncher.unit.test.ts b/src/kernels/raw/launcher/kernelLauncher.unit.test.ts index a0e53607565..2284e1ea8c8 100644 --- a/src/kernels/raw/launcher/kernelLauncher.unit.test.ts +++ b/src/kernels/raw/launcher/kernelLauncher.unit.test.ts @@ -17,11 +17,18 @@ import { dispose } from '../../../platform/common/utils/lifecycle'; import { anything, instance, mock, when } from 'ts-mockito'; import { IFileSystemNode } from '../../../platform/common/platform/types.node'; import { PythonKernelConnectionMetadata } from '../../types'; -import { CancellationTokenSource, Disposable, EventEmitter, PortAutoForwardAction, Uri } from 'vscode'; +import { + CancellationTokenSource, + Disposable, + EventEmitter, + PortAutoForwardAction, + Uri, + type PortAttributesProvider +} from 'vscode'; import { KernelProcess } from './kernelProcess.node'; -import { PortAttributesProviders } from '../../portAttributeProvider.node'; import { IPythonExecutionFactory, IPythonExecutionService } from '../../../platform/interpreter/types.node'; import { UsedPorts } from '../../common/usedPorts'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; suite('kernel Launcher', () => { let disposables: IDisposable[] = []; @@ -67,7 +74,10 @@ suite('kernel Launcher', () => { instance(platform) ); }); - teardown(() => (disposables = dispose(disposables))); + teardown(() => { + disposables = dispose(disposables); + resetVSCodeMocks(); + }); async function launchKernel() { const kernelSpec = PythonKernelConnectionMetadata.create({ id: '1', @@ -103,17 +113,25 @@ suite('kernel Launcher', () => { assert.notDeepEqual(Array.from(UsedPorts), oldPorts, 'Ports not updated'); }); test('Verify Kernel ports are not forwarded', async () => { - await launchKernel(); + const oldPorts = new Set(UsedPorts); + const providers: PortAttributesProvider[] = []; + when(mockedVSCodeNamespaces.workspace.registerPortAttributesProvider(anything(), anything())).thenCall( + (_, provider) => providers.push(provider) + ); - const portAttributeProvider = new PortAttributesProviders(disposables); + await launchKernel(); const cancellation = new CancellationTokenSource(); + disposables.push(cancellation); for (const port of UsedPorts) { - assert.equal( - portAttributeProvider.providePortAttributes({ port }, cancellation.token)?.autoForwardAction, - PortAutoForwardAction.Ignore, - 'Kernel Port should not be forwarded' + if (oldPorts.has(port)) { + continue; + } + const results = await Promise.all( + providers.map((p) => Promise.resolve(p.providePortAttributes({ port }, cancellation.token))) ); + const portForwardingIgnored = results.some((r) => r?.autoForwardAction === PortAutoForwardAction.Ignore); + assert.isTrue(portForwardingIgnored, `Kernel Port ${port} should not be forwarded`); } }); }); diff --git a/src/kernels/raw/launcher/kernelProcess.node.ts b/src/kernels/raw/launcher/kernelProcess.node.ts index 6b8dddb0314..c503dbdd63f 100644 --- a/src/kernels/raw/launcher/kernelProcess.node.ts +++ b/src/kernels/raw/launcher/kernelProcess.node.ts @@ -61,6 +61,7 @@ import { IPythonExecutionFactory } from '../../../platform/interpreter/types.nod import { getDisplayPath } from '../../../platform/common/platform/fs-paths'; import { StopWatch } from '../../../platform/common/utils/stopWatch'; import { ServiceContainer } from '../../../platform/ioc/container'; +import { ObservableDisposable } from '../../../platform/common/utils/lifecycle'; const kernelOutputWithConnectionFile = 'To connect another client to this kernel, use:'; const kernelOutputToNotLog = @@ -68,7 +69,7 @@ const kernelOutputToNotLog = // Launches and disposes a kernel process given a kernelspec and a resource or python interpreter. // Exposes connection information and the process itself. -export class KernelProcess implements IKernelProcess { +export class KernelProcess extends ObservableDisposable implements IKernelProcess { private _pid?: number; private _disposingPromise?: Promise; public get pid() { @@ -100,7 +101,6 @@ export class KernelProcess implements IKernelProcess { private _process?: ChildProcess; private exitEvent = new EventEmitter<{ exitCode?: number; reason?: string }>(); private launchedOnce?: boolean; - private disposed?: boolean; private connectionFile?: Uri; private _launchKernelSpec?: IJupyterKernelSpec; private interrupter?: Interrupter; @@ -122,6 +122,7 @@ export class KernelProcess implements IKernelProcess { private readonly pythonKernelInterruptDaemon: PythonKernelInterruptDaemon, private readonly platform: IPlatformService ) { + super(); this._kernelConnectionMetadata = kernelConnectionMetadata; } public async interrupt(): Promise { @@ -176,7 +177,7 @@ export class KernelProcess implements IKernelProcess { if (proc) { proc.on('exit', (exitCode) => { exitCode = exitCode || providedExitCode; - if (this.disposed) { + if (this.isDisposed) { traceVerbose(`KernelProcess Exited, Exit Code - ${exitCode}`); return; } @@ -248,7 +249,7 @@ export class KernelProcess implements IKernelProcess { this.sendToOutput(output.out); }); exeObs.out.done.catch((error) => { - if (this.disposed) { + if (this.isDisposed) { traceWarning('Kernel died', error, stderr); return; } @@ -337,14 +338,13 @@ export class KernelProcess implements IKernelProcess { } } - public async dispose(): Promise { + public override async dispose(): Promise { if (this._disposingPromise) { return this._disposingPromise; } - if (this.disposed) { + if (this.isDisposed) { return; } - this.disposed = true; const pid = this._process?.pid; traceInfo(`Dispose Kernel process ${pid}.`); this._disposingPromise = (async () => { @@ -370,6 +370,7 @@ export class KernelProcess implements IKernelProcess { }); traceVerbose(`Disposed Kernel process ${pid}.`); })(); + super.dispose(); } private async killChildProcesses(pid?: number) { diff --git a/src/kernels/serviceRegistry.node.ts b/src/kernels/serviceRegistry.node.ts index 6facfce400c..6af15185f92 100644 --- a/src/kernels/serviceRegistry.node.ts +++ b/src/kernels/serviceRegistry.node.ts @@ -22,7 +22,6 @@ import { KernelRefreshIndicator } from './kernelRefreshIndicator.node'; import { KernelStartupCodeProviders } from './kernelStartupCodeProviders.node'; import { KernelStartupTelemetry } from './kernelStartupTelemetry.node'; import { KernelStatusProvider } from './kernelStatusProvider'; -import { PortAttributesProviders } from './portAttributeProvider.node'; import { ContributedLocalKernelSpecFinder } from './raw/finder/contributedLocalKernelSpecFinder.node'; import { JupyterPaths } from './raw/finder/jupyterPaths.node'; import { LocalKnownPathKernelSpecFinder } from './raw/finder/localKnownPathKernelSpecFinder.node'; @@ -54,10 +53,6 @@ import { KernelChatStartupCodeProvider } from './chat/kernelStartupCodeProvider' export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) { serviceManager.addSingleton(IExtensionSyncActivationService, Activation); serviceManager.addSingleton(IExtensionSyncActivationService, ServerPreload); - serviceManager.addSingleton( - IExtensionSyncActivationService, - PortAttributesProviders - ); serviceManager.addSingleton( IRawNotebookSupportedService, RawNotebookSupportedService