1 //===- SanitizerBinaryMetadata.cpp
2 //----------------------------------------------===//
3 //
4 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
5 // See https://llvm.org/LICENSE.txt for license information.
6 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // This file is a part of SanitizerBinaryMetadata.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "llvm/Transforms/Instrumentation/SanitizerBinaryMetadata.h"
15 #include "llvm/CodeGen/MachineFrameInfo.h"
16 #include "llvm/CodeGen/MachineFunction.h"
17 #include "llvm/CodeGen/MachineFunctionPass.h"
18 #include "llvm/CodeGen/Passes.h"
19 #include "llvm/IR/IRBuilder.h"
20 #include "llvm/IR/MDBuilder.h"
21 #include "llvm/InitializePasses.h"
22 #include "llvm/Pass.h"
23 #include <algorithm>
24 
25 using namespace llvm;
26 
27 namespace {
28 class MachineSanitizerBinaryMetadata : public MachineFunctionPass {
29 public:
30   static char ID;
31 
32   MachineSanitizerBinaryMetadata();
33   bool runOnMachineFunction(MachineFunction &F) override;
34 };
35 } // namespace
36 
37 INITIALIZE_PASS(MachineSanitizerBinaryMetadata, "machine-sanmd",
38                 "Machine Sanitizer Binary Metadata", false, false)
39 
40 char MachineSanitizerBinaryMetadata::ID = 0;
41 char &llvm::MachineSanitizerBinaryMetadataID =
42     MachineSanitizerBinaryMetadata::ID;
43 
44 MachineSanitizerBinaryMetadata::MachineSanitizerBinaryMetadata()
45     : MachineFunctionPass(ID) {
46   initializeMachineSanitizerBinaryMetadataPass(
47       *PassRegistry::getPassRegistry());
48 }
49 
50 bool MachineSanitizerBinaryMetadata::runOnMachineFunction(MachineFunction &MF) {
51   MDNode *MD = MF.getFunction().getMetadata(LLVMContext::MD_pcsections);
52   if (!MD)
53     return false;
54   const auto &Section = *cast<MDString>(MD->getOperand(0));
55   if (!Section.getString().starts_with(kSanitizerBinaryMetadataCoveredSection))
56     return false;
57   auto &AuxMDs = *cast<MDTuple>(MD->getOperand(1));
58   // Assume it currently only has features.
59   assert(AuxMDs.getNumOperands() == 1);
60   Constant *Features =
61       cast<ConstantAsMetadata>(AuxMDs.getOperand(0))->getValue();
62   if (!Features->getUniqueInteger()[kSanitizerBinaryMetadataUARBit])
63     return false;
64   // Calculate size of stack args for the function.
65   int64_t Size = 0;
66   uint64_t Align = 0;
67   const MachineFrameInfo &MFI = MF.getFrameInfo();
68   for (int i = -1; i >= (int)-MFI.getNumFixedObjects(); --i) {
69     Size = std::max(Size, MFI.getObjectOffset(i) + MFI.getObjectSize(i));
70     Align = std::max(Align, MFI.getObjectAlign(i).value());
71   }
72   Size = (Size + Align - 1) & ~(Align - 1);
73   if (!Size)
74     return false;
75   // Non-zero size, update metadata.
76   auto &F = MF.getFunction();
77   IRBuilder<> IRB(F.getContext());
78   MDBuilder MDB(F.getContext());
79   // Keep the features and append size of stack args to the metadata.
80   APInt NewFeatures = Features->getUniqueInteger();
81   NewFeatures.setBit(kSanitizerBinaryMetadataUARHasSizeBit);
82   F.setMetadata(
83       LLVMContext::MD_pcsections,
84       MDB.createPCSections({{Section.getString(),
85                              {IRB.getInt(NewFeatures), IRB.getInt32(Size)}}}));
86   return false;
87 }
88