Transform NTSTATUS functions into Result<()> (#994)

microsoft · Jul 22, 2021 · cef626a · cef626a
1 parent 7455af8
commit cef626a
Show file tree

Hide file tree

Showing 8 changed files with 143 additions and 23 deletions.
diff --git a/crates/gen/src/types/function.rs b/crates/gen/src/types/function.rs
@@ -108,36 +108,57 @@ impl Function {
                 }
             }
         } else if let Some(return_type) = &signature.return_type {
-            if return_type.kind == ElementType::HRESULT {
-                quote! {
-                    pub unsafe fn #name<#constraints>(#params) -> ::windows::Result<()> {
-                        #[cfg(windows)]
-                        {
-                            #link_attr
-                            extern "system" {
-                                fn #name(#(#abi_params),*) -> ::windows::HRESULT;
+            match &return_type.kind {
+                ElementType::HRESULT => {
+                    quote! {
+                        pub unsafe fn #name<#constraints>(#params) -> ::windows::Result<()> {
+                            #[cfg(windows)]
+                            {
+                                #link_attr
+                                extern "system" {
+                                    fn #name(#(#abi_params),*) -> ::windows::HRESULT;
+                                }
+                                #name(#(#args),*).ok()
                             }
-                            #name(#(#args),*).ok()
+                            #[cfg(not(windows))]
+                            unimplemented!("Unsupported target OS");
                         }
-                        #[cfg(not(windows))]
-                        unimplemented!("Unsupported target OS");
                     }
                 }
-            } else {
-                let return_type = return_type.gen_win32(gen);
+                ElementType::TypeDef(def)
+                    if def.full_name() == ("Windows.Win32.Foundation", "NTSTATUS") =>
+                {
+                    quote! {
+                        pub unsafe fn #name<#constraints>(#params) -> ::windows::Result<()> {
+                            #[cfg(windows)]
+                            {
+                                #link_attr
+                                extern "system" {
+                                    fn #name(#(#abi_params),*) #abi_return_type;
+                                }
+                                #name(#(#args),*).ok()
+                            }
+                            #[cfg(not(windows))]
+                            unimplemented!("Unsupported target OS");
+                        }
+                    }
+                }
+                _ => {
+                    let return_type = return_type.gen_win32(gen);
 
-                quote! {
-                    pub unsafe fn #name<#constraints>(#params) -> #return_type {
-                        #[cfg(windows)]
-                        {
-                            #link_attr
-                            extern "system" {
-                                fn #name(#(#abi_params),*) #abi_return_type;
+                    quote! {
+                        pub unsafe fn #name<#constraints>(#params) -> #return_type {
+                            #[cfg(windows)]
+                            {
+                                #link_attr
+                                extern "system" {
+                                    fn #name(#(#abi_params),*) #abi_return_type;
+                                }
+                                #name(#(#args),*)
                             }
-                            #name(#(#args),*)
+                            #[cfg(not(windows))]
+                            unimplemented!("Unsupported target OS");
                         }
-                        #[cfg(not(windows))]
-                        unimplemented!("Unsupported target OS");
                     }
                 }
             }

diff --git a/crates/gen/src/types/mod.rs b/crates/gen/src/types/mod.rs
@@ -13,6 +13,7 @@ mod handle;
 mod interface;
 mod matrix3x2;
 mod matrix4x4;
+mod ntstatus;
 mod pstr;
 mod pwstr;
 mod r#struct;
@@ -33,6 +34,7 @@ pub use handle::*;
 pub use interface::*;
 pub use matrix3x2::*;
 pub use matrix4x4::*;
+pub use ntstatus::*;
 pub use pstr::*;
 pub use pwstr::*;
 pub use r#enum::*;

diff --git a/crates/gen/src/types/ntstatus.rs b/crates/gen/src/types/ntstatus.rs
@@ -0,0 +1,34 @@
+use super::*;
+
+pub fn gen_ntstatus() -> TokenStream {
+    quote! {
+        #[repr(transparent)]
+        #[derive(::std::default::Default, ::std::clone::Clone, ::std::marker::Copy, ::std::cmp::PartialEq, ::std::cmp::Eq, ::std::fmt::Debug)]
+        pub struct NTSTATUS(pub u32);
+
+        impl NTSTATUS {
+            #[inline]
+            pub const fn is_ok(self) -> bool {
+                self.0 == 0
+            }
+
+            #[inline]
+            pub const fn is_err(self) -> bool {
+                !self.is_ok()
+            }
+
+            #[inline]
+            pub fn ok(self) -> ::windows::Result<()> {
+                if self.is_ok() {
+                    Ok(())
+                } else {
+                    Err(::windows::Error::fast_error(::windows::HRESULT(self.0 | 0x10000000)))
+                }
+            }
+        }
+
+        unsafe impl ::windows::Abi for NTSTATUS {
+            type Abi = Self;
+        }
+    }
+}
diff --git a/crates/gen/src/types/struct.rs b/crates/gen/src/types/struct.rs
@@ -389,6 +389,7 @@ impl Struct {
             ("Windows.Win32.Foundation", "PWSTR") => Some(gen_pwstr()),
             ("Windows.Win32.Foundation", "PSTR") => Some(gen_pstr()),
             ("Windows.Win32.Foundation", "BSTR") => Some(gen_bstr()),
+            ("Windows.Win32.Foundation", "NTSTATUS") => Some(gen_ntstatus()),
             _ => None,
         }
     }

diff --git a/tests/ntstatus/Cargo.toml b/tests/ntstatus/Cargo.toml
@@ -0,0 +1,11 @@
+[package]
+name = "test_ntstatus"
+version = "0.0.0"
+authors = ["Microsoft"]
+edition = "2018"
+
+[dependencies]
+windows = { path = "../.." }
+
+[build-dependencies]
+windows = { path = "../.." }
diff --git a/tests/ntstatus/build.rs b/tests/ntstatus/build.rs
@@ -0,0 +1,8 @@
+fn main() {
+    windows::build! {
+        Windows::Win32::Foundation::STATUS_NOT_FOUND,
+        Windows::Win32::Security::Cryptography::Core::{
+            BCryptGenRandom, BCryptOpenAlgorithmProvider,
+        },
+    };
+}
diff --git a/tests/ntstatus/src/lib.rs b/tests/ntstatus/src/lib.rs
@@ -0,0 +1 @@
+windows::include_bindings!();
diff --git a/tests/ntstatus/tests/ntstatus.rs b/tests/ntstatus/tests/ntstatus.rs
@@ -0,0 +1,42 @@
+use test_ntstatus::{
+    Windows::Win32::Foundation::{NTSTATUS, STATUS_NOT_FOUND},
+    Windows::Win32::Security::Cryptography::Core::{BCryptGenRandom, BCryptOpenAlgorithmProvider},
+};
+
+use windows::{Guid, Result, HRESULT};
+
+#[test]
+fn test() -> Result<()> {
+    let status = NTSTATUS::default();
+    assert_eq!(status.0, 0);
+    assert_eq!(status.is_ok(), true);
+    assert_eq!(status.is_err(), false);
+    assert_eq!(status.ok().is_ok(), true);
+
+    let status = STATUS_NOT_FOUND;
+    assert_eq!(status.0, 0xC000_0225);
+    assert_eq!(status.is_ok(), false);
+    assert_eq!(status.is_err(), true);
+    assert_eq!(status.ok().is_ok(), false);
+
+    let error = status.ok().unwrap_err();
+    assert_eq!(error.code(), HRESULT(0xD000_0225));
+
+    unsafe {
+        let mut provider = std::ptr::null_mut();
+        BCryptOpenAlgorithmProvider(&mut provider, "RNG", None, Default::default())?;
+
+        let mut random = Guid::zeroed();
+
+        BCryptGenRandom(
+            provider,
+            &mut random as *mut _ as _,
+            std::mem::size_of::<Guid>() as _,
+            0,
+        )?;
+
+        assert_ne!(random, Guid::zeroed());
+    }
+
+    Ok(())
+}