From a955fd68d6451bd42199110c978e99b3d2959db2 Mon Sep 17 00:00:00 2001
From: AN Long <aisk@users.noreply.github.com>
Date: Fri, 8 Dec 2023 01:26:29 +0800
Subject: [PATCH] gh-112278: Disable WMI queries on Windows after they time out
 (GH-112658)

---
 Lib/platform.py                               | 32 +++++++++++--------
 ...-12-03-19-22-37.gh-issue-112278.FiloCE.rst |  2 ++
 PC/_wmimodule.cpp                             | 25 +++++++++++++--
 3 files changed, 44 insertions(+), 15 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Windows/2023-12-03-19-22-37.gh-issue-112278.FiloCE.rst

diff --git a/Lib/platform.py b/Lib/platform.py
index 7bb222088d5061..75aa55510858fd 100755
--- 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
new file mode 100644
index 00000000000000..0350d105d97375
--- /dev/null
+++ 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
index fdf09ec6ec6f63..215350acfb0d8e 100644
--- 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