From 7d9d6b53bcee58252e38732a1fa1ca786eba4aab Mon Sep 17 00:00:00 2001
From: Steve Dower <steve.dower@python.org>
Date: Mon, 15 Apr 2024 15:43:11 +0100
Subject: [PATCH] gh-112278: Improve error handling in wmi module and tests
 (GH-117818)

---
 Lib/test/test_wmi.py |  6 ++++--
 PC/_wmimodule.cpp    | 45 +++++++++++++++++++++++++-------------------
 2 files changed, 30 insertions(+), 21 deletions(-)

diff --git a/Lib/test/test_wmi.py b/Lib/test/test_wmi.py
index bf8c52e646d..f667926d1f8 100644
--- a/Lib/test/test_wmi.py
+++ b/Lib/test/test_wmi.py
@@ -14,11 +14,13 @@ def wmi_exec_query(query):
     # gh-112278: WMI maybe slow response when first call.
     try:
         return _wmi.exec_query(query)
+    except BrokenPipeError:
+        pass
     except WindowsError as e:
         if e.winerror != 258:
             raise
-        time.sleep(LOOPBACK_TIMEOUT)
-        return _wmi.exec_query(query)
+    time.sleep(LOOPBACK_TIMEOUT)
+    return _wmi.exec_query(query)
 
 
 class WmiTests(unittest.TestCase):
diff --git a/PC/_wmimodule.cpp b/PC/_wmimodule.cpp
index 5ab6dcb0325..22ed05276e6 100644
--- a/PC/_wmimodule.cpp
+++ b/PC/_wmimodule.cpp
@@ -279,9 +279,11 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
     // a timeout.  The initEvent will be set after COM initialization, it will
     // take a longer time when first initialized.  The connectEvent will be set
     // after connected to WMI.
-    err = wait_event(data.initEvent, 1000);
     if (!err) {
-        err = wait_event(data.connectEvent, 100);
+        err = wait_event(data.initEvent, 1000);
+        if (!err) {
+            err = wait_event(data.connectEvent, 100);
+        }
     }
 
     while (!err) {
@@ -305,28 +307,33 @@ _wmi_exec_query_impl(PyObject *module, PyObject *query)
         CloseHandle(data.readPipe);
     }
 
-    // Allow the thread some time to clean up
-    switch (WaitForSingleObject(hThread, 100)) {
-    case WAIT_OBJECT_0:
-        // Thread ended cleanly
-        if (!GetExitCodeThread(hThread, (LPDWORD)&err)) {
-            err = GetLastError();
+    if (hThread) {
+        // Allow the thread some time to clean up
+        int thread_err;
+        switch (WaitForSingleObject(hThread, 100)) {
+        case WAIT_OBJECT_0:
+            // Thread ended cleanly
+            if (!GetExitCodeThread(hThread, (LPDWORD)&thread_err)) {
+                thread_err = GetLastError();
+            }
+            break;
+        case WAIT_TIMEOUT:
+            // Probably stuck - there's not much we can do, unfortunately
+            thread_err = WAIT_TIMEOUT;
+            break;
+        default:
+            thread_err = GetLastError();
+            break;
         }
-        break;
-    case WAIT_TIMEOUT:
-        // Probably stuck - there's not much we can do, unfortunately
+        // An error on our side is more likely to be relevant than one from
+        // the thread, but if we don't have one on our side we'll take theirs.
         if (err == 0 || err == ERROR_BROKEN_PIPE) {
-            err = WAIT_TIMEOUT;
+            err = thread_err;
         }
-        break;
-    default:
-        if (err == 0 || err == ERROR_BROKEN_PIPE) {
-            err = GetLastError();
-        }
-        break;
+
+        CloseHandle(hThread);
     }
 
-    CloseHandle(hThread);
     CloseHandle(data.initEvent);
     CloseHandle(data.connectEvent);
     hThread = NULL;