MemoryTaggingSupport.cpp

Go to the documentation of this file.

//== MemoryTaggingSupport.cpp - helpers for memory tagging implementations ===//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file declares common infrastructure for HWAddressSanitizer and
// Aarch64StackTagging.
//
//===----------------------------------------------------------------------===//
 
#include "llvm/Transforms/Utils/MemoryTaggingSupport.h"
 
#include "llvm/ADT/STLExtras.h"
#include "llvm/Analysis/CFG.h"
#include "llvm/Analysis/PostDominators.h"
#include "llvm/Analysis/StackSafetyAnalysis.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/BinaryFormat/Dwarf.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/TargetParser/Triple.h"
#include "llvm/Transforms/Utils/PromoteMemToReg.h"
#include <utility>
 
namespace llvm {
namespace memtag {
 
bool forAllReachableExits(const DominatorTree &DT, const PostDominatorTree &PDT,
                          const LoopInfo &LI, const AllocaInfo &AInfo,
                          const SmallVectorImpl<Instruction *> &RetVec,
                          llvm::function_ref<void(Instruction *)> Callback) {
  if (AInfo.LifetimeEnd.size() == 1 && AInfo.LifetimeStart.size() == 1 &&
      PDT.dominates(AInfo.LifetimeEnd[0], AInfo.LifetimeStart[0])) {
    Callback(AInfo.LifetimeEnd[0]);
    return true;
  }
  SmallPtrSet<BasicBlock *, 2> EndBlocks;
  SmallVector<BasicBlock *, 2> StartBlocks;
  for (const auto &[BB, BBInfo] : AInfo.BBInfos) {
    if (BBInfo.Last == Intrinsic::lifetime_end)
      EndBlocks.insert(BB);
    else
      StartBlocks.push_back(BB);
  }
  bool UncoveredRets = false;
 
  if (!StartBlocks.empty()) {
    for (auto *RI : RetVec) {
      auto WL = StartBlocks;
      // If the block with the return is an EndBlock (i.e. a block where the
      // last relevant lifetime intrinsic is an end), we don't have to run a
      // complicated algorithm to know that the RetInst is never reachable
      // without going through an end.
      if (!EndBlocks.contains(RI->getParent()) &&
          isPotentiallyReachableFromMany(WL, RI->getParent(), &EndBlocks, &DT,
                                         &LI)) {
        Callback(RI);
        UncoveredRets = true;
      }
    }
  }
  for_each(AInfo.LifetimeEnd, Callback);
  // We may have inserted untag outside of the lifetime interval.
  // Signal the caller to remove the lifetime end call for this alloca.
  return !UncoveredRets;
}
 
bool isSupportedLifetime(const AllocaInfo &AInfo, const DominatorTree *DT,
                         const LoopInfo *LI) {
  if (AInfo.LifetimeStart.empty())
    return false;
  SmallVector<BasicBlock *, 2> LastEndBlocks;
  SmallPtrSet<const BasicBlock *, 2> FirstEndBlocks;
  SmallPtrSet<BasicBlock *, 2> StartBlocks;
  if (any_of(AInfo.BBInfos, [&](const auto &It) {
        const auto &[BB, BBI] = It;
        if (BBI.Last == Intrinsic::lifetime_end)
          LastEndBlocks.append(succ_begin(BB), succ_end(BB));
        else
          StartBlocks.insert(BB);
        if (BBI.First == Intrinsic::lifetime_end)
          FirstEndBlocks.insert(BB);
        return BBI.DoubleEnd;
      }))
    return false;
  if (LastEndBlocks.empty() || FirstEndBlocks.empty())
    return true;
  return !isManyPotentiallyReachableFromMany(LastEndBlocks, FirstEndBlocks,
                                             &StartBlocks, DT, LI);
}
 
Instruction *getUntagLocationIfFunctionExit(Instruction &Inst) {
  if (isa<ReturnInst>(Inst)) {
    if (CallInst *CI = Inst.getParent()->getTerminatingMustTailCall())
      return CI;
    return &Inst;
  }
  if (isa<ResumeInst, CleanupReturnInst>(Inst)) {
    return &Inst;
  }
  return nullptr;
}
 
void StackInfoBuilder::visit(OptimizationRemarkEmitter &ORE,
                             Instruction &Inst) {
  // Visit non-intrinsic debug-info records attached to Inst.
  for (DbgVariableRecord &DVR : filterDbgVars(Inst.getDbgRecordRange())) {
    auto AddIfInteresting = [&](Value *V) {
      if (auto *AI = dyn_cast_or_null<AllocaInst>(V)) {
        if (getAllocaInterestingness(*AI) !=
            AllocaInterestingness::kInteresting)
          return;
        AllocaInfo &AInfo = Info.AllocasToInstrument[AI];
        auto &DVRVec = AInfo.DbgVariableRecords;
        if (DVRVec.empty() || DVRVec.back() != &DVR)
          DVRVec.push_back(&DVR);
      }
    };
 
    for_each(DVR.location_ops(), AddIfInteresting);
    if (DVR.isDbgAssign())
      AddIfInteresting(DVR.getAddress());
  }
 
  if (CallInst *CI = dyn_cast<CallInst>(&Inst)) {
    if (CI->canReturnTwice()) {
      Info.CallsReturnTwice = true;
    }
  }
  if (AllocaInst *AI = dyn_cast<AllocaInst>(&Inst)) {
    switch (getAllocaInterestingness(*AI)) {
    case AllocaInterestingness::kInteresting:
      Info.AllocasToInstrument[AI].AI = AI;
      ORE.emit([&]() {
        return OptimizationRemarkMissed(DebugType, "safeAlloca", &Inst);
      });
      break;
    case AllocaInterestingness::kSafe:
      ORE.emit(
          [&]() { return OptimizationRemark(DebugType, "safeAlloca", &Inst); });
      break;
    case AllocaInterestingness::kUninteresting:
      break;
    }
    return;
  }
  if (auto *II = dyn_cast<LifetimeIntrinsic>(&Inst)) {
    AllocaInst *AI = dyn_cast<AllocaInst>(II->getArgOperand(0));
    if (!AI ||
        getAllocaInterestingness(*AI) != AllocaInterestingness::kInteresting)
      return;
    auto &AInfo = Info.AllocasToInstrument[AI];
    auto &BBInfo = AInfo.BBInfos[II->getParent()];
 
    if (II->getIntrinsicID() == Intrinsic::lifetime_start) {
      AInfo.LifetimeStart.push_back(II);
    } else {
      AInfo.LifetimeEnd.push_back(II);
      if (BBInfo.Last == Intrinsic::lifetime_end)
        BBInfo.DoubleEnd = true;
    }
 
    BBInfo.Last = II->getIntrinsicID();
    if (BBInfo.First == Intrinsic::not_intrinsic)
      BBInfo.First = II->getIntrinsicID();
 
    return;
  }
 
  Instruction *ExitUntag = getUntagLocationIfFunctionExit(Inst);
  if (ExitUntag)
    Info.RetVec.push_back(ExitUntag);
}
 
AllocaInterestingness
StackInfoBuilder::getAllocaInterestingness(const AllocaInst &AI) {
  std::optional<TypeSize> Size = AI.getAllocationSize(AI.getDataLayout());
  if (Size &&
      // FIXME: support vscale.
      !Size->isScalable() &&
      // FIXME: instrument dynamic allocas, too
      AI.isStaticAlloca() &&
      // alloca() may be called with 0 size, ignore it.
      !Size->isZero() &&
      // We are only interested in allocas not promotable to registers.
      // Promotable allocas are common under -O0.
      !isAllocaPromotable(&AI) &&
      // inalloca allocas are not treated as static, and we don't want
      // dynamic alloca instrumentation for them as well.
      !AI.isUsedWithInAlloca() &&
      // swifterror allocas are register promoted by ISel
      !AI.isSwiftError()) {
    if (!(SSI && SSI->isSafe(AI))) {
      return AllocaInterestingness::kInteresting;
    }
    // safe allocas are not interesting
    return AllocaInterestingness::kSafe;
  }
  return AllocaInterestingness::kUninteresting;
}
 
uint64_t getAllocaSizeInBytes(const AllocaInst &AI) {
  auto DL = AI.getDataLayout();
  return *AI.getAllocationSize(DL);
}
 
void alignAndPadAlloca(memtag::AllocaInfo &Info, llvm::Align Alignment) {
  const Align NewAlignment = std::max(Info.AI->getAlign(), Alignment);
  Info.AI->setAlignment(NewAlignment);
  auto &Ctx = Info.AI->getFunction()->getContext();
 
  uint64_t Size = getAllocaSizeInBytes(*Info.AI);
  uint64_t AlignedSize = alignTo(Size, Alignment);
  if (Size == AlignedSize)
    return;
 
  // Add padding to the alloca.
  Type *AllocatedType =
      Info.AI->isArrayAllocation()
          ? ArrayType::get(
                Info.AI->getAllocatedType(),
                cast<ConstantInt>(Info.AI->getArraySize())->getZExtValue())
          : Info.AI->getAllocatedType();
  Type *PaddingType = ArrayType::get(Type::getInt8Ty(Ctx), AlignedSize - Size);
  Type *TypeWithPadding = StructType::get(AllocatedType, PaddingType);
  auto *NewAI = new AllocaInst(TypeWithPadding, Info.AI->getAddressSpace(),
                               nullptr, "", Info.AI->getIterator());
  NewAI->takeName(Info.AI);
  NewAI->setAlignment(Info.AI->getAlign());
  NewAI->setUsedWithInAlloca(Info.AI->isUsedWithInAlloca());
  NewAI->setSwiftError(Info.AI->isSwiftError());
  NewAI->copyMetadata(*Info.AI);
 
  Info.AI->replaceAllUsesWith(NewAI);
  Info.AI->eraseFromParent();
  Info.AI = NewAI;
}
 
Value *readRegister(IRBuilder<> &IRB, StringRef Name) {
  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
  MDNode *MD =
      MDNode::get(M->getContext(), {MDString::get(M->getContext(), Name)});
  Value *Args[] = {MetadataAsValue::get(M->getContext(), MD)};
  return IRB.CreateIntrinsic(Intrinsic::read_register,
                             IRB.getIntPtrTy(M->getDataLayout()), Args);
}
 
Value *getPC(const Triple &TargetTriple, IRBuilder<> &IRB) {
  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
  if (TargetTriple.getArch() == Triple::aarch64)
    return memtag::readRegister(IRB, "pc");
  return IRB.CreatePtrToInt(IRB.GetInsertBlock()->getParent(),
                            IRB.getIntPtrTy(M->getDataLayout()));
}
 
Value *getFP(IRBuilder<> &IRB) {
  Function *F = IRB.GetInsertBlock()->getParent();
  Module *M = F->getParent();
  return IRB.CreatePtrToInt(
      IRB.CreateIntrinsic(Intrinsic::frameaddress,
                          IRB.getPtrTy(M->getDataLayout().getAllocaAddrSpace()),
                          {Constant::getNullValue(IRB.getInt32Ty())}),
      IRB.getIntPtrTy(M->getDataLayout()));
}
 
Value *getDarwinSlotPtr(IRBuilder<> &IRB, int Slot) {
  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
  // FIXME: This should use the thread_pointer intrinsic. However, the
  // intrinsic is not currently implemented on Darwin correctly. The
  // TPIDRRO_EL0 register became part of the ABI to access TSD on recent
  // versions of OS and so it is safe to use here
  // Darwin provides a fixed slot for sanitizers at offset 231.
  MDNode *MD = MDNode::get(M->getContext(),
                           {MDString::get(M->getContext(), "TPIDRRO_EL0")});
  Value *Args[] = {MetadataAsValue::get(M->getContext(), MD)};
  return IRB.CreateConstGEP1_32(
      IRB.getInt8Ty(),
      IRB.CreateIntToPtr(
          IRB.CreateIntrinsic(Intrinsic::read_register,
                              {IRB.getIntPtrTy(M->getDataLayout())}, Args),
          IRB.getPtrTy()),
      8 * Slot);
}
 
Value *getAndroidSlotPtr(IRBuilder<> &IRB, int Slot) {
  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
  // Android provides a fixed TLS slot for sanitizers. See TLS_SLOT_SANITIZER
  // in Bionic's libc/private/bionic_tls.h.
  Function *ThreadPointerFunc = Intrinsic::getOrInsertDeclaration(
      M, Intrinsic::thread_pointer,
      IRB.getPtrTy(M->getDataLayout().getDefaultGlobalsAddressSpace()));
  return IRB.CreateConstGEP1_32(IRB.getInt8Ty(),
                                IRB.CreateCall(ThreadPointerFunc), 8 * Slot);
}
 
static DbgVariableRecord *DynCastToDbgAssign(DbgVariableRecord *DVR) {
  return DVR->isDbgAssign() ? DVR : nullptr;
}
 
void annotateDebugRecords(AllocaInfo &Info, unsigned int Tag) {
  auto AnnotateDbgRecord = [&](DbgVariableRecord *DPtr) {
    // Prepend "tag_offset, N" to the dwarf expression.
    // Tag offset logically applies to the alloca pointer, and it makes sense
    // to put it at the beginning of the expression.
    SmallVector<uint64_t, 8> NewOps = {dwarf::DW_OP_LLVM_tag_offset, Tag};
    for (size_t LocNo = 0; LocNo < DPtr->getNumVariableLocationOps(); ++LocNo)
      if (DPtr->getVariableLocationOp(LocNo) == Info.AI)
        DPtr->setExpression(
            DIExpression::appendOpsToArg(DPtr->getExpression(), NewOps, LocNo));
    if (auto *DAI = DynCastToDbgAssign(DPtr)) {
      if (DAI->getAddress() == Info.AI)
        DAI->setAddressExpression(
            DIExpression::prependOpcodes(DAI->getAddressExpression(), NewOps));
    }
  };
 
  llvm::for_each(Info.DbgVariableRecords, AnnotateDbgRecord);
}
 
Value *incrementThreadLong(IRBuilder<> &IRB, Value *ThreadLong,
                           unsigned int Inc, bool IsMemtagDarwin) {
  // Update the ring buffer. Top byte of ThreadLong defines the size of the
  // buffer in pages, it must be a power of two, and the start of the buffer
  // must be aligned by twice that much. Therefore wrap around of the ring
  // buffer is simply Addr &= ~((ThreadLong >> 56) << 12).
  // The use of AShr instead of LShr is due to
  //   https://bugs.llvm.org/show_bug.cgi?id=39030
  // Runtime library makes sure not to use the highest bit.
  //
  // Mechanical proof of this address calculation can be found at:
  // https://github.com/google/sanitizers/blob/master/hwaddress-sanitizer/prove_hwasanwrap.smt2
  //
  // Example of the wrap case for N = 1
  // Pointer:   0x01AAAAAAAAAAAFF8
  //                     +
  //            0x0000000000000008
  //                     =
  //            0x01AAAAAAAAAAB000
  //                     &
  // WrapMask:  0xFFFFFFFFFFFFF000
  //                     =
  //            0x01AAAAAAAAAAA000
  //
  // Then the WrapMask will be a no-op until the next wrap case.
  //
  // Darwin relies on bit 60-62 to store the size of the buffer in pages. This
  // limits N to [0,2] while the rest of the proof remains unchanged. Bits
  // 56-59 are avoided in order to prevent MTE Canonical Tag Faults while
  // accessing the ring buffer. Bit 63 is avoided to prevent unintentional
  // signed extension by AShr.
  assert((4096 % Inc) == 0);
  Value *WrapMask = IRB.CreateXor(
      IRB.CreateShl(IRB.CreateAShr(ThreadLong, IsMemtagDarwin ? 60 : 56), 12,
                    "", true, true),
      ConstantInt::get(ThreadLong->getType(), (uint64_t)-1));
  return IRB.CreateAnd(
      IRB.CreateAdd(ThreadLong, ConstantInt::get(ThreadLong->getType(), Inc)),
      WrapMask);
}
 
} // namespace memtag
} // namespace llvm