From de4e8bc3ace3eb64ac92166617056424f5bd5996 Mon Sep 17 00:00:00 2001
From: Matt Morehouse <mascasa@google.com>
Date: Tue, 1 Feb 2022 11:23:36 -0800
Subject: [PATCH] [HWASan] Properly handle musttail calls.

Fixes a compile error when the `clang::musttail` attribute is used.

Reviewed By: eugenis

Differential Revision: https://reviews.llvm.org/D118712
---
 compiler-rt/test/hwasan/TestCases/musttail.cpp             | 13 +++++++++++++
 llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp |  9 +++++++--
 2 files changed, 20 insertions(+), 2 deletions(-)
 create mode 100644 compiler-rt/test/hwasan/TestCases/musttail.cpp

diff --git a/compiler-rt/test/hwasan/TestCases/musttail.cpp b/compiler-rt/test/hwasan/TestCases/musttail.cpp
new file mode 100644
index 0000000..89b58e6
--- /dev/null
+++ b/compiler-rt/test/hwasan/TestCases/musttail.cpp
@@ -0,0 +1,13 @@
+// RUN: %clangxx_hwasan -mllvm -hwasan-use-stack-safety=0 %s -o %t
+// RUN: %run %t
+//
+// REQUIRES: pointer-tagging
+
+__attribute__((noinline)) int bar(int X) { return X; }
+
+__attribute__((noinline)) int foo(int X) {
+  volatile int A = 5;
+  [[clang::musttail]] return bar(X + A);
+}
+
+int main(int Argc, char *Argv[]) { return foo(Argc) != 6; }
diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index fb10a99..70b3944 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -1531,9 +1531,14 @@ bool HWAddressSanitizer::sanitizeFunction(
         }
       }
 
-      if (isa<ReturnInst>(Inst) || isa<ResumeInst>(Inst) ||
-          isa<CleanupReturnInst>(Inst))
+      if (isa<ReturnInst>(Inst)) {
+        if (CallInst *CI = Inst.getParent()->getTerminatingMustTailCall())
+          RetVec.push_back(CI);
+        else
+          RetVec.push_back(&Inst);
+      } else if (isa<ResumeInst, CleanupReturnInst>(Inst)) {
         RetVec.push_back(&Inst);
+      }
 
       if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
         for (Value *V : DVI->location_ops()) {
-- 
2.7.4