Implement pagerando wrapper functions to initialize POT register

[rinon]

Summary

Pagerando with PIP (Position Independent Pages) segments requires that all
externally visible (including address-taken) functions be wrapped with a
function that initializes the POT (Page Offset Table) register for use in the
protected PIP segment. Randomly located PIP segments cannot load the address of
the POT directly since that would require either being located at a fixed
PC-relative offset to a data segment or mutable code, both of which are not
acceptable options.

The PagerandoWrappersPass creates pagerando wrappers for all functions that
require a wrapper and marks all original functions with the PagerandoBinned
attribute so they will be placed into PIP bins when processing machine IR.

This patch set (D37580, D37581, D37582, D37583, D37584, D37585, D37586, D37587)
is a first draft of the pagerando implementation described in
http://lists.llvm.org/pipermail/llvm-dev/2017-June/113794.html.

[rinon]

This patch set (D37580, D37581, D37582, D37583, D37584, D37585, D37586, D37587) is a first draft of the pagerando implementation described in http://lists.llvm.org/pipermail/llvm-dev/2017-June/113794.html.

[vlad.tsyrklevich]

small style nit: omit curly brackets for single-statement if/for/etc. statements.

lib/Transforms/IPO/PagerandoWrappers.cpp

77

Could you explain this condition?

129

What's the purpose of setting the Function section to the default here?

295

nit: https://llvm.org/docs/CodingStandards.html#prefer-preincrement

[rinon]

Thanks for the comments. I'll fix the nits right away.

lib/Transforms/IPO/PagerandoWrappers.cpp

77

The last condition handles functions that consist of a single unreachable instruction. We don't every want to create a wrapper for these functions, since they shouldn't be emitted. I'll add a better comment here.

129

We want to assign these functions to unique section "bins" later in pagerando. Because we want to control the section that pagerando'd functions end up in to make sure it can be randomly placed, we don't respect the assigned section for the wrapped function. However, we do preserve any section assignment for the wrapper function that takes the given function's place (since we copy attributes (line 156).

[rinon]

• Fix style nits

[nlewycky]

include/llvm/IR/Attributes.td

182–186

I don't see a matching LangRef change.

I don't see a matching Verifier change. These are only applicable to functions (as opposed to return values or arguments)? Are these two attributes mutually exclusive?

lib/Transforms/IPO/PagerandoWrappers.cpp

61

Consider choosing SmallVectorImpl<Use *>?
http://llvm.org/docs/ProgrammersManual.html#sequential-containers-std-vector-std-list-etc

77

That's not quite what this tests for. This would return true for a function that does work but ends in unreachable, say printing a fatal error message and calling abort. If you want to check for the 1-instruction function with unreachable, use isa<UnreachableInst>(&F.getEntryBlock().front()).

141

Please move the comment to before the code instead of using two-column code // comments.

151

Use llvm::Twine which is "a lightweight data structure for efficiently representing the concatenation of temporary values as strings." Value::setName takes a Twine argument. Something like:

F.setName(Twine(F.getName(), F.isVarArg() ? OrigVASuffix : OrigSuffix));

[rinon]

• Remove pagerando_wrapper attribute and rename pagerando_binned
• Use SmallVector for address useses
• Fix comment style
• Do not pagerando trivial functions with only trap+unreachable
• Use Twine for function name

[rinon]

Thanks for taking a look at this, Nick. I think this revision should address all comments so far.

include/llvm/IR/Attributes.td

182–186

I added a LangRef description.

After giving attributes a second look, I removed the PagerandoWrapper attribute as we don't need it at this point. The Pagerando attribute (formerly PagerandoBinned) is checked in the Verifier under isFuncOnlyAttr. I don't think we need to verify any other properties of the attribute, now that PagerandoWrapper is gone.

lib/Transforms/IPO/PagerandoWrappers.cpp

77

Good point. I intended to catch functions like the following (generated for abstract, non-base destructors):

call void @llvm.dbg.value...
tail call void @llvm.trap()
unreachable

Randomizing the page address of such functions is pointless, so we can save code size by not enabling pagerando for functions like these and thus not need wrappers. I've now made this optimization more explicit. If you think we should just leave this out for the sake of simplicity, I'm fine with that as well.

[rinon]

• Rebase
• Don't apply pagerando to naked functions

A naked function should not get a pagerando wrapper, so cannot be placed in a
bin if it would require a wrapper. Naked functions are also used for CFI jump
tables (on non-windows platforms), which cannot be placed in pagerando bins
because they are PC-relative direct jumps in inline ASM.

[peter.smith]

docs/LangRef.rst

1557 ↗

(On Diff #118511)

Is section the right word to use here? The compiler will place the function in a section, the linker will place that in an output section that will presumably be described, at least in ELF terms as a segment.

lib/Transforms/IPO/PagerandoWrappers.cpp

102

Are the issues surrounding comdat that at code-generation time that we can't guarantee at link time which comdat group is selected and it may not be the one with wrappers? If I'm write it would be useful to know, although not necessarily in this review, what ideas you have to resolve this?

[rinon]

docs/LangRef.rst

1557 ↗

(On Diff #118511)

I'm open to either section or segment. I wasn't sure which one made more sense, since the compiler only cares about sections, and in this case the section and segment need to be the same.

lib/Transforms/IPO/PagerandoWrappers.cpp

102

Not quite. We don't really care if non-pagerando comdat groups are selected since this implementation should only refer to wrapped, external functions and cannot call functions in pagerando bins directly. To make sure that the wrappers are selected if their corresponding implementations are selected, wrappers need to go in the same comdat group as their implementation function, which is easy. We would need to treat all comdat functions as externally visible and potentially replaced to handle the case where the external implementation is selected.

However, if we mix comdat groups and pagerando bins, it turns out that we cannot control the layout of the resulting section precisely. The linker may (and in my testing did) reorder resolved comdat sections when merging them into our named output section (pagerando bin), which ruins the section offsets we computed during code generation for referencing binned functions.

Looking closer at the gold plugin interface specifically, the gold documentation seems to promise that all comdat groups will be resolved to the IR input file version, and no comdat symbols should be in the LTO output file (see COMDAT Sections in https://gcc.gnu.org/wiki/whopr/driver). So, as long as we're doing full LTO with gold, we should be able to resolve comdat before assigning to bins and the whole issue goes away. However, as far as I can tell, LLVM doesn't actually do that right now. I built some simple C++ with full (gold) LTO, saved temps, and the LTO output still included COMDAT sections.

Basically, I think we either need cooperation from the linker in how comdat sections are merged together, or we need to resolve and remove comdat during LTO. Alternatively, if we add linker support for new relocations (section->symbol offset and POT entry index), I think this issue goes away, but that is more involved and requires upstreaming to multiple projects.

[rinon]

Ensure that pagerando wrapper functions are placed before binned functions

Gold assumes that the first executable section should be laid out first in the
file, but we want the normal .text section to come first and all pagerando bins
to be placed at the end of the file. To ensure that the first function emitted
will be in the normal .text section we prepend wrapper functions (which are
emitted to the normal .text) onto the Module function list.

[rinon]

Explicitly propagate applicable function attributes to pagerando wrappers.

[rinon]

• Do not wrap thunk functions
• Restore propagation of function attributes not in attribute list

Attributes are now propagated by default to pagerando wrappers and attributes
that do not apply to wrappers are blacklisted explicitly.

[rinon]

• Only wrap functions marked with pagerando attribute.
  
  The frontend should now explicitly opt in to pagerando with a function attribute, allowing control of pagerando on a per-function granularity via the existing sanitizer flag infrastructure.
• Refactor replacement of function references with wrappers
• Don't skip aliases when determining pagerando wrappers
• Don't apply pagerando to trivial functions
  
  Wrapping trivial (single forwarding call) functions results in a large proportional increase in overhead for these functions. Not applying pagerando in this case results in effectively equivalent security, assuming that the trivial function does not contain useful gadgets, since an attacker could reuse the whole function. This optimization can be enabled via a command-line option.
• Rebase

[rinon]

Rebase

[rinon]

Rebase

• Update pagerando wrapper tests

[rinon]

Rebase

include/llvm/Bitcode/LLVMBitCodes.h

enum AttributeKindCodes {
   ATTR_KIND_SWIFT_ERROR = 47,
   ATTR_KIND_NO_RECURSE = 48,
   ATTR_KIND_INACCESSIBLEMEM_ONLY = 49,
   ATTR_KIND_INACCESSIBLEMEM_OR_ARGMEMONLY = 50,
   ATTR_KIND_ALLOC_SIZE = 51,
   ATTR_KIND_WRITEONLY = 52,
   ATTR_KIND_SPECULATABLE = 53,
   ATTR_KIND_STRICT_FP = 54,
+  ATTR_KIND_PAGERANDO_WRAPPER = 55,
+  ATTR_KIND_PAGERANDO_BINNED = 56,
 };
 
 enum ComdatSelectionKindCodes {
   COMDAT_SELECTION_KIND_ANY = 1,
   COMDAT_SELECTION_KIND_EXACT_MATCH = 2,
   COMDAT_SELECTION_KIND_LARGEST = 3,
   COMDAT_SELECTION_KIND_NO_DUPLICATES = 4,
   COMDAT_SELECTION_KIND_SAME_SIZE = 5,

include/llvm/IR/Attributes.td

 def UWTable : EnumAttr<"uwtable">;
 
 /// Function only writes to memory.
 def WriteOnly : EnumAttr<"writeonly">;
 
 /// Zero extended before/after call.
 def ZExt : EnumAttr<"zeroext">;
 
+/// Function is a Pagerando wrapper.
+def PagerandoWrapper : EnumAttr<"pagerando_wrapper">;
+
+/// Function is placed into a Pagerando bin.
+def PagerandoBinned : EnumAttr<"pagerando_binned">;
+
 /// Target-independent string attributes.
 def LessPreciseFPMAD : StrBoolAttr<"less-precise-fpmad">;
 def NoInfsFPMath : StrBoolAttr<"no-infs-fp-math">;
 def NoNansFPMath : StrBoolAttr<"no-nans-fp-math">;
 def UnsafeFPMath : StrBoolAttr<"unsafe-fp-math">;
 def NoJumpTables : StrBoolAttr<"no-jump-tables">;
 def ProfileSampleAccurate : StrBoolAttr<"profile-sample-accurate">;
 

include/llvm/IR/Function.h

public:
 
   /// @brief Determine if the function returns a structure through first
   /// or second pointer argument.
   bool hasStructRetAttr() const {
     return AttributeSets.hasParamAttribute(0, Attribute::StructRet) ||
            AttributeSets.hasParamAttribute(1, Attribute::StructRet);
   }
 
+  /// @brief True if the function is a Pagerando wrapper.
+  bool isPagerandoWrapper() const {
+    return hasFnAttribute(Attribute::PagerandoWrapper);
+  }
+
+  /// @brief True if the function is compatible with Pagerando and should be
+  /// placed into a randomized bin.
+  bool isPagerando() const {
+    return hasFnAttribute(Attribute::PagerandoBinned);
+  }
+
   /// @brief Determine if the parameter or return value is marked with NoAlias
   /// attribute.
   bool returnDoesNotAlias() const {
     return AttributeSets.hasAttribute(AttributeList::ReturnIndex,
                                       Attribute::NoAlias);
   }
   void setReturnDoesNotAlias() {
     addAttribute(AttributeList::ReturnIndex, Attribute::NoAlias);

include/llvm/InitializePasses.h

 void initializeObjCARCAAWrapperPassPass(PassRegistry&);
 void initializeObjCARCAPElimPass(PassRegistry&);
 void initializeObjCARCContractPass(PassRegistry&);
 void initializeObjCARCExpandPass(PassRegistry&);
 void initializeObjCARCOptPass(PassRegistry&);
 void initializeOptimizationRemarkEmitterWrapperPassPass(PassRegistry&);
 void initializeOptimizePHIsPass(PassRegistry&);
 void initializePAEvalPass(PassRegistry&);
+void initializePagerandoWrappersPass(PassRegistry&);
 void initializePEIPass(PassRegistry&);
 void initializePGOIndirectCallPromotionLegacyPassPass(PassRegistry&);
 void initializePGOInstrumentationGenLegacyPassPass(PassRegistry&);
 void initializePGOInstrumentationUseLegacyPassPass(PassRegistry&);
 void initializePGOMemOPSizeOptLegacyPassPass(PassRegistry&);
 void initializePHIEliminationPass(PassRegistry&);
 void initializePartialInlinerLegacyPassPass(PassRegistry&);
 void initializePartiallyInlineLibCallsLegacyPassPass(PassRegistry&);

include/llvm/LinkAllPasses.h

ForcePassLinking() {
       (void) llvm::createLowerSwitchPass();
       (void) llvm::createNaryReassociatePass();
       (void) llvm::createObjCARCAAWrapperPass();
       (void) llvm::createObjCARCAPElimPass();
       (void) llvm::createObjCARCExpandPass();
       (void) llvm::createObjCARCContractPass();
       (void) llvm::createObjCARCOptPass();
       (void) llvm::createPAEvalPass();
+      (void) llvm::createPagerandoWrappersPass();
       (void) llvm::createPromoteMemoryToRegisterPass();
       (void) llvm::createDemoteRegisterToMemoryPass();
       (void) llvm::createPruneEHPass();
       (void) llvm::createPostDomOnlyPrinterPass();
       (void) llvm::createPostDomPrinterPass();
       (void) llvm::createPostDomOnlyViewerPass();
       (void) llvm::createPostDomViewerPass();
       (void) llvm::createReassociatePass();

include/llvm/Transforms/IPO.h

 // IR metadata to reflect the profile.
 ModulePass *createSampleProfileLoaderPass();
 ModulePass *createSampleProfileLoaderPass(StringRef Name);
 
 /// Write ThinLTO-ready bitcode to Str.
 ModulePass *createWriteThinLTOBitcodePass(raw_ostream &Str,
                                           raw_ostream *ThinLinkOS = nullptr);
 
+/// \brief This pass creates Pagerando entry wrappers that properly set up the
+/// POT base register.
+ModulePass *createPagerandoWrappersPass();
+
 } // End llvm namespace
 
 #endif

lib/AsmParser/LLLexer.cpp

#define KEYWORD(STR)                                                           \
   KEYWORD(sanitize_address);
   KEYWORD(sanitize_thread);
   KEYWORD(sanitize_memory);
   KEYWORD(swifterror);
   KEYWORD(swiftself);
   KEYWORD(uwtable);
   KEYWORD(writeonly);
   KEYWORD(zeroext);
+  KEYWORD(pagerando_wrapper);
+  KEYWORD(pagerando_binned);
 
   KEYWORD(type);
   KEYWORD(opaque);
 
   KEYWORD(comdat);
 
   // Comdat types
   KEYWORD(any);

lib/AsmParser/LLParser.cpp

case lltok::kw_sanitize_address:
       B.addAttribute(Attribute::SanitizeAddress); break;
     case lltok::kw_sanitize_thread:
       B.addAttribute(Attribute::SanitizeThread); break;
     case lltok::kw_sanitize_memory:
       B.addAttribute(Attribute::SanitizeMemory); break;
     case lltok::kw_strictfp: B.addAttribute(Attribute::StrictFP); break;
     case lltok::kw_uwtable: B.addAttribute(Attribute::UWTable); break;
     case lltok::kw_writeonly: B.addAttribute(Attribute::WriteOnly); break;
+    case lltok::kw_pagerando_wrapper:
+      B.addAttribute(Attribute::PagerandoWrapper); break;
+    case lltok::kw_pagerando_binned:
+      B.addAttribute(Attribute::PagerandoBinned); break;
 
     // Error handling.
     case lltok::kw_inreg:
     case lltok::kw_signext:
     case lltok::kw_zeroext:
       HaveError |=
         Error(Lex.getLoc(),
               "invalid use of attribute on a function");

lib/AsmParser/LLToken.h

enum Kind {
   kw_sanitize_thread,
   kw_sanitize_memory,
   kw_strictfp,
   kw_swifterror,
   kw_swiftself,
   kw_uwtable,
   kw_writeonly,
   kw_zeroext,
+  kw_pagerando_wrapper,
+  kw_pagerando_binned,
 
   kw_type,
   kw_opaque,
 
   kw_comdat,
 
   // Comdat types
   kw_any,

lib/Bitcode/Reader/BitcodeReader.cpp

static uint64_t getRawAttributeMask(Attribute::AttrKind Val) {
   case Attribute::NoRecurse:       return 1ULL << 48;
   case Attribute::InaccessibleMemOnly:         return 1ULL << 49;
   case Attribute::InaccessibleMemOrArgMemOnly: return 1ULL << 50;
   case Attribute::SwiftSelf:       return 1ULL << 51;
   case Attribute::SwiftError:      return 1ULL << 52;
   case Attribute::WriteOnly:       return 1ULL << 53;
   case Attribute::Speculatable:    return 1ULL << 54;
   case Attribute::StrictFP:        return 1ULL << 55;
+  case Attribute::PagerandoWrapper:return 1ULL << 56;
+  case Attribute::PagerandoBinned: return 1ULL << 57;
   case Attribute::Dereferenceable:
     llvm_unreachable("dereferenceable attribute not supported in raw format");
     break;
   case Attribute::DereferenceableOrNull:
     llvm_unreachable("dereferenceable_or_null attribute not supported in raw "
                      "format");
     break;
   case Attribute::ArgMemOnly:
static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
   case bitc::ATTR_KIND_SWIFT_SELF:
     return Attribute::SwiftSelf;
   case bitc::ATTR_KIND_UW_TABLE:
     return Attribute::UWTable;
   case bitc::ATTR_KIND_WRITEONLY:
     return Attribute::WriteOnly;
   case bitc::ATTR_KIND_Z_EXT:
     return Attribute::ZExt;
+  case bitc::ATTR_KIND_PAGERANDO_WRAPPER:
+    return Attribute::PagerandoWrapper;
+  case bitc::ATTR_KIND_PAGERANDO_BINNED:
+    return Attribute::PagerandoBinned;
   }
 }
 
 Error BitcodeReader::parseAlignmentValue(uint64_t Exponent,
                                          unsigned &Alignment) {
   // Note: Alignment in bitcode files is incremented by 1, so that zero
   // can be used for default alignment.
   if (Exponent > Value::MaxAlignmentExponent + 1)

lib/Bitcode/Writer/BitcodeWriter.cpp

static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
   case Attribute::SwiftSelf:
     return bitc::ATTR_KIND_SWIFT_SELF;
   case Attribute::UWTable:
     return bitc::ATTR_KIND_UW_TABLE;
   case Attribute::WriteOnly:
     return bitc::ATTR_KIND_WRITEONLY;
   case Attribute::ZExt:
     return bitc::ATTR_KIND_Z_EXT;
+  case Attribute::PagerandoWrapper:
+    return bitc::ATTR_KIND_PAGERANDO_WRAPPER;
+  case Attribute::PagerandoBinned:
+    return bitc::ATTR_KIND_PAGERANDO_BINNED;
   case Attribute::EndAttrKinds:
     llvm_unreachable("Can not encode end-attribute kinds marker.");
   case Attribute::None:
     llvm_unreachable("Can not encode none-attribute.");
   }
 
   llvm_unreachable("Trying to encode unknown attribute");
 }

lib/IR/Attributes.cpp

std::string Attribute::getAsString(bool InAttrGrp) const {
   if (hasAttribute(Attribute::SanitizeMemory))
     return "sanitize_memory";
   if (hasAttribute(Attribute::UWTable))
     return "uwtable";
   if (hasAttribute(Attribute::ZExt))
     return "zeroext";
   if (hasAttribute(Attribute::Cold))
     return "cold";
+  if (hasAttribute(Attribute::PagerandoWrapper))
+    return "pagerando_wrapper";
+  if (hasAttribute(Attribute::PagerandoBinned))
+    return "pagerando_binned";
 
   // FIXME: These should be output like this:
   //
   //   align=4
   //   alignstack=8
   //
   if (hasAttribute(Attribute::Alignment)) {
     std::string Result;

lib/IR/Verifier.cpp

static bool isFuncOnlyAttr(Attribute::AttrKind Kind) {
   case Attribute::Convergent:
   case Attribute::ArgMemOnly:
   case Attribute::NoRecurse:
   case Attribute::InaccessibleMemOnly:
   case Attribute::InaccessibleMemOrArgMemOnly:
   case Attribute::AllocSize:
   case Attribute::Speculatable:
   case Attribute::StrictFP:
+  case Attribute::PagerandoWrapper:
+  case Attribute::PagerandoBinned:
     return true;
   default:
     break;
   }
   return false;
 }
 
 /// Return true if this is a function attribute that can also appear on

lib/LTO/LTOBackend.cpp

static void runOldPMPasses(Config &Conf, Module &Mod, TargetMachine *TM,
   PMB.LoopVectorize = true;
   PMB.SLPVectorize = true;
   PMB.OptLevel = Conf.OptLevel;
   PMB.PGOSampleUse = Conf.SampleProfile;
   if (IsThinLTO)
     PMB.populateThinLTOPassManager(passes);
   else
     PMB.populateLTOPassManager(passes);
+  if (TM->isPagerando())
+    passes.add(createPagerandoWrappersPass());
   passes.run(Mod);
 }
 
 bool opt(Config &Conf, TargetMachine *TM, unsigned Task, Module &Mod,
          bool IsThinLTO, ModuleSummaryIndex *ExportSummary,
          const ModuleSummaryIndex *ImportSummary) {
   // FIXME: Plumb the combined index into the new pass manager.
   if (!Conf.OptPipeline.empty())

lib/Transforms/IPO/CMakeLists.txt

add_llvm_library(LLVMipo
   InlineSimple.cpp
   Inliner.cpp
   Internalize.cpp
   LoopExtractor.cpp
   LowerTypeTests.cpp
   MergeFunctions.cpp
   PartialInlining.cpp
   PassManagerBuilder.cpp
+  PagerandoWrappers.cpp
   PruneEH.cpp
   SampleProfile.cpp
   StripDeadPrototypes.cpp
   StripSymbols.cpp
   ThinLTOBitcodeWriter.cpp
   WholeProgramDevirt.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/Transforms/IPO
 
   DEPENDS
   intrinsics_gen
   )

lib/Transforms/IPO/IPO.cpp

void llvm::initializeIPO(PassRegistry &Registry) {
   initializeSimpleInlinerPass(Registry);
   initializeInferFunctionAttrsLegacyPassPass(Registry);
   initializeInternalizeLegacyPassPass(Registry);
   initializeLoopExtractorPass(Registry);
   initializeBlockExtractorPassPass(Registry);
   initializeSingleLoopExtractorPass(Registry);
   initializeLowerTypeTestsPass(Registry);
   initializeMergeFunctionsPass(Registry);
+  initializePagerandoWrappersPass(Registry);
   initializePartialInlinerLegacyPassPass(Registry);
   initializePostOrderFunctionAttrsLegacyPassPass(Registry);
   initializeReversePostOrderFunctionAttrsLegacyPassPass(Registry);
   initializePruneEHPass(Registry);
   initializeStripDeadPrototypesLegacyPassPass(Registry);
   initializeStripSymbolsPass(Registry);
   initializeStripDebugDeclarePass(Registry);
   initializeStripDeadDebugInfoPass(Registry);

lib/Transforms/IPO/PagerandoWrappers.cpp

+//===-- PagerandoWrappers.cpp - Pagerando entry wrappers ------------------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass creates wrappers for pagerando-enabled functions. A function needs
+// a wrapper if it has non-local linkage or its address taken, i.e., if it can
+// be used from outside the module. (As an optimization we could use pointer
+// escape analysis for address-taken functions instead of creating wrappers for
+// all of them.)
+// Vararg functions require special treatment: their variable arguments on the
+// stack need to be preserved even when indirecting through the POT. We replace
+// the original function with a new function that takes an explicit 'va_list'
+// parameter:  foo(int, ...) -> foo$$origva(int, *va_list). The wrapper captures
+// its variable arguments and explicitly passes it to the adapted function to
+// preserve the variable arguments passed by the caller.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/IR/CallSite.h"
+#include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/Transforms/IPO.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+
+using namespace llvm;
+
+#define DEBUG_TYPE "pagerando"
+
+namespace {
+class PagerandoWrappers : public ModulePass {
+public:
+  static char ID;
+  explicit PagerandoWrappers() : ModulePass(ID) {
+    initializePagerandoWrappersPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnModule(Module &M) override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    // Requires nothing, preserves nothing
+    ModulePass::getAnalysisUsage(AU);
+  }
+
+private:
+  static constexpr const char *OrigSuffix = "$$orig";
+  static constexpr const char *OrigVASuffix = "$$origva";
+
+  void processFunction(Function *F);
+  Function *rewriteVarargs(Function &F, Type *&VAListTy);
+  Function *createWrapper(Function &F, const std::vector<Use *> &AddressUses);
+  void createWrapperBody(Function *Wrapper, Function *Callee, Type *VAListTy);
+};
+} // end anonymous namespace
+
+char PagerandoWrappers::ID = 0;
+INITIALIZE_PASS(PagerandoWrappers, "pagerando-wrappers",
+                "Pagerando entry wrappers", false, false)
+
+ModulePass *llvm::createPagerandoWrappersPass() {
+  return new PagerandoWrappers();
+}
+
+static bool skipFunction(const Function &F) {
+  return F.isDeclaration()
+      || F.hasAvailableExternallyLinkage()
+      || F.hasComdat()  // TODO: Support COMDAT
+      || isa<UnreachableInst>(F.getEntryBlock().getTerminator());
+      // Above condition is different from F.doesNotReturn(), which we do not
+      // include (at least for now).
+}
+
+bool PagerandoWrappers::runOnModule(Module &M) {
+  std::vector<Function*> Worklist;
+  for (auto &F : M) {
+    if (!skipFunction(F)) Worklist.push_back(&F);
+  }
+
+  for (auto F : Worklist) {
+    processFunction(F);
+  }
+
+  return !Worklist.empty();
+}
+
+static bool skipFunctionUse(const Use &U);
+static bool IsDirectCallOfBitcast(User *Usr) {
+  auto CE = dyn_cast<ConstantExpr>(Usr);
+  return CE && CE->getOpcode() == Instruction::BitCast
+      && std::all_of(CE->use_begin(), CE->use_end(), skipFunctionUse);
+}
+
+static bool skipFunctionUse(const Use &U) {
+  auto User = U.getUser();
+  auto UserFn = dyn_cast<Function>(User);
+  ImmutableCallSite CS(User);
+
+  return (CS && CS.isCallee(&U))  // Used as the callee
+      || isa<GlobalAlias>(User)   // No need to indirect
+      || isa<BlockAddress>(User)  // Handled in AsmPrinter::EmitBasicBlockStart
+      || (UserFn && UserFn->getPersonalityFn() == U.get()) // Skip pers. fn uses
+      || IsDirectCallOfBitcast(User); // Calls to bitcasted functions end up as direct calls
+}
+
+void PagerandoWrappers::processFunction(Function *F) {
+  std::vector<Use*> AddressUses;
+  for (Use &U : F->uses()) {
+    if (!skipFunctionUse(U)) AddressUses.push_back(&U);
+  }
+
+  if (!F->hasLocalLinkage() || !AddressUses.empty()) {
+    auto Wrapper = createWrapper(*F, AddressUses);
+    Type *VAListTy = nullptr;
+    if (F->isVarArg()) {
+      F = rewriteVarargs(*F, /* out */ VAListTy); // Reassign F, it might have been deleted
+    }
+    createWrapperBody(Wrapper, F, VAListTy);
+  }
+
+  F->setSection("");
+  F->addFnAttr(Attribute::PagerandoBinned);
+}
+
+static void replaceAddressTakenUse(Use *U, Function *F, Function *Wrapper,
+                                   SmallSet<Constant*, 8> &Constants) {
+  if (!U->get()) return; // Already replaced this use?
+
+  if (auto GV = dyn_cast<GlobalVariable>(U->getUser())) {
+    assert(GV->getInitializer() == F);
+    GV->setInitializer(Wrapper);
+  } else if (auto C = dyn_cast<Constant>(U->getUser())) {
+    if (Constants.insert(C).second) {     // Constant::handleOperandChange must
+      C->handleOperandChange(F, Wrapper); // not be called more than once per user
+    }
+  } else {
+    U->set(Wrapper);
+  }
+}
+
+Function *PagerandoWrappers::createWrapper(Function &F,
+                                           const std::vector<Use *> &AddressUses) {
+  std::string Name = F.getName();
+  F.setName(Name + (F.isVarArg() ? OrigVASuffix : OrigSuffix));
+
+  auto Wrapper = Function::Create(F.getFunctionType(), F.getLinkage(), Name,
+                                  F.getParent());
+  Wrapper->copyAttributesFrom(&F);
+  Wrapper->setComdat(F.getComdat());
+
+  Wrapper->addFnAttr(Attribute::PagerandoWrapper);
+  Wrapper->addFnAttr(Attribute::NoInline);
+  Wrapper->addFnAttr(Attribute::OptimizeForSize);
+
+  // +) Calls to a non-local function must go through the wrapper since they
+  //    could be redirected by the dynamic linker (i.e, LD_PRELOAD).
+  // +) Calls to vararg functions must go through the wrapper to ensure that we
+  //    preserve the arguments on the stack when we indirect through the POT.
+  // -) Address-taken uses of local functions might escape, hence we must also
+  //    replace them.
+  if (!F.hasLocalLinkage() || F.isVarArg()) {
+    F.replaceAllUsesWith(Wrapper);
+    if (!F.hasLocalLinkage()) {
+      F.setVisibility(GlobalValue::HiddenVisibility);
+    }
+  } else {
+    assert(!AddressUses.empty());
+    SmallSet<Constant*, 8> Constants;
+    for (auto U : AddressUses) {
+      replaceAddressTakenUse(U, &F, Wrapper, Constants);
+    }
+  }
+
+  return Wrapper;
+}
+
+static SmallVector<VAStartInst*, 1> findVAStarts(Function &F) {
+  SmallVector<VAStartInst*, 1> Insts;
+  for (auto &I : instructions(F)) {
+    if (isa<VAStartInst>(&I)) {
+      Insts.push_back(cast<VAStartInst>(&I));
+    }
+  }
+  return Insts;
+}
+
+static AllocaInst *findAlloca(VAStartInst *VAStart) {
+  Instruction *Inst = VAStart;
+  while (Inst && !isa<AllocaInst>(Inst)) {
+    Inst = dyn_cast<Instruction>(Inst->op_begin());
+  }
+  assert(Inst && "Could not find va_list alloca in var args function");
+  return cast<AllocaInst>(Inst);
+}
+
+static AllocaInst *createVAList(Module *M, IRBuilder<> &Builder, Type *VAListTy) {
+  auto VAListAlloca = Builder.CreateAlloca(VAListTy);
+  Builder.CreateCall(  // @llvm.va_start(i8* <arglist>)
+      Intrinsic::getDeclaration(M, Intrinsic::vastart),
+      {Builder.CreateBitCast(VAListAlloca, Builder.getInt8PtrTy())});
+
+  return VAListAlloca;
+}
+
+static void createVAEndCall(IRBuilder<> &Builder, AllocaInst *VAListAlloca) {
+  Builder.CreateCall(  // @llvm.va_end(i8* <arglist>)
+      Intrinsic::getDeclaration(VAListAlloca->getModule(), Intrinsic::vaend),
+      {Builder.CreateBitCast(VAListAlloca, Builder.getInt8PtrTy())});
+}
+
+static void replaceWithVACopyCall(IRBuilder<> &Builder, VAStartInst *VAStart,
+                                  Argument *VAListArg) {
+  Builder.SetInsertPoint(VAStart);
+  Builder.CreateCall(  // @llvm.va_copy(i8* <destarglist>, i8* <srcarglist>)
+      Intrinsic::getDeclaration(VAStart->getModule(), Intrinsic::vacopy),
+      {VAStart->getArgOperand(0),
+       Builder.CreateBitCast(VAListArg, Builder.getInt8PtrTy())});
+  VAStart->eraseFromParent();
+}
+
+void PagerandoWrappers::createWrapperBody(Function *Wrapper, Function *Callee,
+                                          Type *VAListTy) {
+  auto BB = BasicBlock::Create(Wrapper->getContext(), "", Wrapper);
+  IRBuilder<> Builder(BB);
+
+  // Arguments
+  SmallVector<Value*, 8> Args;
+  for (auto &A : Wrapper->args()) {
+    Args.push_back(&A);
+  }
+
+  // va_list alloca and va_start
+  AllocaInst* VAListAlloca;
+  if (VAListTy) {
+    VAListAlloca = createVAList(Wrapper->getParent(), Builder, VAListTy);
+    Args.push_back(VAListAlloca);
+  }
+
+  // Call
+  auto Call = Builder.CreateCall(Callee, Args);
+  Call->setCallingConv(Callee->getCallingConv());
+
+  // va_end
+  if (VAListTy) {
+    createVAEndCall(Builder, VAListAlloca);
+  }
+
+  // Return
+  if (Wrapper->getReturnType()->isVoidTy()) {
+    Builder.CreateRetVoid();
+  } else {
+    Builder.CreateRet(Call);
+  }
+}
+
+// Replaces the original function with a new function that takes a va_list
+// parameter but is not varargs:  foo(int, ...) -> foo$$origva(int, *va_list)
+Function *PagerandoWrappers::rewriteVarargs(Function &F, Type *&VAListTy) {
+  auto VAStarts = findVAStarts(F);
+  if (VAStarts.empty()) return &F;
+
+  // Determine va_list type
+  auto VAListAlloca = findAlloca(VAStarts[0]);
+  VAListTy = VAListAlloca->getAllocatedType();
+
+  // Adapt function type
+  auto FTy = F.getFunctionType();
+  SmallVector<Type*, 8> Params(FTy->param_begin(), FTy->param_end());
+  Params.push_back(VAListTy->getPointerTo());
+  auto NonVAFty = FunctionType::get(FTy->getReturnType(), Params, false);
+
+  // Create new function definition
+  auto NF = Function::Create(NonVAFty, F.getLinkage(), "", F.getParent());
+  NF->takeName(&F);
+  NF->copyAttributesFrom(&F);
+  NF->setComdat(F.getComdat());
+  NF->setSubprogram(F.getSubprogram());
+
+  // Move basic blocks into new function; F is now dysfunctional
+  NF->getBasicBlockList().splice(NF->begin(), F.getBasicBlockList());
+
+  // Adapt arguments (FN's additional 'va_list' arg does not need adaption)
+  auto DestArg = NF->arg_begin();
+  for (auto &A : F.args()) {
+    A.replaceAllUsesWith(DestArg);
+    DestArg->takeName(&A);
+    DestArg++;
+  }
+
+  // Replace va_start with va_copy
+  IRBuilder<> Builder(NF->getContext());
+  auto VAListArg = NF->arg_end() - 1;
+  for (auto VAStart : VAStarts) {
+    replaceWithVACopyCall(Builder, VAStart, VAListArg);
+  }
+
+  // Delete original function
+  F.eraseFromParent();
+
+  return NF;
+}

test/Transforms/PagerandoWrappers/address-taken.ll

+; RUN: opt < %s -pagerando-wrappers -S | FileCheck %s
+
+; CHECK-LABEL: @fn_ptr1 = global void ()* @global
+; CHECK-LABEL: @fn_ptr2 = global void ()* @internal
+
+@fn_ptr1 = global void ()* @global
+@fn_ptr2 = global void ()* @internal
+
+; CHECK-LABEL: define hidden void @"global$$orig"() #0
+; CHECK-LABEL: define internal void @"internal$$orig"() #0
+
+define void @global() { ret void }
+define internal void @internal() { ret void }
+
+; CHECK-LABEL: define void @global() #1
+; CHECK-LABEL: define internal void @internal() #1

test/Transforms/PagerandoWrappers/calls.ll

+; RUN: opt < %s -pagerando-wrappers -S | FileCheck %s
+
+; CHECK-LABEL: define hidden void @"global$$orig"() #0 {
+; CHECK-LABEL: define internal void @internal() #0 {
+
+define void @global() { ret void }
+define internal void @internal() { ret void }
+
+; CHECK-LABEL: define internal void @user() #0 {
+define internal void @user() {
+; CHECK-NEXT:    call void @global()
+; CHECK-NEXT:    call void @internal()
+; CHECK-NEXT:    ret void
+  call void @global()
+  call void @internal()
+  ret void
+}
+
+; CHECK-LABEL: define void @global() #1 {
+
+; CHECK-LABEL: attributes #0 = { pagerando_binned }
+; CHECK-LABEL: attributes #1 = { noinline optsize pagerando_wrapper }

test/Transforms/PagerandoWrappers/varargs.ll

+; RUN: opt < %s -pagerando-wrappers -S | FileCheck %s
+
+declare void @llvm.va_start(i8*)
+declare void @llvm.va_end(i8*)
+declare void @llvm.va_copy(i8*, i8*)
+
+; CHECK-LABEL: define internal void @user()
+; CHECK-NEXT:    call i32 (...) @varags(i32 13, i32 37)
+
+define internal void @user() {
+  call i32 (...) @varags(i32 13, i32 37)
+  ret void
+}
+
+; CHECK-LABEL: define i32 @varags(...)
+; CHECK-NEXT:    %1 = alloca i8
+; CHECK-NEXT:    call void @llvm.va_start(i8* %1)
+; CHECK-NEXT:    %2 = call i32 @"varags$$origva"(i8* %1)
+; CHECK-NEXT:    call void @llvm.va_end(i8* %1)
+; CHECK-NEXT:    ret i32 %2
+
+; CHECK-LABEL: define hidden i32 @"varags$$origva"(i8*)
+; CHECK-NEXT:    %va = alloca i8
+; CHECK-NEXT:    call void @llvm.va_copy(i8* %va, i8* %0)
+; CHECK-NEXT:    %ret = va_arg i8* %va, i32
+; CHECK-NEXT:    call void @llvm.va_end(i8* %va)
+; CHECK-NEXT:    ret i32 %ret
+
+define i32 @varags(...) {
+  %va = alloca i8
+  call void @llvm.va_start(i8* %va)
+  %ret = va_arg i8* %va, i32
+  call void @llvm.va_end(i8* %va)
+  ret i32 %ret
+}
+
+; CHECK-LABEL: define void @multiple_starts(...)
+; CHECK-NEXT:    %1 = alloca i8
+; CHECK-NEXT:    call void @llvm.va_start(i8* %1)
+; CHECK-NEXT:    call void @"multiple_starts$$origva"(i8* %1)
+; CHECK-NEXT:    call void @llvm.va_end(i8* %1)
+; CHECK-NEXT:    ret void
+
+; CHECK-LABEL: define hidden void @"multiple_starts$$origva"(i8*)
+; CHECK-NEXT:    %va1 = alloca i8
+; CHECK-NEXT:    %va2 = alloca i8
+; CHECK-NEXT:    call void @llvm.va_copy(i8* %va1, i8* %0)
+; CHECK-NEXT:    call void @llvm.va_copy(i8* %va2, i8* %0)
+; CHECK-NEXT:    call void @llvm.va_end(i8* %va1)
+; CHECK-NEXT:    call void @llvm.va_end(i8* %va2)
+; CHECK-NEXT:    ret void
+
+define void @multiple_starts(...) {
+  %va1 = alloca i8
+  %va2 = alloca i8
+  call void @llvm.va_start(i8* %va1)
+  call void @llvm.va_start(i8* %va2)
+  call void @llvm.va_end(i8* %va1)
+  call void @llvm.va_end(i8* %va2)
+  ret void
+}
+
+; CHECK-LABEL: define void @copy(...)
+; CHECK-NEXT:    %1 = alloca i8
+; CHECK-NEXT:    call void @llvm.va_start(i8* %1)
+; CHECK-NEXT:    call void @"copy$$origva"(i8* %1)
+; CHECK-NEXT:    call void @llvm.va_end(i8* %1)
+; CHECK-NEXT:    ret void
+
+; CHECK-LABEL: define hidden void @"copy$$origva"(i8*)
+; CHECK-NEXT:    %va1 = alloca i8
+; CHECK-NEXT:    %va2 = alloca i8
+; CHECK-NEXT:    call void @llvm.va_copy(i8* %va1, i8* %0)
+; CHECK-NEXT:    call void @llvm.va_copy(i8* %va2, i8* %va1)
+; CHECK-NEXT:    call void @llvm.va_end(i8* %va1)
+; CHECK-NEXT:    call void @llvm.va_end(i8* %va2)
+; CHECK-NEXT:    ret void
+
+define void @copy(...) {
+  %va1 = alloca i8
+  %va2 = alloca i8
+  call void @llvm.va_start(i8* %va1)
+  call void @llvm.va_copy(i8* %va2, i8* %va1)
+  call void @llvm.va_end(i8* %va1)
+  call void @llvm.va_end(i8* %va2)
+  ret void
+}
+
+; CHECK-LABEL: define void @real_va_type(...)
+; CHECK-NEXT:    %1 = alloca %struct.va_list
+; CHECK-NEXT:    %2 = bitcast %struct.va_list* %1 to i8*
+; CHECK-NEXT:    call void @llvm.va_start(i8* %2)
+; CHECK-NEXT:    call void @"real_va_type$$origva"(%struct.va_list* %1)
+; CHECK-NEXT:    %3 = bitcast %struct.va_list* %1 to i8*
+; CHECK-NEXT:    call void @llvm.va_end(i8* %3)
+; CHECK-NEXT:    ret void
+
+; CHECK-LABEL: define hidden void @"real_va_type$$origva"(%struct.va_list*)
+; CHECK-NEXT:    %va = alloca %struct.va_list
+; CHECK-NEXT:    %va_ptr = bitcast %struct.va_list* %va to i8*
+; CHECK-NEXT:    %2 = bitcast %struct.va_list* %0 to i8*
+; CHECK-NEXT:    call void @llvm.va_copy(i8* %va_ptr, i8* %2)
+; CHECK-NEXT:    call void @llvm.va_end(i8* %va_ptr)
+; CHECK-NEXT:    ret void
+
+%struct.va_list = type { i8*, i8*, i8*, i32, i32 }
+
+define void @real_va_type(...) {
+  %va = alloca %struct.va_list
+  %va_ptr = bitcast %struct.va_list* %va to i8*
+  call void @llvm.va_start(i8* %va_ptr)
+  call void @llvm.va_end(i8* %va_ptr)
+  ret void
+}

test/Transforms/PagerandoWrappers/wrapper-body.ll

+; RUN: opt < %s -pagerando-wrappers -S | FileCheck %s
+
+; CHECK-LABEL: define hidden i32 @"foo$$orig"(i32 %arg) #0
+; CHECK-NEXT:    ret i32 %arg
+
+define i32 @foo(i32 %arg) {
+  ret i32 %arg
+}
+
+; CHECK-LABEL: define i32 @foo(i32) #1
+; CHECK-NEXT:    %2 = call i32 @"foo$$orig"(i32 %0)
+; CHECK-NEXT:    ret i32 %2
+
+; CHECK-LABEL: attributes #0 = { pagerando_binned }
+; CHECK-LABEL: attributes #1 = { noinline optsize pagerando_wrapper }