From e5c8c5a1bcecff7e2aa60672be6af2062ad63e6a Mon Sep 17 00:00:00 2001
From: Evgeniy Stepanov <eugeni.stepanov@gmail.com>
Date: Mon, 16 Sep 2013 13:24:32 +0000
Subject: [PATCH] [msan] Check return value of main().

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@190782 91177308-0d34-0410-b5e6-96231b3b80d8
---
 .../Instrumentation/MemorySanitizer.cpp        | 17 +++++++++++++----
 .../MemorySanitizer/return_from_main.ll        | 18 ++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 test/Instrumentation/MemorySanitizer/return_from_main.ll

diff --git a/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 075a72f78b1..9a67cfc5a75 100644
--- a/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -427,6 +427,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   bool LoadShadow;
   bool PoisonStack;
   bool PoisonUndef;
+  bool CheckReturnValue;
   OwningPtr<VarArgHelper> VAHelper;
 
   struct ShadowOriginAndInsertPoint {
@@ -449,6 +450,9 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     LoadShadow = SanitizeFunction;
     PoisonStack = SanitizeFunction && ClPoisonStack;
     PoisonUndef = SanitizeFunction && ClPoisonUndef;
+    // FIXME: Consider using SpecialCaseList to specify a list of functions that
+    // must always return fully initialized values. For now, we hardcode "main".
+    CheckReturnValue = SanitizeFunction && (F.getName() == "main");
 
     DEBUG(if (!InsertChecks)
           dbgs() << "MemorySanitizer is not inserting checks into '"
@@ -1686,12 +1690,17 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
 
   void visitReturnInst(ReturnInst &I) {
     IRBuilder<> IRB(&I);
-    if (Value *RetVal = I.getReturnValue()) {
-      // Set the shadow for the RetVal.
+    Value *RetVal = I.getReturnValue();
+    if (!RetVal) return;
+    Value *ShadowPtr = getShadowPtrForRetval(RetVal, IRB);
+    if (CheckReturnValue) {
+      insertCheck(RetVal, &I);
+      Value *Shadow = getCleanShadow(RetVal);
+      IRB.CreateAlignedStore(Shadow, ShadowPtr, kShadowTLSAlignment);
+    } else {
       Value *Shadow = getShadow(RetVal);
-      Value *ShadowPtr = getShadowPtrForRetval(RetVal, IRB);
-      DEBUG(dbgs() << "Return: " << *Shadow << "\n" << *ShadowPtr << "\n");
       IRB.CreateAlignedStore(Shadow, ShadowPtr, kShadowTLSAlignment);
+      // FIXME: make it conditional if ClStoreCleanOrigin==0
       if (MS.TrackOrigins)
         IRB.CreateStore(getOrigin(RetVal), getOriginPtrForRetval(IRB));
     }
diff --git a/test/Instrumentation/MemorySanitizer/return_from_main.ll b/test/Instrumentation/MemorySanitizer/return_from_main.ll
new file mode 100644
index 00000000000..81dc88834db
--- /dev/null
+++ b/test/Instrumentation/MemorySanitizer/return_from_main.ll
@@ -0,0 +1,18 @@
+; RUN: opt < %s -msan -msan-check-access-address=0 -S | FileCheck %s
+target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @main() sanitize_memory {
+entry:
+  %call = tail call i32 @f()
+  ret i32 %call
+}
+
+declare i32 @f() sanitize_memory
+
+; CHECK: @main
+; CHECK: call i32 @f()
+; CHECK: store i32 0, {{.*}} @__msan_retval_tls
+; CHECK: br i1
+; CHECK: call void @__msan_warning_noreturn()
+; CHECK: ret i32
-- 
2.34.1