WTSEnumerateProcesses with socket and data struct

Question

I'm new with C++ Windows API and I've got few questions about using the WTSEnumerateProcesses() API.

My target: use WTSEnumerateProcesses(), take only the process list, put it in any kind of "list" or struct, convert the widechar to ANSI, and send it to a socket using Winsock.

That's my code below:

Winsock:

int connect(char buffer[])
{
    PCSTR ip = "localhost";
    WSADATA wsaData;
    SOCKET ConnectSocket = INVALID_SOCKET;
    struct addrinfo *result = NULL,
        *ptr = NULL,
        hints;
    const char *sendbuf =buffer;
    char recvbuf[DEFAULT_BUFLEN];
    int iResult;
    int recvbuflen = DEFAULT_BUFLEN;

    // Initialize Winsock
    iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != 0) {
        printf("WSAStartup failed with error: %d\n", iResult);
        return 1;
    }

    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    // Resolve the server address and port
    iResult = getaddrinfo(ip, DEFAULT_PORT, &hints, &result);
    if (iResult != 0) {
        printf("getaddrinfo failed with error: %d\n", iResult);
        WSACleanup();
        return 1;
    }

    // Attempt to connect to an address until one succeeds
    for (ptr = result; ptr != NULL; ptr = ptr->ai_next) {
        // Create a SOCKET for connecting to server
        ConnectSocket = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (ConnectSocket == INVALID_SOCKET) {
            printf("socket failed with error: %ld\n", WSAGetLastError());
            WSACleanup();
            return 1;
        }

        // Connect to server.
        iResult = connect(ConnectSocket, ptr->ai_addr, (int)ptr->ai_addrlen);
        if (iResult == SOCKET_ERROR) {
            closesocket(ConnectSocket);
            ConnectSocket = INVALID_SOCKET;
            continue;
        }
        break;
    }

    freeaddrinfo(result);

    if (ConnectSocket == INVALID_SOCKET) {
        printf("Unable to connect to server!\n");
        WSACleanup();
        return 1;
    }

    // Send an initial buffer
    iResult = send(ConnectSocket, sendbuf, (int)strlen(sendbuf), 0);
    if (iResult == SOCKET_ERROR) {
        printf("send failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }

    printf("Bytes Sent: %ld\n", iResult);

    // shutdown the connection since no more data will be sent
    iResult = shutdown(ConnectSocket, SD_SEND);
    if (iResult == SOCKET_ERROR) {
        printf("shutdown failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }

    // Receive until the peer closes the connection
    do {
        iResult = recv(ConnectSocket, recvbuf, recvbuflen, 0);
        if (iResult > 0)
            printf("Bytes received: %d\n", iResult);
        else if (iResult == 0)
            printf("Connection closed\n");
        else
            printf("recv failed with error: %d\n", WSAGetLastError());
    }
    while (iResult > 0);

    // cleanup
    closesocket(ConnectSocket);
    WSACleanup();

    return 0;
}

and the WTSEnumerateProcesses():

void GetProcesslist()
{
    vector<LPWSTR> v;         // no need to prepend std:: any more
    LPWSTR pProcessName;
    WTS_PROCESS_INFO* pWPIs = NULL;
    DWORD dwProcCount = 0;
    if (WTSEnumerateProcesses(WTS_CURRENT_SERVER_HANDLE, NULL, 1, &pWPIs, &dwProcCount))
    {
        //Go through all processes retrieved
        for (DWORD i = 0; i < dwProcCount; i++)
        {
            pProcessName = pWPIs[i].pProcessName;
            v[i] = pProcessName;
            //contains(pProcessName);
            // process file name only, no path!
            //data = pProcessName;
            std::wcout << pProcessName << endl;
            //pWPIs[i].ProcessId = process ID
            //pWPIs[i].SessionId = session ID, if you need to limit it to the logged in user processes
            //pWPIs[i].pUserSid = user SID that started the process
        }
    }

    //Free memory
    if (pWPIs)
    {
        WTSFreeMemory(pWPIs);
        pWPIs = NULL;
    }
}

Can you please help me to understand how I can take the process list only, push it to any struct (like a list?) and send it correctly over the TCP socket, that is waiting to get a const char[] and not wide characters?

Answer 1

You are not using std::vector correctly in your GetProcesslist(). You are trying to access elements you have not allocated any memory for. And worse, you are storing pointers to data that you free before GetProcesslist() exits, leaving the vector containing dangling pointers (were you to access that vector outside of GetProcesslist(), which you are not doing yet).

Try something more like this:

std::vector<std::wstring> GetProcesslist()
{
    std::vector<std::wstring> v;
    WTS_PROCESS_INFO* pWPIs = NULL;
    DWORD dwProcCount = 0;

    if (WTSEnumerateProcesses(WTS_CURRENT_SERVER_HANDLE, NULL, 1, &pWPIs, &dwProcCount))
    {
        //Go through all processes retrieved
        for (DWORD i = 0; i < dwProcCount; ++i)
        {
            LPWSTR pProcessName = pWPIs[i].pProcessName;
            v.push_back(pProcessName);
            // process file name only, no path!
            //data = pProcessName;
            std::wcout << pProcessName << endl;
            //pWPIs[i].ProcessId = process ID
            //pWPIs[i].SessionId = session ID, if you need to limit it to the logged in user processes
            //pWPIs[i].pUserSid = user SID that started the process
        }
        //Free memory
        WTSFreeMemory(pWPIs);
    }
    return v;
}

Then, I would suggest sending the data over the socket like this:

int connect(const char buffer[], size_t buffersize)
{
    PCSTR ip = "localhost";
    WSADATA wsaData;
    SOCKET ConnectSocket = INVALID_SOCKET;
    struct addrinfo *result = NULL,
        *ptr = NULL,
        hints;
    const char *sendbuf = buffer;
    char recvbuf[DEFAULT_BUFLEN];
    int iResult;

    // Initialize Winsock
    iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != 0) {
        printf("WSAStartup failed with error: %d\n", iResult);
        return 1;
    }

    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    // Resolve the server address and port
    iResult = getaddrinfo(ip, DEFAULT_PORT, &hints, &result);
    if (iResult != 0) {
        printf("getaddrinfo failed with error: %d\n", iResult);
        WSACleanup();
        return 1;
    }

    // Attempt to connect to an address until one succeeds
    for (ptr = result; ptr != NULL; ptr = ptr->ai_next) {
        // Create a SOCKET for connecting to server
        ConnectSocket = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (ConnectSocket == INVALID_SOCKET) {
            printf("socket failed with error: %ld\n", WSAGetLastError());
            WSACleanup();
            return 1;
        }

        // Connect to server.
        iResult = connect(ConnectSocket, ptr->ai_addr, (int)ptr->ai_addrlen);
        if (iResult == SOCKET_ERROR) {
            closesocket(ConnectSocket);
            ConnectSocket = INVALID_SOCKET;
            continue;
        }

        break;
    }

    freeaddrinfo(result);

    if (ConnectSocket == INVALID_SOCKET) {
        printf("Unable to connect to server!\n");
        WSACleanup();
        return 1;
    }

    // Send an initial buffer
    while (buffersize > 0) {
        iResult = send(ConnectSocket, sendbuf, (int)buffersize, 0);
        if (iResult == SOCKET_ERROR) {
            printf("send failed with error: %d\n", WSAGetLastError());
            closesocket(ConnectSocket);
            WSACleanup();
            return 1;
        }

        printf("Bytes Sent: %ld\n", iResult);

        sendbuf += iResult;
        buffersize -= iResult;
    }

    // shutdown the connection since no more data will be sent
    iResult = shutdown(ConnectSocket, SD_SEND);
    if (iResult == SOCKET_ERROR) {
        printf("shutdown failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }

    // Receive until the peer closes the connection
    do {
        iResult = recv(ConnectSocket, recvbuf, sizeof(recvbuf), 0);
        if (iResult > 0)
            printf("Bytes received: %d\n", iResult);
        else if (iResult == 0)
            printf("Connection closed\n");
        else
            printf("recv failed with error: %d\n", WSAGetLastError());
    }
    while (iResult > 0);

    // cleanup
    closesocket(ConnectSocket);
    WSACleanup();

    return 0;
}

std::vector<std::wstring> procs = GetProcesslist();
std::vector<char> buffer;

size_t size = sizeof(uint32_t);
for(const std::wstring &name : procs) {
    size += (sizeof(uint32_t) + WideCharToMultiByte(CP_UTF8, 0, name.c_str(), name.size(), NULL, 0, NULL, NULL));
}

buffer.resize(size);
char *ptr = buffer.data();

*reinterpret_cast<uint32_t*>(ptr) = htonl(procs.size());
ptr += sizeof(uint32_t);
for(const std::wstring &name : procs) {
    int len = WideCharToMultiByte(CP_UTF8, 0, name.c_str(), name.size(), NULL, 0, NULL, NULL));
    *reinterpret_cast<uint32_t*>(ptr) = htonl(len);
    ptr += sizeof(uint32_t);
    if (len > 0) {
        ptr += WideCharToMultiByte(CP_UTF8, 0, name.c_str(), name.size(), ptr, len, NULL, NULL));
    }
}

connect(buffer.data(), buffer.size());

Or, if you absolutely cannot change the signature of connect() and must pass it a null-terminated C string, then try this:

int connect(const char buffer[])
{
    PCSTR ip = "localhost";
    WSADATA wsaData;
    SOCKET ConnectSocket = INVALID_SOCKET;
    struct addrinfo *result = NULL,
        *ptr = NULL,
        hints;
    const char *sendbuf = buffer;
    sendbuflen = strlen(buffer);
    char recvbuf[DEFAULT_BUFLEN];
    int iResult;

    // Initialize Winsock
    iResult = WSAStartup(MAKEWORD(2, 2), &wsaData);
    if (iResult != 0) {
        printf("WSAStartup failed with error: %d\n", iResult);
        return 1;
    }

    ZeroMemory(&hints, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    // Resolve the server address and port
    iResult = getaddrinfo(ip, DEFAULT_PORT, &hints, &result);
    if (iResult != 0) {
        printf("getaddrinfo failed with error: %d\n", iResult);
        WSACleanup();
        return 1;
    }

    // Attempt to connect to an address until one succeeds
    for (ptr = result; ptr != NULL; ptr = ptr->ai_next) {
        // Create a SOCKET for connecting to server
        ConnectSocket = socket(ptr->ai_family, ptr->ai_socktype, ptr->ai_protocol);
        if (ConnectSocket == INVALID_SOCKET) {
            printf("socket failed with error: %ld\n", WSAGetLastError());
            WSACleanup();
            return 1;
        }

        // Connect to server.
        iResult = connect(ConnectSocket, ptr->ai_addr, (int)ptr->ai_addrlen);
        if (iResult == SOCKET_ERROR) {
            closesocket(ConnectSocket);
            ConnectSocket = INVALID_SOCKET;
            continue;
        }

        break;
    }

    freeaddrinfo(result);

    if (ConnectSocket == INVALID_SOCKET) {
        printf("Unable to connect to server!\n");
        WSACleanup();
        return 1;
    }

    // Send an initial buffer
    while (sendbuflen > 0) {
        iResult = send(ConnectSocket, sendbuf, sendbuflen, 0);
        if (iResult == SOCKET_ERROR) {
            printf("send failed with error: %d\n", WSAGetLastError());
            closesocket(ConnectSocket);
            WSACleanup();
            return 1;
        }

        printf("Bytes Sent: %ld\n", iResult);

        sendbuf += iResult;
        sendbuflen -= iResult;
    }

    // shutdown the connection since no more data will be sent
    iResult = shutdown(ConnectSocket, SD_SEND);
    if (iResult == SOCKET_ERROR) {
        printf("shutdown failed with error: %d\n", WSAGetLastError());
        closesocket(ConnectSocket);
        WSACleanup();
        return 1;
    }

    // Receive until the peer closes the connection
    do {
        iResult = recv(ConnectSocket, recvbuf, sizeof(recvbuf), 0);
        if (iResult > 0)
            printf("Bytes received: %d\n", iResult);
        else if (iResult == 0)
            printf("Connection closed\n");
        else
            printf("recv failed with error: %d\n", WSAGetLastError());
    }
    while (iResult > 0);

    // cleanup
    closesocket(ConnectSocket);
    WSACleanup();

    return 0;
}

std::vector<std::wstring> procs = GetProcesslist();
std::ostringstream oss;

oss << procs.size() << "\x1";
for(const std::wstring &name : procs) {
    int len = WideCharToMultiByte(CP_UTF8, 0, name.c_str(), name.size(), NULL, 0, NULL, NULL);
    if (len > 0) {
        std::string s;
        s.resize(len);
        WideCharToMultiByte(CP_UTF8, 0, name.c_str(), name.size(), &s[0], len, NULL, NULL);
        oss << s;
    }
    oss << "\x1";
}

connect(oss.str().c_str());