From 54b0ce9b5650040fade22b2e8a9fcd775110361b Mon Sep 17 00:00:00 2001
From: Yanbing <yawang@microsoft.com>
Date: Fri, 16 Nov 2018 15:36:31 -0800
Subject: [PATCH] fix of issue 1290 (#358)

Fix processid assignment, memory leak and handle leaks on conpty session.
---
 contrib/win32/win32compat/win32_pty.c | 110 +++++++++-----------------
 1 file changed, 38 insertions(+), 72 deletions(-)

diff --git a/contrib/win32/win32compat/win32_pty.c b/contrib/win32/win32compat/win32_pty.c
index bb1d10e6d..61b1bba00 100644
--- a/contrib/win32/win32compat/win32_pty.c
+++ b/contrib/win32/win32compat/win32_pty.c
@@ -35,63 +35,6 @@
 #include "misc_internal.h"
 #include "signal_internal.h"
 
-// Return Value: 0 for success, -1 for failure
-int
-CreateConPty(const wchar_t *cmdline,
-	const unsigned short width,
-	const unsigned short height,
-	HANDLE const hInput,
-	HANDLE const hOutput,	
-	HANDLE const tty_sighandle,
-	PROCESS_INFORMATION* const piPty)
-{
-	wchar_t system32_path[PATH_MAX] = { 0, };
-
-	SetHandleInformation(tty_sighandle, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT);
-
-	wchar_t conhostCmdline[8191] = { 0, }; // msdn	
-	wchar_t *cmd_fmt = L"%ls\\conhost.exe --headless --width %d --height %d --signal 0x%x -- %ls";
-	
-	if (!GetSystemDirectoryW(system32_path, PATH_MAX))
-		fatal("unable to retrieve system32 path");
-	
-	_snwprintf_s(conhostCmdline,
-		_countof(conhostCmdline),
-		_countof(conhostCmdline),
-		cmd_fmt,
-		system32_path,
-		width,
-		height,
-		tty_sighandle,
-		cmdline);
-	
-	STARTUPINFOW si;
-	memset(&si, 0, sizeof(STARTUPINFOW));
-	si.cb = sizeof(STARTUPINFOW);
-	si.hStdInput = hInput;
-	si.hStdOutput = hOutput;	
-	si.hStdError = hOutput;
-	si.dwFlags = STARTF_USESTDHANDLES | STARTF_USESIZE | STARTF_USECOUNTCHARS;
-
-	debug3("pty commandline: %ls", conhostCmdline);
-
-	/* 
-	 * process CTRL+C input.
-	 * Child processes will inherit this behavior.
-	 */
-	SetConsoleCtrlHandler(NULL, FALSE);
-	if (0 == CreateProcessW(NULL, conhostCmdline, NULL, NULL, TRUE, 0, NULL, NULL, &si, piPty)) {
-		debug("%s - failed to execute %ls, error:%d", __func__, conhostCmdline, GetLastError());
-		errno = EOTHER;
-		return -1;
-	}
-
-	/* disable Ctrl+C hander in this process*/
-	SetConsoleCtrlHandler(NULL, TRUE);
-
-	return 0;
-}
-
 int
 is_conpty_supported()
 {
@@ -140,33 +83,54 @@ int exec_command_with_pty(int * pid, char* cmd, int in, int out, int err, unsign
 	HANDLE ttyh = (HANDLE)w32_fd_to_handle(ttyfd);
 	wchar_t * cmd_w = NULL;
 
-	if ((cmd_w = utf8_to_utf16(cmd)) == NULL)
+	if ((cmd_w = utf8_to_utf16(cmd)) == NULL) {
+		errno = ENOMEM;
 		return ret;
+	}
 
 	memset(&si, 0, sizeof(STARTUPINFO));
 	si.cb = sizeof(STARTUPINFO);
-	si.dwXCountChars = col;
-	si.dwYCountChars = row;
 	si.dwFlags = STARTF_USESTDHANDLES | STARTF_USESIZE | STARTF_USECOUNTCHARS;
-
 	si.hStdInput = (HANDLE)w32_fd_to_handle(in);
 	si.hStdOutput = (HANDLE)w32_fd_to_handle(out);
-	si.hStdError = (HANDLE)w32_fd_to_handle(err);
 	si.lpDesktop = NULL;
 
-	if (is_conpty_supported())
-		return CreateConPty(cmd_w, (short)si.dwXCountChars, (short)si.dwYCountChars, si.hStdInput, si.hStdOutput, ttyh, &pi);
+	if (is_conpty_supported()) {
+		wchar_t system32_path[PATH_MAX] = { 0, };
+		SetHandleInformation(ttyh, HANDLE_FLAG_INHERIT, HANDLE_FLAG_INHERIT);
+		wchar_t *cmd_fmt = L"%ls\\conhost.exe --headless --width %d --height %d --signal 0x%x -- %ls";
 
-	/* launch via  "ssh-shellhost" -p command*/
+		if (!GetSystemDirectoryW(system32_path, PATH_MAX))
+			fatal("unable to retrieve system32 path");
+
+		_snwprintf_s(pty_cmdline,
+			MAX_CMD_LEN,
+			MAX_CMD_LEN,
+			cmd_fmt,
+			system32_path,
+			col,
+			row,
+			ttyh,
+			cmd_w);
+
+		si.hStdError = si.hStdOutput;
+		/* process CTRL+C input. Child processes will inherit this behavior. */
+		SetConsoleCtrlHandler(NULL, FALSE);
+	}
+	else {
+		/* launch via  "ssh-shellhost" -p command*/
+		_snwprintf_s(pty_cmdline, MAX_CMD_LEN, MAX_CMD_LEN, L"\"%ls\\ssh-shellhost.exe\" ---pty %ls", __wprogdir, cmd_w);
+		si.dwXCountChars = col;
+		si.dwYCountChars = row;
+
+		/*
+		 * In PTY mode, ssh-shellhost takes stderr as control channel
+		 * TODO - fix this and pass control channel pipe as a command line parameter
+		 */
+		si.hStdError = ttyh;
+	}
 
-	_snwprintf_s(pty_cmdline, MAX_CMD_LEN, MAX_CMD_LEN, L"\"%ls\\ssh-shellhost.exe\" ---pty %ls", __wprogdir, cmd_w);
-	/* 
-	 * In PTY mode, ssh-shellhost takes stderr as control channel
-	 * TODO - fix this and pass control channel pipe as a command line parameter
-	 */
-	si.hStdError = ttyh;
 	debug3("pty commandline: %ls", pty_cmdline);
-
 	if (CreateProcessW(NULL, pty_cmdline, NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) {
 		if (register_child(pi.hProcess, pi.dwProcessId) == -1) {
 			TerminateProcess(pi.hProcess, 0);
@@ -184,6 +148,8 @@ int exec_command_with_pty(int * pid, char* cmd, int in, int out, int err, unsign
 	ret = 0;
 
 done:
+	/* disable Ctrl+C hander in this process*/
+	SetConsoleCtrlHandler(NULL, TRUE);
 	if (cmd_w)
 		free(cmd_w);
 	return ret;