diff --git a/src/agent/Agent.cc b/src/agent/Agent.cc index efc3e7f8..a3ab21b9 100644 --- a/src/agent/Agent.cc +++ b/src/agent/Agent.cc @@ -27,6 +27,7 @@ #include #include +#include #include #include #include @@ -332,6 +333,9 @@ void Agent::handlePacket(ReadBuffer &packet) // at once, we can ignore the early ones. handleSetSizePacket(packet); break; + case AgentMsg::GetConsoleProcessList: + handleGetConsoleProcessListPacket(packet); + break; default: trace("Unrecognized message, id:%d", type); } @@ -426,6 +430,33 @@ void Agent::handleSetSizePacket(ReadBuffer &packet) writePacket(reply); } +void Agent::handleGetConsoleProcessListPacket(ReadBuffer &packet) +{ + packet.assertEof(); + + auto processList = std::vector(64); + auto processCount = GetConsoleProcessList(&processList[0], processList.size()); + + // The process list can change while we're trying to read it + while (processList.size() < processCount) { + // Multiplying by two caps the number of iterations + const auto newSize = std::max(processList.size() * 2, processCount); + processList.resize(newSize); + processCount = GetConsoleProcessList(&processList[0], processList.size()); + } + + if (processCount == 0) { + trace("GetConsoleProcessList failed"); + } + + auto reply = newPacket(); + reply.putInt32(processCount); + for (DWORD i = 0; i < processCount; i++) { + reply.putInt32(processList[i]); + } + writePacket(reply); +} + void Agent::pollConinPipe() { const std::string newData = m_coninPipe->readAllToString(); diff --git a/src/agent/Agent.h b/src/agent/Agent.h index 5775e086..1dde48fe 100644 --- a/src/agent/Agent.h +++ b/src/agent/Agent.h @@ -59,6 +59,7 @@ class Agent : public EventLoop, public DsrSender void writePacket(WriteBuffer &packet); void handleStartProcessPacket(ReadBuffer &packet); void handleSetSizePacket(ReadBuffer &packet); + void handleGetConsoleProcessListPacket(ReadBuffer &packet); void pollConinPipe(); protected: diff --git a/src/include/winpty.h b/src/include/winpty.h index 2a196a5d..fdfe4bca 100644 --- a/src/include/winpty.h +++ b/src/include/winpty.h @@ -218,6 +218,11 @@ WINPTY_API BOOL winpty_set_size(winpty_t *wp, int cols, int rows, winpty_error_ptr_t *err /*OPTIONAL*/); +/* Gets a list of processes attached to the console. */ +WINPTY_API int +winpty_get_console_process_list(winpty_t *wp, int *processList, const int processCount, + winpty_error_ptr_t *err /*OPTIONAL*/); + /* Frees the winpty_t object and the OS resources contained in it. This * call breaks the connection with the agent, which should then close its * console, terminating the processes attached to it. diff --git a/src/libwinpty/winpty.cc b/src/libwinpty/winpty.cc index e89e3481..3d977498 100644 --- a/src/libwinpty/winpty.cc +++ b/src/libwinpty/winpty.cc @@ -935,6 +935,33 @@ winpty_set_size(winpty_t *wp, int cols, int rows, } API_CATCH(FALSE) } +WINPTY_API int +winpty_get_console_process_list(winpty_t *wp, int *processList, const int processCount, + winpty_error_ptr_t *err /*OPTIONAL*/) { + API_TRY { + ASSERT(wp != nullptr); + ASSERT(processList != nullptr); + LockGuard lock(wp->mutex); + RpcOperation rpc(*wp); + auto packet = newPacket(); + packet.putInt32(AgentMsg::GetConsoleProcessList); + writePacket(*wp, packet); + auto reply = readPacket(*wp); + + auto actualProcessCount = reply.getInt32(); + + if (actualProcessCount <= processCount) { + for (auto i = 0; i < actualProcessCount; i++) { + processList[i] = reply.getInt32(); + } + } + + reply.assertEof(); + rpc.success(); + return actualProcessCount; + } API_CATCH(0) +} + WINPTY_API void winpty_free(winpty_t *wp) { // At least in principle, CloseHandle can fail, so this deletion can // fail. It won't throw an exception, but maybe there's an error that diff --git a/src/shared/AgentMsg.h b/src/shared/AgentMsg.h index 31859e09..ab60c6b9 100644 --- a/src/shared/AgentMsg.h +++ b/src/shared/AgentMsg.h @@ -26,6 +26,7 @@ struct AgentMsg enum Type { StartProcess, SetSize, + GetConsoleProcessList, }; };