1 files changed, 80 insertions, 0 deletions
diff --git a/llvm/lib/CodeGen/SanitizerBinaryMetadata.cpp b/llvm/lib/CodeGen/SanitizerBinaryMetadata.cpp
new file mode 100644
index 000000000000..dd70a2f23e45
--- /dev/null
+++ b/llvm/lib/CodeGen/SanitizerBinaryMetadata.cpp
@@ -0,0 +1,80 @@
+//===- SanitizerBinaryMetadata.cpp
+//----------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of SanitizerBinaryMetadata.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Instrumentation/SanitizerBinaryMetadata.h"
+#include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/MDBuilder.h"
+#include "llvm/InitializePasses.h"
+#include "llvm/Pass.h"
+#include <algorithm>
+
+using namespace llvm;
+
+namespace {
+class MachineSanitizerBinaryMetadata : public MachineFunctionPass {
+public:
+  static char ID;
+
+  MachineSanitizerBinaryMetadata();
+  bool runOnMachineFunction(MachineFunction &F) override;
+};
+} // namespace
+
+INITIALIZE_PASS(MachineSanitizerBinaryMetadata, "machine-sanmd",
+                "Machine Sanitizer Binary Metadata", false, false)
+
+char MachineSanitizerBinaryMetadata::ID = 0;
+char &llvm::MachineSanitizerBinaryMetadataID =
+    MachineSanitizerBinaryMetadata::ID;
+
+MachineSanitizerBinaryMetadata::MachineSanitizerBinaryMetadata()
+    : MachineFunctionPass(ID) {
+  initializeMachineSanitizerBinaryMetadataPass(
+      *PassRegistry::getPassRegistry());
+}
+
+bool MachineSanitizerBinaryMetadata::runOnMachineFunction(MachineFunction &MF) {
+  MDNode *MD = MF.getFunction().getMetadata(LLVMContext::MD_pcsections);
+  if (!MD)
+    return false;
+  const auto &Section = *cast<MDString>(MD->getOperand(0));
+  if (!Section.getString().equals(kSanitizerBinaryMetadataCoveredSection))
+    return false;
+  auto &AuxMDs = *cast<MDTuple>(MD->getOperand(1));
+  // Assume it currently only has features.
+  assert(AuxMDs.getNumOperands() == 1);
+  auto *Features = cast<ConstantAsMetadata>(AuxMDs.getOperand(0))->getValue();
+  if (!Features->getUniqueInteger()[kSanitizerBinaryMetadataUARBit])
+    return false;
+  // Calculate size of stack args for the function.
+  int64_t Size = 0;
+  uint64_t Align = 0;
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+  for (int i = -1; i >= (int)-MFI.getNumFixedObjects(); --i) {
+    Size = std::max(Size, MFI.getObjectOffset(i) + MFI.getObjectSize(i));
+    Align = std::max(Align, MFI.getObjectAlign(i).value());
+  }
+  Size = (Size + Align - 1) & ~(Align - 1);
+  auto &F = MF.getFunction();
+  IRBuilder<> IRB(F.getContext());
+  MDBuilder MDB(F.getContext());
+  // Keep the features and append size of stack args to the metadata.
+  F.setMetadata(LLVMContext::MD_pcsections,
+                MDB.createPCSections(
+                    {{Section.getString(), {Features, IRB.getInt32(Size)}}}));
+  return false;
+}