Add support for tracking origins of uninitialized memory

4 files changed, 53 insertions, 2 deletions

diff --git a/src/librustc/session/config.rs b/src/librustc/session/config.rs
index afcbafd8a00..66c73d138e2 100644
--- a/src/librustc/session/config.rs
+++ b/src/librustc/session/config.rs
@@ -834,6 +834,7 @@ macro_rules! options {
             Some("one of: `address`, `leak`, `memory` or `thread`");
         pub const parse_sanitizer_list: Option<&str> =
             Some("comma separated list of sanitizers");
+        pub const parse_sanitizer_memory_track_origins: Option<&str> = None;
         pub const parse_linker_flavor: Option<&str> =
             Some(::rustc_target::spec::LinkerFlavor::one_of());
         pub const parse_optimization_fuel: Option<&str> =
@@ -1054,6 +1055,22 @@ macro_rules! options {
             }
         }
 
+        fn parse_sanitizer_memory_track_origins(slot: &mut usize, v: Option<&str>) -> bool {
+            match v.map(|s| s.parse()) {
+                None => {
+                    *slot = 2;
+                    true
+                }
+                Some(Ok(i)) if i <= 2 => {
+                    *slot = i;
+                    true
+                }
+                _ => {
+                    false
+                }
+            }
+        }
+
         fn parse_linker_flavor(slote: &mut Option<LinkerFlavor>, v: Option<&str>) -> bool {
             match v.and_then(LinkerFlavor::from_str) {
                 Some(lf) => *slote = Some(lf),
@@ -1411,6 +1428,8 @@ options! {DebuggingOptions, DebuggingSetter, basic_debugging_options,
                                     "use a sanitizer"),
     sanitizer_recover: Vec<Sanitizer> = (vec![], parse_sanitizer_list, [TRACKED],
         "Enable recovery for selected sanitizers"),
+    sanitizer_memory_track_origins: usize = (0, parse_sanitizer_memory_track_origins, [TRACKED],
+        "Enable origins tracking in MemorySanitizer"),
     fuel: Option<(String, u64)> = (None, parse_optimization_fuel, [TRACKED],
         "set the optimization fuel quota for a crate"),
     print_fuel: Option<String> = (None, parse_opt_string, [TRACKED],
diff --git a/src/librustc_codegen_llvm/back/write.rs b/src/librustc_codegen_llvm/back/write.rs
index 1b94cc605a5..b664c78f658 100644
--- a/src/librustc_codegen_llvm/back/write.rs
+++ b/src/librustc_codegen_llvm/back/write.rs
@@ -29,7 +29,7 @@ use std::path::{Path, PathBuf};
 use std::str;
 use std::sync::Arc;
 use std::slice;
-use libc::{c_uint, c_void, c_char, size_t};
+use libc::{c_int, c_uint, c_void, c_char, size_t};
 
 pub const RELOC_MODEL_ARGS : [(&str, llvm::RelocMode); 7] = [
     ("pic", llvm::RelocMode::PIC),
@@ -373,7 +373,7 @@ pub(crate) unsafe fn optimize(cgcx: &CodegenContext<LlvmCodegenBackend>,
                                 recover));
                     }
                     Sanitizer::Memory => {
-                        let track_origins = 0;
+                        let track_origins = config.sanitizer_memory_track_origins as c_int;
                         extra_passes.push(llvm::LLVMRustCreateMemorySanitizerPass(
                                 track_origins, recover));
                     }
diff --git a/src/librustc_codegen_ssa/back/write.rs b/src/librustc_codegen_ssa/back/write.rs
index f80869132c0..9d3e57449f8 100644
--- a/src/librustc_codegen_ssa/back/write.rs
+++ b/src/librustc_codegen_ssa/back/write.rs
@@ -61,6 +61,7 @@ pub struct ModuleConfig {
 
     pub sanitizer: Option<Sanitizer>,
     pub sanitizer_recover: Vec<Sanitizer>,
+    pub sanitizer_memory_track_origins: usize,
 
     // Flags indicating which outputs to produce.
     pub emit_pre_lto_bc: bool,
@@ -102,6 +103,7 @@ impl ModuleConfig {
 
             sanitizer: None,
             sanitizer_recover: Default::default(),
+            sanitizer_memory_track_origins: 0,
 
             emit_no_opt_bc: false,
             emit_pre_lto_bc: false,
@@ -359,6 +361,8 @@ pub fn start_async_codegen<B: ExtraBackendMethods>(
     modules_config.pgo_use = sess.opts.cg.profile_use.clone();
     modules_config.sanitizer = sess.opts.debugging_opts.sanitizer.clone();
     modules_config.sanitizer_recover = sess.opts.debugging_opts.sanitizer_recover.clone();
+    modules_config.sanitizer_memory_track_origins =
+        sess.opts.debugging_opts.sanitizer_memory_track_origins;
     modules_config.opt_level = Some(sess.opts.optimize);
     modules_config.opt_size = Some(sess.opts.optimize);
 
diff --git a/src/test/codegen/sanitizer-memory-track-orgins.rs b/src/test/codegen/sanitizer-memory-track-orgins.rs
new file mode 100644
index 00000000000..fd8be0bced7
--- /dev/null
+++ b/src/test/codegen/sanitizer-memory-track-orgins.rs
@@ -0,0 +1,28 @@
+// Verifies that MemorySanitizer track-origins level can be controlled
+// with -Zsanitizer-memory-track-origins option.
+//
+// needs-sanitizer-support
+// only-linux
+// only-x86_64
+// revisions:MSAN-0 MSAN-1 MSAN-2
+//
+//[MSAN-0] compile-flags: -Zsanitizer=memory
+//[MSAN-1] compile-flags: -Zsanitizer=memory -Zsanitizer-memory-track-origins=1
+//[MSAN-2] compile-flags: -Zsanitizer=memory -Zsanitizer-memory-track-origins
+
+#![crate_type="lib"]
+
+// MSAN-0-NOT: @__msan_track_origins
+// MSAN-1:     @__msan_track_origins = weak_odr local_unnamed_addr constant i32 1
+// MSAN-2:     @__msan_track_origins = weak_odr local_unnamed_addr constant i32 2
+//
+// MSAN-0-LABEL: define void @copy(
+// MSAN-1-LABEL: define void @copy(
+// MSAN-2-LABEL: define void @copy(
+#[no_mangle]
+pub fn copy(dst: &mut i32, src: &i32) {
+    // MSAN-0-NOT: call i32 @__msan_chain_origin(
+    // MSAN-1-NOT: call i32 @__msan_chain_origin(
+    // MSAN-2:     call i32 @__msan_chain_origin(
+    *dst = *src;
+}