Apply LLVM sanitize attributes to generated entry wrapper

3 files changed, 43 insertions, 32 deletions

diff --git a/src/librustc_codegen_llvm/attributes.rs b/src/librustc_codegen_llvm/attributes.rs
index 3e23df09c66..a9e4fdba030 100644
--- a/src/librustc_codegen_llvm/attributes.rs
+++ b/src/librustc_codegen_llvm/attributes.rs
@@ -46,6 +46,31 @@ fn inline(cx: &CodegenCx<'ll, '_>, val: &'ll Value, inline: InlineAttr) {
     };
 }
 
+/// Apply LLVM sanitize attributes.
+#[inline]
+pub fn sanitize(cx: &CodegenCx<'ll, '_>, codegen_fn_flags: CodegenFnAttrFlags, llfn: &'ll Value) {
+    if let Some(ref sanitizer) = cx.tcx.sess.opts.debugging_opts.sanitizer {
+        match *sanitizer {
+            Sanitizer::Address => {
+                if !codegen_fn_flags.contains(CodegenFnAttrFlags::NO_SANITIZE_ADDRESS) {
+                    llvm::Attribute::SanitizeAddress.apply_llfn(Function, llfn);
+                }
+            }
+            Sanitizer::Memory => {
+                if !codegen_fn_flags.contains(CodegenFnAttrFlags::NO_SANITIZE_MEMORY) {
+                    llvm::Attribute::SanitizeMemory.apply_llfn(Function, llfn);
+                }
+            }
+            Sanitizer::Thread => {
+                if !codegen_fn_flags.contains(CodegenFnAttrFlags::NO_SANITIZE_THREAD) {
+                    llvm::Attribute::SanitizeThread.apply_llfn(Function, llfn);
+                }
+            }
+            Sanitizer::Leak => {}
+        }
+    }
+}
+
 /// Tell LLVM to emit or not emit the information necessary to unwind the stack for the function.
 #[inline]
 pub fn emit_uwtable(val: &'ll Value, emit: bool) {
@@ -288,26 +313,7 @@ pub fn from_fn_attrs(
     if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::ALLOCATOR) {
         Attribute::NoAlias.apply_llfn(llvm::AttributePlace::ReturnValue, llfn);
     }
-    if let Some(ref sanitizer) = cx.tcx.sess.opts.debugging_opts.sanitizer {
-        match *sanitizer {
-            Sanitizer::Address => {
-                if !codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_SANITIZE_ADDRESS) {
-                    llvm::Attribute::SanitizeAddress.apply_llfn(Function, llfn);
-                }
-            }
-            Sanitizer::Memory => {
-                if !codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_SANITIZE_MEMORY) {
-                    llvm::Attribute::SanitizeMemory.apply_llfn(Function, llfn);
-                }
-            }
-            Sanitizer::Thread => {
-                if !codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_SANITIZE_THREAD) {
-                    llvm::Attribute::SanitizeThread.apply_llfn(Function, llfn);
-                }
-            }
-            Sanitizer::Leak => {}
-        }
-    }
+    sanitize(cx, codegen_fn_attrs.flags, llfn);
 
     unwind(
         llfn,
diff --git a/src/librustc_codegen_llvm/base.rs b/src/librustc_codegen_llvm/base.rs
index d3b524c1a1e..04c084e459e 100644
--- a/src/librustc_codegen_llvm/base.rs
+++ b/src/librustc_codegen_llvm/base.rs
@@ -15,6 +15,7 @@
 
 use super::ModuleLlvm;
 
+use crate::attributes;
 use crate::builder::Builder;
 use crate::common;
 use crate::context::CodegenCx;
@@ -23,7 +24,7 @@ use crate::metadata;
 use crate::value::Value;
 
 use rustc::dep_graph;
-use rustc::middle::codegen_fn_attrs::CodegenFnAttrs;
+use rustc::middle::codegen_fn_attrs::{CodegenFnAttrFlags, CodegenFnAttrs};
 use rustc::middle::cstore::EncodedMetadata;
 use rustc::middle::exported_symbols;
 use rustc::mir::mono::{Linkage, Visibility};
@@ -131,7 +132,9 @@ pub fn compile_codegen_unit(
 
             // If this codegen unit contains the main function, also create the
             // wrapper here
-            maybe_create_entry_wrapper::<Builder<'_, '_, '_>>(&cx);
+            if let Some(entry) = maybe_create_entry_wrapper::<Builder<'_, '_, '_>>(&cx) {
+                attributes::sanitize(&cx, CodegenFnAttrFlags::empty(), entry);
+            }
 
             // Run replace-all-uses-with for statics that need it
             for &(old_g, new_g) in cx.statics_to_rauw().borrow().iter() {
diff --git a/src/librustc_codegen_ssa/base.rs b/src/librustc_codegen_ssa/base.rs
index 1f43a4027c5..90015091384 100644
--- a/src/librustc_codegen_ssa/base.rs
+++ b/src/librustc_codegen_ssa/base.rs
@@ -391,10 +391,12 @@ pub fn codegen_instance<'a, 'tcx: 'a, Bx: BuilderMethods<'a, 'tcx>>(
 
 /// Creates the `main` function which will initialize the rust runtime and call
 /// users main function.
-pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(cx: &'a Bx::CodegenCx) {
+pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
+    cx: &'a Bx::CodegenCx,
+) -> Option<Bx::Function> {
     let (main_def_id, span) = match cx.tcx().entry_fn(LOCAL_CRATE) {
         Some((def_id, _)) => (def_id, cx.tcx().def_span(def_id)),
-        None => return,
+        None => return None,
     };
 
     let instance = Instance::mono(cx.tcx(), main_def_id);
@@ -402,17 +404,15 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(cx: &'
     if !cx.codegen_unit().contains_item(&MonoItem::Fn(instance)) {
         // We want to create the wrapper in the same codegen unit as Rust's main
         // function.
-        return;
+        return None;
     }
 
     let main_llfn = cx.get_fn_addr(instance);
 
-    let et = cx.tcx().entry_fn(LOCAL_CRATE).map(|e| e.1);
-    match et {
-        Some(EntryFnType::Main) => create_entry_fn::<Bx>(cx, span, main_llfn, main_def_id, true),
-        Some(EntryFnType::Start) => create_entry_fn::<Bx>(cx, span, main_llfn, main_def_id, false),
-        None => {} // Do nothing.
-    }
+    return cx.tcx().entry_fn(LOCAL_CRATE).map(|(_, et)| {
+        let use_start_lang_item = EntryFnType::Start != et;
+        create_entry_fn::<Bx>(cx, span, main_llfn, main_def_id, use_start_lang_item)
+    });
 
     fn create_entry_fn<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         cx: &'a Bx::CodegenCx,
@@ -420,7 +420,7 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(cx: &'
         rust_main: Bx::Value,
         rust_main_def_id: DefId,
         use_start_lang_item: bool,
-    ) {
+    ) -> Bx::Function {
         // The entry function is either `int main(void)` or `int main(int argc, char **argv)`,
         // depending on whether the target needs `argc` and `argv` to be passed in.
         let llfty = if cx.sess().target.target.options.main_needs_argc_argv {
@@ -481,6 +481,8 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(cx: &'
         let result = bx.call(start_fn, &args, None);
         let cast = bx.intcast(result, cx.type_int(), true);
         bx.ret(cast);
+
+        llfn
     }
 }