Merge pull request #130 from the-ress/get-console-process-list

Add winpty_get_console_process_list

src/agent/Agent.cc

@@ -27,6 +27,7 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include <algorithm>
 #include <string>
 #include <utility>
 #include <vector>
@@ -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<DWORD>(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<DWORD>(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();

src/agent/Agent.h

@@ -59,6 +59,7 @@ private:
     void writePacket(WriteBuffer &packet);
     void handleStartProcessPacket(ReadBuffer &packet);
     void handleSetSizePacket(ReadBuffer &packet);
+    void handleGetConsoleProcessListPacket(ReadBuffer &packet);
     void pollConinPipe();
 
 protected:

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.

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<Mutex> 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

src/shared/AgentMsg.h

@@ -26,6 +26,7 @@ struct AgentMsg
     enum Type {
         StartProcess,
         SetSize,
+        GetConsoleProcessList,
     };
 };