From 0f0aa4f8a20f6c7c851f81dc653cc1fb8d292cc4 Mon Sep 17 00:00:00 2001
From: Daniel Kiss <daniel.kiss@arm.com>
Date: Wed, 10 Jul 2024 10:06:43 +0200
Subject: [PATCH] [llvm][ARM][AArch64] Add attributes to synthetic functions.
 (#83153)

Module flags represent the original intention.

Depends on #82819
---
 llvm/lib/IR/Function.cpp                      | 29 ++++++++++++++++
 .../AddressSanitizer/module-flags-aarch64.ll  | 33 +++++++++++++++++++
 .../AArch64/module-flags-aarch64.ll           | 33 +++++++++++++++++++
 3 files changed, 95 insertions(+)
 create mode 100644 llvm/test/Instrumentation/AddressSanitizer/module-flags-aarch64.ll
 create mode 100644 llvm/test/Instrumentation/MemorySanitizer/AArch64/module-flags-aarch64.ll

diff --git a/llvm/lib/IR/Function.cpp b/llvm/lib/IR/Function.cpp
index 5fb348a8bbcd4..20871982afb06 100644
--- a/llvm/lib/IR/Function.cpp
+++ b/llvm/lib/IR/Function.cpp
@@ -407,6 +407,35 @@ Function *Function::createWithDefaultAttr(FunctionType *Ty,
   StringRef DefaultFeatures = F->getContext().getDefaultTargetFeatures();
   if (!DefaultFeatures.empty())
     B.addAttribute("target-features", DefaultFeatures);
+
+  // Check if the module attribute is present and not zero.
+  auto isModuleAttributeSet = [&](const StringRef &ModAttr) -> bool {
+    const auto *Attr =
+        mdconst::extract_or_null<ConstantInt>(M->getModuleFlag(ModAttr));
+    return Attr && !Attr->isZero();
+  };
+
+  auto AddAttributeIfSet = [&](const StringRef &ModAttr) {
+    if (isModuleAttributeSet(ModAttr))
+      B.addAttribute(ModAttr);
+  };
+
+  StringRef SignType = "none";
+  if (isModuleAttributeSet("sign-return-address"))
+    SignType = "non-leaf";
+  if (isModuleAttributeSet("sign-return-address-all"))
+    SignType = "all";
+  if (SignType != "none") {
+    B.addAttribute("sign-return-address", SignType);
+    B.addAttribute("sign-return-address-key",
+                   isModuleAttributeSet("sign-return-address-with-bkey")
+                       ? "b_key"
+                       : "a_key");
+  }
+  AddAttributeIfSet("branch-target-enforcement");
+  AddAttributeIfSet("branch-protection-pauth-lr");
+  AddAttributeIfSet("guarded-control-stack");
+
   F->addFnAttrs(B);
   return F;
 }
diff --git a/llvm/test/Instrumentation/AddressSanitizer/module-flags-aarch64.ll b/llvm/test/Instrumentation/AddressSanitizer/module-flags-aarch64.ll
new file mode 100644
index 0000000000000..91eaa360d18df
--- /dev/null
+++ b/llvm/test/Instrumentation/AddressSanitizer/module-flags-aarch64.ll
@@ -0,0 +1,33 @@
+;; Verify that the synthetic functions inherit their flags from the corresponding
+;; BTE and return address signing module flags.
+; RUN: opt < %s -passes=asan -S | FileCheck %s
+; REQUIRES: aarch64-registered-target
+
+target triple = "aarch64-unknown-linux-gnu"
+
+@g = dso_local global i32 0, align 4
+
+define i32 @test_load() sanitize_address {
+entry:
+  %tmp = load i32, ptr @g, align 4
+  ret i32 %tmp
+}
+
+!llvm.module.flags = !{!0, !1}
+
+;; Due to -fasynchronous-unwind-tables.
+!0 = !{i32 7, !"uwtable", i32 2}
+
+;; Due to -fno-omit-frame-pointer.
+!1 = !{i32 7, !"frame-pointer", i32 2}
+
+!llvm.module.flags = !{!2, !3, !4}
+
+!2 = !{i32 8, !"branch-target-enforcement", i32 1}
+!3 = !{i32 8, !"sign-return-address", i32 1}
+!4 = !{i32 8, !"sign-return-address-all", i32 0}
+
+;; Set the uwtable attribute on ctor/dtor.
+; CHECK: define internal void @asan.module_ctor() #[[#ATTR:]]
+; CHECK: define internal void @asan.module_dtor() #[[#ATTR]]
+; CHECK: attributes #[[#ATTR]] = { nounwind uwtable "branch-target-enforcement" "frame-pointer"="all" "sign-return-address"="non-leaf" "sign-return-address-key"="a_key" }
diff --git a/llvm/test/Instrumentation/MemorySanitizer/AArch64/module-flags-aarch64.ll b/llvm/test/Instrumentation/MemorySanitizer/AArch64/module-flags-aarch64.ll
new file mode 100644
index 0000000000000..f6a1009c932a0
--- /dev/null
+++ b/llvm/test/Instrumentation/MemorySanitizer/AArch64/module-flags-aarch64.ll
@@ -0,0 +1,33 @@
+;; Verify that the synthetic functions inherit their flags from the corresponding
+;; BTE and return address signing module flags.
+; RUN: opt < %s -passes=asan -S | FileCheck %s
+; REQUIRES: aarch64-registered-target
+
+target triple = "aarch64-unknown-linux-gnu"
+
+@g = dso_local global i32 0, align 4
+
+define i32 @test_load() sanitize_memory {
+entry:
+  %tmp = load i32, ptr @g, align 4
+  ret i32 %tmp
+}
+
+!llvm.module.flags = !{!0, !1}
+
+;; Due to -fasynchronous-unwind-tables.
+!0 = !{i32 7, !"uwtable", i32 2}
+
+;; Due to -fno-omit-frame-pointer.
+!1 = !{i32 7, !"frame-pointer", i32 2}
+
+!llvm.module.flags = !{!2, !3, !4}
+
+!2 = !{i32 8, !"branch-target-enforcement", i32 1}
+!3 = !{i32 8, !"sign-return-address", i32 1}
+!4 = !{i32 8, !"sign-return-address-all", i32 0}
+
+;; Set the uwtable attribute on ctor/dtor.
+; CHECK: define internal void @asan.module_ctor() #[[#ATTR:]]
+; CHECK: define internal void @asan.module_dtor() #[[#ATTR]]
+; CHECK: attributes #[[#ATTR]] = { nounwind uwtable "branch-target-enforcement" "frame-pointer"="all" "sign-return-address"="non-leaf" "sign-return-address-key"="a_key" }