llvm/lib/Transforms/IPO/CrossDSOCFI.cpp

   1 //===-- CrossDSOCFI.cpp - Externalize this module's CFI checks ------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This pass exports all llvm.bitset's found in the module in the form of a
  10 // __cfi_check function, which can be used to verify cross-DSO call targets.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/Transforms/IPO/CrossDSOCFI.h"
  15 #include "llvm/ADT/SetVector.h"
  16 #include "llvm/ADT/Statistic.h"
  17 #include "llvm/IR/Constants.h"
  18 #include "llvm/IR/Function.h"
  19 #include "llvm/IR/GlobalObject.h"
  20 #include "llvm/IR/IRBuilder.h"
  21 #include "llvm/IR/Instructions.h"
  22 #include "llvm/IR/Intrinsics.h"
  23 #include "llvm/IR/MDBuilder.h"
  24 #include "llvm/IR/Module.h"
  25 #include "llvm/TargetParser/Triple.h"
  26 #include "llvm/Transforms/IPO.h"
  27
  28 using namespace llvm;
  29
  30 #define DEBUG_TYPE "cross-dso-cfi"
  31
  32 STATISTIC(NumTypeIds, "Number of unique type identifiers");
  33
  34 namespace {
  35
  36 struct CrossDSOCFI {
  37   MDNode *VeryLikelyWeights;
  38
  39   ConstantInt *extractNumericTypeId(MDNode *MD);
  40   void buildCFICheck(Module &M);
  41   bool runOnModule(Module &M);
  42 };
  43
  44 } // anonymous namespace
  45
  46 /// Extracts a numeric type identifier from an MDNode containing type metadata.
  47 ConstantInt *CrossDSOCFI::extractNumericTypeId(MDNode *MD) {
  48   // This check excludes vtables for classes inside anonymous namespaces.
  49   auto TM = dyn_cast<ValueAsMetadata>(MD->getOperand(1));
  50   if (!TM)
  51     return nullptr;
  52   auto C = dyn_cast_or_null<ConstantInt>(TM->getValue());
  53   if (!C) return nullptr;
  54   // We are looking for i64 constants.
  55   if (C->getBitWidth() != 64) return nullptr;
  56
  57   return C;
  58 }
  59
  60 /// buildCFICheck - emits __cfi_check for the current module.
  61 void CrossDSOCFI::buildCFICheck(Module &M) {
  62   // FIXME: verify that __cfi_check ends up near the end of the code section,
  63   // but before the jump slots created in LowerTypeTests.
  64   SetVector<uint64_t> TypeIds;
  65   SmallVector<MDNode *, 2> Types;
  66   for (GlobalObject &GO : M.global_objects()) {
  67     Types.clear();
  68     GO.getMetadata(LLVMContext::MD_type, Types);
  69     for (MDNode *Type : Types)
  70       if (ConstantInt *TypeId = extractNumericTypeId(Type))
  71         TypeIds.insert(TypeId->getZExtValue());
  72   }
  73
  74   NamedMDNode *CfiFunctionsMD = M.getNamedMetadata("cfi.functions");
  75   if (CfiFunctionsMD) {
  76     for (auto *Func : CfiFunctionsMD->operands()) {
  77       assert(Func->getNumOperands() >= 2);
  78       for (unsigned I = 2; I < Func->getNumOperands(); ++I)
  79         if (ConstantInt *TypeId =
  80                 extractNumericTypeId(cast<MDNode>(Func->getOperand(I).get())))
  81           TypeIds.insert(TypeId->getZExtValue());
  82     }
  83   }
  84
  85   LLVMContext &Ctx = M.getContext();
  86   FunctionCallee C = M.getOrInsertFunction(
  87       "__cfi_check", Type::getVoidTy(Ctx), Type::getInt64Ty(Ctx),
  88       PointerType::getUnqual(Ctx), PointerType::getUnqual(Ctx));
  89   Function *F = cast<Function>(C.getCallee());
  90   // Take over the existing function. The frontend emits a weak stub so that the
  91   // linker knows about the symbol; this pass replaces the function body.
  92   F->deleteBody();
  93   F->setAlignment(Align(4096));
  94
  95   Triple T(M.getTargetTriple());
  96   if (T.isARM() || T.isThumb())
  97     F->addFnAttr("target-features", "+thumb-mode");
  98
  99   auto args = F->arg_begin();
 100   Value &CallSiteTypeId = *(args++);
 101   CallSiteTypeId.setName("CallSiteTypeId");
 102   Value &Addr = *(args++);
 103   Addr.setName("Addr");
 104   Value &CFICheckFailData = *(args++);
 105   CFICheckFailData.setName("CFICheckFailData");
 106   assert(args == F->arg_end());
 107
 108   BasicBlock *BB = BasicBlock::Create(Ctx, "entry", F);
 109   BasicBlock *ExitBB = BasicBlock::Create(Ctx, "exit", F);
 110
 111   BasicBlock *TrapBB = BasicBlock::Create(Ctx, "fail", F);
 112   IRBuilder<> IRBFail(TrapBB);
 113   FunctionCallee CFICheckFailFn = M.getOrInsertFunction(
 114       "__cfi_check_fail", Type::getVoidTy(Ctx), PointerType::getUnqual(Ctx),
 115       PointerType::getUnqual(Ctx));
 116   IRBFail.CreateCall(CFICheckFailFn, {&CFICheckFailData, &Addr});
 117   IRBFail.CreateBr(ExitBB);
 118
 119   IRBuilder<> IRBExit(ExitBB);
 120   IRBExit.CreateRetVoid();
 121
 122   IRBuilder<> IRB(BB);
 123   SwitchInst *SI = IRB.CreateSwitch(&CallSiteTypeId, TrapBB, TypeIds.size());
 124   for (uint64_t TypeId : TypeIds) {
 125     ConstantInt *CaseTypeId = ConstantInt::get(Type::getInt64Ty(Ctx), TypeId);
 126     BasicBlock *TestBB = BasicBlock::Create(Ctx, "test", F);
 127     IRBuilder<> IRBTest(TestBB);
 128     Function *BitsetTestFn = Intrinsic::getDeclaration(&M, Intrinsic::type_test);
 129
 130     Value *Test = IRBTest.CreateCall(
 131         BitsetTestFn, {&Addr, MetadataAsValue::get(
 132                                   Ctx, ConstantAsMetadata::get(CaseTypeId))});
 133     BranchInst *BI = IRBTest.CreateCondBr(Test, ExitBB, TrapBB);
 134     BI->setMetadata(LLVMContext::MD_prof, VeryLikelyWeights);
 135
 136     SI->addCase(CaseTypeId, TestBB);
 137     ++NumTypeIds;
 138   }
 139 }
 140
 141 bool CrossDSOCFI::runOnModule(Module &M) {
 142   VeryLikelyWeights =
 143     MDBuilder(M.getContext()).createBranchWeights((1U << 20) - 1, 1);
 144   if (M.getModuleFlag("Cross-DSO CFI") == nullptr)
 145     return false;
 146   buildCFICheck(M);
 147   return true;
 148 }
 149
 150 PreservedAnalyses CrossDSOCFIPass::run(Module &M, ModuleAnalysisManager &AM) {
 151   CrossDSOCFI Impl;
 152   bool Changed = Impl.runOnModule(M);
 153   if (!Changed)
 154     return PreservedAnalyses::all();
 155   return PreservedAnalyses::none();
 156 }