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/CrossDSOCFI.h"
16 #include "llvm/ADT/EquivalenceClasses.h"
17 #include "llvm/ADT/SetVector.h"
18 #include "llvm/ADT/Statistic.h"
19 #include "llvm/ADT/Triple.h"
20 #include "llvm/IR/Constant.h"
21 #include "llvm/IR/Constants.h"
22 #include "llvm/IR/Function.h"
23 #include "llvm/IR/GlobalObject.h"
24 #include "llvm/IR/GlobalVariable.h"
25 #include "llvm/IR/IRBuilder.h"
26 #include "llvm/IR/Instructions.h"
27 #include "llvm/IR/Intrinsics.h"
28 #include "llvm/IR/MDBuilder.h"
29 #include "llvm/IR/Module.h"
30 #include "llvm/IR/Operator.h"
31 #include "llvm/Pass.h"
32 #include "llvm/Support/Debug.h"
33 #include "llvm/Support/raw_ostream.h"
34 #include "llvm/Transforms/IPO.h"
35 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
36 
37 using namespace llvm;
38 
39 #define DEBUG_TYPE "cross-dso-cfi"
40 
41 STATISTIC(NumTypeIds, "Number of unique type identifiers");
42 
43 namespace {
44 
45 struct CrossDSOCFI : public ModulePass {
46   static char ID;
47   CrossDSOCFI() : ModulePass(ID) {
48     initializeCrossDSOCFIPass(*PassRegistry::getPassRegistry());
49   }
50 
51   MDNode *VeryLikelyWeights;
52 
53   ConstantInt *extractNumericTypeId(MDNode *MD);
54   void buildCFICheck(Module &M);
55   bool runOnModule(Module &M) override;
56 };
57 
58 } // anonymous namespace
59 
60 INITIALIZE_PASS_BEGIN(CrossDSOCFI, "cross-dso-cfi", "Cross-DSO CFI", false,
61                       false)
62 INITIALIZE_PASS_END(CrossDSOCFI, "cross-dso-cfi", "Cross-DSO CFI", false, false)
63 char CrossDSOCFI::ID = 0;
64 
65 ModulePass *llvm::createCrossDSOCFIPass() { return new CrossDSOCFI; }
66 
67 /// Extracts a numeric type identifier from an MDNode containing type metadata.
68 ConstantInt *CrossDSOCFI::extractNumericTypeId(MDNode *MD) {
69   // This check excludes vtables for classes inside anonymous namespaces.
70   auto TM = dyn_cast<ValueAsMetadata>(MD->getOperand(1));
71   if (!TM)
72     return nullptr;
73   auto C = dyn_cast_or_null<ConstantInt>(TM->getValue());
74   if (!C) return nullptr;
75   // We are looking for i64 constants.
76   if (C->getBitWidth() != 64) return nullptr;
77 
78   return C;
79 }
80 
81 /// buildCFICheck - emits __cfi_check for the current module.
82 void CrossDSOCFI::buildCFICheck(Module &M) {
83   // FIXME: verify that __cfi_check ends up near the end of the code section,
84   // but before the jump slots created in LowerTypeTests.
85   SetVector<uint64_t> TypeIds;
86   SmallVector<MDNode *, 2> Types;
87   for (GlobalObject &GO : M.global_objects()) {
88     Types.clear();
89     GO.getMetadata(LLVMContext::MD_type, Types);
90     for (MDNode *Type : Types) {
91       // Sanity check. GO must not be a function declaration.
92       assert(!isa<Function>(&GO) || !cast<Function>(&GO)->isDeclaration());
93 
94       if (ConstantInt *TypeId = extractNumericTypeId(Type))
95         TypeIds.insert(TypeId->getZExtValue());
96     }
97   }
98 
99   NamedMDNode *CfiFunctionsMD = M.getNamedMetadata("cfi.functions");
100   if (CfiFunctionsMD) {
101     for (auto Func : CfiFunctionsMD->operands()) {
102       assert(Func->getNumOperands() >= 2);
103       for (unsigned I = 2; I < Func->getNumOperands(); ++I)
104         if (ConstantInt *TypeId =
105                 extractNumericTypeId(cast<MDNode>(Func->getOperand(I).get())))
106           TypeIds.insert(TypeId->getZExtValue());
107     }
108   }
109 
110   LLVMContext &Ctx = M.getContext();
111   Constant *C = M.getOrInsertFunction(
112       "__cfi_check", Type::getVoidTy(Ctx), Type::getInt64Ty(Ctx),
113       Type::getInt8PtrTy(Ctx), Type::getInt8PtrTy(Ctx));
114   Function *F = dyn_cast<Function>(C);
115   // Take over the existing function. The frontend emits a weak stub so that the
116   // linker knows about the symbol; this pass replaces the function body.
117   F->deleteBody();
118   F->setAlignment(4096);
119 
120   Triple T(M.getTargetTriple());
121   if (T.isARM() || T.isThumb())
122     F->addFnAttr("target-features", "+thumb-mode");
123 
124   auto args = F->arg_begin();
125   Value &CallSiteTypeId = *(args++);
126   CallSiteTypeId.setName("CallSiteTypeId");
127   Value &Addr = *(args++);
128   Addr.setName("Addr");
129   Value &CFICheckFailData = *(args++);
130   CFICheckFailData.setName("CFICheckFailData");
131   assert(args == F->arg_end());
132 
133   BasicBlock *BB = BasicBlock::Create(Ctx, "entry", F);
134   BasicBlock *ExitBB = BasicBlock::Create(Ctx, "exit", F);
135 
136   BasicBlock *TrapBB = BasicBlock::Create(Ctx, "fail", F);
137   IRBuilder<> IRBFail(TrapBB);
138   Constant *CFICheckFailFn = M.getOrInsertFunction(
139       "__cfi_check_fail", Type::getVoidTy(Ctx), Type::getInt8PtrTy(Ctx),
140       Type::getInt8PtrTy(Ctx));
141   IRBFail.CreateCall(CFICheckFailFn, {&CFICheckFailData, &Addr});
142   IRBFail.CreateBr(ExitBB);
143 
144   IRBuilder<> IRBExit(ExitBB);
145   IRBExit.CreateRetVoid();
146 
147   IRBuilder<> IRB(BB);
148   SwitchInst *SI = IRB.CreateSwitch(&CallSiteTypeId, TrapBB, TypeIds.size());
149   for (uint64_t TypeId : TypeIds) {
150     ConstantInt *CaseTypeId = ConstantInt::get(Type::getInt64Ty(Ctx), TypeId);
151     BasicBlock *TestBB = BasicBlock::Create(Ctx, "test", F);
152     IRBuilder<> IRBTest(TestBB);
153     Function *BitsetTestFn = Intrinsic::getDeclaration(&M, Intrinsic::type_test);
154 
155     Value *Test = IRBTest.CreateCall(
156         BitsetTestFn, {&Addr, MetadataAsValue::get(
157                                   Ctx, ConstantAsMetadata::get(CaseTypeId))});
158     BranchInst *BI = IRBTest.CreateCondBr(Test, ExitBB, TrapBB);
159     BI->setMetadata(LLVMContext::MD_prof, VeryLikelyWeights);
160 
161     SI->addCase(CaseTypeId, TestBB);
162     ++NumTypeIds;
163   }
164 }
165 
166 bool CrossDSOCFI::runOnModule(Module &M) {
167   if (skipModule(M))
168     return false;
169 
170   VeryLikelyWeights =
171     MDBuilder(M.getContext()).createBranchWeights((1U << 20) - 1, 1);
172   if (M.getModuleFlag("Cross-DSO CFI") == nullptr)
173     return false;
174   buildCFICheck(M);
175   return true;
176 }
177 
178 PreservedAnalyses CrossDSOCFIPass::run(Module &M, ModuleAnalysisManager &AM) {
179   CrossDSOCFI Impl;
180   bool Changed = Impl.runOnModule(M);
181   if (!Changed)
182     return PreservedAnalyses::all();
183   return PreservedAnalyses::none();
184 }
185