From 033f79d038423b59f1392abd7b01dda082629e5b Mon Sep 17 00:00:00 2001
From: Lisanna Dettwyler <lisanna.dettwyler@intel.com>
Date: Mon, 23 Sep 2024 12:56:52 -0400
Subject: [PATCH] windows: do not exclude DN_NEED_RESTART drivers if
 LevelZeroStagedDriverPath is set

Signed-off-by: Lisanna Dettwyler <lisanna.dettwyler@intel.com>
---
 .../loader/windows/driver_discovery_win.cpp   | 33 ++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/source/loader/windows/driver_discovery_win.cpp b/source/loader/windows/driver_discovery_win.cpp
index 4f619ca3..6f264f78 100644
--- a/source/loader/windows/driver_discovery_win.cpp
+++ b/source/loader/windows/driver_discovery_win.cpp
@@ -61,7 +61,9 @@ bool isDeviceAvailable(DEVINST devnode) {
 
     bool isInInvalidState = (devStatus & DN_HAS_PROBLEM)
                             && (devProblem == CM_PROB_NEED_RESTART);
-    isInInvalidState |= (DN_NEED_RESTART == (devStatus & DN_NEED_RESTART));
+    if (! isStagedDriverPathSetForDisplayAdapter(devnode)) {
+        isInInvalidState |= (DN_NEED_RESTART == (devStatus & DN_NEED_RESTART));
+    }
     isInInvalidState |= (devStatus & DN_HAS_PROBLEM)
                              && (devProblem == CM_PROB_DISABLED);
 
@@ -69,6 +71,35 @@ bool isDeviceAvailable(DEVINST devnode) {
     return false == isInInvalidState;
 }
 
+bool isStagedDriverPathSetForDisplayAdapter(DEVINST dnDevNode) {
+    static constexpr char stagedDriverPathKey[] = "LevelZeroStagedDriverPath";
+
+    HKEY hkey = {};
+    CONFIGRET configErr = CM_Open_DevNode_Key(dnDevNode, KEY_QUERY_VALUE, 0,
+                                              RegDisposition_OpenExisting, &hkey,
+                                              CM_REGISTRY_SOFTWARE);
+
+    if (CR_SUCCESS != configErr) {
+        assert(false && "CM_Open_DevNode_Key failed");
+        return false;
+    }
+
+    DWORD regValueType = {};
+    DWORD pathSize = {};
+    LSTATUS regOpStatus = RegQueryValueExA(hkey, stagedDriverPathKey, NULL,
+                                           &regValueType, NULL, &pathSize);
+
+    bool isStagedDriverPathSet = false;
+    if (ERROR_SUCCESS == regOpStatus) {
+        isStagedDriverPathSet = true;
+    }
+
+    regOpStatus = RegCloseKey(hkey);
+    assert((ERROR_SUCCESS == regOpStatus) && "RegCloseKey failed");
+
+    return isStagedDriverPathSet;
+}
+
 DriverLibraryPath readDriverPathForDisplayAdapter(DEVINST dnDevNode) {
     static constexpr char levelZeroDriverPathKey[] = "LevelZeroDriverPath";