Rooba
/
cpython
mirror of https://github.com/python/cpython.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

gh-112278: In _wmi, treat initialization timeout separately from connection timeout (GH-112878)

This commit is contained in:
AN Long 2023-12-09 00:52:22 +08:00 committed by GitHub
parent 4d1eea59bd
commit 5a0137ca34
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 34 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

46
PC/_wmimodule.cpp
View File

@ -44,6 +44,7 @@ struct _query_data {
LPCWSTR query; LPCWSTR query;
HANDLE writePipe; HANDLE writePipe;
HANDLE readPipe; HANDLE readPipe;
HANDLE initEvent;
HANDLE connectEvent; HANDLE connectEvent;
}; };
@ -81,13 +82,16 @@ _query_thread(LPVOID param)
IID_IWbemLocator, (LPVOID *)&locator IID_IWbemLocator, (LPVOID *)&locator
); );
} }
if (SUCCEEDED(hr) && !SetEvent(data->initEvent)) {
hr = HRESULT_FROM_WIN32(GetLastError());
}
if (SUCCEEDED(hr)) { if (SUCCEEDED(hr)) {
hr = locator->ConnectServer( hr = locator->ConnectServer(
bstr_t(L"ROOT\\CIMV2"), bstr_t(L"ROOT\\CIMV2"),
NULL, NULL, 0, NULL, 0, 0, &services NULL, NULL, 0, NULL, 0, 0, &services
); );
} }
if (!SetEvent(data->connectEvent)) { if (SUCCEEDED(hr) && !SetEvent(data->connectEvent)) {
hr = HRESULT_FROM_WIN32(GetLastError()); hr = HRESULT_FROM_WIN32(GetLastError());
} }
if (SUCCEEDED(hr)) { if (SUCCEEDED(hr)) {
@ -193,6 +197,24 @@ _query_thread(LPVOID param)
} }
static DWORD
wait_event(HANDLE event, DWORD timeout)
{
DWORD err = 0;
switch (WaitForSingleObject(event, timeout)) {
case WAIT_OBJECT_0:
break;
case WAIT_TIMEOUT:
err = WAIT_TIMEOUT;
break;
default:
err = GetLastError();
break;
}
return err;
}
/*[clinic input] /*[clinic input]
_wmi.exec_query _wmi.exec_query
@ -235,8 +257,11 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
Py_BEGIN_ALLOW_THREADS Py_BEGIN_ALLOW_THREADS
data.initEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
data.connectEvent = CreateEvent(NULL, TRUE, FALSE, NULL); data.connectEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
if (!data.connectEvent || !CreatePipe(&data.readPipe, &data.writePipe, NULL, 0)) { if (!data.initEvent || !data.connectEvent ||
!CreatePipe(&data.readPipe, &data.writePipe, NULL, 0))
{
err = GetLastError(); err = GetLastError();
} else { } else {
hThread = CreateThread(NULL, 0, _query_thread, (LPVOID*)&data, 0, NULL); hThread = CreateThread(NULL, 0, _query_thread, (LPVOID*)&data, 0, NULL);
@ -251,16 +276,12 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
// gh-112278: If current user doesn't have permission to query the WMI, the // gh-112278: If current user doesn't have permission to query the WMI, the
// function IWbemLocator::ConnectServer will hang for 5 seconds, and there // function IWbemLocator::ConnectServer will hang for 5 seconds, and there
// is no way to specify the timeout. So we use an Event object to simulate // is no way to specify the timeout. So we use an Event object to simulate
// a timeout. // a timeout. The initEvent will be set after COM initialization, it will
switch (WaitForSingleObject(data.connectEvent, 100)) { // take a longer time when first initialized. The connectEvent will be set
case WAIT_OBJECT_0: // after connected to WMI.
break; err = wait_event(data.initEvent, 1000);
case WAIT_TIMEOUT: if (!err) {
err = WAIT_TIMEOUT; err = wait_event(data.connectEvent, 100);
break;
default:
err = GetLastError();
break;
} }
while (!err) { while (!err) {
@ -306,6 +327,7 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
} }
CloseHandle(hThread); CloseHandle(hThread);
CloseHandle(data.initEvent);
CloseHandle(data.connectEvent); CloseHandle(data.connectEvent);
hThread = NULL; hThread = NULL;