[ASan] Added stack safety support in address sanitizer.

[kstoimenov]

Summary

Added and implemented -asan-use-stack-safety flag, which control if ASan would use the Stack Safety results to emit less code for operations which are marked as 'safe' by the static analysis.

[kstoimenov]

Fixed the test to pass.

[fmayer]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

1496–1499

Doesn't need nested if.

I would order this as ClUseStackSafety && SSI && findAllocaForValue(Ptr) && ... (cheapest first).

[vitalybuka]

clang/include/clang/Basic/CodeGenOptions.def

227

We should avoid this flag.
this is internal optimization. if it works we need to make it on
for experiments -mllvm is enough.

Also clang stuff better to extract to a separate patch.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

813

functionality of the pass should be the same, so I don't see why this is the fatal error. so just ignore it?

btw, why don't you want to support it?

1262

You can make this module pass to calculate StackSafetyGlobalAnalysis and use cached one below.
This way we will avoid exposure of this logic to PM

[fmayer]

clang/test/CodeGen/asan-stack-safety-analysis.c

2

Should this file be in llvm/test/Instrumentation/ instead? Also consider porting some of the tests from HWASan (https://github.com/llvm/llvm-project/blob/main/llvm/test/Instrumentation/HWAddressSanitizer/stack-safety-analysis.ll).

[fmayer]

clang/test/CodeGen/asan-stack-safety-analysis.c

2

Uhm nevermind my first sentence, sorry about that. This is of course the right location.

[kstoimenov]

Removed the top level flag and addressed comments.

[fmayer]

I think you also want to use SSI->isSafe(AllocaInst*) in isInterestingAlloca to prevent use-after-scope instrumentation if all accesses are safe.

[kstoimenov]

clang/test/CodeGen/asan-stack-safety-analysis.c

2

I've removed this file. Also added a .ll test for the internal flag.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

813

I would like to support it, but I am not sure how to get an instance of StackSafetyGlobalAnalysis, because in the legacy pass manger I don't have access to AnalysisManager<Function> &AM. If you know how to make this work, please let me know.

As far as the report_fatal_error my logic is that if someone sets ClUseStackSafety flag they expect it to work and if it silently doesn't it is a bad user experience.

[fmayer]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

813

Maybe it would help to look at the HWASan instrumentation pass, which uses the stack safety analysis for both old and new pass managers.

[kstoimenov]

Added support for legacy pass manager and a test for it.

[kstoimenov]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

813

Thanks for the hint. I've used the HWAsan approach.

[kstoimenov]

This should be good for re-review now. PTAL.

[vitalybuka]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

814

usually we don't use {} for one liners.
But I'd prefer here:
const StackSafetyGlobalInfo *SSI = ClUseStackSafety ? get...: nullptr

1262

It should probably be in ModuleAddressSanitizerPass, so it wll return ::all() and we don't care about invalidation.

1337

INITIALIZE_PASS_DEPENDENCY(StackSafetyGlobalInfoWrapperPass)

[kstoimenov]

Fixed test.

[kstoimenov]

Added a comment.

[kstoimenov]

Removed {} after if.

[vitalybuka]

Can you please fix in a separate patch llvm::runPassPipeline: it runs ModulePass after the function pass for -passes=asan-pipeline

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

1496

1516

using poiter after the cast looks cleaner to me
"ignoreAccess(LI, LI->"

[kstoimenov]

PTAL.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

1262

getCachedResult uses the invalidation so I couldn't get rid of it.

[vitalybuka]

llvm/include/llvm/Analysis/StackSafetyAnalysis.h

93

we still can't do that, some passes can make results irrelevant

[kstoimenov]

s/I/LI/.

[vitalybuka]

llvm/include/llvm/Analysis/StackSafetyAnalysis.h

93

Looks like only immutable analysis can be used through proxy, and this analysis cannot be immutable.
Maybe we have to convert Asan into ModulePass, like HWAsan.
Any other ideas? @eugenis @pcc @morehouse

[eugenis]

llvm/include/llvm/Analysis/StackSafetyAnalysis.h

93

Is the problem here that an invalidated analysis needs to be rerun, and the module-to-function proxy does not know how to do that?

TBH, I'm not sure how big the performance benefits of a function pass are. With the amount of time we've sunk in working around function pass requirements and I'd say just go for it.

[kstoimenov]

llvm/include/llvm/Analysis/StackSafetyAnalysis.h

93

See https://reviews.llvm.org/D112732. I changed it for the new pass manager only.

[kstoimenov]

After removing function pass.

[kstoimenov]

Restored StackSafetyAnalysis.h.

[vitalybuka]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

1496

check is not needed, we SSGI should be already null in this case

1496–1499

findAllocaForValue is DFS (expensive), stackAccessIsSafe is hash table lookup (cheep)
so I would recommend to switch them

[kstoimenov]

Fixed tests.

[vitalybuka]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

1317

There is ::ProcessedAllocas which will be unnececary overpopulated with irrelevand allocas
Maybe add FunctionSanitizer.Reset()

But please do that in a separate patch, this optimization is unrelated to SSGI

[kstoimenov]

Moved AddressSanitizer back to the loop.

[RKSimon]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

1521

@kstoimenov You're using the LI pointer for all IgnoreAccess calls which is causing nullptr dereference warnings in static analyzer.

Should we just be using I or the dyn_cast<> pointers in each case?

https://llvm.org/reports/scan-build/report-AddressSanitizer.cpp-ignoreAccess-21-f37ec0.html#EndPath

[RKSimon]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

1521

@kstoimenov Have you been able to check this at all please?

[kstoimenov]

Sorry i just noticed this email. I will take a look next week.

clang/include/clang/Basic/CodeGenOptions.def

ENUM_CODEGENOPT(SanitizeAddressUseAfterReturn,
                 llvm::AsanDetectStackUseAfterReturnMode::Runtime
                 ) ///< Set detection mode for stack-use-after-return.
 CODEGENOPT(SanitizeAddressPoisonCustomArrayCookie, 1,
            0) ///< Enable poisoning operator new[] which is not a replaceable
               ///< global allocation function in AddressSanitizer
 CODEGENOPT(SanitizeAddressGlobalsDeadStripping, 1, 0) ///< Enable linker dead stripping
                                                       ///< of globals in AddressSanitizer
 CODEGENOPT(SanitizeAddressUseOdrIndicator, 1, 0) ///< Enable ODR indicator globals
+CODEGENOPT(SanitizeAddressUseStackSafety, 1, 0) ///< Use stack safety result.
 CODEGENOPT(SanitizeMemoryTrackOrigins, 2, 0) ///< Enable tracking origins in
                                              ///< MemorySanitizer
 ENUM_CODEGENOPT(SanitizeAddressDtor, llvm::AsanDtorKind, 2,
                 llvm::AsanDtorKind::Global)  ///< Set how ASan global
                                              ///< destructors are emitted.
 CODEGENOPT(SanitizeMemoryUseAfterDtor, 1, 0) ///< Enable use-after-delete detection
                                              ///< in MemorySanitizer
 CODEGENOPT(SanitizeCfiCrossDso, 1, 0) ///< Enable cross-dso support in CFI.

clang/include/clang/Driver/Options.td

def fno_sanitize_memory_track_origins : Flag<["-"], "fno-sanitize-memory-track-origins">,
                                         Flags<[CoreOption, NoXarchOption]>,
                                         HelpText<"Disable origins tracking in MemorySanitizer">;
 def fsanitize_address_outline_instrumentation : Flag<["-"], "fsanitize-address-outline-instrumentation">,
                                                 Group<f_clang_Group>,
                                                 HelpText<"Always generate function calls for address sanitizer instrumentation">;
 def fno_sanitize_address_outline_instrumentation : Flag<["-"], "fno-sanitize-address-outline-instrumentation">,
                                                    Group<f_clang_Group>,
                                                    HelpText<"Use default code inlining logic for the address sanitizer">;
+defm sanitize_address_use_stack_safety : BoolOption<"f", "sanitize-address-use-stack-safety",
+  CodeGenOpts<"SanitizeAddressUseStackSafety">, DefaultFalse,
+  PosFlag<SetTrue, [CC1Option], "Enable">, NegFlag<SetFalse, [], "Disable">,
+  BothFlags<[], "Use stack safety result to optimize instrumentation">>,
+  Group<f_clang_Group>;
 def fsanitize_hwaddress_experimental_aliasing
   : Flag<["-"], "fsanitize-hwaddress-experimental-aliasing">,
     Group<f_clang_Group>,
     HelpText<"Enable aliasing mode in HWAddressSanitizer">;
 def fno_sanitize_hwaddress_experimental_aliasing
   : Flag<["-"], "fno-sanitize-hwaddress-experimental-aliasing">,
     Group<f_clang_Group>, Flags<[CoreOption, NoXarchOption]>,
     HelpText<"Disable aliasing mode in HWAddressSanitizer">;

 def fopenmp_cuda_blocks_per_sm_EQ : Joined<["-"], "fopenmp-cuda-blocks-per-sm=">, Group<f_Group>,
   Flags<[CC1Option, NoArgumentUnused, HelpHidden]>;
 def fopenmp_cuda_teams_reduction_recs_num_EQ : Joined<["-"], "fopenmp-cuda-teams-reduction-recs-num=">, Group<f_Group>,
   Flags<[CC1Option, NoArgumentUnused, HelpHidden]>;
 def fopenmp_target_debug : Flag<["-"], "fopenmp-target-debug">, Group<f_Group>, Flags<[CC1Option, NoArgumentUnused]>,
   HelpText<"Enable debugging in the OpenMP offloading device RTL">;
 def fno_openmp_target_debug : Flag<["-"], "fno-openmp-target-debug">, Group<f_Group>, Flags<[NoArgumentUnused]>;
 def fopenmp_target_debug_EQ : Joined<["-"], "fopenmp-target-debug=">, Group<f_Group>, Flags<[CC1Option, NoArgumentUnused, HelpHidden]>;
-def fopenmp_assume_teams_oversubscription : Flag<["-"], "fopenmp-assume-teams-oversubscription">, 
+def fopenmp_assume_teams_oversubscription : Flag<["-"], "fopenmp-assume-teams-oversubscription">,
   Group<f_Group>, Flags<[CC1Option, NoArgumentUnused, HelpHidden]>;
-def fopenmp_assume_threads_oversubscription : Flag<["-"], "fopenmp-assume-threads-oversubscription">, 
+def fopenmp_assume_threads_oversubscription : Flag<["-"], "fopenmp-assume-threads-oversubscription">,
   Group<f_Group>, Flags<[CC1Option, NoArgumentUnused, HelpHidden]>;
-def fno_openmp_assume_teams_oversubscription : Flag<["-"], "fno-openmp-assume-teams-oversubscription">, 
+def fno_openmp_assume_teams_oversubscription : Flag<["-"], "fno-openmp-assume-teams-oversubscription">,
   Group<f_Group>, Flags<[CC1Option, NoArgumentUnused, HelpHidden]>;
-def fno_openmp_assume_threads_oversubscription : Flag<["-"], "fno-openmp-assume-threads-oversubscription">, 
+def fno_openmp_assume_threads_oversubscription : Flag<["-"], "fno-openmp-assume-threads-oversubscription">,
   Group<f_Group>, Flags<[CC1Option, NoArgumentUnused, HelpHidden]>;
 defm openmp_target_new_runtime: BoolFOption<"openmp-target-new-runtime",
   LangOpts<"OpenMPTargetNewRuntime">, DefaultFalse,
   PosFlag<SetTrue, [CC1Option], "Use the new bitcode library for OpenMP offloading">,
   NegFlag<SetFalse>>;
 defm openmp_optimistic_collapse : BoolFOption<"openmp-optimistic-collapse",
   LangOpts<"OpenMPOptimisticCollapse">, DefaultFalse,
   PosFlag<SetTrue, [CC1Option]>, NegFlag<SetFalse>, BothFlags<[NoArgumentUnused, HelpHidden]>>;

clang/include/clang/Driver/SanitizerArgs.h

class SanitizerArgs {
   bool SharedRuntime = false;
   bool AsanUseAfterScope = true;
   bool AsanPoisonCustomArrayCookie = false;
   bool AsanGlobalsDeadStripping = false;
   bool AsanUseOdrIndicator = false;
   bool AsanInvalidPointerCmp = false;
   bool AsanInvalidPointerSub = false;
   bool AsanOutlineInstrumentation = false;
+  bool AsanUseStackSafety = false;
   llvm::AsanDtorKind AsanDtorKind = llvm::AsanDtorKind::Invalid;
   std::string HwasanAbi;
   bool LinkRuntimes = true;
   bool LinkCXXRuntimes = false;
   bool NeedPIE = false;
   bool SafeStackRuntime = false;
   bool Stats = false;
   bool TsanMemoryAccess = true;

clang/lib/CodeGen/BackendUtil.cpp

auto ASanPass = [&](SanitizerMask Mask, bool CompileKernel) {
         bool UseAfterScope = CodeGenOpts.SanitizeAddressUseAfterScope;
         bool ModuleUseAfterScope = asanUseGlobalsGC(TargetTriple, CodeGenOpts);
         bool UseOdrIndicator = CodeGenOpts.SanitizeAddressUseOdrIndicator;
         llvm::AsanDtorKind DestructorKind =
             CodeGenOpts.getSanitizeAddressDtor();
         llvm::AsanDetectStackUseAfterReturnMode UseAfterReturn =
             CodeGenOpts.getSanitizeAddressUseAfterReturn();
         MPM.addPass(RequireAnalysisPass<ASanGlobalsMetadataAnalysis, Module>());
+        if (CodeGenOpts.SanitizeAddressUseStackSafety) {
+          MPM.addPass(RequireAnalysisPass<StackSafetyGlobalAnalysis, Module>());
+        }
         MPM.addPass(ModuleAddressSanitizerPass(
             CompileKernel, Recover, ModuleUseAfterScope, UseOdrIndicator,
             DestructorKind));
         MPM.addPass(createModuleToFunctionPassAdaptor(AddressSanitizerPass(
             {CompileKernel, Recover, UseAfterScope, UseAfterReturn})));
       }
     };
     ASanPass(SanitizerKind::Address, false);

clang/lib/Driver/SanitizerArgs.cpp

AsanPoisonCustomArrayCookie = Args.hasFlag(
         options::OPT_fno_sanitize_address_poison_custom_array_cookie,
         AsanPoisonCustomArrayCookie);
 
     AsanOutlineInstrumentation =
         Args.hasFlag(options::OPT_fsanitize_address_outline_instrumentation,
                      options::OPT_fno_sanitize_address_outline_instrumentation,
                      AsanOutlineInstrumentation);
 
+    AsanUseStackSafety = Args.hasFlag(
+        options::OPT_fsanitize_address_use_stack_safety,
+        options::OPT_fno_sanitize_address_use_stack_safety, AsanUseStackSafety);
+
     // As a workaround for a bug in gold 2.26 and earlier, dead stripping of
     // globals in ASan is disabled by default on ELF targets.
     // See https://sourceware.org/bugzilla/show_bug.cgi?id=19002
     AsanGlobalsDeadStripping =
         !TC.getTriple().isOSBinFormatELF() || TC.getTriple().isOSFuchsia() ||
         TC.getTriple().isPS4() ||
         Args.hasArg(options::OPT_fsanitize_address_globals_dead_stripping);
 
if (AsanInvalidPointerSub) {
     CmdArgs.push_back("-asan-detect-invalid-pointer-sub");
   }
 
   if (AsanOutlineInstrumentation) {
     CmdArgs.push_back("-mllvm");
     CmdArgs.push_back("-asan-instrumentation-with-call-threshold=0");
   }
 
+  if (AsanUseStackSafety) {
+    CmdArgs.push_back("-fsanitize-address-use-stack-safety");
+    CmdArgs.push_back("-mllvm");
+    CmdArgs.push_back("-asan-use-stack-safety");
+  }
+
   // Only pass the option to the frontend if the user requested,
   // otherwise the frontend will just use the codegen default.
   if (AsanDtorKind != llvm::AsanDtorKind::Invalid) {
     CmdArgs.push_back(Args.MakeArgString("-fsanitize-address-destructor=" +
                                          AsanDtorKindToString(AsanDtorKind)));
   }
 
   if (AsanUseAfterReturn != llvm::AsanDetectStackUseAfterReturnMode::Invalid) {

clang/test/CodeGen/asan-stack-safety-analysis.c

+// REQUIRES: x86_64-pc-linux-gnu
+
+// RUN: %clang -fno-legacy-pass-manager -fsanitize=address -fsanitize-address-outline-instrumentation -target x86_64-pc-linux-gnu -S -emit-llvm -mllvm -asan-use-stack-safety=true -O2 %s -o - | FileCheck %s --check-prefix=SAFETY
+// RUN: %clang -fno-legacy-pass-manager -fsanitize=address -fsanitize-address-outline-instrumentation -target x86_64-pc-linux-gnu -S -emit-llvm -mllvm -asan-use-stack-safety=false  -O2 %s -o - | FileCheck %s --check-prefix=NOSAFETY
+
+// RUN: %clang -flegacy-pass-manager -fsanitize=address -target x86_64-pc-linux-gnu -S -emit-llvm -mllvm -asan-use-stack-safety=true  -O2 %s -o - | FileCheck %s --check-prefix=SAFETY
+// RUN: %clang -flegacy-pass-manager -fsanitize=address -target x86_64-pc-linux-gnu -S -emit-llvm -mllvm -asan-use-stack-safety=false  -O2 %s -o - | FileCheck %s --check-prefix=NOSAFETY
+
+int main(int argc, char **argv) {
+  char buf[10];
+  volatile char *x = buf;
+  *x = 0;
+  return buf[0];
+  // NOSAFETY: call void @__asan_load1
+  // SAFETY-NOT: call void @__asan_load1
+}

clang/test/Driver/fsanitize.c

 // RUN:     FileCheck %s --check-prefix=CHECK-ASAN-OUTLINE-OK
 // CHECK-ASAN-OUTLINE-OK: "-mllvm" "-asan-instrumentation-with-call-threshold=0"
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=address -fno-sanitize-address-outline-instrumentation %s -### 2>&1 | \
 // RUN:     FileCheck %s --check-prefix=CHECK-NO-CHECK-ASAN-CALLBACK
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=address -fsanitize-address-outline-instrumentation -fno-sanitize-address-outline-instrumentation %s -### 2>&1 | \
 // RUN:     FileCheck %s --check-prefix=CHECK-NO-CHECK-ASAN-CALLBACK
 // CHECK-NO-CHECK-ASAN-CALLBACK-NOT: "-mllvm" "-asan-instrumentation-with-call-threshold=0"
 
+// RUN: %clang -target x86_64-linux-gnu -fsanitize-address-use-stack-safety %s -### 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=CHECK-ASAN-OUTLINE-STACK-SAFETY-WARN
+// CHECK-ASAN-OUTLINE-STACK-SAFETY-WARN: warning: argument unused during compilation: '-fsanitize-address-use-stack-safety'
+// RUN: %clang -target x86_64-linux-gnu -fsanitize=address -fsanitize-address-use-stack-safety %s -### 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=CHECK-ASAN-OUTLINE-STACK-SAFETY-OK
+// RUN: %clang -target x86_64-linux-gnu -fsanitize=address -fno-sanitize-address-use-stack-safety \
+// RUN:     -fsanitize-address-use-stack-safety %s -### 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=CHECK-ASAN-OUTLINE-STACK-SAFETY-OK
+// CHECK-ASAN-OUTLINE-STACK-SAFETY-OK: "-fsanitize-address-use-stack-safety" "-mllvm" "-asan-use-stack-safety"
+// RUN: %clang -target x86_64-linux-gnu -fsanitize=address -fno-sanitize-address-use-stack-safety %s -### 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=CHECK-NO-CHECK-ASAN-CALLBACK
+// RUN: %clang -target x86_64-linux-gnu -fsanitize=address -fsanitize-address-use-stack-safety \
+// RUN:     -fno-sanitize-address-use-stack-safety %s -### 2>&1 | \
+// RUN:     FileCheck %s --check-prefix=CHECK-NO-CHECK-ASAN-CALLBACK
+// CHECK-NO-CHECK-ASAN-CALLBACK-NOT: "-fsanitize-address-use-stack-safety" "-mllvm" "-asan-use-stack-safety"
+
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=address -fsanitize-address-use-odr-indicator %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-ASAN-ODR-INDICATOR
 // RUN: %clang_cl --target=x86_64-windows -fsanitize=address -fsanitize-address-use-odr-indicator -### -- %s 2>&1 | FileCheck %s --check-prefix=CHECK-ASAN-ODR-INDICATOR
 // CHECK-ASAN-ODR-INDICATOR: -cc1{{.*}}-fsanitize-address-use-odr-indicator
 
 // RUN: %clang -target x86_64-linux-gnu -fsanitize=address -fno-sanitize-address-use-odr-indicator %s -### 2>&1 | FileCheck %s --check-prefix=CHECK-ASAN-ODR-INDICATOR-OFF
 // RUN: %clang_cl --target=x86_64-windows -fsanitize=address -fno-sanitize-address-use-odr-indicator -### -- %s 2>&1 | FileCheck %s --check-prefix=CHECK-ASAN-ODR-INDICATOR-OFF
 // CHECK-ASAN-ODR-INDICATOR-OFF-NOT: -cc1{{.*}}address-generate-odr-globals
 

llvm/include/llvm/Analysis/StackSafetyAnalysis.h

public:
 
   // Returns true if the instruction can be proven to do only two types of
   // memory accesses:
   //  (1) live stack locations in-bounds or
   //  (2) non-stack locations.
   bool stackAccessIsSafe(const Instruction &I) const;
   void print(raw_ostream &O) const;
   void dump() const;
+
+  /// Handle invalidation from the pass manager.
+  /// These results are never invalidated.
+  bool invalidate(Module &, const PreservedAnalyses &,
+                  ModuleAnalysisManager::Invalidator &) {
+    return false;
+  }
+  bool invalidate(Function &, const PreservedAnalyses &,
+                  FunctionAnalysisManager::Invalidator &) {
+    return false;
+  }
 };
 
 /// StackSafetyInfo wrapper for the new pass manager.
 class StackSafetyAnalysis : public AnalysisInfoMixin<StackSafetyAnalysis> {
   friend AnalysisInfoMixin<StackSafetyAnalysis>;
   static AnalysisKey Key;
 
 public:

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Analysis/MemoryBuiltins.h"
+#include "llvm/Analysis/StackSafetyAnalysis.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/BinaryFormat/MachO.h"
 #include "llvm/IR/Argument.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/BasicBlock.h"
 #include "llvm/IR/Comdat.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalAlias.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InlineAsm.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/MDBuilder.h"

 static cl::opt<bool> ClInstrumentReads("asan-instrument-reads",
                                        cl::desc("instrument read instructions"),
                                        cl::Hidden, cl::init(true));
 
 static cl::opt<bool> ClInstrumentWrites(
     "asan-instrument-writes", cl::desc("instrument write instructions"),
     cl::Hidden, cl::init(true));
 
+static cl::opt<bool>
+    ClUseStackSafety("asan-use-stack-safety", cl::Hidden, cl::init(false),
+                     cl::Hidden, cl::desc("Use Stack Safety analysis results"),
+                     cl::Optional);
+
 static cl::opt<bool> ClInstrumentAtomics(
     "asan-instrument-atomics",
     cl::desc("instrument atomic instructions (rmw, cmpxchg)"), cl::Hidden,
     cl::init(true));
 
 static cl::opt<bool>
     ClInstrumentByval("asan-instrument-byval",
                       cl::desc("instrument byval call arguments"), cl::Hidden,
private:
   GlobalsMetadata GlobalsMD;
 };
 
 char ASanGlobalsMetadataWrapperPass::ID = 0;
 
 /// AddressSanitizer: instrument the code in module to find memory bugs.
 struct AddressSanitizer {
   AddressSanitizer(Module &M, const GlobalsMetadata *GlobalsMD,
-                   bool CompileKernel = false, bool Recover = false,
-                   bool UseAfterScope = false,
+                   const StackSafetyGlobalInfo *SSI, bool CompileKernel = false,
+                   bool Recover = false, bool UseAfterScope = false,
                    AsanDetectStackUseAfterReturnMode UseAfterReturn =
                        AsanDetectStackUseAfterReturnMode::Runtime)
       : CompileKernel(ClEnableKasan.getNumOccurrences() > 0 ? ClEnableKasan
                                                             : CompileKernel),
         Recover(ClRecover.getNumOccurrences() > 0 ? ClRecover : Recover),
         UseAfterScope(UseAfterScope || ClUseAfterScope),
         UseAfterReturn(ClUseAfterReturn.getNumOccurrences() ? ClUseAfterReturn
                                                             : UseAfterReturn),
-        GlobalsMD(*GlobalsMD) {
+        GlobalsMD(*GlobalsMD), SSI(SSI) {
+    assert(GlobalsMD != nullptr);
     C = &(M.getContext());
     LongSize = M.getDataLayout().getPointerSizeInBits();
     IntptrTy = Type::getIntNTy(*C, LongSize);
     Int8PtrTy = Type::getInt8PtrTy(*C);
     Int32Ty = Type::getInt32Ty(*C);
     TargetTriple = Triple(M.getTargetTriple());
 
     Mapping = getShadowMapping(TargetTriple, LongSize, this->CompileKernel);
uint64_t getAllocaSizeInBytes(const AllocaInst &AI) const {
     uint64_t SizeInBytes =
         AI.getModule()->getDataLayout().getTypeAllocSize(Ty);
     return SizeInBytes * ArraySize;
   }
 
   /// Check if we want (and can) handle this alloca.
   bool isInterestingAlloca(const AllocaInst &AI);
 
-  bool ignoreAccess(Value *Ptr);
+  bool ignoreAccess(Instruction *Inst, Value *Ptr);
   void getInterestingMemoryOperands(
       Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting);
 
   void instrumentMop(ObjectSizeOffsetVisitor &ObjSizeVis,
                      InterestingMemoryOperand &O, bool UseCalls,
                      const DataLayout &DL);
   void instrumentPointerComparisonOrSubtraction(Instruction *I);
   void instrumentAddress(Instruction *OrigIns, Instruction *InsertBefore,
private:
 
   // These arrays is indexed by AccessIsWrite and Experiment.
   FunctionCallee AsanErrorCallbackSized[2][2];
   FunctionCallee AsanMemoryAccessCallbackSized[2][2];
 
   FunctionCallee AsanMemmove, AsanMemcpy, AsanMemset;
   Value *LocalDynamicShadow = nullptr;
   const GlobalsMetadata &GlobalsMD;
+  const StackSafetyGlobalInfo *SSI;
   DenseMap<const AllocaInst *, bool> ProcessedAllocas;
 
   FunctionCallee AMDGPUAddressShared;
   FunctionCallee AMDGPUAddressPrivate;
 };
 
 class AddressSanitizerLegacyPass : public FunctionPass {
 public:
public:
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<ASanGlobalsMetadataWrapperPass>();
     AU.addRequired<TargetLibraryInfoWrapperPass>();
   }
 
   bool runOnFunction(Function &F) override {
     GlobalsMetadata &GlobalsMD =
         getAnalysis<ASanGlobalsMetadataWrapperPass>().getGlobalsMD();
     const TargetLibraryInfo *TLI =
         &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
-    AddressSanitizer ASan(*F.getParent(), &GlobalsMD, CompileKernel, Recover,
-                          UseAfterScope, UseAfterReturn);
+    AddressSanitizer ASan(*F.getParent(), &GlobalsMD, nullptr, CompileKernel,
+                          Recover, UseAfterScope, UseAfterReturn);
     return ASan.instrumentFunction(F, TLI);
   }
 
 private:
   bool CompileKernel;
   bool Recover;
   bool UseAfterScope;
   AsanDetectStackUseAfterReturnMode UseAfterReturn;
ConstantInt *IsExcluded =
         mdconst::extract<ConstantInt>(MDN->getOperand(4));
     E.IsExcluded |= IsExcluded->isOne();
   }
 }
 
 AnalysisKey ASanGlobalsMetadataAnalysis::Key;
 
 GlobalsMetadata ASanGlobalsMetadataAnalysis::run(Module &M,
                                                  ModuleAnalysisManager &AM) {
   return GlobalsMetadata(M);
 }
 
 PreservedAnalyses AddressSanitizerPass::run(Function &F,
                                             AnalysisManager<Function> &AM) {
   auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
   Module &M = *F.getParent();
   if (auto *R = MAMProxy.getCachedResult<ASanGlobalsMetadataAnalysis>(M)) {
     const TargetLibraryInfo *TLI = &AM.getResult<TargetLibraryAnalysis>(F);
-    AddressSanitizer Sanitizer(M, R, Options.CompileKernel, Options.Recover,
-                               Options.UseAfterScope, Options.UseAfterReturn);
+    const StackSafetyGlobalInfo *SSI =
+        MAMProxy.getCachedResult<StackSafetyGlobalAnalysis>(M);
+    AddressSanitizer Sanitizer(M, R, SSI, Options.CompileKernel,
+                               Options.Recover, Options.UseAfterScope,
+                               Options.UseAfterReturn);
     if (Sanitizer.instrumentFunction(F, TLI))
       return PreservedAnalyses::none();
     return PreservedAnalyses::all();
   }
 
   report_fatal_error(
       "The ASanGlobalsMetadataAnalysis is required to run before "
       "AddressSanitizer can run");
ModuleAddressSanitizerPass::ModuleAddressSanitizerPass(
     bool CompileKernel, bool Recover, bool UseGlobalGC, bool UseOdrIndicator,
     AsanDtorKind DestructorKind)
     : CompileKernel(CompileKernel), Recover(Recover), UseGlobalGC(UseGlobalGC),
       UseOdrIndicator(UseOdrIndicator), DestructorKind(DestructorKind) {}
 
 PreservedAnalyses ModuleAddressSanitizerPass::run(Module &M,
                                                   AnalysisManager<Module> &AM) {
   GlobalsMetadata &GlobalsMD = AM.getResult<ASanGlobalsMetadataAnalysis>(M);
   ModuleAddressSanitizer Sanitizer(M, &GlobalsMD, CompileKernel, Recover,
                                    UseGlobalGC, UseOdrIndicator,
                                    DestructorKind);
   if (Sanitizer.instrumentModule(M))
     return PreservedAnalyses::none();
   return PreservedAnalyses::all();
 }
 
 INITIALIZE_PASS(ASanGlobalsMetadataWrapperPass, "asan-globals-md",
                 "Read metadata to mark which globals should be instrumented "
                 "when running ASan.",
                 false, true)
 
 char AddressSanitizerLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(
     AddressSanitizerLegacyPass, "asan",
     "AddressSanitizer: detects use-after-free and out-of-bounds bugs.", false,
     false)
 INITIALIZE_PASS_DEPENDENCY(ASanGlobalsMetadataWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_END(
     AddressSanitizerLegacyPass, "asan",
     "AddressSanitizer: detects use-after-free and out-of-bounds bugs.", false,
     false)
 
 FunctionPass *llvm::createAddressSanitizerFunctionPass(
     bool CompileKernel, bool Recover, bool UseAfterScope,
     AsanDetectStackUseAfterReturnMode UseAfterReturn) {
bool IsInteresting =
        !AI.isUsedWithInAlloca() &&
        // swifterror allocas are register promoted by ISel
        !AI.isSwiftError());
 
   ProcessedAllocas[&AI] = IsInteresting;
   return IsInteresting;
 }
 
-bool AddressSanitizer::ignoreAccess(Value *Ptr) {
+bool AddressSanitizer::ignoreAccess(Instruction *Inst, Value *Ptr) {
   // Instrument acesses from different address spaces only for AMDGPU.
   Type *PtrTy = cast<PointerType>(Ptr->getType()->getScalarType());
   if (PtrTy->getPointerAddressSpace() != 0 &&
       !(TargetTriple.isAMDGPU() && !isUnsupportedAMDGPUAddrspace(Ptr)))
     return true;
 
   // Ignore swifterror addresses.
   // swifterror memory addresses are mem2reg promoted by instruction
   // selection. As such they cannot have regular uses like an instrumentation
   // function and it makes no sense to track them as memory.
   if (Ptr->isSwiftError())
     return true;
 
   // Treat memory accesses to promotable allocas as non-interesting since they
   // will not cause memory violations. This greatly speeds up the instrumented
   // executable at -O0.
   if (auto AI = dyn_cast_or_null<AllocaInst>(Ptr))
     if (ClSkipPromotableAllocas && !isInterestingAlloca(*AI))
       return true;
 
+  if (ClUseStackSafety && findAllocaForValue(Ptr)) {
+    if (SSI && SSI->stackAccessIsSafe(*Inst)) {
+      return true;
+    }
+  }
+
   return false;
 }
 
 void AddressSanitizer::getInterestingMemoryOperands(
     Instruction *I, SmallVectorImpl<InterestingMemoryOperand> &Interesting) {
   // Skip memory accesses inserted by another instrumentation.
   if (I->hasMetadata("nosanitize"))
     return;
 
   // Do not instrument the load fetching the dynamic shadow address.
   if (LocalDynamicShadow == I)
     return;
 
   if (LoadInst *LI = dyn_cast<LoadInst>(I)) {
-    if (!ClInstrumentReads || ignoreAccess(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(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(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(XCHG->getPointerOperand()))
+    if (!ClInstrumentAtomics || ignoreAccess(I, XCHG->getPointerOperand()))
       return;
     Interesting.emplace_back(I, XCHG->getPointerOperandIndex(), true,
                              XCHG->getCompareOperand()->getType(), None);
   } else if (auto CI = dyn_cast<CallInst>(I)) {
     auto *F = CI->getCalledFunction();
     if (F && (F->getName().startswith("llvm.masked.load.") ||
               F->getName().startswith("llvm.masked.store."))) {
       bool IsWrite = F->getName().startswith("llvm.masked.store.");
       // Masked store has an initial operand for the value.
       unsigned OpOffset = IsWrite ? 1 : 0;
       if (IsWrite ? !ClInstrumentWrites : !ClInstrumentReads)
         return;
 
       auto BasePtr = CI->getOperand(OpOffset);
-      if (ignoreAccess(BasePtr))
+      if (ignoreAccess(I, BasePtr))
         return;
       auto Ty = cast<PointerType>(BasePtr->getType())->getElementType();
       MaybeAlign Alignment = Align(1);
       // Otherwise no alignment guarantees. We probably got Undef.
       if (auto *Op = dyn_cast<ConstantInt>(CI->getOperand(1 + OpOffset)))
         Alignment = Op->getMaybeAlignValue();
       Value *Mask = CI->getOperand(2 + OpOffset);
       Interesting.emplace_back(I, OpOffset, IsWrite, Ty, Alignment, Mask);
     } else {
       for (unsigned ArgNo = 0; ArgNo < CI->arg_size(); ArgNo++) {
         if (!ClInstrumentByval || !CI->isByValArgument(ArgNo) ||
-            ignoreAccess(CI->getArgOperand(ArgNo)))
+            ignoreAccess(I, CI->getArgOperand(ArgNo)))
           continue;
         Type *Ty = CI->getParamByValType(ArgNo);
         Interesting.emplace_back(I, ArgNo, false, Ty, Align(1));
       }
     }
   }
 }