Fix kill under conpty

Author: Tyriar

binding.gyp

@@ -15,6 +15,15 @@
             'shlwapi.lib'
           ]
         },
+        {
+          'target_name': 'conpty_console_list',
+          'include_dirs' : [
+            '<!(node -e "require(\'nan\')")'
+          ],
+          'sources' : [
+            'src/win/conpty_console_list.cc'
+          ]
+        },
         {
           'target_name': 'pty',
           'include_dirs' : [

package.json

@@ -34,6 +34,7 @@
   ],
   "scripts": {
     "tsc": "tsc",
+    "watch": "tsc -w",
     "lint": "tslint 'src/**/*.ts'",
     "install": "node scripts/install.js",
     "postinstall": "node scripts/post-install.js",
@@ -46,11 +47,14 @@
   "devDependencies": {
     "@types/mocha": "^5.0.0",
     "@types/node": "^6.0.104",
+    "@types/ps-list": "^6.0.0",
     "cross-env": "^5.1.4",
     "mocha": "^5.0.5",
     "pollUntil": "^1.0.3",
+    "ps-list": "^6.0.0",
     "tslint": "^5.12.1",
     "tslint-consistent-codestyle": "^1.15.0",
-    "typescript": "^3.2.2"
+    "typescript": "^3.2.2",
+    "windows-process-tree": "^0.2.3"
   }
 }

scripts/post-install.js

@@ -5,6 +5,8 @@ var RELEASE_DIR = path.join(__dirname, '..', 'build', 'Release');
 var BUILD_FILES = [
   path.join(RELEASE_DIR, 'conpty.node'),
   path.join(RELEASE_DIR, 'conpty.pdb'),
+  path.join(RELEASE_DIR, 'conpty_console_list.node'),
+  path.join(RELEASE_DIR, 'conpty_console_list.pdb'),
   path.join(RELEASE_DIR, 'pty.node'),
   path.join(RELEASE_DIR, 'pty.pdb'),
   path.join(RELEASE_DIR, 'winpty-agent.exe'),

src/conpty_console_list_agent.ts

@@ -0,0 +1,15 @@
+/**
+ * Copyright (c) 2019, Microsoft Corporation (MIT License).
+ *
+ * This module fetches the console process list for a particular PID. It must be
+ * called from a different process (child_process.fork) as there can only be a
+ * single console attached to a process.
+ */
+
+import { loadNative } from './utils';
+
+const getConsoleProcessList = loadNative('conpty_console_list').getConsoleProcessList;
+const shellPid = parseInt(process.argv[2], 10);
+const consoleProcessList = getConsoleProcessList(shellPid);
+process.send({ consoleProcessList });
+process.exit(0);

src/native.d.ts

@@ -5,7 +5,6 @@
 interface IConptyNative {
   startProcess(file: string, cols: number, rows: number, debug: boolean, pipeName: string): IConptyProcess;
   connect(ptyId: number, commandLine: string, cwd: string, env: string[], onExitCallback: (exitCode: number) => void): { pid: number };
-  getProcessList(shellPid: number): number[];
   resize(ptyId: number, cols: number, rows: number): void;
   kill(ptyId: number): void;
 }

src/win/conpty.cc

@@ -419,52 +419,11 @@ static NAN_METHOD(PtyKill) {
     }
   }
 
-  // Fetch the console process tree
-  // auto processList = std::vector<DWORD>(64);
-  // auto processCount = GetConsoleProcessList(&processList[0], processList.size());
-  // if (processList.size() < processCount) {
-  //     processList.resize(processCount);
-  //     processCount = GetConsoleProcessList(&processList[0], processList.size());
-  // }
-
-  // // Kill the console process tree
-  // for (DWORD i = 0; i < processCount; i++) {
-  //   TerminateProcess(processList[i], 1);
-  // }
-
-  // Close the shell handle
   CloseHandle(handle->hShell);
 
   return info.GetReturnValue().SetUndefined();
 }
 
-static NAN_METHOD(PtyGetProcessList) {
-  Nan::HandleScope scope;
-
-  if (info.Length() != 1 ||
-      !info[0]->IsNumber()) {
-    Nan::ThrowError("Usage: pty.getProcessList(shellPid)");
-    return;
-  }
-
-  int shellPid = info[0]->Int32Value();
-
-  AttachConsole(shellPid);
-  auto processList = std::vector<DWORD>(64);
-  auto processCount = GetConsoleProcessList(&processList[0], processList.size());
-  if (processList.size() < processCount) {
-      processList.resize(processCount);
-      processCount = GetConsoleProcessList(&processList[0], processList.size());
-  }
-  FreeConsole();
-
-  auto result = Nan::New<v8::Array>(processCount);
-  for (DWORD i = 0; i < processCount; i++) {
-    result->Set(i, Nan::New<v8::Number>(processList[i]));
-  }
-  info.GetReturnValue().Set(result);
-}
-
 /**
 * Init
 */
@@ -475,7 +434,6 @@ extern "C" void init(v8::Handle<v8::Object> target) {
   Nan::SetMethod(target, "connect", PtyConnect);
   Nan::SetMethod(target, "resize", PtyResize);
   Nan::SetMethod(target, "kill", PtyKill);
-  Nan::SetMethod(target, "getProcessList", PtyGetProcessList);
 };
 
 NODE_MODULE(pty, init);

src/win/conpty_console_list.cc

@@ -0,0 +1,43 @@
+/**
+ * Copyright (c) 2019, Microsoft Corporation (MIT License).
+ */
+
+#include <nan.h>
+#include <windows.h>
+
+static NAN_METHOD(ApiConsoleProcessList) {
+  if (info.Length() != 1 ||
+      !info[0]->IsNumber()) {
+    Nan::ThrowError("Usage: getConsoleProcessList(shellPid)");
+    return;
+  }
+
+  const SHORT pid = info[0]->Uint32Value();
+
+  if (!FreeConsole()) {
+    Nan::ThrowError("FreeConsole failed");
+  }
+  if (!AttachConsole(pid)) {
+    Nan::ThrowError("AttachConsole failed");
+  }
+  auto processList = std::vector<DWORD>(64);
+  auto processCount = GetConsoleProcessList(&processList[0], processList.size());
+  if (processList.size() < processCount) {
+      processList.resize(processCount);
+      processCount = GetConsoleProcessList(&processList[0], processList.size());
+  }
+  FreeConsole();
+
+  v8::Local<v8::Array> result = Nan::New<v8::Array>();
+  for (DWORD i = 0; i < processCount; i++) {
+    result->Set(i, Nan::New<v8::Number>(processList[i]));
+  }
+  info.GetReturnValue().Set(result);
+}
+
+extern "C" void init(v8::Handle<v8::Object> target) {
+  Nan::HandleScope scope;
+  Nan::SetMethod(target, "getConsoleProcessList", ApiConsoleProcessList);
+};
+
+NODE_MODULE(pty, init);

src/windowsPtyAgent.ts

@@ -9,6 +9,7 @@ import * as path from 'path';
 import { Socket } from 'net';
 import { ArgvOrCommandLine } from './types';
 import { loadNative } from './utils';
+import { fork } from 'child_process';
 
 let conptyNative: IConptyNative;
 let winptyNative: IWinptyNative;
@@ -130,24 +131,48 @@ export class WindowsPtyAgent {
     this._outSocket.writable = false;
     // Tell the agent to kill the pty, this releases handles to the process
     if (this._useConpty) {
-      (this._ptyNative as IConptyNative).kill(this._pty);
+      this._getConsoleProcessList().then(consoleProcessList => {
+        consoleProcessList.forEach((pid: number) => {
+          try {
+            process.kill(pid);
+          } catch (e) {
+            // Ignore if process cannot be found (kill ESRCH error)
+          }
+        });
+        (this._ptyNative as IConptyNative).kill(this._pty);
+      });
     } else {
       (this._ptyNative as IWinptyNative).kill(this._pid, this._innerPidHandle);
+      // Since pty.kill closes the handle it will kill most processes by itself
+      // and process IDs can be reused as soon as all handles to them are
+      // dropped, we want to immediately kill the entire console process list.
+      // If we do not force kill all processes here, node servers in particular
+      // seem to become detached and remain running (see
+      // Microsoft/vscode#26807).
+      const processList: number[] = (this._ptyNative as IWinptyNative).getProcessList(this._pid);
+      processList.forEach(pid => {
+        try {
+          process.kill(pid);
+        } catch (e) {
+          // Ignore if process cannot be found (kill ESRCH error)
+        }
+      });
     }
-    // Since pty.kill closes the handle it will kill most processes by itself
-    // and process IDs can be reused as soon as all handles to them are
-    // dropped, we want to immediately kill the entire console process list.
-    // If we do not force kill all processes here, node servers in particular
-    // seem to become detached and remain running (see
-    // Microsoft/vscode#26807).
-    const processList: number[] = this._ptyNative.getProcessList(this._useConpty ? this._innerPid : this._pid);
-    console.log('processList', processList.join(', '));
-    processList.forEach(pid => {
-      try {
-        process.kill(pid);
-      } catch (e) {
-        // Ignore if process cannot be found (kill ESRCH error)
-      }
+  }
+
+  private _getConsoleProcessList(): Promise<number[]> {
+    return new Promise<number[]>(resolve => {
+      const agent = fork(path.join(__dirname, 'conpty_console_list_agent'), [ this._innerPid.toString() ]);
+      agent.on('message', message => {
+        clearTimeout(timeout);
+        resolve(message.consoleProcessList);
+      });
+      const timeout = setTimeout(() => {
+        // Something went wrong, just send back the shell PID
+        console.error('Could not fetch console process list');
+        agent.kill();
+        resolve([ this._innerPid ]);
+      }, 5000);
     });
   }
 

src/windowsTerminal.test.ts

@@ -8,6 +8,48 @@ import * as fs from 'fs';
 import * as assert from 'assert';
 import { WindowsTerminal } from './windowsTerminal';
 import * as path from 'path';
+import { getProcessList } from 'windows-process-tree';
+import * as psList from 'ps-list';
+
+interface IProcessState {
+  // Whether the PID must exist or must not exist
+  [pid: number]: boolean;
+}
+
+function pollForProcessState(desiredState: IProcessState, intervalMs: number = 100, timeoutMs: number = 2000): Promise<void> {
+  return new Promise<void>(resolve => {
+    let tries = 0;
+    const interval = setInterval(() => {
+      psList({ all: true }).then(ps => {
+        let success = true;
+        const pids = Object.keys(desiredState).map(k => parseInt(k, 10));
+        pids.forEach(pid => {
+          if (desiredState[pid]) {
+            if (!ps.some(p => p.pid === pid)) {
+              success = false;
+            }
+          } else {
+            if (ps.some(p => p.pid === pid)) {
+              success = false;
+            }
+          }
+        });
+        if (success) {
+          clearInterval(interval);
+          resolve();
+          return;
+        }
+        tries++;
+        if (tries * intervalMs >= timeoutMs) {
+          clearInterval(interval);
+          const processListing = pids.map(k => `${k}: ${desiredState[k]}`).join('\n');
+          assert.fail(`Bad process state, expected:\n${processListing}`);
+          resolve();
+        }
+      });
+    }, intervalMs);
+  });
+}
 
 if (process.platform === 'win32') {
   describe('WindowsTerminal', () => {
@@ -18,15 +60,34 @@ if (process.platform === 'win32') {
         // Add done call to deferred function queue to ensure the kill call has completed
         (<any>term)._defer(done);
       });
-      it('should kill the process tree', (done) => {
+      it('should kill the process tree', function (done: Mocha.Done): void {
+        this.timeout(5000);
+
         const term = new WindowsTerminal('cmd.exe', [], {});
         // Start a sub-process
-        term.write('powershell.exe');
-        const proc = cp.execSync(`tasklist /fi "PID eq ${term.pid}"`);
-        const index = proc.toString().indexOf(term.pid.toString());
-        console.log('******* ' + index + ', ' + proc.toString());
-        // term.pid
-        // term.kill();
+        term.write('powershell.exe\r');
+        term.write('notepad.exe\r');
+        term.write('node.exe\r');
+        setTimeout(() => {
+          getProcessList(term.pid, list => {
+            assert.equal(list.length, 4);
+            assert.equal(list[0].name, 'cmd.exe');
+            assert.equal(list[1].name, 'powershell.exe');
+            assert.equal(list[2].name, 'notepad.exe');
+            assert.equal(list[3].name, 'node.exe');
+            term.kill();
+            const desiredState: IProcessState = {};
+            desiredState[list[0].pid] = false;
+            desiredState[list[1].pid] = false;
+            desiredState[list[2].pid] = true;
+            desiredState[list[3].pid] = false;
+            pollForProcessState(desiredState).then(() => {
+              // Kill notepad before done
+              process.kill(list[2].pid);
+              done();
+            });
+          });
+        }, 1000);
       });
     });
 

tsconfig.json

@@ -6,7 +6,7 @@
     "outDir": "lib",
     "sourceMap": true,
     "lib": [
-      "es5"
+      "es2015"
     ],
     "alwaysStrict": true,
     "noImplicitAny": true,