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

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.arch_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 ucx_header_path = self
+                .latest_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",
+                ucx_header_path,
                 "ude/1.1/UdeCx.h",
                 "ufx/1.1/ufxbase.h",
                 "ufxproprietarycharger.h",
@@ -1120,6 +1111,67 @@ 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.
+    ///
+    /// 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 arch_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)
+    }
+
+    fn latest_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
+            .arch_sdk_library_path(sdk_version)?
+            .join("ucx")
+            .canonicalize()?
+            .strip_extended_length_path_prefix()?;
+
+        // Find the latest version of UCX headers in the directory
+        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.0, max_version.1);
+        Ok(Box::leak(path.into_boxed_str()))
+    }
 }
 
 impl From<DeserializableDriverConfig> for DriverConfig {

crates/wdk-build/src/utils.rs

@@ -366,6 +366,43 @@ fn read_registry_key_string_value(
     None
 }
 
+/// Finds the maximum version in a directory where subdirectories are named with
+/// version format "x.y"
+///
+/// # Arguments
+/// * `directory_path` - The path to the directory to search for version
+///   subdirectories
+///
+/// # Returns
+/// * `Some((major, minor))` - The maximum version found as a tuple of (major,
+///   minor)
+/// * `None` - If no valid version directories are found or if the directory
+///   cannot be read
+///
+/// # Examples
+/// Given a directory structure:
+/// ```text
+/// /some/path/
+/// ├── 1.4/
+/// ├── 1.5/
+/// └── 1.6/
+/// ```
+/// This function will return `Some((1, 6))`
+pub fn find_max_version_in_directory<P: AsRef<Path>>(directory_path: P) -> Option<(u32, u32)> {
+    std::fs::read_dir(directory_path.as_ref())
+        .ok()?
+        .flatten()
+        .filter_map(|entry| {
+            let dir_name = entry.file_name();
+            let dir_name = dir_name.to_str()?;
+            let (major_str, minor_str) = dir_name.split_once('.')?;
+            let major = major_str.parse::<u32>().ok()?;
+            let minor = minor_str.parse::<u32>().ok()?;
+            Some((major, minor))
+        })
+        .max()
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -515,4 +552,24 @@ mod tests {
             )
         );
     }
+
+    #[test]
+    fn test_find_max_version_in_directory() {
+        use assert_fs::prelude::*;
+        let temp_dir = assert_fs::TempDir::new().unwrap();
+        temp_dir.child("1.2").create_dir_all().unwrap();
+        temp_dir.child("1.10").create_dir_all().unwrap();
+        temp_dir.child("2.0").create_dir_all().unwrap();
+        temp_dir.child("not_a_version").create_dir_all().unwrap();
+        assert_eq!(find_max_version_in_directory(temp_dir.path()), Some((2, 0)));
+    }
+
+    #[test]
+    fn test_find_max_version_in_directory_no_versions() {
+        use assert_fs::prelude::*;
+        let temp_dir = assert_fs::TempDir::new().unwrap();
+        temp_dir.child("folder1").create_dir_all().unwrap();
+        temp_dir.child("folder2").create_dir_all().unwrap();
+        assert_eq!(find_max_version_in_directory(temp_dir.path()), None);
+    }
 }