From 98ca5853f2f3eab2a39d6152f5eb278e756a45e1 Mon Sep 17 00:00:00 2001
From: Kenny Kerr <kenny@kennykerr.ca>
Date: Wed, 16 Oct 2024 07:14:34 -0500
Subject: [PATCH] Add `unwrap` helper for `NTSTATUS` (#3324)

---
 .../extensions/Win32/Foundation/NTSTATUS.rs   | 11 +++
 crates/tests/misc/result/Cargo.toml           |  6 ++
 crates/tests/misc/result/tests/ntstatus.rs    | 68 +++++++++++++++++++
 3 files changed, 85 insertions(+)
 create mode 100644 crates/tests/misc/result/tests/ntstatus.rs

diff --git a/crates/libs/windows/src/extensions/Win32/Foundation/NTSTATUS.rs b/crates/libs/windows/src/extensions/Win32/Foundation/NTSTATUS.rs
index 2d40cf001f..62867fe273 100644
--- a/crates/libs/windows/src/extensions/Win32/Foundation/NTSTATUS.rs
+++ b/crates/libs/windows/src/extensions/Win32/Foundation/NTSTATUS.rs
@@ -5,14 +5,23 @@ impl NTSTATUS {
     pub const fn is_ok(self) -> bool {
         self.0 >= 0
     }
+
     #[inline]
     pub const fn is_err(self) -> bool {
         !self.is_ok()
     }
+
     #[inline]
     pub const fn to_hresult(self) -> windows_core::HRESULT {
         windows_core::HRESULT::from_nt(self.0)
     }
+
+    #[inline]
+    #[track_caller]
+    pub fn unwrap(self) {
+        assert!(self.is_ok(), "NTSTATUS 0x{:X}", self.0);
+    }
+
     #[inline]
     pub fn ok(self) -> windows_core::Result<()> {
         if self.is_ok() {
@@ -22,11 +31,13 @@ impl NTSTATUS {
         }
     }
 }
+
 impl From<NTSTATUS> for windows_core::HRESULT {
     fn from(value: NTSTATUS) -> Self {
         value.to_hresult()
     }
 }
+
 impl From<NTSTATUS> for windows_core::Error {
     fn from(value: NTSTATUS) -> Self {
         value.to_hresult().into()
diff --git a/crates/tests/misc/result/Cargo.toml b/crates/tests/misc/result/Cargo.toml
index 5e4ebb00e6..9e029f0333 100644
--- a/crates/tests/misc/result/Cargo.toml
+++ b/crates/tests/misc/result/Cargo.toml
@@ -8,6 +8,12 @@ publish = false
 doc = false
 doctest = false
 
+[dependencies.windows]
+workspace = true
+features = [
+    "Win32_Foundation",
+]
+
 [dependencies.windows-result]
 workspace = true
 
diff --git a/crates/tests/misc/result/tests/ntstatus.rs b/crates/tests/misc/result/tests/ntstatus.rs
new file mode 100644
index 0000000000..3a02f72984
--- /dev/null
+++ b/crates/tests/misc/result/tests/ntstatus.rs
@@ -0,0 +1,68 @@
+use windows::Win32::Foundation::{NTSTATUS, STATUS_INVALID_ACL, STATUS_SUCCESS, S_OK};
+use windows_result::Result as WindowsResult;
+use windows_result::HRESULT;
+
+#[test]
+fn test() {
+    assert!(STATUS_SUCCESS.is_ok());
+    assert!(!STATUS_INVALID_ACL.is_ok());
+
+    assert!(!STATUS_SUCCESS.is_err());
+    assert!(STATUS_INVALID_ACL.is_err());
+
+    STATUS_SUCCESS.unwrap();
+    let result: WindowsResult<()> = STATUS_SUCCESS.ok();
+    assert!(result.is_ok());
+
+    assert_eq!(STATUS_SUCCESS.to_hresult(), S_OK);
+
+    // Tests convertibility.
+    a().unwrap();
+    b().unwrap();
+}
+
+#[test]
+#[should_panic(expected = "NTSTATUS 0xC0000077")]
+fn test_panic() {
+    STATUS_INVALID_ACL.unwrap();
+}
+
+fn a() -> WindowsResult<()> {
+    fn a() -> WindowsResult<()> {
+        Ok(())
+    }
+
+    fn b() -> Result<(), HRESULT> {
+        Ok(())
+    }
+
+    fn c() -> Result<(), NTSTATUS> {
+        Ok(())
+    }
+
+    a()?;
+    b()?;
+    c()?;
+
+    Ok(())
+}
+
+fn b() -> Result<(), HRESULT> {
+    fn a() -> WindowsResult<()> {
+        Ok(())
+    }
+
+    fn b() -> Result<(), HRESULT> {
+        Ok(())
+    }
+
+    fn c() -> Result<(), NTSTATUS> {
+        Ok(())
+    }
+
+    a()?;
+    b()?;
+    c()?;
+
+    Ok(())
+}