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

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; LPCWSTR query;
HANDLE writePipe; HANDLE writePipe;
HANDLE readPipe; HANDLE readPipe;
HANDLE connectEvent;
}; };
@ -86,6 +87,9 @@ _query_thread(LPVOID param)
NULL, NULL, 0, NULL, 0, 0, &services NULL, NULL, 0, NULL, 0, 0, &services
); );
} }
if (!SetEvent(data->connectEvent)) {
hr = HRESULT_FROM_WIN32(GetLastError());
}
if (SUCCEEDED(hr)) { if (SUCCEEDED(hr)) {
hr = CoSetProxyBlanket( hr = CoSetProxyBlanket(
services, RPC_C_AUTHN_WINNT, RPC_C_AUTHZ_NONE, NULL, 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 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(); err = GetLastError();
} else { } else {
hThread = CreateThread(NULL, 0, _query_thread, (LPVOID*)&data, 0, NULL); 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) { while (!err) {
if (ReadFile( if (ReadFile(
data.readPipe, data.readPipe,
@ -265,7 +285,7 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
} }
// Allow the thread some time to clean up // Allow the thread some time to clean up
switch (WaitForSingleObject(hThread, 1000)) { switch (WaitForSingleObject(hThread, 100)) {
case WAIT_OBJECT_0: case WAIT_OBJECT_0:
// Thread ended cleanly // Thread ended cleanly
if (!GetExitCodeThread(hThread, (LPDWORD)&err)) { if (!GetExitCodeThread(hThread, (LPDWORD)&err)) {
@ -286,6 +306,7 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
} }
CloseHandle(hThread); CloseHandle(hThread);
CloseHandle(data.connectEvent);
hThread = NULL; hThread = NULL;
Py_END_ALLOW_THREADS Py_END_ALLOW_THREADS