xref: /llvm-project/llvm/lib/Transforms/Utils/MemoryTaggingSupport.cpp (revision 85c17e40926132575d1b98ca1a36b8394fe511cd) 
1 //== MemoryTaggingSupport.cpp - helpers for memory tagging implementations ===//
2 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
3 // See https://llvm.org/LICENSE.txt for license information.
4 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
5 //
6 //===----------------------------------------------------------------------===//
7 //
8 // This file declares common infrastructure for HWAddressSanitizer and
9 // Aarch64StackTagging.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
14 
15 #include "llvm/ADT/STLExtras.h"
16 #include "llvm/Analysis/CFG.h"
17 #include "llvm/Analysis/PostDominators.h"
18 #include "llvm/Analysis/StackSafetyAnalysis.h"
19 #include "llvm/Analysis/ValueTracking.h"
20 #include "llvm/BinaryFormat/Dwarf.h"
21 #include "llvm/IR/BasicBlock.h"
22 #include "llvm/IR/IRBuilder.h"
23 #include "llvm/IR/IntrinsicInst.h"
24 #include "llvm/TargetParser/Triple.h"
25 #include "llvm/Transforms/Utils/PromoteMemToReg.h"
26 
27 namespace llvm {
28 namespace memtag {
29 namespace {
30 bool maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts,
31                                  const DominatorTree *DT, const LoopInfo *LI,
32                                  size_t MaxLifetimes) {
33   // If we have too many lifetime ends, give up, as the algorithm below is N^2.
34   if (Insts.size() > MaxLifetimes)
35     return true;
36   for (size_t I = 0; I < Insts.size(); ++I) {
37     for (size_t J = 0; J < Insts.size(); ++J) {
38       if (I == J)
39         continue;
40       if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, DT, LI))
41         return true;
42     }
43   }
44   return false;
45 }
46 } // namespace
47 
48 bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT,
49                           const LoopInfo &LI, const Instruction *Start,
50                           const SmallVectorImpl<IntrinsicInst *> &Ends,
51                           const SmallVectorImpl<Instruction *> &RetVec,
52                           llvm::function_ref<void(Instruction *)> Callback) {
53   if (Ends.size() == 1 && PDT.dominates(Ends[0], Start)) {
54     Callback(Ends[0]);
55     return true;
56   }
57   SmallPtrSet<BasicBlock *, 2> EndBlocks;
58   for (auto *End : Ends) {
59     EndBlocks.insert(End->getParent());
60   }
61   SmallVector<Instruction *, 8> ReachableRetVec;
62   unsigned NumCoveredExits = 0;
63   for (auto *RI : RetVec) {
64     if (!isPotentiallyReachable(Start, RI, nullptr, &DT, &LI))
65       continue;
66     ReachableRetVec.push_back(RI);
67     // If there is an end in the same basic block as the return, we know for
68     // sure that the return is covered. Otherwise, we can check whether there
69     // is a way to reach the RI from the start of the lifetime without passing
70     // through an end.
71     if (EndBlocks.contains(RI->getParent()) ||
72         !isPotentiallyReachable(Start, RI, &EndBlocks, &DT, &LI)) {
73       ++NumCoveredExits;
74     }
75   }
76   if (NumCoveredExits == ReachableRetVec.size()) {
77     for_each(Ends, Callback);
78   } else {
79     // If there's a mix of covered and non-covered exits, just put the untag
80     // on exits, so we avoid the redundancy of untagging twice.
81     for_each(ReachableRetVec, Callback);
82     // We may have inserted untag outside of the lifetime interval.
83     // Signal the caller to remove the lifetime end call for this alloca.
84     return false;
85   }
86   return true;
87 }
88 
89 bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart,
90                         const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd,
91                         const DominatorTree *DT, const LoopInfo *LI,
92                         size_t MaxLifetimes) {
93   // An alloca that has exactly one start and end in every possible execution.
94   // If it has multiple ends, they have to be unreachable from each other, so
95   // at most one of them is actually used for each execution of the function.
96   return LifetimeStart.size() == 1 &&
97          (LifetimeEnd.size() == 1 ||
98           (LifetimeEnd.size() > 0 &&
99            !maybeReachableFromEachOther(LifetimeEnd, DT, LI, MaxLifetimes)));
100 }
101 
102 Instruction *getUntagLocationIfFunctionExit(Instruction &Inst) {
103   if (isa<ReturnInst>(Inst)) {
104     if (CallInst *CI = Inst.getParent()->getTerminatingMustTailCall())
105       return CI;
106     return &Inst;
107   }
108   if (isa<ResumeInst, CleanupReturnInst>(Inst)) {
109     return &Inst;
110   }
111   return nullptr;
112 }
113 
114 void StackInfoBuilder::visit(OptimizationRemarkEmitter &ORE,
115                              Instruction &Inst) {
116   // Visit non-intrinsic debug-info records attached to Inst.
117   for (DbgVariableRecord &DVR : filterDbgVars(Inst.getDbgRecordRange())) {
118     auto AddIfInteresting = [&](Value *V) {
119       if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
120         if (getAllocaInterestingness(*AI) !=
121             AllocaInterestingness::kInteresting)
122           return;
123         AllocaInfo &AInfo = Info.AllocasToInstrument[AI];
124         auto &DVRVec = AInfo.DbgVariableRecords;
125         if (DVRVec.empty() || DVRVec.back() != &DVR)
126           DVRVec.push_back(&DVR);
127       }
128     };
129 
130     for_each(DVR.location_ops(), AddIfInteresting);
131     if (DVR.isDbgAssign())
132       AddIfInteresting(DVR.getAddress());
133   }
134 
135   if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
136     if (CI->canReturnTwice()) {
137       Info.CallsReturnTwice = true;
138     }
139   }
140   if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
141     switch (getAllocaInterestingness(*AI)) {
142     case AllocaInterestingness::kInteresting:
143       Info.AllocasToInstrument[AI].AI = AI;
144       ORE.emit([&]() {
145         return OptimizationRemarkMissed(DebugType, "safeAlloca", &Inst);
146       });
147       break;
148     case AllocaInterestingness::kSafe:
149       ORE.emit(
150           [&]() { return OptimizationRemark(DebugType, "safeAlloca", &Inst); });
151       break;
152     case AllocaInterestingness::kUninteresting:
153       break;
154     }
155     return;
156   }
157   auto *II = dyn_cast<IntrinsicInst>(&Inst);
158   if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start ||
159              II->getIntrinsicID() == Intrinsic::lifetime_end)) {
160     AllocaInst *AI = findAllocaForValue(II->getArgOperand(1));
161     if (!AI) {
162       Info.UnrecognizedLifetimes.push_back(&Inst);
163       return;
164     }
165     if (getAllocaInterestingness(*AI) != AllocaInterestingness::kInteresting)
166       return;
167     if (II->getIntrinsicID() == Intrinsic::lifetime_start)
168       Info.AllocasToInstrument[AI].LifetimeStart.push_back(II);
169     else
170       Info.AllocasToInstrument[AI].LifetimeEnd.push_back(II);
171     return;
172   }
173   if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) {
174     auto AddIfInteresting = [&](Value *V) {
175       if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
176         if (getAllocaInterestingness(*AI) !=
177             AllocaInterestingness::kInteresting)
178           return;
179         AllocaInfo &AInfo = Info.AllocasToInstrument[AI];
180         auto &DVIVec = AInfo.DbgVariableIntrinsics;
181         if (DVIVec.empty() || DVIVec.back() != DVI)
182           DVIVec.push_back(DVI);
183       }
184     };
185     for_each(DVI->location_ops(), AddIfInteresting);
186     if (auto *DAI = dyn_cast<DbgAssignIntrinsic>(DVI))
187       AddIfInteresting(DAI->getAddress());
188   }
189 
190   Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
191   if (ExitUntag)
192     Info.RetVec.push_back(ExitUntag);
193 }
194 
195 AllocaInterestingness
196 StackInfoBuilder::getAllocaInterestingness(const AllocaInst &AI) {
197   if (AI.getAllocatedType()->isSized() &&
198       // FIXME: support vscale.
199       !AI.getAllocatedType()->isScalableTy() &&
200       // FIXME: instrument dynamic allocas, too
201       AI.isStaticAlloca() &&
202       // alloca() may be called with 0 size, ignore it.
203       memtag::getAllocaSizeInBytes(AI) > 0 &&
204       // We are only interested in allocas not promotable to registers.
205       // Promotable allocas are common under -O0.
206       !isAllocaPromotable(&AI) &&
207       // inalloca allocas are not treated as static, and we don't want
208       // dynamic alloca instrumentation for them as well.
209       !AI.isUsedWithInAlloca() &&
210       // swifterror allocas are register promoted by ISel
211       !AI.isSwiftError()) {
212     if (!(SSI && SSI->isSafe(AI))) {
213       return AllocaInterestingness::kInteresting;
214     }
215     // safe allocas are not interesting
216     return AllocaInterestingness::kSafe;
217   }
218   return AllocaInterestingness::kUninteresting;
219 }
220 
221 uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
222   auto DL = AI.getDataLayout();
223   return *AI.getAllocationSize(DL);
224 }
225 
226 void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Alignment) {
227   const Align NewAlignment = std::max(Info.AI->getAlign(), Alignment);
228   Info.AI->setAlignment(NewAlignment);
229   auto &Ctx = Info.AI->getFunction()->getContext();
230 
231   uint64_t Size = getAllocaSizeInBytes(*Info.AI);
232   uint64_t AlignedSize = alignTo(Size, Alignment);
233   if (Size == AlignedSize)
234     return;
235 
236   // Add padding to the alloca.
237   Type *AllocatedType =
238       Info.AI->isArrayAllocation()
239           ? ArrayType::get(
240                 Info.AI->getAllocatedType(),
241                 cast<ConstantInt>(Info.AI->getArraySize())->getZExtValue())
242           : Info.AI->getAllocatedType();
243   Type *PaddingType = ArrayType::get(Type::getInt8Ty(Ctx), AlignedSize - Size);
244   Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType);
245   auto *NewAI = new AllocaInst(TypeWithPadding, Info.AI->getAddressSpace(),
246                                nullptr, "", Info.AI->getIterator());
247   NewAI->takeName(Info.AI);
248   NewAI->setAlignment(Info.AI->getAlign());
249   NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca());
250   NewAI->setSwiftError(Info.AI->isSwiftError());
251   NewAI->copyMetadata(*Info.AI);
252 
253   Value *NewPtr = NewAI;
254 
255   // TODO: Remove when typed pointers dropped
256   if (Info.AI->getType() != NewAI->getType())
257     NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI->getIterator());
258 
259   Info.AI->replaceAllUsesWith(NewPtr);
260   Info.AI->eraseFromParent();
261   Info.AI = NewAI;
262 }
263 
264 bool isLifetimeIntrinsic(Value *V) {
265   auto *II = dyn_cast<IntrinsicInst>(V);
266   return II && II->isLifetimeStartOrEnd();
267 }
268 
269 Value *readRegister(IRBuilder<> &IRB, StringRef Name) {
270   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
271   MDNode *MD =
272       MDNode::get(M->getContext(), {MDString::get(M->getContext(), Name)});
273   Value *Args[] = {MetadataAsValue::get(M->getContext(), MD)};
274   return IRB.CreateIntrinsic(Intrinsic::read_register,
275                              IRB.getIntPtrTy(M->getDataLayout()), Args);
276 }
277 
278 Value *getPC(const Triple &TargetTriple, IRBuilder<> &IRB) {
279   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
280   if (TargetTriple.getArch() == Triple::aarch64)
281     return memtag::readRegister(IRB, "pc");
282   return IRB.CreatePtrToInt(IRB.GetInsertBlock()->getParent(),
283                             IRB.getIntPtrTy(M->getDataLayout()));
284 }
285 
286 Value *getFP(IRBuilder<> &IRB) {
287   Function *F = IRB.GetInsertBlock()->getParent();
288   Module *M = F->getParent();
289   return IRB.CreatePtrToInt(
290       IRB.CreateIntrinsic(Intrinsic::frameaddress,
291                           IRB.getPtrTy(M->getDataLayout().getAllocaAddrSpace()),
292                           {Constant::getNullValue(IRB.getInt32Ty())}),
293       IRB.getIntPtrTy(M->getDataLayout()));
294 }
295 
296 Value *getAndroidSlotPtr(IRBuilder<> &IRB, int Slot) {
297   Module *M = IRB.GetInsertBlock()->getParent()->getParent();
298   // Android provides a fixed TLS slot for sanitizers. See TLS_SLOT_SANITIZER
299   // in Bionic's libc/private/bionic_tls.h.
300   Function *ThreadPointerFunc =
301       Intrinsic::getOrInsertDeclaration(M, Intrinsic::thread_pointer);
302   return IRB.CreateConstGEP1_32(IRB.getInt8Ty(),
303                                 IRB.CreateCall(ThreadPointerFunc), 8 * Slot);
304 }
305 
306 static DbgAssignIntrinsic *DynCastToDbgAssign(DbgVariableIntrinsic *DVI) {
307   return dyn_cast<DbgAssignIntrinsic>(DVI);
308 }
309 
310 static DbgVariableRecord *DynCastToDbgAssign(DbgVariableRecord *DVR) {
311   return DVR->isDbgAssign() ? DVR : nullptr;
312 }
313 
314 void annotateDebugRecords(AllocaInfo &Info, unsigned int Tag) {
315   // Helper utility for adding DW_OP_LLVM_tag_offset to debug-info records,
316   // abstracted over whether they're intrinsic-stored or DbgVariableRecord
317   // stored.
318   auto AnnotateDbgRecord = [&](auto *DPtr) {
319     // Prepend "tag_offset, N" to the dwarf expression.
320     // Tag offset logically applies to the alloca pointer, and it makes sense
321     // to put it at the beginning of the expression.
322     SmallVector<uint64_t, 8> NewOps = {dwarf::DW_OP_LLVM_tag_offset, Tag};
323     for (size_t LocNo = 0; LocNo < DPtr->getNumVariableLocationOps(); ++LocNo)
324       if (DPtr->getVariableLocationOp(LocNo) == Info.AI)
325         DPtr->setExpression(
326             DIExpression::appendOpsToArg(DPtr->getExpression(), NewOps, LocNo));
327     if (auto *DAI = DynCastToDbgAssign(DPtr)) {
328       if (DAI->getAddress() == Info.AI)
329         DAI->setAddressExpression(
330             DIExpression::prependOpcodes(DAI->getAddressExpression(), NewOps));
331     }
332   };
333 
334   llvm::for_each(Info.DbgVariableIntrinsics, AnnotateDbgRecord);
335   llvm::for_each(Info.DbgVariableRecords, AnnotateDbgRecord);
336 }
337 
338 Value *incrementThreadLong(IRBuilder<> &IRB, Value *ThreadLong,
339                            unsigned int Inc) {
340   // Update the ring buffer. Top byte of ThreadLong defines the size of the
341   // buffer in pages, it must be a power of two, and the start of the buffer
342   // must be aligned by twice that much. Therefore wrap around of the ring
343   // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
344   // The use of AShr instead of LShr is due to
345   //   https://bugs.llvm.org/show_bug.cgi?id=39030
346   // Runtime library makes sure not to use the highest bit.
347   //
348   // Mechanical proof of this address calculation can be found at:
349   // https://github.com/google/sanitizers/blob/master/hwaddress-sanitizer/prove_hwasanwrap.smt2
350   //
351   // Example of the wrap case for N = 1
352   // Pointer:   0x01AAAAAAAAAAAFF8
353   //                     +
354   //            0x0000000000000008
355   //                     =
356   //            0x01AAAAAAAAAAB000
357   //                     &
358   // WrapMask:  0xFFFFFFFFFFFFF000
359   //                     =
360   //            0x01AAAAAAAAAAA000
361   //
362   // Then the WrapMask will be a no-op until the next wrap case.
363   assert((4096 % Inc) == 0);
364   Value *WrapMask = IRB.CreateXor(
365       IRB.CreateShl(IRB.CreateAShr(ThreadLong, 56), 12, "", true, true),
366       ConstantInt::get(ThreadLong->getType(), (uint64_t)-1));
367   return IRB.CreateAnd(
368       IRB.CreateAdd(ThreadLong, ConstantInt::get(ThreadLong->getType(), Inc)),
369       WrapMask);
370 }
371 
372 } // namespace memtag
373 } // namespace llvm
374