Revert "revert due to buffer overflow on subprocess (microsoft#21762)"

This reverts commit dd20561.

src/client/testing/testController/common/server.ts

@@ -3,10 +3,11 @@
 
 import * as net from 'net';
 import * as crypto from 'crypto';
-import { Disposable, Event, EventEmitter } from 'vscode';
+import { Disposable, Event, EventEmitter, TestRun } from 'vscode';
 import * as path from 'path';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
+    ExecutionResult,
     IPythonExecutionFactory,
     SpawnOptions,
 } from '../../../common/process/types';
@@ -15,6 +16,7 @@ import { DataReceivedEvent, ITestServer, TestCommandOptions } from './types';
 import { ITestDebugLauncher, LaunchOptions } from '../../common/types';
 import { UNITTEST_PROVIDER } from '../../common/constants';
 import { jsonRPCHeaders, jsonRPCContent, JSONRPC_UUID_HEADER } from './utils';
+import { createDeferred } from '../../../common/utils/async';
 
 export class PythonTestServer implements ITestServer, Disposable {
     private _onDataReceived: EventEmitter<DataReceivedEvent> = new EventEmitter<DataReceivedEvent>();
@@ -140,7 +142,12 @@ export class PythonTestServer implements ITestServer, Disposable {
         return this._onDataReceived.event;
     }
 
-    async sendCommand(options: TestCommandOptions, runTestIdPort?: string, callback?: () => void): Promise<void> {
+    async sendCommand(
+        options: TestCommandOptions,
+        runTestIdPort?: string,
+        runInstance?: TestRun,
+        callback?: () => void,
+    ): Promise<void> {
         const { uuid } = options;
 
         const pythonPathParts: string[] = process.env.PYTHONPATH?.split(path.delimiter) ?? [];
@@ -154,7 +161,7 @@ export class PythonTestServer implements ITestServer, Disposable {
         };
 
         if (spawnOptions.extraVariables) spawnOptions.extraVariables.RUN_TEST_IDS_PORT = runTestIdPort;
-        const isRun = !options.testIds;
+        const isRun = runTestIdPort !== undefined;
         // Create the Python environment in which to execute the command.
         const creationOptions: ExecutionFactoryCreateWithEnvironmentOptions = {
             allowEnvironmentFetchExceptions: false,
@@ -195,7 +202,19 @@ export class PythonTestServer implements ITestServer, Disposable {
                     // This means it is running discovery
                     traceLog(`Discovering unittest tests with arguments: ${args}\r\n`);
                 }
-                await execService.exec(args, spawnOptions);
+                const deferred = createDeferred<ExecutionResult<string>>();
+
+                const result = execService.execObservable(args, spawnOptions);
+
+                runInstance?.token.onCancellationRequested(() => {
+                    result?.proc?.kill();
+                });
+                result?.proc?.on('close', () => {
+                    traceLog('Exec server closed.', uuid);
+                    deferred.resolve({ stdout: '', stderr: '' });
+                    callback?.();
+                });
+                await deferred.promise;
             }
         } catch (ex) {
             this.uuids = this.uuids.filter((u) => u !== uuid);

src/client/testing/testController/common/types.ts

@@ -174,7 +174,12 @@ export interface ITestServer {
     readonly onDataReceived: Event<DataReceivedEvent>;
     readonly onRunDataReceived: Event<DataReceivedEvent>;
     readonly onDiscoveryDataReceived: Event<DataReceivedEvent>;
-    sendCommand(options: TestCommandOptions, runTestIdsPort?: string, callback?: () => void): Promise<void>;
+    sendCommand(
+        options: TestCommandOptions,
+        runTestIdsPort?: string,
+        runInstance?: TestRun,
+        callback?: () => void,
+    ): Promise<void>;
     serverReady(): Promise<void>;
     getPort(): number;
     createUUID(cwd: string): string;

src/client/testing/testController/pytest/pytestDiscoveryAdapter.ts

@@ -4,13 +4,14 @@ import * as path from 'path';
 import { Uri } from 'vscode';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
+    ExecutionResult,
     IPythonExecutionFactory,
     SpawnOptions,
 } from '../../../common/process/types';
 import { IConfigurationService, ITestOutputChannel } from '../../../common/types';
 import { createDeferred } from '../../../common/utils/async';
 import { EXTENSION_ROOT_DIR } from '../../../constants';
-import { traceError, traceVerbose } from '../../../logging';
+import { traceVerbose } from '../../../logging';
 import {
     DataReceivedEvent,
     DiscoveredTestPayload,
@@ -48,7 +49,7 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         return discoveryPayload;
     }
 
-    async runPytestDiscovery(uri: Uri, executionFactory?: IPythonExecutionFactory): Promise<DiscoveredTestPayload> {
+    async runPytestDiscovery(uri: Uri, executionFactory?: IPythonExecutionFactory): Promise<void> {
         const deferred = createDeferred<DiscoveredTestPayload>();
         const relativePathToPytest = 'pythonFiles';
         const fullPluginPath = path.join(EXTENSION_ROOT_DIR, relativePathToPytest);
@@ -78,17 +79,15 @@ export class PytestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         };
         const execService = await executionFactory?.createActivatedEnvironment(creationOptions);
         // delete UUID following entire discovery finishing.
-        execService
-            ?.exec(['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs), spawnOptions)
-            .then(() => {
-                this.testServer.deleteUUID(uuid);
-                return deferred.resolve();
-            })
-            .catch((err) => {
-                traceError(`Error while trying to run pytest discovery, \n${err}\r\n\r\n`);
-                this.testServer.deleteUUID(uuid);
-                return deferred.reject(err);
-            });
-        return deferred.promise;
+        const deferredExec = createDeferred<ExecutionResult<string>>();
+        const execArgs = ['-m', 'pytest', '-p', 'vscode_pytest', '--collect-only'].concat(pytestArgs);
+        const result = execService?.execObservable(execArgs, spawnOptions);
+
+        result?.proc?.on('close', () => {
+            deferredExec.resolve({ stdout: '', stderr: '' });
+            this.testServer.deleteUUID(uuid);
+            deferred.resolve();
+        });
+        await deferredExec.promise;
     }
 }

src/client/testing/testController/pytest/pytestExecutionAdapter.ts

@@ -15,20 +15,15 @@ import {
 } from '../common/types';
 import {
     ExecutionFactoryCreateWithEnvironmentOptions,
+    ExecutionResult,
     IPythonExecutionFactory,
     SpawnOptions,
 } from '../../../common/process/types';
 import { removePositionalFoldersAndFiles } from './arguments';
 import { ITestDebugLauncher, LaunchOptions } from '../../common/types';
 import { PYTEST_PROVIDER } from '../../common/constants';
 import { EXTENSION_ROOT_DIR } from '../../../common/constants';
-import { startTestIdServer } from '../common/utils';
-
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-// (global as any).EXTENSION_ROOT_DIR = EXTENSION_ROOT_DIR;
-/**
- * Wrapper Class for pytest test execution..
- */
+import * as utils from '../common/utils';
 
 export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
     constructor(
@@ -48,18 +43,20 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
     ): Promise<ExecutionTestPayload> {
         const uuid = this.testServer.createUUID(uri.fsPath);
         traceVerbose(uri, testIds, debugBool);
-        const disposable = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
+        const disposedDataReceived = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
             if (runInstance) {
                 this.resultResolver?.resolveExecution(JSON.parse(e.data), runInstance);
             }
         });
-        try {
-            await this.runTestsNew(uri, testIds, uuid, debugBool, executionFactory, debugLauncher);
-        } finally {
-            this.testServer.deleteUUID(uuid);
-            disposable.dispose();
-            // confirm with testing that this gets called (it must clean this up)
-        }
+        const dispose = function (testServer: ITestServer) {
+            testServer.deleteUUID(uuid);
+            disposedDataReceived.dispose();
+        };
+        runInstance?.token.onCancellationRequested(() => {
+            dispose(this.testServer);
+        });
+        await this.runTestsNew(uri, testIds, uuid, runInstance, debugBool, executionFactory, debugLauncher);
+
         // placeholder until after the rewrite is adopted
         // TODO: remove after adoption.
         const executionPayload: ExecutionTestPayload = {
@@ -74,6 +71,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
         uri: Uri,
         testIds: string[],
         uuid: string,
+        runInstance?: TestRun,
         debugBool?: boolean,
         executionFactory?: IPythonExecutionFactory,
         debugLauncher?: ITestDebugLauncher,
@@ -124,7 +122,7 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
             }
             traceLog(`Running PYTEST execution for the following test ids: ${testIds}`);
 
-            const pytestRunTestIdsPort = await startTestIdServer(testIds);
+            const pytestRunTestIdsPort = await utils.startTestIdServer(testIds);
             if (spawnOptions.extraVariables)
                 spawnOptions.extraVariables.RUN_TEST_IDS_PORT = pytestRunTestIdsPort.toString();
 
@@ -143,14 +141,27 @@ export class PytestTestExecutionAdapter implements ITestExecutionAdapter {
                 traceInfo(`Running DEBUG pytest with arguments: ${testArgs.join(' ')}\r\n`);
                 await debugLauncher!.launchDebugger(launchOptions, () => {
                     deferred.resolve();
+                    this.testServer.deleteUUID(uuid);
                 });
             } else {
                 // combine path to run script with run args
                 const scriptPath = path.join(fullPluginPath, 'vscode_pytest', 'run_pytest_script.py');
                 const runArgs = [scriptPath, ...testArgs];
                 traceInfo(`Running pytests with arguments: ${runArgs.join(' ')}\r\n`);
 
-                await execService?.exec(runArgs, spawnOptions);
+                const deferredExec = createDeferred<ExecutionResult<string>>();
+                const result = execService?.execObservable(runArgs, spawnOptions);
+
+                runInstance?.token.onCancellationRequested(() => {
+                    result?.proc?.kill();
+                });
+
+                result?.proc?.on('close', () => {
+                    deferredExec.resolve({ stdout: '', stderr: '' });
+                    this.testServer.deleteUUID(uuid);
+                    deferred.resolve();
+                });
+                await deferredExec.promise;
             }
         } catch (ex) {
             traceError(`Error while running tests: ${testIds}\r\n${ex}\r\n\r\n`);

src/client/testing/testController/unittest/testDiscoveryAdapter.ts

@@ -46,12 +46,11 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         const disposable = this.testServer.onDiscoveryDataReceived((e: DataReceivedEvent) => {
             this.resultResolver?.resolveDiscovery(JSON.parse(e.data));
         });
-        try {
-            await this.callSendCommand(options);
-        } finally {
+
+        await this.callSendCommand(options, () => {
             this.testServer.deleteUUID(uuid);
             disposable.dispose();
-        }
+        });
         // placeholder until after the rewrite is adopted
         // TODO: remove after adoption.
         const discoveryPayload: DiscoveredTestPayload = {
@@ -61,8 +60,8 @@ export class UnittestTestDiscoveryAdapter implements ITestDiscoveryAdapter {
         return discoveryPayload;
     }
 
-    private async callSendCommand(options: TestCommandOptions): Promise<DiscoveredTestPayload> {
-        await this.testServer.sendCommand(options);
+    private async callSendCommand(options: TestCommandOptions, callback: () => void): Promise<DiscoveredTestPayload> {
+        await this.testServer.sendCommand(options, undefined, undefined, callback);
         const discoveryPayload: DiscoveredTestPayload = { cwd: '', status: 'success' };
         return discoveryPayload;
     }

src/client/testing/testController/unittest/testExecutionAdapter.ts

@@ -37,18 +37,19 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         runInstance?: TestRun,
     ): Promise<ExecutionTestPayload> {
         const uuid = this.testServer.createUUID(uri.fsPath);
-        const disposable = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
+        const disposedDataReceived = this.testServer.onRunDataReceived((e: DataReceivedEvent) => {
             if (runInstance) {
                 this.resultResolver?.resolveExecution(JSON.parse(e.data), runInstance);
             }
         });
-        try {
-            await this.runTestsNew(uri, testIds, uuid, debugBool);
-        } finally {
+        const dispose = function () {
+            disposedDataReceived.dispose();
+        };
+        runInstance?.token.onCancellationRequested(() => {
             this.testServer.deleteUUID(uuid);
-            disposable.dispose();
-            // confirm with testing that this gets called (it must clean this up)
-        }
+            dispose();
+        });
+        await this.runTestsNew(uri, testIds, uuid, runInstance, debugBool, dispose);
         const executionPayload: ExecutionTestPayload = { cwd: uri.fsPath, status: 'success', error: '' };
         return executionPayload;
     }
@@ -57,7 +58,9 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
         uri: Uri,
         testIds: string[],
         uuid: string,
+        runInstance?: TestRun,
         debugBool?: boolean,
+        dispose?: () => void,
     ): Promise<ExecutionTestPayload> {
         const settings = this.configSettings.getSettings(uri);
         const { unittestArgs } = settings.testing;
@@ -80,8 +83,10 @@ export class UnittestTestExecutionAdapter implements ITestExecutionAdapter {
 
         const runTestIdsPort = await startTestIdServer(testIds);
 
-        await this.testServer.sendCommand(options, runTestIdsPort.toString(), () => {
+        await this.testServer.sendCommand(options, runTestIdsPort.toString(), runInstance, () => {
+            this.testServer.deleteUUID(uuid);
             deferred.resolve();
+            dispose?.();
         });
         // placeholder until after the rewrite is adopted
         // TODO: remove after adoption.