Just Let It Flow

#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <wtsapi32.h>
#include <cstdlib>
#include <cassert>
#include <vector>
#include <malloc.h> // for _alloca
 
#pragma comment(lib, "wtsapi32.lib")
#pragma comment(lib, "advapi32.lib")
#pragma comment(lib, "kernel32.lib")
 
// RAII wrapper around Windows handles and misc other data
template<class Data = HANDLE, class ReleaseRet = BOOL>
struct WinData
{
    typedef ReleaseRet (WINAPI*Releaser)(Data);
private:
    Data data;
    Releaser release;
 
public:
    WinData(Data d, Releaser r) : data(d), release(r) {}
    WinData() : data(NULL), release(NULL) {}
 
    ~WinData()
    {
        if(data)
        {
            release(data);
        }
    }
 
    void Init(Data d, Releaser r)
    {
        assert((d && r) && !(data || release));
        data = d;
        release = r;
    }
 
    Data operator*() const
    {
        return data;
    }
 
    Data* operator&()
    {
        // only allow this if we have a releaser and no valid data
        assert(!data && release);
        return &data;
    }
 
    Data Release()
    {
        Data temp = data;
        data = NULL;
        release = NULL;
        return temp;
    }
};
 
#define FAIL_NULL_RETURN(x) \
    do { \
        if(!(x)) \
        { \
            return NULL; \
        } \
    } while(0)
 
// for easy use with the WinData wrapper above
PVOID WINAPI Alloc(SIZE_T size)
{
    return malloc(size);
}
 
void WINAPI Free(PVOID p)
{
    free(p);
}
 
HANDLE GetProcessToken(DWORD procId, DWORD perms)
{
    WinData<> hProc(OpenProcess(PROCESS_QUERY_INFORMATION, FALSE, procId), &CloseHandle);
    FAIL_NULL_RETURN(*hProc);
    HANDLE hProcToken = NULL;
    FAIL_NULL_RETURN(OpenProcessToken(*hProc, perms, &hProcToken));
    return hProcToken;
}
 
PVOID GetTokenInformation(HANDLE hToken, TOKEN_INFORMATION_CLASS info, DWORD* pSize)
{
    PVOID pData = NULL;
    DWORD required = 0;
    GetTokenInformation(hToken, info, NULL, 0, &required);
    if(required)
    {
        pData = Alloc(required);
        if(!pData) return NULL;
        if(!GetTokenInformation(hToken, info, pData, required, &required))
        {
            Free(pData);
            pData = NULL;
        }
        else if(pSize)
        {
            *pSize = required;
        }
    }
    return pData;
}
 
BOOL GetProcessLogonSid(DWORD procID, PSID* ppLogonSid)
{
    *ppLogonSid = NULL;
    WinData<> hProcToken(GetProcessToken(procID, TOKEN_QUERY), &CloseHandle);
    if(!*hProcToken)
    {
        return FALSE;
    }
    // doesn't just return a sid as the docs would have you believe
    PTOKEN_GROUPS pLogonGroups = static_cast<PTOKEN_GROUPS>(GetTokenInformation(*hProcToken, TokenLogonSid, NULL));
    if(!pLogonGroups)
    {
        return FALSE;
    }
    // move the actual sid to the start of the memory block
    // and return that instead of doing a second allocation
    PSID pLogonSid = pLogonGroups->Groups[0].Sid;
    MoveMemory(pLogonGroups, pLogonSid, GetLengthSid(pLogonSid));
    *ppLogonSid = pLogonGroups;
    return TRUE;
}
 
HANDLE FindTokenOfProcessWithEqualSids(PSID pUserSid, PSID pLogonSid)
{
    PWTS_PROCESS_INFO pProcesses = NULL;
    DWORD numProcesses = 0;
    // get data on all running processes
    WTSEnumerateProcesses(WTS_CURRENT_SERVER, 0, 1, &pProcesses, &numProcesses);
    HANDLE hProcToken = NULL;
    for(DWORD i = 0; i < numProcesses; ++i)
    {
        // if we have access to the process' user sid
        PSID pProcessUserSid = pProcesses[i].pUserSid;
        if(pProcessUserSid)
        {
            // and can get its logon sid
            WinData<PVOID, void> pProcessLogonSid(NULL, &Free);
            if(GetProcessLogonSid(pProcesses[i].ProcessId, &pProcessLogonSid))
            {
                // and they're equal to what we're looking for
                if(EqualSid(pUserSid, pProcessUserSid) && EqualSid(pLogonSid, *pProcessLogonSid))
                {
                    // grab it...
                    hProcToken = GetProcessToken(pProcesses[i].ProcessId, TOKEN_DUPLICATE | TOKEN_QUERY), &CloseHandle);
                    // ... if it's valid that is
                    if(hProcToken)
                    {
                        break;
                    }
                }
            }
        }
    }
    WTSFreeMemory(pProcesses);
    return hProcToken;
}
 
PSID GetAccountSid(LPCWSTR domain, LPCWSTR userName)
{
    PSID pSid = NULL;
    DWORD sidSize = 0, domSize = 0;
    SID_NAME_USE snu;
    // get the required buffer sizes first
    LookupAccountName(domain, userName, NULL, &sidSize, NULL, &domSize, &snu);
    if(sidSize)
    {
        pSid = Alloc(sidSize);
        if(pSid)
        {
            WCHAR* unneededInfo = static_cast<WCHAR*>(_alloca(domSize * sizeof(WCHAR)));
            if(!LookupAccountName(domain, userName, pSid, &sidSize, unneededInfo, &domSize, &snu))
            {
                Free(pSid);
                pSid = NULL;
            }
        }
    }
    return pSid;
}
 
HANDLE GetLoggedOnUserToken(DWORD perms, TOKEN_TYPE typeOfToken)
{
    HWINSTA hWinsta = GetProcessWindowStation();
    DWORD required = 0;
    GetUserObjectInformation(hWinsta, UOI_USER_SID, NULL, 0, &required);
    FAIL_NULL_RETURN(required);
    // this is a logon sid, contrary to the flag name
    std::vector<BYTE> logonSid(required);
    FAIL_NULL_RETURN(GetUserObjectInformation(hWinsta, UOI_USER_SID, &logonSid[0], required, &required));
    // get session user's details
    PWTSINFO pClientInfo = NULL;
    ULONG clientInfoSize = 0;
    FAIL_NULL_RETURN(
        WTSQuerySessionInformation(
            WTS_CURRENT_SERVER_HANDLE,
            WTS_CURRENT_SESSION,
            WTSSessionInfo, 
            reinterpret_cast<PWSTR*>(&pClientInfo),
            &clientInfoSize
        )
    );
    // grab their sid
    WinData<PVOID, void> pUserSid(GetAccountSid(pClientInfo->Domain, pClientInfo->UserName), &Free);
    WTSFreeMemory(pClientInfo);
    FAIL_NULL_RETURN(*pUserSid);
    // find a process with the same user and logon sid, and retreive its token
    WinData<> hToken(FindTokenOfProcessWithEqualSids(*pUserSid, &logonSid[0]), &CloseHandle);
    HANDLE hReturnedToken = NULL;
    // turn the token into what the caller wants
    DuplicateTokenEx(*hToken, perms, NULL, SecurityImpersonation, typeOfToken, &hReturnedToken);
    return hReturnedToken;
}