Skip to content

Commit

Permalink
Merge pull request #130 from the-ress/get-console-process-list
Browse files Browse the repository at this point in the history
Add winpty_get_console_process_list
  • Loading branch information
rprichard authored Oct 8, 2017
2 parents ce9239a + f678472 commit 850661d
Show file tree
Hide file tree
Showing 5 changed files with 65 additions and 0 deletions.
31 changes: 31 additions & 0 deletions src/agent/Agent.cc
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
#include <stdlib.h>
#include <string.h>

#include <algorithm>
#include <string>
#include <utility>
#include <vector>
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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();
Expand Down
1 change: 1 addition & 0 deletions src/agent/Agent.h
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
5 changes: 5 additions & 0 deletions src/include/winpty.h
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
27 changes: 27 additions & 0 deletions src/libwinpty/winpty.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/shared/AgentMsg.h
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ struct AgentMsg
enum Type {
StartProcess,
SetSize,
GetConsoleProcessList,
};
};

Expand Down

0 comments on commit 850661d

Please sign in to comment.