fix: ensure proper cleanup of console process on shutdown

src/native.d.ts

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

src/win/conpty.cc

@@ -37,42 +37,70 @@ typedef void (__stdcall *PFNCLOSEPSEUDOCONSOLE)(HPCON hpc);
 
 #endif
 
+typedef struct _PseudoConsole
+{
+    HANDLE hSignal;
+    HANDLE hPtyReference;
+    HANDLE hConPtyProcess;
+} PseudoConsole;
+
 struct pty_baton {
   int id;
   HANDLE hIn;
   HANDLE hOut;
   HPCON hpc;
 
   HANDLE hShell;
+  HANDLE pty_job_handle;
 
   pty_baton(int _id, HANDLE _hIn, HANDLE _hOut, HPCON _hpc) : id(_id), hIn(_hIn), hOut(_hOut), hpc(_hpc) {};
 };
 
-static std::vector<pty_baton*> ptyHandles;
+static std::vector<std::unique_ptr<pty_baton>> ptyHandles;
 static volatile LONG ptyCounter;
 
 static pty_baton* get_pty_baton(int id) {
-  for (size_t i = 0; i < ptyHandles.size(); ++i) {
-    pty_baton* ptyHandle = ptyHandles[i];
-    if (ptyHandle->id == id) {
-      return ptyHandle;
-    }
+  auto it = std::find_if(ptyHandles.begin(), ptyHandles.end(), [id](const auto& ptyHandle) {
+    return ptyHandle->id == id;
+  });
+  if (it != ptyHandles.end()) {
+    return it->get();
   }
   return nullptr;
 }
 
 static bool remove_pty_baton(int id) {
-  for (size_t i = 0; i < ptyHandles.size(); ++i) {
-    pty_baton* ptyHandle = ptyHandles[i];
-    if (ptyHandle->id == id) {
-      ptyHandles.erase(ptyHandles.begin() + i);
-      ptyHandle = nullptr;
-      return true;
-    }
+  auto it = std::remove_if(ptyHandles.begin(), ptyHandles.end(), [id](const auto& ptyHandle) {
+    return ptyHandle->id == id;
+  });
+  if (it != ptyHandles.end()) {
+    ptyHandles.erase(it);
+    return true;
   }
   return false;
 }
 
+static BOOL InitPtyJobHandle(HANDLE* pty_job_handle) {
+  SECURITY_ATTRIBUTES attr;
+  memset(&attr, 0, sizeof(attr));
+  attr.bInheritHandle = FALSE;
+  *pty_job_handle = CreateJobObjectW(&attr, nullptr);
+  if (*pty_job_handle == NULL) {
+    return FALSE;
+  }
+  JOBOBJECT_EXTENDED_LIMIT_INFORMATION limits = {};
+  limits.BasicLimitInformation.LimitFlags =
+      JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION |
+      JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
+  if (!SetInformationJobObject(*pty_job_handle,
+                               JobObjectExtendedLimitInformation,
+                               &limits,
+                               sizeof(limits))) {
+    return FALSE;
+  }
+  return TRUE;
+}
+
 struct ExitEvent {
   int exit_code = 0;
 };
@@ -107,6 +135,8 @@ void SetupExitCallback(Napi::Env env, Napi::Function cb, pty_baton* baton) {
     // so we only call CloseHandle for now.
     CloseHandle(baton->hIn);
     CloseHandle(baton->hOut);
+    CloseHandle(baton->hShell);
+    assert(remove_pty_baton(baton->id));
 
     auto status = tsfn.BlockingCall(exit_event, callback); // In main thread
     switch (status) {
@@ -304,7 +334,8 @@ static Napi::Value PtyStartProcess(const Napi::CallbackInfo& info) {
     // We were able to instantiate a conpty
     const int ptyId = InterlockedIncrement(&ptyCounter);
     marshal.Set("pty", Napi::Number::New(env, ptyId));
-    ptyHandles.insert(ptyHandles.end(), new pty_baton(ptyId, hIn, hOut, hpc));
+    ptyHandles.emplace_back(
+        std::make_unique<pty_baton>(ptyId, hIn, hOut, hpc));
   } else {
     throw Napi::Error::New(env, "Cannot launch conpty");
   }
@@ -335,20 +366,22 @@ static Napi::Value PtyConnect(const Napi::CallbackInfo& info) {
   std::stringstream errorText;
   BOOL fSuccess = FALSE;
 
-  if (info.Length() != 5 ||
+  if (info.Length() != 6 ||
       !info[0].IsNumber() ||
       !info[1].IsString() ||
       !info[2].IsString() ||
       !info[3].IsArray() ||
-      !info[4].IsFunction()) {
-    throw Napi::Error::New(env, "Usage: pty.connect(id, cmdline, cwd, env, exitCallback)");
+      !info[4].IsBoolean() ||
+      !info[5].IsFunction()) {
+    throw Napi::Error::New(env, "Usage: pty.connect(id, cmdline, cwd, env, useConptyDll, exitCallback)");
   }
 
   const int id = info[0].As<Napi::Number>().Int32Value();
   const std::wstring cmdline(path_util::to_wstring(info[1].As<Napi::String>()));
   const std::wstring cwd(path_util::to_wstring(info[2].As<Napi::String>()));
   const Napi::Array envValues = info[3].As<Napi::Array>();
-  Napi::Function exitCallback = info[4].As<Napi::Function>();
+  const bool useConptyDll = info[4].As<Napi::Boolean>().Value();
+  Napi::Function exitCallback = info[5].As<Napi::Function>();
 
   // Fetch pty handle from ID and start process
   pty_baton* handle = get_pty_baton(id);
@@ -426,6 +459,17 @@ static Napi::Value PtyConnect(const Napi::CallbackInfo& info) {
     throw errorWithCode(info, "Cannot create process");
   }
 
+  if (useConptyDll) {
+    if (!InitPtyJobHandle(&handle->pty_job_handle)) {
+      throw errorWithCode(info, "Initialize pty job handle failed");
+    }
+
+    PseudoConsole* server = reinterpret_cast<PseudoConsole*>(handle->hpc);
+    if (!AssignProcessToJobObject(handle->pty_job_handle, server->hConPtyProcess)) {
+      throw errorWithCode(info, "AssignProcessToJobObject for server failed");
+    }
+  }
+
   // Update handle
   handle->hShell = piClient.hProcess;
 
@@ -545,8 +589,10 @@ static Napi::Value PtyKill(const Napi::CallbackInfo& info) {
       }
     }
 
-    CloseHandle(handle->hShell);
-    assert(remove_pty_baton(id));
+    TerminateProcess(handle->hShell, 1);
+    if (useConptyDll) {
+      CloseHandle(handle->pty_job_handle);
+    }
   }
 
   return env.Undefined();

src/windowsPtyAgent.ts

@@ -7,9 +7,9 @@
 import * as fs from 'fs';
 import * as os from 'os';
 import * as path from 'path';
+import { fork } from 'child_process';
 import { Socket } from 'net';
 import { ArgvOrCommandLine } from './types';
-import { fork } from 'child_process';
 import { ConoutConnection } from './windowsConoutConnection';
 
 let conptyNative: IConptyNative;
@@ -135,7 +135,7 @@ export class WindowsPtyAgent {
     this._inSocket.setEncoding('utf8');
 
     if (this._useConpty) {
-      const connect = (this._ptyNative as IConptyNative).connect(this._pty, commandLine, cwd, env, c => this._$onProcessExit(c));
+      const connect = (this._ptyNative as IConptyNative).connect(this._pty, commandLine, cwd, env, this._useConptyDll, c => this._$onProcessExit(c));
       this._innerPid = connect.pid;
     }
   }
@@ -162,16 +162,18 @@ export class WindowsPtyAgent {
     this._outSocket.readable = false;
     // Tell the agent to kill the pty, this releases handles to the process
     if (this._useConpty) {
-      this._getConsoleProcessList().then(consoleProcessList => {
-        consoleProcessList.forEach((pid: number) => {
-          try {
-            process.kill(pid);
-          } catch (e) {
-            // Ignore if process cannot be found (kill ESRCH error)
-          }
+      if (!this._useConptyDll) {
+        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, this._useConptyDll);
-      });
+      }
+      (this._ptyNative as IConptyNative).kill(this._pty, this._useConptyDll);
     } else {
       // Because pty.kill closes the handle, it will kill most processes by itself.
       // Process IDs can be reused as soon as all handles to them are

src/windowsTerminal.test.ts

@@ -88,18 +88,18 @@ function pollForProcessTreeSize(pid: number, size: number, intervalMs: number =
 }
 
 if (process.platform === 'win32') {
-  [true, false].forEach((useConpty) => {
-    describe(`WindowsTerminal (useConpty = ${useConpty})`, () => {
+  [[true, false], [true, true], [false, false]].forEach(([useConpty, useConptyDll]) => {
+    describe(`WindowsTerminal (useConpty = ${useConpty}, useConptyDll = ${useConptyDll})`, () => {
       describe('kill', () => {
         it('should not crash parent process', (done) => {
-          const term = new WindowsTerminal('cmd.exe', [], { useConpty });
+          const term = new WindowsTerminal('cmd.exe', [], { useConpty, useConptyDll });
           term.kill();
           // Add done call to deferred function queue to ensure the kill call has completed
           (<any>term)._defer(done);
         });
         it('should kill the process tree', function (done: Mocha.Done): void {
           this.timeout(10000);
-          const term = new WindowsTerminal('cmd.exe', [], { useConpty });
+          const term = new WindowsTerminal('cmd.exe', [], { useConpty, useConptyDll });
           // Start sub-processes
           term.write('powershell.exe\r');
           term.write('notepad.exe\r');
@@ -115,24 +115,26 @@ if (process.platform === 'win32') {
             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();
+            term.on('exit', () => {
+              pollForProcessState(desiredState).then(() => {
+                // Kill notepad before done
+                process.kill(list[2].pid);
+                done();
+              });
             });
           });
         });
       });
 
       describe('resize', () => {
         it('should throw a non-native exception when resizing an invalid value', () => {
-          const term = new WindowsTerminal('cmd.exe', [], { useConpty });
+          const term = new WindowsTerminal('cmd.exe', [], { useConpty, useConptyDll });
           assert.throws(() => term.resize(-1, -1));
           assert.throws(() => term.resize(0, 0));
           assert.doesNotThrow(() => term.resize(1, 1));
         });
         it('should throw a non-native exception when resizing a killed terminal', (done) => {
-          const term = new WindowsTerminal('cmd.exe', [], { useConpty });
+          const term = new WindowsTerminal('cmd.exe', [], { useConpty, useConptyDll });
           (<any>term)._defer(() => {
             term.once('exit', () => {
               assert.throws(() => term.resize(1, 1));
@@ -158,7 +160,7 @@ if (process.platform === 'win32') {
             // Skip test if git bash isn't installed
             return;
           }
-          const term = new WindowsTerminal(cmdCopiedPath, '/c echo "hello world"', { useConpty });
+          const term = new WindowsTerminal(cmdCopiedPath, '/c echo "hello world"', { useConpty, useConptyDll });
           let result = '';
           term.on('data', (data) => {
             result += data;
@@ -186,15 +188,15 @@ if (process.platform === 'win32') {
 
       describe('On close', () => {
         it('should return process zero exit codes', (done) => {
-          const term = new WindowsTerminal('cmd.exe', '/C exit', { useConpty });
+          const term = new WindowsTerminal('cmd.exe', '/C exit', { useConpty, useConptyDll });
           term.on('exit', (code) => {
             assert.strictEqual(code, 0);
             done();
           });
         });
 
         it('should return process non-zero exit codes', (done) => {
-          const term = new WindowsTerminal('cmd.exe', '/C exit 2', { useConpty });
+          const term = new WindowsTerminal('cmd.exe', '/C exit 2', { useConpty, useConptyDll });
           term.on('exit', (code) => {
             assert.strictEqual(code, 2);
             done();
@@ -204,7 +206,7 @@ if (process.platform === 'win32') {
 
       describe('Write', () => {
         it('should accept input', (done) => {
-          const term = new WindowsTerminal('cmd.exe', '', { useConpty });
+          const term = new WindowsTerminal('cmd.exe', '', { useConpty, useConptyDll });
           term.write('exit\r');
           term.on('exit', () => {
             done();