lib/Transforms/IPO/CrossDSOCFI.cpp

   1 //===-- CrossDSOCFI.cpp - Externalize this module's CFI checks ------------===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // This pass exports all llvm.bitset's found in the module in the form of a
  11 // __cfi_check function, which can be used to verify cross-DSO call targets.
  12 //
  13 //===----------------------------------------------------------------------===//
  14
  15 #include "llvm/Transforms/IPO.h"
  16 #include "llvm/ADT/DenseSet.h"
  17 #include "llvm/ADT/EquivalenceClasses.h"
  18 #include "llvm/ADT/Statistic.h"
  19 #include "llvm/IR/Constant.h"
  20 #include "llvm/IR/Constants.h"
  21 #include "llvm/IR/Function.h"
  22 #include "llvm/IR/GlobalObject.h"
  23 #include "llvm/IR/GlobalVariable.h"
  24 #include "llvm/IR/IRBuilder.h"
  25 #include "llvm/IR/Instructions.h"
  26 #include "llvm/IR/Intrinsics.h"
  27 #include "llvm/IR/MDBuilder.h"
  28 #include "llvm/IR/Module.h"
  29 #include "llvm/IR/Operator.h"
  30 #include "llvm/Pass.h"
  31 #include "llvm/Support/Debug.h"
  32 #include "llvm/Support/raw_ostream.h"
  33 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
  34
  35 using namespace llvm;
  36
  37 #define DEBUG_TYPE "cross-dso-cfi"
  38
  39 STATISTIC(TypeIds, "Number of unique type identifiers");
  40
  41 namespace {
  42
  43 struct CrossDSOCFI : public ModulePass {
  44   static char ID;
  45   CrossDSOCFI() : ModulePass(ID) {
  46     initializeCrossDSOCFIPass(*PassRegistry::getPassRegistry());
  47   }
  48
  49   Module *M;
  50   MDNode *VeryLikelyWeights;
  51
  52   ConstantInt *extractBitSetTypeId(MDNode *MD);
  53   void buildCFICheck();
  54
  55   bool doInitialization(Module &M) override;
  56   bool runOnModule(Module &M) override;
  57 };
  58
  59 } // anonymous namespace
  60
  61 INITIALIZE_PASS_BEGIN(CrossDSOCFI, "cross-dso-cfi", "Cross-DSO CFI", false,
  62                       false)
  63 INITIALIZE_PASS_END(CrossDSOCFI, "cross-dso-cfi", "Cross-DSO CFI", false, false)
  64 char CrossDSOCFI::ID = 0;
  65
  66 ModulePass *llvm::createCrossDSOCFIPass() { return new CrossDSOCFI; }
  67
  68 bool CrossDSOCFI::doInitialization(Module &Mod) {
  69   M = &Mod;
  70   VeryLikelyWeights =
  71       MDBuilder(M->getContext()).createBranchWeights((1U << 20) - 1, 1);
  72
  73   return false;
  74 }
  75
  76 /// extractBitSetTypeId - Extracts TypeId from a hash-based bitset MDNode.
  77 ConstantInt *CrossDSOCFI::extractBitSetTypeId(MDNode *MD) {
  78   // This check excludes vtables for classes inside anonymous namespaces.
  79   auto TM = dyn_cast<ValueAsMetadata>(MD->getOperand(0));
  80   if (!TM)
  81     return nullptr;
  82   auto C = dyn_cast_or_null<ConstantInt>(TM->getValue());
  83   if (!C) return nullptr;
  84   // We are looking for i64 constants.
  85   if (C->getBitWidth() != 64) return nullptr;
  86
  87   // Sanity check.
  88   auto FM = dyn_cast_or_null<ValueAsMetadata>(MD->getOperand(1));
  89   // Can be null if a function was removed by an optimization.
  90   if (FM) {
  91     auto F = dyn_cast<Function>(FM->getValue());
  92     (void)F;
  93     // But can never be a function declaration.
  94     assert(!F || !F->isDeclaration());
  95     (void)F; // Suppress unused variable warning in the no-asserts build.
  96   }
  97   return C;
  98 }
  99
 100 /// buildCFICheck - emits __cfi_check for the current module.
 101 void CrossDSOCFI::buildCFICheck() {
 102   // FIXME: verify that __cfi_check ends up near the end of the code section,
 103   // but before the jump slots created in LowerBitSets.
 104   llvm::DenseSet<uint64_t> BitSetIds;
 105   NamedMDNode *BitSetNM = M->getNamedMetadata("llvm.bitsets");
 106
 107   if (BitSetNM)
 108     for (unsigned I = 0, E = BitSetNM->getNumOperands(); I != E; ++I)
 109       if (ConstantInt *TypeId = extractBitSetTypeId(BitSetNM->getOperand(I)))
 110         BitSetIds.insert(TypeId->getZExtValue());
 111
 112   LLVMContext &Ctx = M->getContext();
 113   Constant *C = M->getOrInsertFunction(
 114       "__cfi_check",
 115       FunctionType::get(
 116           Type::getVoidTy(Ctx),
 117           {Type::getInt64Ty(Ctx), PointerType::getUnqual(Type::getInt8Ty(Ctx))},
 118           false));
 119   Function *F = dyn_cast<Function>(C);
 120   F->setAlignment(4096);
 121   auto args = F->arg_begin();
 122   Argument &CallSiteTypeId = *(args++);
 123   CallSiteTypeId.setName("CallSiteTypeId");
 124   Argument &Addr = *(args++);
 125   Addr.setName("Addr");
 126   assert(args == F->arg_end());
 127
 128   BasicBlock *BB = BasicBlock::Create(Ctx, "entry", F);
 129
 130   BasicBlock *TrapBB = BasicBlock::Create(Ctx, "trap", F);
 131   IRBuilder<> IRBTrap(TrapBB);
 132   Function *TrapFn = Intrinsic::getDeclaration(M, Intrinsic::trap);
 133   llvm::CallInst *TrapCall = IRBTrap.CreateCall(TrapFn);
 134   TrapCall->setDoesNotReturn();
 135   TrapCall->setDoesNotThrow();
 136   IRBTrap.CreateUnreachable();
 137
 138   BasicBlock *ExitBB = BasicBlock::Create(Ctx, "exit", F);
 139   IRBuilder<> IRBExit(ExitBB);
 140   IRBExit.CreateRetVoid();
 141
 142   IRBuilder<> IRB(BB);
 143   SwitchInst *SI = IRB.CreateSwitch(&CallSiteTypeId, TrapBB, BitSetIds.size());
 144   for (uint64_t TypeId : BitSetIds) {
 145     ConstantInt *CaseTypeId = ConstantInt::get(Type::getInt64Ty(Ctx), TypeId);
 146     BasicBlock *TestBB = BasicBlock::Create(Ctx, "test", F);
 147     IRBuilder<> IRBTest(TestBB);
 148     Function *BitsetTestFn =
 149         Intrinsic::getDeclaration(M, Intrinsic::bitset_test);
 150
 151     Value *Test = IRBTest.CreateCall(
 152         BitsetTestFn, {&Addr, MetadataAsValue::get(
 153                                   Ctx, ConstantAsMetadata::get(CaseTypeId))});
 154     BranchInst *BI = IRBTest.CreateCondBr(Test, ExitBB, TrapBB);
 155     BI->setMetadata(LLVMContext::MD_prof, VeryLikelyWeights);
 156
 157     SI->addCase(CaseTypeId, TestBB);
 158     ++TypeIds;
 159   }
 160 }
 161
 162 bool CrossDSOCFI::runOnModule(Module &M) {
 163   if (M.getModuleFlag("Cross-DSO CFI") == nullptr)
 164     return false;
 165   buildCFICheck();
 166   return true;
 167 }