From 1ab5e5d9bb27548a774ad8c0b6995ad126bcd05f Mon Sep 17 00:00:00 2001
From: dylni <46035563+dylni@users.noreply.github.com>
Date: Thu, 6 Jan 2022 17:36:34 -0500
Subject: [PATCH] Fix windows bug (#5)

---
 src/windows.rs   | 26 ++++++++++++++++++++++----
 tests/common.rs  | 10 ++++++----
 tests/windows.rs | 18 ++++++++++++++++++
 3 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/src/windows.rs b/src/windows.rs
index dd7dc5f..dd8944e 100644
--- a/src/windows.rs
+++ b/src/windows.rs
@@ -109,18 +109,36 @@ pub(super) fn normalize_virtually(
         // This assertion should never fail.
         static_assert!(mem::size_of::<DWORD>() <= mem::size_of::<usize>());
 
-        let capacity = capacity as usize;
-        if let Some(additional_capacity) =
-            capacity.checked_sub(buffer.capacity())
+        let length = capacity as usize;
+        if let Some(mut additional_capacity) =
+            length.checked_sub(buffer.capacity())
         {
             assert_ne!(0, additional_capacity);
+
+            // WinAPI can recommend an insufficient capacity that causes it to
+            // return incorrect results, so extra space is reserved as a
+            // workaround.
+            macro_rules! extra_capacity {
+                () => {
+                    2
+                };
+            }
+            capacity =
+                capacity.checked_add(extra_capacity!()).ok_or_else(|| {
+                    io::Error::new(
+                        io::ErrorKind::Other,
+                        "required path length is too large for WinAPI",
+                    )
+                })?;
+            additional_capacity += extra_capacity!();
+
             buffer.reserve(additional_capacity);
             continue;
         }
 
         // SAFETY: These characters were initialized by the system call.
         unsafe {
-            buffer.set_len(capacity);
+            buffer.set_len(length);
         }
         break Ok(BasePathBuf(OsString::from_wide(&buffer)));
     }
diff --git a/tests/common.rs b/tests/common.rs
index 69bf0f8..04a1bb0 100644
--- a/tests/common.rs
+++ b/tests/common.rs
@@ -15,7 +15,7 @@ use normpath::PathExt;
 #[rustversion::attr(since(1.46.0), track_caller)]
 pub(crate) fn assert_eq<P>(expected: &Path, result: io::Result<P>)
 where
-    P: AsRef<BasePath>,
+    P: AsRef<Path>,
 {
     struct Wrapper<'a>(&'a Path);
 
@@ -25,9 +25,11 @@ where
         }
     }
 
-    impl PartialEq<Result<&BasePath, &io::Error>> for Wrapper<'_> {
-        fn eq(&self, other: &Result<&BasePath, &io::Error>) -> bool {
-            other.map(|x| self.0 == x).unwrap_or(false)
+    impl PartialEq<Result<&Path, &io::Error>> for Wrapper<'_> {
+        fn eq(&self, other: &Result<&Path, &io::Error>) -> bool {
+            other
+                .map(|x| self.0.as_os_str() == x.as_os_str())
+                .unwrap_or(false)
         }
     }
 
diff --git a/tests/windows.rs b/tests/windows.rs
index aad0483..edfa388 100644
--- a/tests/windows.rs
+++ b/tests/windows.rs
@@ -2,6 +2,8 @@
 
 #![cfg(windows)]
 
+use std::env;
+use std::io;
 use std::path::Path;
 
 use normpath::PathExt;
@@ -128,3 +130,19 @@ fn test_edge_cases() {
     // test!(r"/??/X:/", r"\??\X:/", SAME);
     // test!(r"/??/X:", r"\??\X:", SAME);
 }
+
+// https://github.com/dylni/normpath/issues/5
+#[test]
+fn test_windows_bug() -> io::Result<()> {
+    let initial_current_dir = env::current_dir()?;
+
+    for current_dir in &[r"C:\", r"C:\Users"] {
+        let current_dir = Path::new(current_dir);
+        env::set_current_dir(current_dir)?;
+        common::assert_eq(current_dir, env::current_dir());
+
+        common::assert_eq(current_dir, Path::new(".").normalize());
+    }
+
+    env::set_current_dir(initial_current_dir)
+}