gh-112278: Disable WMI queries on Windows after they time out (GH-112658)

This commit is contained in:
AN Long 2023-12-08 01:26:29 +08:00 committed by GitHub
parent b2923a61a1
commit a955fd68d6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 44 additions and 15 deletions

View File

@ -118,6 +118,10 @@ import re
import sys
import functools
import itertools
try:
import _wmi
except ImportError:
_wmi = None
### Globals & Constants
@ -312,21 +316,23 @@ def _syscmd_ver(system='', release='', version='',
version = _norm_version(version)
return system, release, version
try:
import _wmi
except ImportError:
def _wmi_query(*keys):
raise OSError("not supported")
else:
def _wmi_query(table, *keys):
global _wmi
if not _wmi:
raise OSError("not supported")
table = {
"OS": "Win32_OperatingSystem",
"CPU": "Win32_Processor",
}[table]
try:
data = _wmi.exec_query("SELECT {} FROM {}".format(
",".join(keys),
table,
)).split("\0")
except OSError:
_wmi = None
raise OSError("not supported")
split_data = (i.partition("=") for i in data)
dict_data = {i[0]: i[2] for i in split_data}
return (dict_data[k] for k in keys)

View File

@ -0,0 +1,2 @@
Reduce the time cost for some functions in :mod:`platform` on Windows if
current user has no permission to the WMI.

View File

@ -44,6 +44,7 @@ struct _query_data {
LPCWSTR query;
HANDLE writePipe;
HANDLE readPipe;
HANDLE connectEvent;
};
@ -86,6 +87,9 @@ _query_thread(LPVOID param)
NULL, NULL, 0, NULL, 0, 0, &services
);
}
if (!SetEvent(data->connectEvent)) {
hr = HRESULT_FROM_WIN32(GetLastError());
}
if (SUCCEEDED(hr)) {
hr = CoSetProxyBlanket(
services, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL,
@ -231,7 +235,8 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
Py_BEGIN_ALLOW_THREADS
if (!CreatePipe(&data.readPipe, &data.writePipe, NULL, 0)) {
data.connectEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
if (!data.connectEvent || !CreatePipe(&data.readPipe, &data.writePipe, NULL, 0)) {
err = GetLastError();
} else {
hThread = CreateThread(NULL, 0, _query_thread, (LPVOID*)&data, 0, NULL);
@ -243,6 +248,21 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
}
}
// gh-112278: If current user doesn't have permission to query the WMI, the
// 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
// a timeout.
switch (WaitForSingleObject(data.connectEvent, 100)) {
case WAIT_OBJECT_0:
break;
case WAIT_TIMEOUT:
err = WAIT_TIMEOUT;
break;
default:
err = GetLastError();
break;
}
while (!err) {
if (ReadFile(
data.readPipe,
@ -265,7 +285,7 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
}
// Allow the thread some time to clean up
switch (WaitForSingleObject(hThread, 1000)) {
switch (WaitForSingleObject(hThread, 100)) {
case WAIT_OBJECT_0:
// Thread ended cleanly
if (!GetExitCodeThread(hThread, (LPDWORD)&err)) {
@ -286,6 +306,7 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
}
CloseHandle(hThread);
CloseHandle(data.connectEvent);
hThread = NULL;
Py_END_ALLOW_THREADS