1 //===---- IndirectThunks.h - Indirect thunk insertion helpers ---*- C++ -*-===//
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 /// \file
10 /// Contains a base ThunkInserter class that simplifies injection of MI thunks
11 /// as well as a default implementation of MachineFunctionPass wrapping
12 /// several `ThunkInserter`s for targets to extend.
13 ///
14 //===----------------------------------------------------------------------===//
15
16 #ifndef LLVM_CODEGEN_INDIRECTTHUNKS_H
17 #define LLVM_CODEGEN_INDIRECTTHUNKS_H
18
19 #include "llvm/CodeGen/MachineFunction.h"
20 #include "llvm/CodeGen/MachineFunctionPass.h"
21 #include "llvm/CodeGen/MachineModuleInfo.h"
22 #include "llvm/IR/IRBuilder.h"
23 #include "llvm/IR/Module.h"
24
25 namespace llvm {
26
27 /// This class assists in inserting MI thunk functions into the module and
28 /// rewriting the existing machine functions to call these thunks.
29 ///
30 /// One of the common cases is implementing security mitigations that involve
31 /// replacing some machine code patterns with calls to special thunk functions.
32 ///
33 /// Inserting a module pass late in the codegen pipeline may increase memory
34 /// usage, as it serializes the transformations and forces preceding passes to
35 /// produce machine code for all functions before running the module pass.
36 /// For that reason, ThunkInserter can be driven by a MachineFunctionPass by
37 /// passing one MachineFunction at a time to its `run(MMI, MF)` method.
38 /// Then, the derived class should
39 /// * call createThunkFunction from its insertThunks method exactly once for
40 /// each of the thunk functions to be inserted
41 /// * populate the thunk in its populateThunk method
42 ///
43 /// Note that if some other pass is responsible for rewriting the functions,
44 /// the insertThunks method may simply create all possible thunks at once,
45 /// probably postponed until the first occurrence of possibly affected MF.
46 ///
47 /// Alternatively, insertThunks method can rewrite MF by itself and only insert
48 /// the thunks being called. In that case InsertedThunks variable can be used
49 /// to track which thunks were already inserted.
50 ///
51 /// In any case, the thunk function has to be inserted on behalf of some other
52 /// function and then populated on its own "iteration" later - this is because
53 /// MachineFunctionPass will see the newly created functions, but they first
54 /// have to go through the preceding passes from the same pass manager,
55 /// possibly even through the instruction selector.
56 //
57 // FIXME Maybe implement a documented and less surprising way of modifying
58 // the module from a MachineFunctionPass that is restricted to inserting
59 // completely new functions to the module.
60 template <typename Derived, typename InsertedThunksTy = bool>
61 class ThunkInserter {
getDerived()62 Derived &getDerived() { return *static_cast<Derived *>(this); }
63
64 // A variable used to track whether (and possible which) thunks have been
65 // inserted so far. InsertedThunksTy is usually a bool, but can be other types
66 // to represent more than one type of thunk. Requires an |= operator to
67 // accumulate results.
68 InsertedThunksTy InsertedThunks;
69
70 protected:
71 // Interface for subclasses to use.
72
73 /// Create an empty thunk function.
74 ///
75 /// The new function will eventually be passed to populateThunk. If multiple
76 /// thunks are created, populateThunk can distinguish them by their names.
77 void createThunkFunction(MachineModuleInfo &MMI, StringRef Name,
78 bool Comdat = true, StringRef TargetAttrs = "");
79
80 protected:
81 // Interface for subclasses to implement.
82 //
83 // Note: all functions are non-virtual and are called via getDerived().
84 // Note: only doInitialization() has an implementation.
85
86 /// Initializes thunk inserter.
doInitialization(Module & M)87 void doInitialization(Module &M) {}
88
89 /// Returns common prefix for thunk function's names.
90 const char *getThunkPrefix(); // undefined
91
92 /// Checks if MF may use thunks (true - maybe, false - definitely not).
93 bool mayUseThunk(const MachineFunction &MF); // undefined
94
95 /// Rewrites the function if necessary, returns the set of thunks added.
96 InsertedThunksTy insertThunks(MachineModuleInfo &MMI, MachineFunction &MF,
97 InsertedThunksTy ExistingThunks); // undefined
98
99 /// Populate the thunk function with instructions.
100 ///
101 /// If multiple thunks are created, the content that must be inserted in the
102 /// thunk function body should be derived from the MF's name.
103 ///
104 /// Depending on the preceding passes in the pass manager, by the time
105 /// populateThunk is called, MF may have a few target-specific instructions
106 /// (such as a single MBB containing the return instruction).
107 void populateThunk(MachineFunction &MF); // undefined
108
109 public:
init(Module & M)110 void init(Module &M) {
111 InsertedThunks = InsertedThunksTy{};
112 getDerived().doInitialization(M);
113 }
114 // return `true` if `MMI` or `MF` was modified
115 bool run(MachineModuleInfo &MMI, MachineFunction &MF);
116 };
117
118 template <typename Derived, typename InsertedThunksTy>
createThunkFunction(MachineModuleInfo & MMI,StringRef Name,bool Comdat,StringRef TargetAttrs)119 void ThunkInserter<Derived, InsertedThunksTy>::createThunkFunction(
120 MachineModuleInfo &MMI, StringRef Name, bool Comdat,
121 StringRef TargetAttrs) {
122 assert(Name.starts_with(getDerived().getThunkPrefix()) &&
123 "Created a thunk with an unexpected prefix!");
124
125 Module &M = const_cast<Module &>(*MMI.getModule());
126 LLVMContext &Ctx = M.getContext();
127 auto *Type = FunctionType::get(Type::getVoidTy(Ctx), false);
128 Function *F = Function::Create(Type,
129 Comdat ? GlobalValue::LinkOnceODRLinkage
130 : GlobalValue::InternalLinkage,
131 Name, &M);
132 if (Comdat) {
133 F->setVisibility(GlobalValue::HiddenVisibility);
134 F->setComdat(M.getOrInsertComdat(Name));
135 }
136
137 // Add Attributes so that we don't create a frame, unwind information, or
138 // inline.
139 AttrBuilder B(Ctx);
140 B.addAttribute(llvm::Attribute::NoUnwind);
141 B.addAttribute(llvm::Attribute::Naked);
142 if (TargetAttrs != "")
143 B.addAttribute("target-features", TargetAttrs);
144 F->addFnAttrs(B);
145
146 // Populate our function a bit so that we can verify.
147 BasicBlock *Entry = BasicBlock::Create(Ctx, "entry", F);
148 IRBuilder<> Builder(Entry);
149
150 Builder.CreateRetVoid();
151
152 // MachineFunctions aren't created automatically for the IR-level constructs
153 // we already made. Create them and insert them into the module.
154 MachineFunction &MF = MMI.getOrCreateMachineFunction(*F);
155 // A MachineBasicBlock must not be created for the Entry block; code
156 // generation from an empty naked function in C source code also does not
157 // generate one. At least GlobalISel asserts if this invariant isn't
158 // respected.
159
160 // Set MF properties. We never use vregs...
161 MF.getProperties().set(MachineFunctionProperties::Property::NoVRegs);
162 }
163
164 template <typename Derived, typename InsertedThunksTy>
run(MachineModuleInfo & MMI,MachineFunction & MF)165 bool ThunkInserter<Derived, InsertedThunksTy>::run(MachineModuleInfo &MMI,
166 MachineFunction &MF) {
167 // If MF is not a thunk, check to see if we need to insert a thunk.
168 if (!MF.getName().starts_with(getDerived().getThunkPrefix())) {
169 // Only add thunks if one of the functions may use them.
170 if (!getDerived().mayUseThunk(MF))
171 return false;
172
173 // The target can use InsertedThunks to detect whether relevant thunks
174 // have already been inserted.
175 // FIXME: Provide the way for insertThunks to notify us whether it changed
176 // the MF, instead of conservatively assuming it did.
177 InsertedThunks |= getDerived().insertThunks(MMI, MF, InsertedThunks);
178 return true;
179 }
180
181 // If this *is* a thunk function, we need to populate it with the correct MI.
182 getDerived().populateThunk(MF);
183 return true;
184 }
185
186 /// Basic implementation of MachineFunctionPass wrapping one or more
187 /// `ThunkInserter`s passed as type parameters.
188 template <typename... Inserters>
189 class ThunkInserterPass : public MachineFunctionPass {
190 protected:
191 std::tuple<Inserters...> TIs;
192
ThunkInserterPass(char & ID)193 ThunkInserterPass(char &ID) : MachineFunctionPass(ID) {}
194
195 public:
doInitialization(Module & M)196 bool doInitialization(Module &M) override {
197 initTIs(M, TIs);
198 return false;
199 }
200
runOnMachineFunction(MachineFunction & MF)201 bool runOnMachineFunction(MachineFunction &MF) override {
202 auto &MMI = getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
203 return runTIs(MMI, MF, TIs);
204 }
205
206 private:
207 template <typename... ThunkInserterT>
initTIs(Module & M,std::tuple<ThunkInserterT...> & ThunkInserters)208 static void initTIs(Module &M,
209 std::tuple<ThunkInserterT...> &ThunkInserters) {
210 (..., std::get<ThunkInserterT>(ThunkInserters).init(M));
211 }
212
213 template <typename... ThunkInserterT>
runTIs(MachineModuleInfo & MMI,MachineFunction & MF,std::tuple<ThunkInserterT...> & ThunkInserters)214 static bool runTIs(MachineModuleInfo &MMI, MachineFunction &MF,
215 std::tuple<ThunkInserterT...> &ThunkInserters) {
216 return (0 | ... | std::get<ThunkInserterT>(ThunkInserters).run(MMI, MF));
217 }
218 };
219
220 } // namespace llvm
221
222 #endif
223