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