From d8b204c31f07ab98aa0f1a47a00c0b30efebecc2 Mon Sep 17 00:00:00 2001
From: Harry Fei <tiziyuanfang@gmail.com>
Date: Sun, 30 Sep 2018 17:53:14 +0800
Subject: init commit

---
 .gitignore       |  3 +++
 Cargo.toml       |  8 ++++++++
 LICENSE.txt      | 19 ++++++++++++++++++
 build.rs         | 61 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 rustfmt.toml     |  1 +
 src/alloc_sys.rs | 12 +++++++++++
 src/lib.rs       | 20 +++++++++++++++++++
 7 files changed, 124 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 Cargo.toml
 create mode 100644 LICENSE.txt
 create mode 100644 build.rs
 create mode 100644 rustfmt.toml
 create mode 100644 src/alloc_sys.rs
 create mode 100644 src/lib.rs

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..6936990
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+/target
+**/*.rs.bk
+Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..cbda0ca
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,8 @@
+[package]
+name = "win-kmd-alloc"
+version = "0.1.0"
+authors = ["Harry Fei <tiziyuanfang@gmail.com>"]
+
+[build-dependencies]
+winreg = "0.5"
+failure = {version = "0.1.2", default-features = false, features = ["std"]}
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100644
index 0000000..5307124
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,19 @@
+Copyright (c) 2018 Trantect Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
diff --git a/build.rs b/build.rs
new file mode 100644
index 0000000..3b9da80
--- /dev/null
+++ b/build.rs
@@ -0,0 +1,61 @@
+extern crate winreg;
+#[macro_use]
+extern crate failure;
+
+use std::env::var;
+use std::path::{Path, PathBuf};
+
+use winreg::enums::*;
+use winreg::RegKey;
+
+use failure::Error;
+
+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())
+}
+
+fn get_km_dir(windows_kits_dir: &PathBuf) -> Result<PathBuf, 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"))
+}
+
+fn main() {
+    let windows_kits_dir = get_windows_kits_dir().unwrap();
+
+    let km_dir = get_km_dir(&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()
+    );
+}
diff --git a/rustfmt.toml b/rustfmt.toml
new file mode 100644
index 0000000..ab4ccee
--- /dev/null
+++ b/rustfmt.toml
@@ -0,0 +1 @@
+newline_style = "unix"
diff --git a/src/alloc_sys.rs b/src/alloc_sys.rs
new file mode 100644
index 0000000..5d46f02
--- /dev/null
+++ b/src/alloc_sys.rs
@@ -0,0 +1,12 @@
+#[repr(C)]
+pub enum POOL_TYPE {
+    PagedPool,
+}
+
+pub type PVOID = *mut u8;
+
+#[link(name = "ntoskrnl")]
+extern "system" {
+    pub fn ExAllocatePoolWithTag(PoolType: POOL_TYPE, NumberOfBytes: usize, Tag: u32) -> PVOID;
+    pub fn ExFreePoolWithTag(P: PVOID, Tag: u32);
+}
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..678735d
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,20 @@
+#![no_std]
+
+mod alloc_sys;
+
+use self::alloc_sys::*;
+use core::alloc::{GlobalAlloc, Layout};
+
+const KMRS_TAG: u32 = 0x4B4D5253; // 'KMRS'
+
+pub struct KernelAlloc;
+
+unsafe impl GlobalAlloc for KernelAlloc {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        ExAllocatePoolWithTag(POOL_TYPE::PagedPool, layout.size(), KMRS_TAG)
+    }
+
+    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
+        ExFreePoolWithTag(ptr, KMRS_TAG);
+    }
+}
-- 
cgit v1.2.3