From 76e6681b1f79bc7ed1c177bb95a37acead80a08c Mon Sep 17 00:00:00 2001
From: Wesley Wiser <wesleywiser@microsoft.com>
Date: Wed, 29 Nov 2023 17:36:45 -0600
Subject: [PATCH] Dispose llvm::TargetMachines prior to llvm::Context being
 disposed

If the TargetMachine is disposed after the Context is disposed, it can
lead to use after frees in some cases.

I've observed this happening occasionally on code compiled for
aarch64-pc-windows-msvc using `-Zstack-protector=strong` but other users
have reported AVs from host aarch64-pc-windows-msvc compilers as well.

(cherry picked from commit 3323e4dc04e57cc64ac77dbff2f6bf50ac6832f0)
---
 compiler/rustc_codegen_llvm/src/back/lto.rs |  3 ++-
 compiler/rustc_codegen_llvm/src/lib.rs      | 21 ++++++++++++++++-----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/compiler/rustc_codegen_llvm/src/back/lto.rs b/compiler/rustc_codegen_llvm/src/back/lto.rs
index 8655aeec13dd6..db297425b03bf 100644
--- a/compiler/rustc_codegen_llvm/src/back/lto.rs
+++ b/compiler/rustc_codegen_llvm/src/back/lto.rs
@@ -25,6 +25,7 @@ use std::ffi::{CStr, CString};
 use std::fs::File;
 use std::io;
 use std::iter;
+use std::mem::ManuallyDrop;
 use std::path::Path;
 use std::slice;
 use std::sync::Arc;
@@ -734,7 +735,7 @@ pub unsafe fn optimize_thin_module(
     let llcx = llvm::LLVMRustContextCreate(cgcx.fewer_names);
     let llmod_raw = parse_module(llcx, module_name, thin_module.data(), &diag_handler)? as *const _;
     let mut module = ModuleCodegen {
-        module_llvm: ModuleLlvm { llmod_raw, llcx, tm },
+        module_llvm: ModuleLlvm { llmod_raw, llcx, tm: ManuallyDrop::new(tm) },
         name: thin_module.name().to_string(),
         kind: ModuleKind::Regular,
     };
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
index 8a6a5f79b3bb9..4ea2d30ad6d20 100644
--- a/compiler/rustc_codegen_llvm/src/lib.rs
+++ b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -53,6 +53,7 @@ use rustc_span::symbol::Symbol;
 use std::any::Any;
 use std::ffi::CStr;
 use std::io::Write;
+use std::mem::ManuallyDrop;
 
 mod back {
     pub mod archive;
@@ -408,8 +409,9 @@ pub struct ModuleLlvm {
     llcx: &'static mut llvm::Context,
     llmod_raw: *const llvm::Module,
 
-    // independent from llcx and llmod_raw, resources get disposed by drop impl
-    tm: OwnedTargetMachine,
+    // This field is `ManuallyDrop` because it is important that the `TargetMachine`
+    // is disposed prior to the `Context` being disposed otherwise UAFs can occur.
+    tm: ManuallyDrop<OwnedTargetMachine>,
 }
 
 unsafe impl Send for ModuleLlvm {}
@@ -420,7 +422,11 @@ impl ModuleLlvm {
         unsafe {
             let llcx = llvm::LLVMRustContextCreate(tcx.sess.fewer_names());
             let llmod_raw = context::create_module(tcx, llcx, mod_name) as *const _;
-            ModuleLlvm { llmod_raw, llcx, tm: create_target_machine(tcx, mod_name) }
+            ModuleLlvm {
+                llmod_raw,
+                llcx,
+                tm: ManuallyDrop::new(create_target_machine(tcx, mod_name)),
+            }
         }
     }
 
@@ -428,7 +434,11 @@ impl ModuleLlvm {
         unsafe {
             let llcx = llvm::LLVMRustContextCreate(tcx.sess.fewer_names());
             let llmod_raw = context::create_module(tcx, llcx, mod_name) as *const _;
-            ModuleLlvm { llmod_raw, llcx, tm: create_informational_target_machine(tcx.sess) }
+            ModuleLlvm {
+                llmod_raw,
+                llcx,
+                tm: ManuallyDrop::new(create_informational_target_machine(tcx.sess)),
+            }
         }
     }
 
@@ -449,7 +459,7 @@ impl ModuleLlvm {
                 }
             };
 
-            Ok(ModuleLlvm { llmod_raw, llcx, tm })
+            Ok(ModuleLlvm { llmod_raw, llcx, tm: ManuallyDrop::new(tm) })
         }
     }
 
@@ -461,6 +471,7 @@ impl ModuleLlvm {
 impl Drop for ModuleLlvm {
     fn drop(&mut self) {
         unsafe {
+            drop(ManuallyDrop::take(&mut self.tm));
             llvm::LLVMContextDispose(&mut *(self.llcx as *mut _));
         }
     }