From 237bdc6a97c0dd455da99c16e6b3b7cac4c79e86 Mon Sep 17 00:00:00 2001 From: Sean Hall Date: Fri, 2 Sep 2022 16:12:42 -0500 Subject: Add WixCanRestart and make wixstdba elevate for restart when required. Fixes 5499 --- src/libs/dutil/WixToolset.DUtil/inc/procutil.h | 16 +++- src/libs/dutil/WixToolset.DUtil/procutil.cpp | 106 +++++++++++++++++---- src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp | 28 +++++- 3 files changed, 129 insertions(+), 21 deletions(-) (limited to 'src/libs') diff --git a/src/libs/dutil/WixToolset.DUtil/inc/procutil.h b/src/libs/dutil/WixToolset.DUtil/inc/procutil.h index d61d91b5..6a641a5b 100644 --- a/src/libs/dutil/WixToolset.DUtil/inc/procutil.h +++ b/src/libs/dutil/WixToolset.DUtil/inc/procutil.h @@ -23,9 +23,21 @@ HRESULT DAPI ProcSystem( __out BOOL* pfSystem ); -HRESULT DAPI ProcTokenUser( +HRESULT DAPI ProcGetTokenInformation( __in HANDLE hProcess, - __out TOKEN_USER** ppTokenUser + __in TOKEN_INFORMATION_CLASS tokenInformationClass, + __out LPVOID* ppvTokenInformation + ); + +HRESULT DAPI ProcHasPrivilege( + __in HANDLE hProcess, + __in LPCWSTR wzPrivilegeName, + __out BOOL* pfHasPrivilege + ); + +HRESULT DAPI ProcEnablePrivilege( + __in HANDLE hProcess, + __in LPCWSTR wzPrivilegeName ); HRESULT DAPI ProcWow64( diff --git a/src/libs/dutil/WixToolset.DUtil/procutil.cpp b/src/libs/dutil/WixToolset.DUtil/procutil.cpp index 29f575ae..376aec6d 100644 --- a/src/libs/dutil/WixToolset.DUtil/procutil.cpp +++ b/src/libs/dutil/WixToolset.DUtil/procutil.cpp @@ -85,7 +85,7 @@ extern "C" HRESULT DAPI ProcSystem( HRESULT hr = S_OK; TOKEN_USER* pTokenUser = NULL; - hr = ProcTokenUser(hProcess, &pTokenUser); + hr = ProcGetTokenInformation(hProcess, TokenUser, reinterpret_cast(&pTokenUser)); ProcExitOnFailure(hr, "Failed to get TokenUser from process token."); *pfSystem = ::IsWellKnownSid(pTokenUser->User.Sid, WinLocalSystemSid); @@ -96,15 +96,16 @@ LExit: return hr; } -extern "C" HRESULT DAPI ProcTokenUser( +extern "C" HRESULT DAPI ProcGetTokenInformation( __in HANDLE hProcess, - __out TOKEN_USER** ppTokenUser + __in TOKEN_INFORMATION_CLASS tokenInformationClass, + __out LPVOID* ppvTokenInformation ) { HRESULT hr = S_OK; DWORD er = ERROR_SUCCESS; HANDLE hToken = NULL; - TOKEN_USER* pTokenUser = NULL; + LPVOID pvTokenInformation = NULL; DWORD cbToken = 0; if (!::OpenProcessToken(hProcess, TOKEN_QUERY, &hToken)) @@ -112,33 +113,104 @@ extern "C" HRESULT DAPI ProcTokenUser( ProcExitWithLastError(hr, "Failed to open process token."); } - if (::GetTokenInformation(hToken, TokenUser, pTokenUser, 0, &cbToken)) - { - er = ERROR_SUCCESS; - } - else + if (!::GetTokenInformation(hToken, tokenInformationClass, pvTokenInformation, 0, &cbToken)) { er = ::GetLastError(); } if (er != ERROR_INSUFFICIENT_BUFFER) { - ProcExitOnWin32Error(er, hr, "Failed to get user from process token size."); + ProcExitOnWin32Error(er, hr, "Failed to get information from process token size."); } - pTokenUser = reinterpret_cast(MemAlloc(cbToken, TRUE)); - ProcExitOnNull(pTokenUser, hr, E_OUTOFMEMORY, "Failed to allocate token information."); + pvTokenInformation = MemAlloc(cbToken, TRUE); + ProcExitOnNull(pvTokenInformation, hr, E_OUTOFMEMORY, "Failed to allocate token information."); - if (!::GetTokenInformation(hToken, TokenUser, pTokenUser, cbToken, &cbToken)) + if (!::GetTokenInformation(hToken, tokenInformationClass, pvTokenInformation, cbToken, &cbToken)) { - ProcExitWithLastError(hr, "Failed to get user from process token."); + ProcExitWithLastError(hr, "Failed to get information from process token."); } - *ppTokenUser = pTokenUser; - pTokenUser = NULL; + *ppvTokenInformation = pvTokenInformation; + pvTokenInformation = NULL; + +LExit: + ReleaseMem(pvTokenInformation); + ReleaseHandle(hToken); + + return hr; +} + +extern "C" HRESULT DAPI ProcHasPrivilege( + __in HANDLE hProcess, + __in LPCWSTR wzPrivilegeName, + __out BOOL* pfHasPrivilege + ) +{ + HRESULT hr = S_OK; + TOKEN_PRIVILEGES* pTokenPrivileges = NULL; + LUID luidPrivilege = { }; + + *pfHasPrivilege = FALSE; + + if (!::LookupPrivilegeValueW(NULL, wzPrivilegeName, &luidPrivilege)) + { + ProcExitWithLastError(hr, "Failed to get privilege LUID: %ls", wzPrivilegeName); + } + + hr = ProcGetTokenInformation(hProcess, TokenPrivileges, reinterpret_cast(&pTokenPrivileges)); + ProcExitOnFailure(hr, "Failed to get token privilege information."); + + for (DWORD i = 0; i < pTokenPrivileges->PrivilegeCount; ++i) + { + LUID* pTokenLuid = &pTokenPrivileges->Privileges[i].Luid; + + if (luidPrivilege.LowPart == pTokenLuid->LowPart && luidPrivilege.HighPart == pTokenLuid->HighPart) + { + *pfHasPrivilege = TRUE; + break; + } + } + +LExit: + ReleaseMem(pTokenPrivileges); + + return hr; +} + +extern "C" HRESULT DAPI ProcEnablePrivilege( + __in HANDLE hProcess, + __in LPCWSTR wzPrivilegeName + ) +{ + HRESULT hr = S_OK; + HANDLE hToken = NULL; + TOKEN_PRIVILEGES priv = { }; + + priv.PrivilegeCount = 1; + priv.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED; + + if (!::LookupPrivilegeValueW(NULL, wzPrivilegeName, &priv.Privileges[0].Luid)) + { + ProcExitWithLastError(hr, "Failed to get privilege LUID: %ls", wzPrivilegeName); + } + + if (!::OpenProcessToken(hProcess, TOKEN_ADJUST_PRIVILEGES, &hToken)) + { + ProcExitWithLastError(hr, "Failed to get process token to adjust privileges."); + } + + if (!::AdjustTokenPrivileges(hToken, FALSE, &priv, sizeof(TOKEN_PRIVILEGES), NULL, 0)) + { + ProcExitWithLastError(hr, "Failed to adjust token to add privilege: %ls", wzPrivilegeName); + } + + if (ERROR_NOT_ALL_ASSIGNED == ::GetLastError()) + { + hr = S_FALSE; + } LExit: - ReleaseMem(pTokenUser); ReleaseHandle(hToken); return hr; diff --git a/src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp b/src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp index 297d90f4..8672ed38 100644 --- a/src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp +++ b/src/libs/dutil/test/DUtilUnitTest/ProcUtilTest.cpp @@ -6,6 +6,7 @@ using namespace System; using namespace System::Security::Principal; using namespace Xunit; using namespace WixBuildTools::TestSupport; +using namespace WixBuildTools::TestSupport::XunitExtensions; namespace DutilTests { @@ -13,7 +14,7 @@ namespace DutilTests { public: [Fact] - void ProcTokenUserTest() + void ProcGetTokenInformationTest() { HRESULT hr = S_OK; TOKEN_USER* pTokenUser = NULL; @@ -21,7 +22,7 @@ namespace DutilTests try { - hr = ProcTokenUser(::GetCurrentProcess(), &pTokenUser); + hr = ProcGetTokenInformation(::GetCurrentProcess(), TokenUser, reinterpret_cast(&pTokenUser)); NativeAssert::Succeeded(hr, "Failed to get TokenUser for current process."); if (!::ConvertSidToStringSidW(pTokenUser->User.Sid, &sczSid)) @@ -38,5 +39,28 @@ namespace DutilTests ReleaseStr(sczSid); } } + + [SkippableFact] + void ProcHasPrivilegeTest() + { + HRESULT hr = S_OK; + BOOL fHasPrivilege = FALSE; + + hr = ProcHasPrivilege(::GetCurrentProcess(), SE_CREATE_TOKEN_NAME, &fHasPrivilege); + NativeAssert::Succeeded(hr, "Failed to check privilege for current process."); + + if (fHasPrivilege) + { + WixAssert::Skip("Didn't expect process to have SE_CREATE_TOKEN_NAME privilege"); + } + + hr = ProcHasPrivilege(::GetCurrentProcess(), SE_INC_WORKING_SET_NAME, &fHasPrivilege); + NativeAssert::Succeeded(hr, "Failed to check privilege for current process."); + + if (!fHasPrivilege) + { + WixAssert::Skip("Expected process to have SE_INC_WORKING_SET_NAME privilege"); + } + } }; } -- cgit v1.2.3-55-g6feb