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

aisk · web-flow · commit a955fd68d645 · 2023-12-07T17:26:29.000Z
diff --git a/Lib/platform.py b/Lib/platform.py
@@ -118,6 +118,10 @@
 import sys
 import functools
 import itertools
+try:
+    import _wmi
+except ImportError:
+    _wmi = None
 
 ### Globals & Constants
 
@@ -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 = [
diff --git a/Misc/NEWS.d/next/Windows/2023-12-03-19-22-37.gh-issue-112278.FiloCE.rst b/Misc/NEWS.d/next/Windows/2023-12-03-19-22-37.gh-issue-112278.FiloCE.rst
@@ -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.
diff --git a/PC/_wmimodule.cpp b/PC/_wmimodule.cpp
@@ -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

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