Skip to content

Commit

Permalink
pythongh-112278: Disable WMI queries on Windows after they time out (p…
Browse files Browse the repository at this point in the history
  • Loading branch information
aisk committed Feb 11, 2024
1 parent e24ff3d commit f5e64ad
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 15 deletions.
32 changes: 19 additions & 13 deletions Lib/platform.py
Expand Up @@ -118,6 +118,10 @@
import sys
import functools
import itertools
try:
import _wmi
except ImportError:
_wmi = None

### Globals & Constants

Expand Down Expand Up @@ -312,24 +316,26 @@ def _syscmd_ver(system='', release='', version='',
version = _norm_version(version)
return system, release, version

try:
import _wmi
except ImportError:
def _wmi_query(*keys):

def _wmi_query(table, *keys):
global _wmi
if not _wmi:
raise OSError("not supported")
else:
def _wmi_query(table, *keys):
table = {
"OS": "Win32_OperatingSystem",
"CPU": "Win32_Processor",
}[table]
table = {
"OS": "Win32_OperatingSystem",
"CPU": "Win32_Processor",
}[table]
try:
data = _wmi.exec_query("SELECT {} FROM {}".format(
",".join(keys),
table,
)).split("\0")
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)
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)


_WIN32_CLIENT_RELEASES = [
Expand Down
@@ -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.
25 changes: 23 additions & 2 deletions PC/_wmimodule.cpp
Expand Up @@ -44,6 +44,7 @@ struct _query_data {
LPCWSTR query;
HANDLE writePipe;
HANDLE readPipe;
HANDLE connectEvent;
};


Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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);
Expand All @@ -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,
Expand All @@ -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)) {
Expand All @@ -286,6 +306,7 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
}

CloseHandle(hThread);
CloseHandle(data.connectEvent);
hThread = NULL;

Py_END_ALLOW_THREADS
Expand Down

0 comments on commit f5e64ad

Please sign in to comment.