[llvm] [HWASan] optimize AttrInfer fix for selective HWASan (PR #108111)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 10 16:07:09 PDT 2024
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-compiler-rt-sanitizer
Author: Florian Mayer (fmayer)
<details>
<summary>Changes</summary>
---
Full diff: https://github.com/llvm/llvm-project/pull/108111.diff
2 Files Affected:
- (modified) llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp (+107-44)
- (added) llvm/test/Instrumentation/HWAddressSanitizer/mem-attr-declr.ll (+82)
``````````diff
diff --git a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
index a7e7f9a570dac7..bdd0764e1e5e15 100644
--- a/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/HWAddressSanitizer.cpp
@@ -54,6 +54,7 @@
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/MD5.h"
+#include "llvm/Support/ModRef.h"
#include "llvm/Support/RandomNumberGenerator.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/TargetParser/Triple.h"
@@ -301,7 +302,8 @@ class HWAddressSanitizer {
initializeModule();
}
- void sanitizeFunction(Function &F, FunctionAnalysisManager &FAM);
+ bool sanitizeFunction(Function &F, FunctionAnalysisManager &FAM);
+ void unsanitizeFunction(Function &F);
private:
struct ShadowTagCheckInfo {
@@ -316,6 +318,8 @@ class HWAddressSanitizer {
FunctionAnalysisManager &FAM) const;
void initializeModule();
void createHwasanCtorComdat();
+ void removeFnAttributes(Function *F);
+ void restoreFnAttributes(CallInst *CI);
void initializeCallbacks(Module &M);
@@ -450,6 +454,13 @@ class HWAddressSanitizer {
Value *StackBaseTag = nullptr;
Value *CachedFP = nullptr;
GlobalValue *ThreadPtrGlobal = nullptr;
+
+ struct ChangedFn {
+ std::optional<MemoryEffects> OriginalME;
+ SmallVector<size_t, 2> RemovedArgAttrs;
+ };
+
+ SmallMapVector<const Function *, ChangedFn, 2> ChangedFns;
};
} // end anonymous namespace
@@ -466,8 +477,15 @@ PreservedAnalyses HWAddressSanitizerPass::run(Module &M,
HWAddressSanitizer HWASan(M, Options.CompileKernel, Options.Recover, SSI);
auto &FAM = MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
- for (Function &F : M)
- HWASan.sanitizeFunction(F, FAM);
+ SmallVector<Function *, 5> Unsanitize;
+ for (Function &F : M) {
+ if (!HWASan.sanitizeFunction(F, FAM))
+ Unsanitize.emplace_back(&F);
+ }
+
+ for (Function *F : Unsanitize) {
+ HWASan.unsanitizeFunction(*F);
+ }
PreservedAnalyses PA = PreservedAnalyses::none();
// DominatorTreeAnalysis, PostDominatorTreeAnalysis, and LoopAnalysis
@@ -591,47 +609,81 @@ void HWAddressSanitizer::createHwasanCtorComdat() {
appendToCompilerUsed(M, Dummy);
}
-/// Module-level initialization.
-///
-/// inserts a call to __hwasan_init to the module's constructor list.
-void HWAddressSanitizer::initializeModule() {
- LLVM_DEBUG(dbgs() << "Init " << M.getName() << "\n");
- TargetTriple = Triple(M.getTargetTriple());
-
- for (auto &F : M.functions()) {
- // Remove memory attributes that are invalid with HWASan.
- // HWASan checks read from shadow, which invalidates memory(argmem: *)
- // Short granule checks on function arguments read from the argument memory
- // (last byte of the granule), which invalidates writeonly.
- //
- // This is not only true for sanitized functions, because AttrInfer can
- // infer those attributes on libc functions, which is not true if those
- // are instrumented (Android) or intercepted.
-
- // The API is weird. `onlyReadsMemory` actually means "does not write", and
- // `onlyWritesMemory` actually means "does not read". So we reconstruct
- // "accesses memory" && "does not read" <=> "writes".
- bool Changed = false;
- if (!F.doesNotAccessMemory()) {
- bool WritesMemory = !F.onlyReadsMemory();
- bool ReadsMemory = !F.onlyWritesMemory();
- if ((WritesMemory && !ReadsMemory) || F.onlyAccessesArgMemory()) {
- F.removeFnAttr(Attribute::Memory);
- Changed = true;
- }
+void HWAddressSanitizer::removeFnAttributes(Function *F) {
+ if (!F || ChangedFns.contains(F))
+ return;
+ // Remove memory attributes that are invalid with HWASan.
+ // HWASan checks read from shadow, which invalidates memory(argmem: *)
+ // Short granule checks on function arguments read from the argument memory
+ // (last byte of the granule), which invalidates writeonly.
+ //
+ // This is not only true for sanitized functions, because AttrInfer can
+ // infer those attributes on libc functions, which is not true if those
+ // are instrumented (Android) or intercepted.
+ //
+ // We might want to model HWASan shadow memory more opaquely to get rid of
+ // this problem altogether, by hiding the shadow memory write in an
+ // intrinsic, essentially like in the AArch64StackTagging pass. But that's
+ // for another day.
+
+ // The API is weird. `onlyReadsMemory` actually means "does not write", and
+ // `onlyWritesMemory` actually means "does not read". So we reconstruct
+ // "accesses memory" && "does not read" <=> "writes".
+ bool Changed = false;
+ if (!F->doesNotAccessMemory()) {
+ bool WritesMemory = !F->onlyReadsMemory();
+ bool ReadsMemory = !F->onlyWritesMemory();
+ if ((WritesMemory && !ReadsMemory) || F->onlyAccessesArgMemory()) {
+ ChangedFns[F].OriginalME = F->getMemoryEffects();
+ F->removeFnAttr(Attribute::Memory);
+ Changed = true;
}
- for (Argument &A : F.args()) {
+ }
+ if (UseShortGranules) {
+ for (Argument &A : F->args()) {
if (A.hasAttribute(Attribute::WriteOnly)) {
- Changed = true;
A.removeAttr(Attribute::WriteOnly);
+ ChangedFns[F].RemovedArgAttrs.emplace_back(A.getArgNo());
+ Changed = true;
}
}
- if (Changed) {
- // nobuiltin makes sure later passes don't restore assumptions about
- // the function.
- F.addFnAttr(Attribute::NoBuiltin);
- }
}
+ if (Changed) {
+ // nobuiltin makes sure later passes don't restore assumptions about
+ // the function.
+ F->addFnAttr(Attribute::NoBuiltin);
+ }
+}
+
+void HWAddressSanitizer::restoreFnAttributes(CallInst *CI) {
+ auto It = ChangedFns.find(CI->getCalledFunction());
+ if (It == ChangedFns.end())
+ return;
+ // The memory effects are stored as a bitmap where a bit e.g. means
+ // "may write memory". CI->getMemoryEffects is the AND (i.e. intersection)
+ // of the bits of the CallInst and the underlying function. So, we removed
+ // the original ME from the function, which means the bitmap there all 1.
+ // Now we are adding the original mask back to the CallInst, which means
+ // CI->getMemoryEffects is now OriginalME & all 1 = OriginalME.
+ //
+ // Restoring the memory access is slightly dubious, because the called
+ // function will still read the memory. But because the function does not
+ // operate on the shadow memory, there should not be any difficulties with
+ // removing the shadow write, or reordering it.
+ ChangedFn &CFn = It->second;
+ if (CFn.OriginalME.has_value())
+ CI->setMemoryEffects(*CFn.OriginalME);
+ for (size_t I : CFn.RemovedArgAttrs)
+ CI->addAttributeAtIndex(I + AttributeList::FirstArgIndex,
+ Attribute::WriteOnly);
+}
+
+/// Module-level initialization.
+///
+/// inserts a call to __hwasan_init to the module's constructor list.
+void HWAddressSanitizer::initializeModule() {
+ LLVM_DEBUG(dbgs() << "Init " << M.getName() << "\n");
+ TargetTriple = Triple(M.getTargetTriple());
// x86_64 currently has two modes:
// - Intel LAM (default)
@@ -895,6 +947,7 @@ void HWAddressSanitizer::getInterestingMemoryOperands(
Interesting.emplace_back(I, ArgNo, false, Ty, Align(1));
}
maybeMarkSanitizerLibraryCallNoBuiltin(CI, &TLI);
+ removeFnAttributes(CI->getCalledFunction());
}
}
@@ -1593,16 +1646,23 @@ bool HWAddressSanitizer::selectiveInstrumentationShouldSkip(
return Skip;
}
-void HWAddressSanitizer::sanitizeFunction(Function &F,
+void HWAddressSanitizer::unsanitizeFunction(Function &F) {
+ for (auto &Inst : instructions(F)) {
+ if (CallInst *CI = dyn_cast<CallInst>(&Inst))
+ restoreFnAttributes(CI);
+ }
+}
+
+bool HWAddressSanitizer::sanitizeFunction(Function &F,
FunctionAnalysisManager &FAM) {
if (&F == HwasanCtorFunction)
- return;
+ return false;
if (!F.hasFnAttribute(Attribute::SanitizeHWAddress))
- return;
+ return false;
if (F.empty())
- return;
+ return false;
NumTotalFuncs++;
@@ -1610,7 +1670,9 @@ void HWAddressSanitizer::sanitizeFunction(Function &F,
FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
if (selectiveInstrumentationShouldSkip(F, FAM))
- return;
+ return false;
+
+ removeFnAttributes(&F);
NumInstrumentedFuncs++;
@@ -1653,7 +1715,7 @@ void HWAddressSanitizer::sanitizeFunction(Function &F,
if (SInfo.AllocasToInstrument.empty() && OperandsToInstrument.empty() &&
IntrinToInstrument.empty())
- return;
+ return false;
assert(!ShadowBase);
@@ -1702,6 +1764,7 @@ void HWAddressSanitizer::sanitizeFunction(Function &F,
ShadowBase = nullptr;
StackBaseTag = nullptr;
CachedFP = nullptr;
+ return true;
}
void HWAddressSanitizer::instrumentGlobal(GlobalVariable *GV, uint8_t Tag) {
diff --git a/llvm/test/Instrumentation/HWAddressSanitizer/mem-attr-declr.ll b/llvm/test/Instrumentation/HWAddressSanitizer/mem-attr-declr.ll
new file mode 100644
index 00000000000000..d8965730b1e0b3
--- /dev/null
+++ b/llvm/test/Instrumentation/HWAddressSanitizer/mem-attr-declr.ll
@@ -0,0 +1,82 @@
+; Test that HWASan remove writeonly and memory(*) attributes from instrumented functions.
+; RUN: opt -S -passes=hwasan %s | FileCheck %s
+
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128-Fn32"
+target triple = "aarch64-unknown-linux-android30"
+
+; CHECK: declare dso_local void @test_argmemwriteonly(ptr nocapture noundef) local_unnamed_addr #0
+declare dso_local void @test_argmemwriteonly(ptr nocapture noundef writeonly) local_unnamed_addr #0
+
+; CHECK: declare dso_local void @test_argmemwriteonly2(ptr nocapture noundef) local_unnamed_addr #0
+declare dso_local void @test_argmemwriteonly2(ptr nocapture noundef) local_unnamed_addr #0
+
+; CHECK: declare dso_local void @test_writeonly(ptr nocapture noundef) local_unnamed_addr #0
+declare dso_local void @test_writeonly(ptr nocapture noundef) local_unnamed_addr #1
+
+; CHECK: declare dso_local void @test_writeonly2(ptr nocapture noundef) local_unnamed_addr #0
+declare dso_local void @test_writeonly2(ptr nocapture noundef writeonly) local_unnamed_addr #1
+
+; CHECK: declare dso_local void @test_argmemreadonly(ptr nocapture noundef readonly) local_unnamed_addr #0
+declare dso_local void @test_argmemreadonly(ptr nocapture noundef readonly) local_unnamed_addr #2
+
+; CHECK: declare dso_local void @test_argmemreadonly2(ptr nocapture noundef) local_unnamed_addr #0
+declare dso_local void @test_argmemreadonly2(ptr nocapture noundef) local_unnamed_addr #2
+
+; CHECK: declare dso_local void @test_readonly(ptr nocapture noundef) local_unnamed_addr #1
+declare dso_local void @test_readonly(ptr nocapture noundef) local_unnamed_addr #3
+
+; CHECK: declare dso_local void @test_readonly2(ptr nocapture noundef readonly) local_unnamed_addr #1
+declare dso_local void @test_readonly2(ptr nocapture noundef readonly) local_unnamed_addr #3
+
+; CHECK: declare dso_local void @test_writeonly_notcalled_from_sanitized(ptr nocapture noundef readonly) local_unnamed_addr #2
+declare dso_local void @test_writeonly_notcalled_from_sanitized(ptr nocapture noundef readonly) local_unnamed_addr #1
+
+define dso_local void @test_sanitized(ptr %p) sanitize_hwaddress {
+entry:
+ call void @test_argmemwriteonly(ptr %p)
+ call void @test_argmemwriteonly2(ptr %p)
+ call void @test_writeonly(ptr %p)
+ call void @test_writeonly2(ptr %p)
+ call void @test_argmemreadonly(ptr %p)
+ call void @test_argmemreadonly2(ptr %p)
+ call void @test_readonly(ptr %p)
+ call void @test_readonly2(ptr %p)
+ ret void
+}
+
+define dso_local void @test_not_sanitized(ptr %p) {
+entry:
+; CHECK: call void @test_argmemwriteonly(ptr writeonly %p) #5
+ call void @test_argmemwriteonly(ptr %p)
+; CHECK: call void @test_argmemwriteonly2(ptr %p) #5
+ call void @test_argmemwriteonly2(ptr %p)
+; CHECK: call void @test_writeonly(ptr %p) #6
+ call void @test_writeonly(ptr %p)
+; CHECK: call void @test_writeonly2(ptr writeonly %p) #6
+ call void @test_writeonly2(ptr %p)
+; CHECK: call void @test_argmemreadonly(ptr %p) #7
+ call void @test_argmemreadonly(ptr %p)
+; CHECK: call void @test_argmemreadonly2(ptr %p) #7
+ call void @test_argmemreadonly2(ptr %p)
+; CHECK: call void @test_readonly(ptr %p)
+ call void @test_readonly(ptr %p)
+; CHECK: call void @test_readonly2(ptr %p)
+ call void @test_readonly2(ptr %p)
+; CHECK: call void @test_writeonly_notcalled_from_sanitized(ptr %p)
+ call void @test_writeonly_notcalled_from_sanitized(ptr %p)
+ ret void
+}
+
+; CHECK: attributes #0 = { nobuiltin uwtable }
+; CHECK: attributes #1 = { memory(read) uwtable }
+; CHECK: attributes #2 = { memory(write) uwtable }
+; CHECK: attributes #3 = { sanitize_hwaddress }
+; CHECK: attributes #4 = { nounwind }
+; CHECK: attributes #5 = { memory(argmem: write) }
+; CHECK: attributes #6 = { memory(write) }
+; CHECK: attributes #7 = { memory(argmem: read) }
+
+attributes #0 = { memory(argmem: write) uwtable }
+attributes #1 = { memory(write) uwtable }
+attributes #2 = { memory(argmem: read) uwtable }
+attributes #3 = { memory(read) uwtable }
``````````
</details>
https://github.com/llvm/llvm-project/pull/108111
More information about the llvm-commits
mailing list