From fd54f3c5bfacb451d60fec66bd069b6194c44699 Mon Sep 17 00:00:00 2001
From: Chris Denton <chris@chrisdenton.dev>
Date: Sun, 13 Aug 2023 11:11:25 +0100
Subject: [PATCH] Windows: Load DLLs from system32

---
 build.rs               | 40 ++++++++++++++++++++++++++++++++++++++++
 src/bin/rustup-init.rs | 21 +++++++++++++++++++++
 2 files changed, 61 insertions(+)

diff --git a/build.rs b/build.rs
index fae0a48931..407545c869 100644
--- a/build.rs
+++ b/build.rs
@@ -27,4 +27,44 @@ fn main() {
     }
     let target = env::var("TARGET").unwrap();
     println!("cargo:rustc-env=TARGET={target}");
+
+    // Set linker options specific to Windows MSVC.
+    let target_os = env::var("CARGO_CFG_TARGET_OS");
+    let target_env = env::var("CARGO_CFG_TARGET_ENV");
+    if !(target_os.as_deref() == Ok("windows") && target_env.as_deref() == Ok("msvc")) {
+        return;
+    }
+
+    // # Only search system32 for DLLs
+    //
+    // This applies to DLLs loaded at load time. However, this setting is ignored
+    // before Windows 10 RS1 (aka 1601).
+    // https://learn.microsoft.com/en-us/cpp/build/reference/dependentloadflag?view=msvc-170
+    println!("cargo:cargo:rustc-link-arg-bin=rustup-init=/DEPENDENTLOADFLAG:0x800");
+
+    // # Delay load
+    //
+    // Delay load dlls that are not "known DLLs"[1].
+    // Known DLLs are always loaded from the system directory whereas other DLLs
+    // are loaded from the application directory. By delay loading the latter
+    // we can ensure they are instead loaded from the system directory.
+    // [1]: https://learn.microsoft.com/en-us/windows/win32/dlls/dynamic-link-library-search-order#factors-that-affect-searching
+    //
+    // This will work on all supported Windows versions but it relies on
+    // us using `SetDefaultDllDirectories` before any libraries are loaded.
+    // See also: src/bin/rustup-init.rs
+    let delay_load_dlls = ["bcrypt", "powrprof", "secur32"];
+    for dll in delay_load_dlls {
+        println!("cargo:rustc-link-arg-bin=rustup-init=/delayload:{dll}.dll");
+    }
+    // When using delayload, it's necessary to also link delayimp.lib
+    // https://learn.microsoft.com/en-us/cpp/build/reference/dependentloadflag?view=msvc-170
+    println!("cargo:rustc-link-arg-bin=rustup-init=delayimp.lib");
+
+    // # Turn linker warnings into errors
+    //
+    // Rust hides linker warnings meaning mistakes may go unnoticed.
+    // Turning them into errors forces them to be displayed (and the build to fail).
+    // If we do want to ignore specific warnings then `/IGNORE:` should be used.
+    println!("cargo:cargo:rustc-link-arg-bin=rustup-init=/WX");
 }
diff --git a/src/bin/rustup-init.rs b/src/bin/rustup-init.rs
index e486ff6fcf..19c8dd99e9 100644
--- a/src/bin/rustup-init.rs
+++ b/src/bin/rustup-init.rs
@@ -29,6 +29,9 @@ use rustup::is_proxyable_tools;
 use rustup::utils::utils;
 
 fn main() {
+    #[cfg(windows)]
+    pre_rustup_main_init();
+
     let process = OSProcess::default();
     with(process.into(), || match maybe_trace_rustup() {
         Err(e) => {
@@ -163,3 +166,21 @@ fn do_recursion_guard() -> Result<()> {
 
     Ok(())
 }
+
+/// Windows pre-main security mitigations.
+///
+/// This attempts to defend against malicious DLLs that may sit alongside
+/// rustup-init in the user's download folder.
+#[cfg(windows)]
+pub fn pre_rustup_main_init() {
+    use winapi::um::libloaderapi::{SetDefaultDllDirectories, LOAD_LIBRARY_SEARCH_SYSTEM32};
+    // Default to loading delay loaded DLLs from the system directory.
+    // For DLLs loaded at load time, this relies on the `delayload` linker flag.
+    // This is only necessary prior to Windows 10 RS1. See build.rs for details.
+    unsafe {
+        let result = SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_SYSTEM32);
+        // SetDefaultDllDirectories should never fail if given valid arguments.
+        // But just to be safe and to catch mistakes, assert that it succeeded.
+        assert_ne!(result, 0);
+    }
+}