From a5dd6c74195551f0743dce7624683dce8084835f Mon Sep 17 00:00:00 2001
From: Kirill Stoimenov <kstoimenov@google.com>
Date: Mon, 31 Jan 2022 20:51:03 +0000
Subject: [PATCH] [ASan] Fixed null pointer bug introduced in D112098.

Also added some more test to cover the "else if" part.

Reviewed By: RKSimon

Differential Revision: https://reviews.llvm.org/D118645
---
 .../Instrumentation/AddressSanitizer.cpp      | 12 ++++----
 .../AddressSanitizer/asan-stack-safety.ll     | 30 ++++++++++++++++++-
 2 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index 6e72255e51ae..8f94172a6402 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -1527,22 +1527,22 @@ void AddressSanitizer::getInterestingMemoryOperands(
     return;
 
   if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
-    if (!ClInstrumentReads || ignoreAccess(LI, LI->getPointerOperand()))
+    if (!ClInstrumentReads || ignoreAccess(I, LI->getPointerOperand()))
       return;
     Interesting.emplace_back(I, LI->getPointerOperandIndex(), false,
                              LI->getType(), LI->getAlign());
   } else if (StoreInst *SI = dyn_cast<StoreInst>(I)) {
-    if (!ClInstrumentWrites || ignoreAccess(LI, SI->getPointerOperand()))
+    if (!ClInstrumentWrites || ignoreAccess(I, SI->getPointerOperand()))
       return;
     Interesting.emplace_back(I, SI->getPointerOperandIndex(), true,
                              SI->getValueOperand()->getType(), SI->getAlign());
   } else if (AtomicRMWInst *RMW = dyn_cast<AtomicRMWInst>(I)) {
-    if (!ClInstrumentAtomics || ignoreAccess(LI, RMW->getPointerOperand()))
+    if (!ClInstrumentAtomics || ignoreAccess(I, RMW->getPointerOperand()))
       return;
     Interesting.emplace_back(I, RMW->getPointerOperandIndex(), true,
                              RMW->getValOperand()->getType(), None);
   } else if (AtomicCmpXchgInst *XCHG = dyn_cast<AtomicCmpXchgInst>(I)) {
-    if (!ClInstrumentAtomics || ignoreAccess(LI, XCHG->getPointerOperand()))
+    if (!ClInstrumentAtomics || ignoreAccess(I, XCHG->getPointerOperand()))
       return;
     Interesting.emplace_back(I, XCHG->getPointerOperandIndex(), true,
                              XCHG->getCompareOperand()->getType(), None);
@@ -1556,7 +1556,7 @@ void AddressSanitizer::getInterestingMemoryOperands(
         return;
 
       auto BasePtr = CI->getOperand(OpOffset);
-      if (ignoreAccess(LI, BasePtr))
+      if (ignoreAccess(I, BasePtr))
         return;
       Type *Ty = IsWrite ? CI->getArgOperand(0)->getType() : CI->getType();
       MaybeAlign Alignment = Align(1);
@@ -1568,7 +1568,7 @@ void AddressSanitizer::getInterestingMemoryOperands(
     } else {
       for (unsigned ArgNo = 0; ArgNo < CI->arg_size(); ArgNo++) {
         if (!ClInstrumentByval || !CI->isByValArgument(ArgNo) ||
-            ignoreAccess(LI, CI->getArgOperand(ArgNo)))
+            ignoreAccess(I, CI->getArgOperand(ArgNo)))
           continue;
         Type *Ty = CI->getParamByValType(ArgNo);
         Interesting.emplace_back(I, ArgNo, false, Ty, Align(1));
diff --git a/llvm/test/Instrumentation/AddressSanitizer/asan-stack-safety.ll b/llvm/test/Instrumentation/AddressSanitizer/asan-stack-safety.ll
index d05f37db27b1..29fc68660a36 100644
--- a/llvm/test/Instrumentation/AddressSanitizer/asan-stack-safety.ll
+++ b/llvm/test/Instrumentation/AddressSanitizer/asan-stack-safety.ll
@@ -9,11 +9,39 @@
 ; RUN: opt < %s -S -enable-new-pm=1 -asan-instrumentation-with-call-threshold=0 \
 ; RUN:     -passes='asan-pipeline' -asan-use-stack-safety=1 -o - | FileCheck %s --check-prefixes=SAFETY
 ; NOSAFETY: call void @__asan_load1
+; NOSAFETY: call void @__asan_store1
+; NOSAFETY: call void @__asan_store1
+; NOSAFETY: call void @__asan_store1
 ; SAFETY-NOT: call void @__asan_load1
+; SAFETY-NOT: call void @__asan_store1
+; SAFETY-NOT: call void @__asan_store1
+; SAFETY-NOT: call void @__asan_store1
 
-define i32 @stack-safety() sanitize_address {
+define i32 @load() sanitize_address {
   %buf = alloca [10 x i8], align 1
   %arrayidx = getelementptr inbounds [10 x i8], [10 x i8]* %buf, i64 0, i64 0
   %1 = load i8, i8* %arrayidx, align 1
   ret i32 0
 }
+
+define i32 @store() sanitize_address {
+  %buf = alloca [10 x i8], align 1
+  %arrayidx = getelementptr inbounds [10 x i8], [10 x i8]* %buf, i64 0, i64 0
+  store i8 0, i8* %arrayidx
+  ret i32 0
+}
+
+
+define void @atomicrmw() sanitize_address {
+  %buf = alloca [10 x i8], align 1
+  %arrayidx = getelementptr inbounds [10 x i8], [10 x i8]* %buf, i64 0, i64 0
+  %1 = atomicrmw add i8* %arrayidx, i8 1 seq_cst
+  ret void
+}
+
+define void @cmpxchg(i8 %compare_to, i8 %new_value) sanitize_address {
+  %buf = alloca [10 x i8], align 1
+  %arrayidx = getelementptr inbounds [10 x i8], [10 x i8]* %buf, i64 0, i64 0
+  %1 = cmpxchg i8* %arrayidx, i8 %compare_to, i8 %new_value seq_cst seq_cst
+  ret void
+}