Refactor ConsoleProcessList

microsoft · Nov 22, 2022 · 76a1367 · 76a1367
1 parent 437b5ac
commit 76a1367
Show file tree

Hide file tree

Showing 7 changed files with 131 additions and 180 deletions.
diff --git a/src/host/ft_fuzzer/fuzzmain.cpp b/src/host/ft_fuzzer/fuzzmain.cpp
@@ -74,7 +74,6 @@ struct NullDeviceComm : public IDeviceComm
     RETURN_IF_FAILED(gci.ProcessHandleList.AllocProcessData(GetCurrentProcessId(),
                                                             GetCurrentThreadId(),
                                                             0,
-                                                            nullptr,
                                                             &pProcessHandle));
     pProcessHandle->fRootProcess = true;
 

diff --git a/src/host/input.cpp b/src/host/input.cpp
@@ -275,27 +275,21 @@ void ProcessCtrlEvents()
     const auto LimitingProcessId = gci.LimitingProcessId;
     gci.LimitingProcessId = 0;
 
-    ConsoleProcessTerminationRecord* rgProcessHandleList;
-    size_t cProcessHandleList;
-
-    auto hr = gci.ProcessHandleList
-                  .GetTerminationRecordsByGroupId(LimitingProcessId,
-                                                  WI_IsFlagSet(gci.CtrlFlags,
-                                                               CONSOLE_CTRL_CLOSE_FLAG),
-                                                  &rgProcessHandleList,
-                                                  &cProcessHandleList);
-
-    if (FAILED(hr) || cProcessHandleList == 0)
+    std::vector<ConsoleProcessTerminationRecord> termRecords;
+    const auto hr = gci.ProcessHandleList
+                        .GetTerminationRecordsByGroupId(LimitingProcessId,
+                                                        WI_IsFlagSet(gci.CtrlFlags,
+                                                                     CONSOLE_CTRL_CLOSE_FLAG),
+                                                        termRecords);
+
+    if (FAILED(hr) || termRecords.empty())
     {
         gci.UnlockConsole();
         return;
     }
 
     // Copy ctrl flags.
-    auto CtrlFlags = gci.CtrlFlags;
-    FAIL_FAST_IF(!(!((CtrlFlags & (CONSOLE_CTRL_CLOSE_FLAG | CONSOLE_CTRL_BREAK_FLAG | CONSOLE_CTRL_C_FLAG)) && (CtrlFlags & (CONSOLE_CTRL_LOGOFF_FLAG | CONSOLE_CTRL_SHUTDOWN_FLAG)))));
-
-    gci.CtrlFlags = 0;
+    const auto CtrlFlags = std::exchange(gci.CtrlFlags, 0);
 
     gci.UnlockConsole();
 
@@ -330,10 +324,13 @@ void ProcessCtrlEvents()
     case CONSOLE_CTRL_SHUTDOWN_FLAG:
         EventType = CTRL_SHUTDOWN_EVENT;
         break;
+
+    default:
+        return;
     }
 
     auto Status = STATUS_SUCCESS;
-    for (size_t i = 0; i < cProcessHandleList; i++)
+    for (const auto& r : termRecords)
     {
         /*
          * Status will be non-successful if a process attached to this console
@@ -347,20 +344,13 @@ void ProcessCtrlEvents()
         if (NT_SUCCESS(Status))
         {
             Status = ServiceLocator::LocateConsoleControl()
-                         ->EndTask(rgProcessHandleList[i].dwProcessID,
+                         ->EndTask(r.dwProcessID,
                                    EventType,
                                    CtrlFlags);
-            if (rgProcessHandleList[i].hProcess == nullptr)
+            if (!r.hProcess)
             {
                 Status = STATUS_SUCCESS;
             }
         }
-
-        if (rgProcessHandleList[i].hProcess != nullptr)
-        {
-            CloseHandle(rgProcessHandleList[i].hProcess);
-        }
     }
-
-    delete[] rgProcessHandleList;
 }
diff --git a/src/interactivity/win32/windowio.cpp b/src/interactivity/win32/windowio.cpp
@@ -78,7 +78,7 @@ VOID SetConsoleWindowOwner(const HWND hwnd, _Inout_opt_ ConsoleProcessHandle* pP
     else
     {
         // Find a process to own the console window. If there are none then let's use conhost's.
-        pProcessData = gci.ProcessHandleList.GetFirstProcess();
+        pProcessData = gci.ProcessHandleList.GetOldestProcess();
         if (pProcessData != nullptr)
         {
             dwProcessId = pProcessData->dwProcessId;
@@ -988,7 +988,7 @@ LRESULT CALLBACK DialogHookProc(int nCode, WPARAM /*wParam*/, LPARAM lParam)
 NTSTATUS InitWindowsSubsystem(_Out_ HHOOK* phhook)
 {
     auto& gci = ServiceLocator::LocateGlobals().getConsoleInformation();
-    auto ProcessData = gci.ProcessHandleList.FindProcessInList(ConsoleProcessList::ROOT_PROCESS_ID);
+    auto ProcessData = gci.ProcessHandleList.GetRootProcess();
     FAIL_FAST_IF(!(ProcessData != nullptr && ProcessData->fRootProcess));
 
     // Create and activate the main window

diff --git a/src/server/ApiDispatchersInternal.cpp b/src/server/ApiDispatchersInternal.cpp
@@ -90,11 +90,17 @@ using Microsoft::Console::Interactivity::ServiceLocator;
             RETURN_IF_NTSTATUS_FAILED((NtPrivApi::s_GetProcessParentId(&ProcessId)));
             ProcessHandle = gci.ProcessHandleList.FindProcessInList(ProcessId);
             RETURN_HR_IF_NULL(E_INVALIDARG, ProcessHandle);
-            RETURN_IF_FAILED(gci.ProcessHandleList.AllocProcessData(a->ProcessGroupId,
-                                                                    0,
-                                                                    a->ProcessGroupId,
-                                                                    ProcessHandle,
-                                                                    nullptr));
+
+            const auto hr = gci.ProcessHandleList.AllocProcessData(a->ProcessGroupId,
+                                                                   0,
+                                                                   a->ProcessGroupId,
+                                                                   nullptr);
+            // AllocProcessData returns E_FAIL if the given process ID already exists.
+            // The process ID already existing isn't really an error though so we ignore it.
+            if (FAILED(hr) && hr != E_FAIL)
+            {
+                RETURN_HR(hr);
+            }
         }
     }
 

diff --git a/src/server/IoDispatchers.cpp b/src/server/IoDispatchers.cpp
@@ -452,7 +452,6 @@ PCONSOLE_API_MSG IoDispatchers::ConsoleHandleConnectionRequest(_In_ PCONSOLE_API
     Status = NTSTATUS_FROM_HRESULT(gci.ProcessHandleList.AllocProcessData(dwProcessId,
                                                                           dwThreadId,
                                                                           Cac.ProcessGroupId,
-                                                                          nullptr,
                                                                           &ProcessData));
 
     if (!NT_SUCCESS(Status))