Implemented post command and auto terminate batch response

RLoadDll/RLoadDll.cpp

@@ -68,16 +68,13 @@ BOOL CallRemoteFunc(HANDLE hProcess, LPCSTR pProcName, void* pLibRemote, DWORD*
     if (hThread == NULL)
     {
         _ftprintf(stderr, _T("Error: CreateRemoteThread 0x%08x\n"), GetLastError());
-        CloseHandle(hThread);
         return FALSE;
     }
     if (WaitForSingleObject(hThread, INFINITE) == WAIT_FAILED)
         _ftprintf(stderr, _T("Error: WaitForSingleObject 0x%08x\n"), GetLastError());
 
     if (!GetExitCodeThread(hThread, pRet))
-    {
         _ftprintf(stderr, _T("Error: GetExitCodeThread 0x%08x\n"), GetLastError());
-    }
 
     CloseHandle(hThread);
 

RadLine/DynEnv.cpp

@@ -52,8 +52,12 @@ extern "C" {
     )
     {
         DWORD ret = 0;
-        size_t len = lpName != nullptr ? wcslen(lpName) : 0;
-        if (lpName != nullptr && lpName[0] == L'(' && lpName[len - 1] == L')')
+        const size_t len = lpName != nullptr ? wcslen(lpName) : 0;
+
+        if (lpBuffer == nullptr)
+            // TODO Handle the case when lpBuffer == nullptr
+            ret = pOrigGetEnvironmentVariableW(lpName, lpBuffer, nSize);
+        else if (lpName != nullptr && lpName[0] == L'(' && lpName[len - 1] == L')')
         {
             WCHAR strCmdLine[2048];
             wcsncpy_s(strCmdLine, ARRAYSIZE(strCmdLine), lpName + 1, len - 2);

RadLine/RadReadConsole.cpp

@@ -5,6 +5,7 @@
 #include "RadLine.h"
 #include "WinHelpers.h"
 #include "Debug.h"
+#include "bufstring.h"
 
 #include <wchar.h>
 
@@ -21,23 +22,63 @@ extern "C" {
         _In_opt_ PCONSOLE_READCONSOLE_CONTROL pInputControl
     )
     {
-        DebugOut(TEXT("RadLine RadReadConsoleW 0x%08x 0x%p %d 0x%p 0x%p\n"), hConsoleInput, lpBuffer, nNumberOfCharsToRead, lpNumberOfCharsRead, pInputControl);
+        DebugOut(TEXT("RadLine RadReadConsoleW 0x%08p 0x%p %d 0x%p 0x%p\n"), hConsoleInput, lpBuffer, nNumberOfCharsToRead, lpNumberOfCharsRead, pInputControl);
         if (pInputControl != nullptr)
             DebugOut(TEXT("RadLine PCONSOLE_READCONSOLE_CONTROL %d %d 0x%08x 0x%08x\n"), pInputControl->nLength, pInputControl->nInitialChars, pInputControl->dwCtrlWakeupMask, pInputControl->dwControlKeyState);
 
         wchar_t enabled[100] = L"";
         GetEnvironmentVariableW(L"RADLINE", enabled);
 
-        // TODO Check console mode
         const HANDLE hStdOutput = GetStdHandle(STD_OUTPUT_HANDLE);
-        DWORD modeout = 0;
-        GetConsoleMode(hStdOutput, &modeout);
 
-        if (nNumberOfCharsToRead == 1 || GetFileType(hConsoleInput) != FILE_TYPE_CHAR || GetFileType(hStdOutput) != FILE_TYPE_CHAR)
+        if (GetFileType(hConsoleInput) != FILE_TYPE_CHAR || GetFileType(hStdOutput) != FILE_TYPE_CHAR)
         {
-            DebugOut(TEXT("RadLine RadReadConsoleW Orig\n"));
             return pOrigReadConsoleW(hConsoleInput, lpBuffer, nNumberOfCharsToRead, lpNumberOfCharsRead, pInputControl);
         }
+        else if (nNumberOfCharsToRead == 1)
+        {
+            DebugOut(TEXT("RadLine RadReadConsoleW Orig\n"));
+
+            static const WCHAR response[] = L"\n\rY";
+            static int response_count = 0;
+
+            if (response_count == 0)
+            {
+                wchar_t terminate[100] = L"";
+                GetEnvironmentVariableW(L"RADLINE_AUTO_TERMINATE_BATCH", terminate);
+                if (_wcsicmp(terminate, L"1") == 0)
+                {
+                    CONSOLE_SCREEN_BUFFER_INFO csbi = {};
+                    GetConsoleScreenBufferInfo(hStdOutput, &csbi);
+
+                    std::vector<WCHAR> buf(27);
+
+                    COORD pos = GetConsoleCursorPosition(hStdOutput);
+                    pos = Add(pos, -(SHORT) buf.size(), csbi.dwSize.X);
+
+                    DWORD read = 0;
+                    ReadConsoleOutputCharacterW(hStdOutput, buf.data(), (DWORD) buf.size(), pos, &read);
+
+                    if (wcsncmp(buf.data(), L"Terminate batch job (Y/N)? ", buf.size()) == 0)
+                        response_count = ARRAYSIZE(response) - 1;
+                }
+            }
+
+            if (response_count > 0)
+            {
+                WCHAR* pStr = reinterpret_cast<TCHAR*>(lpBuffer);
+                --response_count;
+                pStr[0] = response[response_count];
+                *lpNumberOfCharsRead = 1;
+
+                DWORD written = 0;
+                WriteConsoleW(hStdOutput, lpBuffer, 1, &written, nullptr);
+
+                return TRUE;
+            }
+            else
+                return pOrigReadConsoleW(hConsoleInput, lpBuffer, nNumberOfCharsToRead, lpNumberOfCharsRead, pInputControl);
+        }
         else if ((_wcsicmp(enabled, L"") == 0 || _wcsicmp(enabled, L"1") == 0)
             && pInputControl != nullptr && pInputControl->dwCtrlWakeupMask != 0
             && lpNumberOfCharsRead != nullptr)
@@ -58,7 +99,7 @@ extern "C" {
 
                 if (r)
                 {
-                    WCHAR* const pStr = (WCHAR*) lpBuffer;
+                    WCHAR* const pStr = reinterpret_cast<TCHAR*>(lpBuffer);
                     WCHAR* const pTab = wmemchr(pStr, c, *lpNumberOfCharsRead);
                     if (pTab != nullptr)
                     {
@@ -79,6 +120,21 @@ extern "C" {
                     repeat = false;
             }
 
+            // nNumberOfCharsToRead == 1023 when used for "set /p"
+            if (nNumberOfCharsToRead != 1023 && *lpNumberOfCharsRead > 2)
+            {
+                wchar_t post[100] = L"";
+                GetEnvironmentVariableW(L"RADLINE_POST", post);
+                if (post[0] != TEXT('\0'))
+                {
+                    bufstring cmd(reinterpret_cast<TCHAR*>(lpBuffer), nNumberOfCharsToRead, *lpNumberOfCharsRead);
+                    const wchar_t* w = cmd.begin() + *lpNumberOfCharsRead - 2;
+                    cmd.insert(w, post);
+                    cmd.insert(w, L"& ");
+                    *lpNumberOfCharsRead = static_cast<DWORD>(cmd.length());
+                }
+            }
+
             CleanUpExtra(hStdOutput, &extra);
 
             return r;
@@ -87,14 +143,16 @@ extern "C" {
             && (pInputControl == nullptr || pInputControl->nInitialChars == 0))
         {
             DebugOut(TEXT("RadLine RadReadConsoleW 1\n"));
+            DWORD modeout = 0;
+            GetConsoleMode(hStdOutput, &modeout);
             SetConsoleMode(hStdOutput, modeout & ~ENABLE_PROCESSED_OUTPUT); // Needed so that cursor moves on to next line
 
             size_t length = RadLine(hConsoleInput, hStdOutput, (wchar_t*) lpBuffer, nNumberOfCharsToRead);
 
             SetConsoleMode(hStdOutput, modeout);
 
             if (lpNumberOfCharsRead != nullptr)
-                *lpNumberOfCharsRead = (DWORD) length;
+                *lpNumberOfCharsRead = static_cast<DWORD>(length);
 
             return TRUE;
         }

RadLine/bufstring.h

@@ -68,6 +68,18 @@ class bufstring
         assert(m_buf[m_len] == L'\0');
     }
 
+    void insert(const wchar_t* b, const wchar_t* s)
+    {
+        assert(b >= begin());
+        assert(b <= end());
+        const size_t l = wcslen(s);
+        wchar_t* lb = begin() + (b - begin());
+        wmemmove(lb + l, lb, end() - lb + l);
+        wmemmove(lb, s, l);
+        m_len += l;
+        assert(m_buf[m_len] == L'\0');
+    }
+
 #if 0
     void replace(size_t b, size_t l, const std::wstring& s)
     {