Fix driver export on Windows 11

lostindark · Apr 5, 2022 · 55016e4 · 55016e4
1 parent c8f43b9
commit 55016e4
Showing 1 changed file with 21 additions and 1 deletion.
diff --git a/Rapr/Utils/DismUtil.cs b/Rapr/Utils/DismUtil.cs
@@ -7,6 +7,7 @@
 using System.Runtime.InteropServices;
 
 using Microsoft.Dism;
+using Microsoft.Win32.SafeHandles;
 
 namespace Rapr.Utils
 {
@@ -255,8 +256,14 @@ public bool ExportAllDrivers(string destinationPath)
                 try
                 {
                     using (DismSession session = this.GetSession())
+                    using (SafeWaitHandle safeWaitHandle = new SafeWaitHandle(IntPtr.Zero, false))
                     {
-                        int hresult = NativeMethods._DismExportDriver(session, destinationPath);
+                        int hresult = NativeMethods._DismExportDriver(session, destinationPath, safeWaitHandle, null, IntPtr.Zero);
+
+                        if (hresult == unchecked((int)0x80070006))
+                        {
+                            hresult = NativeMethods._DismExportDriver(session, destinationPath);
+                        }
 
                         if (hresult != 0 && hresult != 1)
                         {
@@ -292,13 +299,26 @@ public bool ExportAllDrivers(string destinationPath)
 
         internal static class NativeMethods
         {
+            public delegate void ProgressCallback(uint Current, uint Total, IntPtr UserData);
+
             [DllImport("DismApi", CharSet = CharSet.Unicode)]
             [return: MarshalAs(UnmanagedType.Error)]
             public static extern int DismInitialize(DismLogLevel logLevel, string logFilePath, string scratchDirectory);
 
+            // Pre Win 11 version.
             [DllImport("DismApi", CharSet = CharSet.Unicode)]
             [System.Diagnostics.CodeAnalysis.SuppressMessage("Style", "IDE1006:Naming Styles", Justification = "<Pending>")]
             public static extern int _DismExportDriver(DismSession Session, string Destination);
+
+            // Win 11 version.
+            [DllImport("DismApi", CharSet = CharSet.Unicode)]
+            [System.Diagnostics.CodeAnalysis.SuppressMessage("Style", "IDE1006:Naming Styles", Justification = "<Pending>")]
+            public static extern int _DismExportDriver(
+                DismSession Session,
+                [MarshalAs(UnmanagedType.LPWStr)] string Destination,
+                SafeWaitHandle CancelHandle,
+                ProgressCallback Progress,
+                IntPtr UserData);
         }
     }
 }