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/Analysis/CFG.h" 16 #include "llvm/Analysis/PostDominators.h" 17 #include "llvm/Analysis/ValueTracking.h" 18 #include "llvm/IR/BasicBlock.h" 19 #include "llvm/IR/IntrinsicInst.h" 20 21 namespace llvm { 22 namespace memtag { 23 namespace { 24 bool maybeReachableFromEachOther(const SmallVectorImpl<IntrinsicInst *> &Insts, 25 const DominatorTree *DT, const LoopInfo *LI, 26 size_t MaxLifetimes) { 27 // If we have too many lifetime ends, give up, as the algorithm below is N^2. 28 if (Insts.size() > MaxLifetimes) 29 return true; 30 for (size_t I = 0; I < Insts.size(); ++I) { 31 for (size_t J = 0; J < Insts.size(); ++J) { 32 if (I == J) 33 continue; 34 if (isPotentiallyReachable(Insts[I], Insts[J], nullptr, DT, LI)) 35 return true; 36 } 37 } 38 return false; 39 } 40 } // namespace 41 42 bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT, 43 const LoopInfo &LI, const Instruction *Start, 44 const SmallVectorImpl<IntrinsicInst *> &Ends, 45 const SmallVectorImpl<Instruction *> &RetVec, 46 llvm::function_ref<void(Instruction *)> Callback) { 47 if (Ends.size() == 1 && PDT.dominates(Ends[0], Start)) { 48 Callback(Ends[0]); 49 return true; 50 } 51 SmallPtrSet<BasicBlock *, 2> EndBlocks; 52 for (auto *End : Ends) { 53 EndBlocks.insert(End->getParent()); 54 } 55 SmallVector<Instruction *, 8> ReachableRetVec; 56 unsigned NumCoveredExits = 0; 57 for (auto *RI : RetVec) { 58 if (!isPotentiallyReachable(Start, RI, nullptr, &DT, &LI)) 59 continue; 60 ReachableRetVec.push_back(RI); 61 // If there is an end in the same basic block as the return, we know for 62 // sure that the return is covered. Otherwise, we can check whether there 63 // is a way to reach the RI from the start of the lifetime without passing 64 // through an end. 65 if (EndBlocks.count(RI->getParent()) > 0 || 66 !isPotentiallyReachable(Start, RI, &EndBlocks, &DT, &LI)) { 67 ++NumCoveredExits; 68 } 69 } 70 // If there's a mix of covered and non-covered exits, just put the untag 71 // on exits, so we avoid the redundancy of untagging twice. 72 if (NumCoveredExits == ReachableRetVec.size()) { 73 for (auto *End : Ends) 74 Callback(End); 75 } else { 76 for (auto *RI : ReachableRetVec) 77 Callback(RI); 78 // We may have inserted untag outside of the lifetime interval. 79 // Signal the caller to remove the lifetime end call for this alloca. 80 return false; 81 } 82 return true; 83 } 84 85 bool isStandardLifetime(const SmallVectorImpl<IntrinsicInst *> &LifetimeStart, 86 const SmallVectorImpl<IntrinsicInst *> &LifetimeEnd, 87 const DominatorTree *DT, const LoopInfo *LI, 88 size_t MaxLifetimes) { 89 // An alloca that has exactly one start and end in every possible execution. 90 // If it has multiple ends, they have to be unreachable from each other, so 91 // at most one of them is actually used for each execution of the function. 92 return LifetimeStart.size() == 1 && 93 (LifetimeEnd.size() == 1 || 94 (LifetimeEnd.size() > 0 && 95 !maybeReachableFromEachOther(LifetimeEnd, DT, LI, MaxLifetimes))); 96 } 97 98 Instruction *getUntagLocationIfFunctionExit(Instruction &Inst) { 99 if (isa<ReturnInst>(Inst)) { 100 if (CallInst *CI = Inst.getParent()->getTerminatingMustTailCall()) 101 return CI; 102 return &Inst; 103 } 104 if (isa<ResumeInst, CleanupReturnInst>(Inst)) { 105 return &Inst; 106 } 107 return nullptr; 108 } 109 110 void StackInfoBuilder::visit(Instruction &Inst) { 111 if (CallInst *CI = dyn_cast<CallInst>(&Inst)) { 112 if (CI->canReturnTwice()) { 113 Info.CallsReturnTwice = true; 114 } 115 } 116 if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) { 117 if (IsInterestingAlloca(*AI)) { 118 Info.AllocasToInstrument[AI].AI = AI; 119 } 120 return; 121 } 122 auto *II = dyn_cast<IntrinsicInst>(&Inst); 123 if (II && (II->getIntrinsicID() == Intrinsic::lifetime_start || 124 II->getIntrinsicID() == Intrinsic::lifetime_end)) { 125 AllocaInst *AI = findAllocaForValue(II->getArgOperand(1)); 126 if (!AI) { 127 Info.UnrecognizedLifetimes.push_back(&Inst); 128 return; 129 } 130 if (!IsInterestingAlloca(*AI)) 131 return; 132 if (II->getIntrinsicID() == Intrinsic::lifetime_start) 133 Info.AllocasToInstrument[AI].LifetimeStart.push_back(II); 134 else 135 Info.AllocasToInstrument[AI].LifetimeEnd.push_back(II); 136 return; 137 } 138 if (auto *DVI = dyn_cast<DbgVariableIntrinsic>(&Inst)) { 139 for (Value *V : DVI->location_ops()) { 140 if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) { 141 if (!IsInterestingAlloca(*AI)) 142 continue; 143 AllocaInfo &AInfo = Info.AllocasToInstrument[AI]; 144 auto &DVIVec = AInfo.DbgVariableIntrinsics; 145 if (DVIVec.empty() || DVIVec.back() != DVI) 146 DVIVec.push_back(DVI); 147 } 148 } 149 } 150 Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst); 151 if (ExitUntag) 152 Info.RetVec.push_back(ExitUntag); 153 } 154 155 uint64_t getAllocaSizeInBytes(const AllocaInst &AI) { 156 auto DL = AI.getModule()->getDataLayout(); 157 return *AI.getAllocationSizeInBits(DL) / 8; 158 } 159 160 void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Alignment) { 161 const Align NewAlignment = std::max(Info.AI->getAlign(), Alignment); 162 Info.AI->setAlignment(NewAlignment); 163 auto &Ctx = Info.AI->getFunction()->getContext(); 164 165 uint64_t Size = getAllocaSizeInBytes(*Info.AI); 166 uint64_t AlignedSize = alignTo(Size, Alignment); 167 if (Size == AlignedSize) 168 return; 169 170 // Add padding to the alloca. 171 Type *AllocatedType = 172 Info.AI->isArrayAllocation() 173 ? ArrayType::get( 174 Info.AI->getAllocatedType(), 175 cast<ConstantInt>(Info.AI->getArraySize())->getZExtValue()) 176 : Info.AI->getAllocatedType(); 177 Type *PaddingType = ArrayType::get(Type::getInt8Ty(Ctx), AlignedSize - Size); 178 Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType); 179 auto *NewAI = 180 new AllocaInst(TypeWithPadding, Info.AI->getType()->getAddressSpace(), 181 nullptr, "", Info.AI); 182 NewAI->takeName(Info.AI); 183 NewAI->setAlignment(Info.AI->getAlign()); 184 NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca()); 185 NewAI->setSwiftError(Info.AI->isSwiftError()); 186 NewAI->copyMetadata(*Info.AI); 187 188 auto *NewPtr = new BitCastInst(NewAI, Info.AI->getType(), "", Info.AI); 189 Info.AI->replaceAllUsesWith(NewPtr); 190 Info.AI->eraseFromParent(); 191 Info.AI = NewAI; 192 } 193 194 } // namespace memtag 195 } // namespace llvm 196