fix: use latest version of ucx in the WDKContent as default

.github/workflows/codeql.yml

@@ -60,12 +60,35 @@ jobs:
 
     - name: Install WDK (${{ matrix.wdk }})
       run: |
-        if ((Get-WinGetPackage -Id ${{ matrix.wdk }} -Source winget -MatchOption Equals).Id -eq '${{ matrix.wdk }}') {
-          Write-Host "${{ matrix.wdk }} is already installed. Attempting to update..."
-          Update-WinGetPackage -Id ${{ matrix.wdk }} -Source winget -MatchOption Equals -Mode Silent -Force
-        } else {
-          Write-Host "Installing ${{ matrix.wdk }}..."
-          Install-WinGetPackage -Id ${{ matrix.wdk }} -Source winget -MatchOption Equals -Mode Silent -Force
+        # Delete the Ge WDK (10.0.26100) Include and Lib folders if present on the system.
+        # Ni WDK (10.0.22621) will be the latest available Kit on the system.
+        # FIXME: Remove once issues with Ge WDK are resolved.
+        $wdkGePaths = @(
+          "C:\Program Files (x86)\Windows Kits\10\Include\10.0.26100.0",
+          "C:\Program Files (x86)\Windows Kits\10\Lib\10.0.26100.0"
+        )
+        foreach ($path in $wdkGePaths) {
+          if (Test-Path $path) {
+            Write-Host "Deleting folder: $path"
+            Remove-Item -Path $path -Recurse -Force -ErrorAction SilentlyContinue
+          } else {
+            Write-Host "Ge WDK folder not found, skipping deletion: $path"
+          }
+        }
+
+        # Install or update the target WDK
+        try {
+          $targetWdk = Get-WinGetPackage -Id ${{ matrix.wdk }} -Source winget -MatchOption Equals -ErrorAction SilentlyContinue
+          if ($targetWdk) {
+            Write-Host "${{ matrix.wdk }} is already installed. Attempting to update..."
+            Update-WinGetPackage -Id ${{ matrix.wdk }} -Source winget -MatchOption Equals -Mode Silent -Force
+          } else {
+            Write-Host "Installing ${{ matrix.wdk }}..."
+            Install-WinGetPackage -Id ${{ matrix.wdk }} -Source winget -MatchOption Equals -Mode Silent -Force
+          }
+        } catch {
+          Write-Host "Error installing ${{ matrix.wdk }}: $($_.Exception.Message)"
+          exit 1
         }
 
     - name: Install Rust Toolchain (${{ matrix.rust_toolchain }})

crates/wdk-build/Cargo.toml

@@ -37,6 +37,7 @@ windows = { workspace = true, features = [
 rustversion.workspace = true
 
 [dev-dependencies]
+assert_fs.workspace = true
 windows = { workspace = true, features = ["Win32_UI_Shell"] }
 
 # Cannot inherit workspace lints since overriding them is not supported yet: https://github.com/rust-lang/cargo/issues/13157

crates/wdk-build/src/lib.rs

@@ -519,29 +519,13 @@ impl Config {
     /// exist.
     pub fn library_paths(&self) -> Result<impl Iterator<Item = PathBuf>, ConfigError> {
         let mut library_paths = vec![];
-
         let library_directory = self.wdk_content_root.join("Lib");
 
         // Add windows sdk library paths
         // Based off of logic from WindowsDriver.KernelMode.props &
         // WindowsDriver.UserMode.props in NI(22H2) WDK
         let sdk_version = utils::get_latest_windows_sdk_version(library_directory.as_path())?;
-        let windows_sdk_library_path =
-            library_directory
-                .join(sdk_version)
-                .join(match self.driver_config {
-                    DriverConfig::Wdm | DriverConfig::Kmdf(_) => {
-                        format!("km/{}", self.cpu_architecture.as_windows_str(),)
-                    }
-                    DriverConfig::Umdf(_) => {
-                        format!("um/{}", self.cpu_architecture.as_windows_str(),)
-                    }
-                });
-        if !windows_sdk_library_path.is_dir() {
-            return Err(ConfigError::DirectoryNotFound {
-                directory: windows_sdk_library_path.to_string_lossy().into(),
-            });
-        }
+        let windows_sdk_library_path = self.sdk_library_path(sdk_version)?;
         library_paths.push(
             windows_sdk_library_path
                 .canonicalize()?
@@ -869,11 +853,18 @@ impl Config {
         }
 
         if matches!(self.driver_config, DriverConfig::Kmdf(_)) {
+            let latest_ucx_header_path = self
+                .ucx_header_path()
+                .map_err(|e| {
+                    tracing::error!("Failed to get latest UCX header: {e}");
+                    e
+                })
+                .expect("Failed to get UCX header. Please verify WDK installation");
             headers.extend([
                 "ucm/1.0/UcmCx.h",
                 "UcmTcpci/1.0/UcmTcpciCx.h",
                 "UcmUcsi/1.0/UcmucsiCx.h",
-                "ucx/1.6/ucxclass.h",
+                latest_ucx_header_path,
                 "ude/1.1/UdeCx.h",
                 "ufx/1.1/ufxbase.h",
                 "ufxproprietarycharger.h",
@@ -1120,6 +1111,63 @@ impl Config {
 
         enabled_cpu_target_features.contains(STATICALLY_LINKED_C_RUNTIME_FEATURE_NAME)
     }
+
+    /// Constructs the architecture-specific Windows SDK library path using the
+    /// provided SDK Version and the driver configuration.
+    ///
+    /// Builds the library path following the Windows SDK convention:
+    /// `{library_directory}/{sdk_version}/{km|um}/{architecture}/`
+    ///
+    /// # Arguments
+    ///
+    /// * `sdk_version` - Windows SDK version string (e.g., "10.0.22621.0")
+    ///
+    /// # Returns
+    ///
+    /// The constructed library path if it exists, otherwise
+    /// `ConfigError::DirectoryNotFound`.
+    ///
+    /// # Examples
+    ///
+    /// KMDF/AMD64: `C:\...\Lib\10.0.22621.0\km\x64`
+    /// UMDF/ARM64: `C:\...\Lib\10.0.22621.0\um\arm64`
+    fn sdk_library_path(&self, sdk_version: String) -> Result<PathBuf, ConfigError> {
+        let windows_sdk_library_path =
+            self.wdk_content_root
+                .join("Lib")
+                .join(sdk_version)
+                .join(match self.driver_config {
+                    DriverConfig::Wdm | DriverConfig::Kmdf(_) => {
+                        format!("km/{}", self.cpu_architecture.as_windows_str(),)
+                    }
+                    DriverConfig::Umdf(_) => {
+                        format!("um/{}", self.cpu_architecture.as_windows_str(),)
+                    }
+                });
+        if !windows_sdk_library_path.is_dir() {
+            return Err(ConfigError::DirectoryNotFound {
+                directory: windows_sdk_library_path.to_string_lossy().into(),
+            });
+        }
+        Ok(windows_sdk_library_path)
+    }
+
+    /// Returns the path to the latest available UCX header file present in the
+    /// Lib folder of the WDK content root
+    // TODO: SDK and UCX version should be configurable
+    fn ucx_header_path(&self) -> Result<&'static str, ConfigError> {
+        let sdk_version =
+            utils::get_latest_windows_sdk_version(&self.wdk_content_root.join("Lib"))?;
+        let ucx_header_root_dir = self.sdk_library_path(sdk_version)?.join("ucx");
+        let max_version =
+            utils::find_max_version_in_directory(&ucx_header_root_dir).ok_or_else(|| {
+                ConfigError::DirectoryNotFound {
+                    directory: ucx_header_root_dir.to_string_lossy().into(),
+                }
+            })?;
+        let path = format!("ucx/{}.{}/ucxclass.h", max_version.major, max_version.minor);
+        Ok(Box::leak(path.into_boxed_str()))
+    }
 }
 
 impl From<DeserializableDriverConfig> for DriverConfig {