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(NumTypeIds, "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 *extractNumericTypeId(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 /// Extracts a numeric type identifier from an MDNode containing type metadata.
77 ConstantInt *CrossDSOCFI::extractNumericTypeId(MDNode *MD) {
78   // This check excludes vtables for classes inside anonymous namespaces.
79   auto TM = dyn_cast<ValueAsMetadata>(MD->getOperand(1));
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   return C;
88 }
89 
90 /// buildCFICheck - emits __cfi_check for the current module.
91 void CrossDSOCFI::buildCFICheck() {
92   // FIXME: verify that __cfi_check ends up near the end of the code section,
93   // but before the jump slots created in LowerTypeTests.
94   llvm::DenseSet<uint64_t> TypeIds;
95   SmallVector<MDNode *, 2> Types;
96   for (GlobalObject &GO : M->global_objects()) {
97     Types.clear();
98     GO.getMetadata(LLVMContext::MD_type, Types);
99     for (MDNode *Type : Types) {
100       // Sanity check. GO must not be a function declaration.
101       assert(!isa<Function>(&GO) || !cast<Function>(&GO)->isDeclaration());
102 
103       if (ConstantInt *TypeId = extractNumericTypeId(Type))
104         TypeIds.insert(TypeId->getZExtValue());
105     }
106   }
107 
108   LLVMContext &Ctx = M->getContext();
109   Constant *C = M->getOrInsertFunction(
110       "__cfi_check", Type::getVoidTy(Ctx), Type::getInt64Ty(Ctx),
111       Type::getInt8PtrTy(Ctx), Type::getInt8PtrTy(Ctx), nullptr);
112   Function *F = dyn_cast<Function>(C);
113   F->setAlignment(4096);
114   auto args = F->arg_begin();
115   Value &CallSiteTypeId = *(args++);
116   CallSiteTypeId.setName("CallSiteTypeId");
117   Value &Addr = *(args++);
118   Addr.setName("Addr");
119   Value &CFICheckFailData = *(args++);
120   CFICheckFailData.setName("CFICheckFailData");
121   assert(args == F->arg_end());
122 
123   BasicBlock *BB = BasicBlock::Create(Ctx, "entry", F);
124   BasicBlock *ExitBB = BasicBlock::Create(Ctx, "exit", F);
125 
126   BasicBlock *TrapBB = BasicBlock::Create(Ctx, "fail", F);
127   IRBuilder<> IRBFail(TrapBB);
128   Constant *CFICheckFailFn = M->getOrInsertFunction(
129       "__cfi_check_fail", Type::getVoidTy(Ctx), Type::getInt8PtrTy(Ctx),
130       Type::getInt8PtrTy(Ctx), nullptr);
131   IRBFail.CreateCall(CFICheckFailFn, {&CFICheckFailData, &Addr});
132   IRBFail.CreateBr(ExitBB);
133 
134   IRBuilder<> IRBExit(ExitBB);
135   IRBExit.CreateRetVoid();
136 
137   IRBuilder<> IRB(BB);
138   SwitchInst *SI = IRB.CreateSwitch(&CallSiteTypeId, TrapBB, TypeIds.size());
139   for (uint64_t TypeId : TypeIds) {
140     ConstantInt *CaseTypeId = ConstantInt::get(Type::getInt64Ty(Ctx), TypeId);
141     BasicBlock *TestBB = BasicBlock::Create(Ctx, "test", F);
142     IRBuilder<> IRBTest(TestBB);
143     Function *BitsetTestFn = Intrinsic::getDeclaration(M, Intrinsic::type_test);
144 
145     Value *Test = IRBTest.CreateCall(
146         BitsetTestFn, {&Addr, MetadataAsValue::get(
147                                   Ctx, ConstantAsMetadata::get(CaseTypeId))});
148     BranchInst *BI = IRBTest.CreateCondBr(Test, ExitBB, TrapBB);
149     BI->setMetadata(LLVMContext::MD_prof, VeryLikelyWeights);
150 
151     SI->addCase(CaseTypeId, TestBB);
152     ++NumTypeIds;
153   }
154 }
155 
156 bool CrossDSOCFI::runOnModule(Module &M) {
157   if (skipModule(M))
158     return false;
159 
160   if (M.getModuleFlag("Cross-DSO CFI") == nullptr)
161     return false;
162   buildCFICheck();
163   return true;
164 }
165