From 6e11c77384989726bb4f412a0e23b59c27222c34 Mon Sep 17 00:00:00 2001
From: Chris Denton <chris@chrisdenton.dev>
Date: Fri, 15 Dec 2023 03:20:45 +0000
Subject: [PATCH] Replace SHGetFolderPathW with SHGetKnownFolderPath

---
 crates/home/Cargo.toml               |  2 +-
 crates/home/src/lib.rs               |  8 ++++----
 crates/home/src/windows.rs           | 23 +++++++++++++++--------
 src/cargo/sources/git/known_hosts.rs |  2 +-
 4 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/crates/home/Cargo.toml b/crates/home/Cargo.toml
index 294bcd56b19..33cd6ba5a95 100644
--- a/crates/home/Cargo.toml
+++ b/crates/home/Cargo.toml
@@ -17,7 +17,7 @@ repository = "https://github.com/rust-lang/cargo"
 description = "Shared definitions of home directories."
 
 [target.'cfg(windows)'.dependencies]
-windows-sys = { workspace = true, features = ["Win32_Foundation", "Win32_UI_Shell"] }
+windows-sys = { workspace = true, features = ["Win32_Foundation", "Win32_UI_Shell", "Win32_System_Com"] }
 
 [lints]
 workspace = true
diff --git a/crates/home/src/lib.rs b/crates/home/src/lib.rs
index 380405e22f1..bbe7c32ca84 100644
--- a/crates/home/src/lib.rs
+++ b/crates/home/src/lib.rs
@@ -44,11 +44,11 @@ use std::path::{Path, PathBuf};
 ///
 /// Returns the value of the `USERPROFILE` environment variable if it is set
 /// **and** it is not an empty string. Otherwise, it tries to determine the
-/// home directory by invoking the [`SHGetFolderPathW`][shgfp] function with
-/// [`CSIDL_PROFILE`][csidl].
+/// home directory by invoking the [`SHGetKnownFolderPath`][shgkfp] function with
+/// [`FOLDERID_Profile`][knownfolderid].
 ///
-/// [shgfp]: https://docs.microsoft.com/en-us/windows/win32/api/shlobj_core/nf-shlobj_core-shgetfolderpathw
-/// [csidl]: https://learn.microsoft.com/en-us/windows/win32/shell/csidl
+/// [shgkfp]: https://learn.microsoft.com/en-us/windows/win32/api/shlobj_core/nf-shlobj_core-shgetknownfolderpath
+/// [knownfolderid]: https://learn.microsoft.com/en-us/windows/win32/shell/knownfolderid
 ///
 /// # Examples
 ///
diff --git a/crates/home/src/windows.rs b/crates/home/src/windows.rs
index a35dc9c5719..c9a63d97b99 100644
--- a/crates/home/src/windows.rs
+++ b/crates/home/src/windows.rs
@@ -2,9 +2,12 @@ use std::env;
 use std::ffi::OsString;
 use std::os::windows::ffi::OsStringExt;
 use std::path::PathBuf;
+use std::ptr;
+use std::slice;
 
-use windows_sys::Win32::Foundation::{MAX_PATH, S_OK};
-use windows_sys::Win32::UI::Shell::{SHGetFolderPathW, CSIDL_PROFILE};
+use windows_sys::Win32::Foundation::S_OK;
+use windows_sys::Win32::System::Com::CoTaskMemFree;
+use windows_sys::Win32::UI::Shell::{FOLDERID_Profile, SHGetKnownFolderPath, KF_FLAG_DONT_VERIFY};
 
 pub fn home_dir_inner() -> Option<PathBuf> {
     env::var_os("USERPROFILE")
@@ -16,15 +19,19 @@ pub fn home_dir_inner() -> Option<PathBuf> {
 #[cfg(not(target_vendor = "uwp"))]
 fn home_dir_crt() -> Option<PathBuf> {
     unsafe {
-        let mut path: Vec<u16> = Vec::with_capacity(MAX_PATH as usize);
-        match SHGetFolderPathW(0, CSIDL_PROFILE as i32, 0, 0, path.as_mut_ptr()) {
+        let mut path = ptr::null_mut();
+        match SHGetKnownFolderPath(&FOLDERID_Profile, KF_FLAG_DONT_VERIFY as u32, 0, &mut path) {
             S_OK => {
-                let len = wcslen(path.as_ptr());
-                path.set_len(len);
-                let s = OsString::from_wide(&path);
+                let path_slice = slice::from_raw_parts(path, wcslen(path));
+                let s = OsString::from_wide(&path_slice);
+                CoTaskMemFree(path.cast());
                 Some(PathBuf::from(s))
             }
-            _ => None,
+            _ => {
+                // Free any allocated memory even on failure. A null ptr is a no-op for `CoTaskMemFree`.
+                CoTaskMemFree(path.cast());
+                None
+            }
         }
     }
 }
diff --git a/src/cargo/sources/git/known_hosts.rs b/src/cargo/sources/git/known_hosts.rs
index 0b0dd3208aa..f316cc25342 100644
--- a/src/cargo/sources/git/known_hosts.rs
+++ b/src/cargo/sources/git/known_hosts.rs
@@ -538,7 +538,7 @@ fn user_known_host_location() -> Option<PathBuf> {
     // - OpenSSH (most unix platforms): Uses `pw->pw_dir` from `getpwuid()`.
     //
     // This doesn't do anything close to that. home_dir's behavior is:
-    // - Windows: $USERPROFILE, or SHGetFolderPathW()
+    // - Windows: $USERPROFILE, or SHGetKnownFolderPath()
     // - Unix: $HOME, or getpwuid_r()
     //
     // Since there is a mismatch here, the location returned here might be