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