IndirectThunks.h

Go to the documentation of this file.

//===---- IndirectThunks.h - Indirect thunk insertion helpers ---*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
///
/// \file
/// Contains a base ThunkInserter class that simplifies injection of MI thunks
/// as well as a default implementation of MachineFunctionPass wrapping
/// several `ThunkInserter`s for targets to extend.
///
//===----------------------------------------------------------------------===//
 
#ifndef LLVM_CODEGEN_INDIRECTTHUNKS_H
#define LLVM_CODEGEN_INDIRECTTHUNKS_H
 
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineModuleInfo.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Module.h"
 
namespace llvm {
 
/// This class assists in inserting MI thunk functions into the module and
/// rewriting the existing machine functions to call these thunks.
///
/// One of the common cases is implementing security mitigations that involve
/// replacing some machine code patterns with calls to special thunk functions.
///
/// Inserting a module pass late in the codegen pipeline may increase memory
/// usage, as it serializes the transformations and forces preceding passes to
/// produce machine code for all functions before running the module pass.
/// For that reason, ThunkInserter can be driven by a MachineFunctionPass by
/// passing one MachineFunction at a time to its `run(MMI, MF)` method.
/// Then, the derived class should
/// * call createThunkFunction from its insertThunks method exactly once for
///   each of the thunk functions to be inserted
/// * populate the thunk in its populateThunk method
///
/// Note that if some other pass is responsible for rewriting the functions,
/// the insertThunks method may simply create all possible thunks at once,
/// probably postponed until the first occurrence of possibly affected MF.
///
/// Alternatively, insertThunks method can rewrite MF by itself and only insert
/// the thunks being called. In that case InsertedThunks variable can be used
/// to track which thunks were already inserted.
///
/// In any case, the thunk function has to be inserted on behalf of some other
/// function and then populated on its own "iteration" later - this is because
/// MachineFunctionPass will see the newly created functions, but they first
/// have to go through the preceding passes from the same pass manager,
/// possibly even through the instruction selector.
//
// FIXME Maybe implement a documented and less surprising way of modifying
//       the module from a MachineFunctionPass that is restricted to inserting
//       completely new functions to the module.
template <typename Derived, typename InsertedThunksTy = bool>
class ThunkInserter {
  Derived &getDerived() { return *static_cast<Derived *>(this); }
 
  // A variable used to track whether (and possible which) thunks have been
  // inserted so far. InsertedThunksTy is usually a bool, but can be other types
  // to represent more than one type of thunk. Requires an |= operator to
  // accumulate results.
  InsertedThunksTy InsertedThunks{};
 
protected:
  // Interface for subclasses to use.
 
  /// Create an empty thunk function.
  ///
  /// The new function will eventually be passed to populateThunk. If multiple
  /// thunks are created, populateThunk can distinguish them by their names.
  void createThunkFunction(MachineModuleInfo &MMI, StringRef Name,
                           bool Comdat = true, StringRef TargetAttrs = "");
 
protected:
  // Interface for subclasses to implement.
  //
  // Note: all functions are non-virtual and are called via getDerived().
 
  /// Returns common prefix for thunk function's names.
  const char *getThunkPrefix(); // undefined
 
  /// Checks if MF may use thunks (true - maybe, false - definitely not).
  bool mayUseThunk(const MachineFunction &MF); // undefined
 
  /// Rewrites the function if necessary, returns the set of thunks added.
  InsertedThunksTy insertThunks(MachineModuleInfo &MMI, MachineFunction &MF,
                                InsertedThunksTy ExistingThunks); // undefined
 
  /// Populate the thunk function with instructions.
  ///
  /// If multiple thunks are created, the content that must be inserted in the
  /// thunk function body should be derived from the MF's name.
  ///
  /// Depending on the preceding passes in the pass manager, by the time
  /// populateThunk is called, MF may have a few target-specific instructions
  /// (such as a single MBB containing the return instruction).
  void populateThunk(MachineFunction &MF); // undefined
 
public:
  // return `true` if `MMI` or `MF` was modified
  bool run(MachineModuleInfo &MMI, MachineFunction &MF);
};
 
template <typename Derived, typename InsertedThunksTy>
void ThunkInserter<Derived, InsertedThunksTy>::createThunkFunction(
    MachineModuleInfo &MMI, StringRef Name, bool Comdat,
    StringRef TargetAttrs) {
  assert(Name.starts_with(getDerived().getThunkPrefix()) &&
         "Created a thunk with an unexpected prefix!");
 
  Module &M = const_cast<Module &>(*MMI.getModule());
  LLVMContext &Ctx = M.getContext();
  auto *Type = FunctionType::get(Type::getVoidTy(Ctx), false);
  Function *F = Function::Create(Type,
                                 Comdat ? GlobalValue::LinkOnceODRLinkage
                                        : GlobalValue::InternalLinkage,
                                 Name, &M);
  if (Comdat) {
    F->setVisibility(GlobalValue::HiddenVisibility);
    F->setComdat(M.getOrInsertComdat(Name));
  }
 
  // Add Attributes so that we don't create a frame, unwind information, or
  // inline.
  AttrBuilder B(Ctx);
  B.addAttribute(llvm::Attribute::NoUnwind);
  B.addAttribute(llvm::Attribute::Naked);
  if (TargetAttrs != "")
    B.addAttribute("target-features", TargetAttrs);
  F->addFnAttrs(B);
 
  // Populate our function a bit so that we can verify.
  BasicBlock *Entry = BasicBlock::Create(Ctx, "entry", F);
  IRBuilder<> Builder(Entry);
 
  Builder.CreateRetVoid();
 
  // MachineFunctions aren't created automatically for the IR-level constructs
  // we already made. Create them and insert them into the module.
  MachineFunction &MF = MMI.getOrCreateMachineFunction(*F);
  // A MachineBasicBlock must not be created for the Entry block; code
  // generation from an empty naked function in C source code also does not
  // generate one.  At least GlobalISel asserts if this invariant isn't
  // respected.
 
  // Set MF properties. We never use vregs.
  MF.getProperties().setNoVRegs();
}
 
template <typename Derived, typename InsertedThunksTy>
bool ThunkInserter<Derived, InsertedThunksTy>::run(MachineModuleInfo &MMI,
                                                   MachineFunction &MF) {
  // If MF is not a thunk, check to see if we need to insert a thunk.
  if (!MF.getName().starts_with(getDerived().getThunkPrefix())) {
    // Only add thunks if one of the functions may use them.
    if (!getDerived().mayUseThunk(MF))
      return false;
 
    // The target can use InsertedThunks to detect whether relevant thunks
    // have already been inserted.
    // FIXME: Provide the way for insertThunks to notify us whether it changed
    //        the MF, instead of conservatively assuming it did.
    InsertedThunks |= getDerived().insertThunks(MMI, MF, InsertedThunks);
    return true;
  }
 
  // If this *is* a thunk function, we need to populate it with the correct MI.
  getDerived().populateThunk(MF);
  return true;
}
 
/// Basic implementation of MachineFunctionPass wrapping one or more
/// `ThunkInserter`s passed as type parameters.
template <typename... Inserters>
class ThunkInserterPass : public MachineFunctionPass {
protected:
  std::tuple<Inserters...> TIs;
 
  ThunkInserterPass(char &ID) : MachineFunctionPass(ID) {}
 
public:
  bool doInitialization(Module &M) override {
    TIs = {};
    return false;
  }
 
  bool runOnMachineFunction(MachineFunction &MF) override {
    auto &MMI = getAnalysis<MachineModuleInfoWrapperPass>().getMMI();
    return runTIs(MMI, MF, TIs);
  }
 
private:
  template <typename... ThunkInserterT>
  static bool runTIs(MachineModuleInfo &MMI, MachineFunction &MF,
                     std::tuple<ThunkInserterT...> &ThunkInserters) {
    return (0 | ... | std::get<ThunkInserterT>(ThunkInserters).run(MMI, MF));
  }
};
 
} // namespace llvm
 
#endif