windows-kernel-build 2.11.19

use std::env::var;
use std::path::Path;
use std::path::PathBuf;
use thiserror::Error;
use winreg::RegKey;
use winreg::enums::*;
use failure::format_err;

/// The error type
#[derive(Debug, Error)]
pub enum Error {
    #[error(transparent)]
    IoError(#[from] std::io::Error),
    #[error("cannot find the directory")]
    DirectoryNotFound,
}

/// Represents the type of directory to retrieve.
pub enum DirectoryType {
    Include,
    Library,
}

/// Retrieves the path to the Windows Kits directory. The default should be
/// `C:\Program Files (x86)\Windows Kits\10`.
pub fn get_windows_kits_dir() -> Result<PathBuf, Error> {
    let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
    let key = r"SOFTWARE\Microsoft\Windows Kits\Installed Roots";
    let dir: String = hklm.open_subkey(key)?.get_value("KitsRoot10")?;

    Ok(dir.into())
}

/// Retrieves the path to the kernel mode libraries. The path may look something like:
/// `C:\Program Files (x86)\Windows Kits\10\lib\10.0.18362.0\km`.
/// in [DirectoryType]
pub fn get_km_dir(dir_type: DirectoryType) -> Result<PathBuf, Error> {
    // We first append lib to the path and read the directory..
    let dir = get_windows_kits_dir()?
        .join(match dir_type {
            DirectoryType::Include => "Include",
            DirectoryType::Library => "Lib",
        })
        .read_dir()?;

    // In the lib directory we may have one or more directories named after the version of Windows,
    // we will be looking for the highest version number.
    let dir = dir
        .filter_map(|dir| dir.ok())
        .map(|dir| dir.path())
        .filter(|dir| {
            dir.components()
                .last()
                .and_then(|c| c.as_os_str().to_str())
                .map(|c| c.starts_with("10.") && dir.join("km").is_dir())
                .unwrap_or(false)
        })
        .max()
        .ok_or_else(|| Error::DirectoryNotFound)?;

    // Finally append km to the path to get the path to the kernel mode libraries.
    Ok(dir.join("km"))
}

/// build 
pub fn build() -> Result<(), Error> {
    // Get the path to the kernel libraries.
    let dir = get_km_dir(DirectoryType::Library).unwrap();

    // Append the architecture based on our target.
    let target = std::env::var("TARGET").unwrap();

    let arch = if target.contains("x86_64") {
        "x64"
    } else if target.contains("i686") {
        "x86"
    } else {
        panic!("The target {} is currently not supported.", target);
    };

    let dir = dir.join(arch);

    // Specify the link path.
    println!("cargo:rustc-link-search=native={}", dir.to_str().unwrap());

    // Ensure the right linker flags are passed for building a driver.
    println!("cargo:rustc-link-arg=/NODEFAULTLIB");
    println!("cargo:rustc-link-arg=/SUBSYSTEM:NATIVE");
    println!("cargo:rustc-link-arg=/DRIVER");
    println!("cargo:rustc-link-arg=/DYNAMICBASE");
    println!("cargo:rustc-link-arg=/MANIFEST:NO");
    println!("cargo:rustc-link-arg=/ENTRY:driver_entry");
    println!("cargo:rustc-link-arg=/MERGE:.edata=.rdata");
    println!("cargo:rustc-link-arg=/MERGE:.rustc=.data");
    println!("cargo:rustc-link-arg=/INTEGRITYCHECK");

    Ok(())
}


/// Retrieves the path to the kernel mode libraries. The path may look something like:
/// `C:\Program Files (x86)\Windows Kits\10\lib\10.0.18362.0\km`.
///  in [`PathBuf`]
pub fn get_km_dir_path(windows_kits_dir: &PathBuf) -> Result<PathBuf, failure::Error> {
    let readdir = Path::new(windows_kits_dir).join("lib").read_dir()?;

    let max_libdir = readdir
        .filter_map(|dir| dir.ok())
        .map(|dir| dir.path())
        .filter(|dir| {
            dir.components()
                .last()
                .and_then(|c| c.as_os_str().to_str())
                .map(|c| c.starts_with("10.") && dir.join("km").is_dir())
                .unwrap_or(false)
        })
        .max()
        .ok_or_else(|| format_err!("Can not find a valid km dir in `{:?}`", windows_kits_dir))?;

    Ok(max_libdir.join("km"))
}


/// search link interna
/// # Examples
/// ```
/// use windows_kernel_build::internal_link_search;
/// use windows_kernel_build::extra_link_search;
/// use std::env::var;
/// pub fn build_driver() {
///    if var(format!(
///        "CARGO_FEATURE_{}",
///        "extra_link_search".to_uppercase()
///    ))
///    .is_ok()
///    {
///        extra_link_search()
///    } else {
///        internal_link_search()
///    }
/// }
/// ```
pub fn internal_link_search() {
    let windows_kits_dir = get_windows_kits_dir().unwrap();
    let km_dir = get_km_dir_path(&windows_kits_dir).unwrap();
    let target = var("TARGET").unwrap();

    let arch = if target.contains("x86_64") {
        "x64"
    } else if target.contains("i686") {
        "x86"
    } else {
        panic!("Only support x86_64 and i686!");
    };

    let lib_dir = km_dir.join(arch);
    println!(
        "cargo:rustc-link-search=native={}",
        lib_dir.to_str().unwrap()
    );
    println!("cargo:rustc-link-lib=msvcrt");
}


/// if extra link is true -> {}
pub fn extra_link_search() {}


/// build driver
/// # Examples
/// ```
/// use windows_kernel_build::build_driver;
/// use windows_kernel_build::extra_link_search;
/// use windows_kernel_build::internal_link_search;
/// fn test() {
///     build_driver();
/// }
/// ```
pub fn build_driver() {
    if var(format!(
        "CARGO_FEATURE_{}",
        "extra_link_search".to_uppercase()
    ))
    .is_ok()
    {
        extra_link_search()
    } else {
        internal_link_search()
    }
}