From 11fe59df64118ef01ecba47a261d09eb4a12f4d1 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               | 32 ++++++++++++++++++++++++++++++++
 src/bin/rustup-init.rs | 19 +++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/build.rs b/build.rs
index fae0a48931e..003151bb0bb 100644
--- a/build.rs
+++ b/build.rs
@@ -27,4 +27,36 @@ fn main() {
     }
     let target = env::var("TARGET").unwrap();
     println!("cargo:rustc-env=TARGET={target}");
+
+    let target_os = env::var("CARGO_CFG_TARGET_OS");
+    let target_env = env::var("CARGO_CFG_TARGET_ENV");
+    if Ok("windows") == target_os.as_deref() && Ok("msvc") == target_env.as_deref() {
+        // # Only search system32 for DLLs
+        //
+        // This applies to DLLs loaded at load time. However, this setting is ignored
+        // before Windows 10 RS1.
+        println!("cargo:cargo:rustc-link-arg-bin=rustup-init=/DEPENDENTLOADFLAG:0x800");
+
+        // # Delay load
+        //
+        // Delay load dlls that are not "known DLLs".
+        // 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.
+        //
+        // This will work on all supported Windows versions but it relies on
+        // using `SetDefaultDllDirectories` before any libraries are loaded.
+        let delay_load_dlls = ["bcrypt", "powrprof", "secur32"];
+        for dll in delay_load_dlls {
+            println!("cargo:rustc-link-arg-bin=rustup-init=/delayload:{dll}.dll");
+        }
+        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 e486ff6fcfd..4cc8f473539 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,19 @@ fn do_recursion_guard() -> Result<()> {
 
     Ok(())
 }
+
+/// Windows pre-main security mitigations.
+///
+/// This is 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.
+    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);
+    }
+}