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       Type::getInt8PtrTy(Ctx), Type::getInt8PtrTy(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 =
114       M.getOrInsertFunction("__cfi_check_fail", Type::getVoidTy(Ctx),
115                             Type::getInt8PtrTy(Ctx), Type::getInt8PtrTy(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 }
157