[llvm] [AArch64] Bugfix when using execute-only and memtag sanitizer together (PR #133084)

[Csanád Hajdú via llvm-commits]

https://github.com/Il-Capitano created https://github.com/llvm/llvm-project/pull/133084

Support for execute-only code generation (#125687) introduced a bug in the case where the memtag sanitizer is used in a module containing a mix of execute-only and non-execute-only functions.

The bug is caused by using `return` instead of `break` to short-circuit a loop, which meant that the rest of the function dealing with memtag sanitizer logic wasn't run.

>From 6e3d9461aa7abb30ffc97e20812b18adab306fe4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Csan=C3=A1d=20Hajd=C3=BA?= <csanad.hajdu at arm.com>
Date: Wed, 26 Mar 2025 14:37:56 +0100
Subject: [PATCH] [AArch64] Bugfix when using execute-only and memtag sanitizer
 together

Support for execute-only code generation (#125687) introduced a bug in
the case where the memtag sanitizer is used in a module containing a mix
of execute-only and non-execute-only functions.

The bug is caused by using `return` instead of `break` to short-circuit
a loop, which meant that the rest of the function dealing with memtag
sanitizer logic wasn't run.
---
 .../MCTargetDesc/AArch64ELFStreamer.cpp        | 16 +++++++++++-----
 llvm/test/MC/AArch64/execute-only-memtag.ll    | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/MC/AArch64/execute-only-memtag.ll

diff --git a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp b/llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp
index 98bd102d8f4c1..b12a12436db81 100644
--- a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp
+++ b/llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp
@@ -511,11 +511,17 @@ void AArch64TargetELFStreamer::finish() {
       })) {
     auto *Text =
         static_cast<MCSectionELF *>(Ctx.getObjectFileInfo()->getTextSection());
-    for (auto &F : *Text)
-      if (auto *DF = dyn_cast<MCDataFragment>(&F))
-        if (!DF->getContents().empty())
-          return;
-    Text->setFlags(Text->getFlags() | ELF::SHF_AARCH64_PURECODE);
+    bool Empty = true;
+    for (auto &F : *Text) {
+      if (auto *DF = dyn_cast<MCDataFragment>(&F)) {
+        if (!DF->getContents().empty()) {
+          Empty = false;
+          break;
+        }
+      }
+    }
+    if (Empty)
+      Text->setFlags(Text->getFlags() | ELF::SHF_AARCH64_PURECODE);
   }
 
   MCSectionELF *MemtagSec = nullptr;
diff --git a/llvm/test/MC/AArch64/execute-only-memtag.ll b/llvm/test/MC/AArch64/execute-only-memtag.ll
new file mode 100644
index 0000000000000..02daf3179101f
--- /dev/null
+++ b/llvm/test/MC/AArch64/execute-only-memtag.ll
@@ -0,0 +1,18 @@
+; RUN: llc %s -mtriple=aarch64-linux-android31 -filetype=obj -o %t.o
+; RUN: llvm-readelf -r %t.o | FileCheck %s
+
+; CHECK:      Relocation section '.rela.memtag.globals.static' at offset {{.*}} contains 1 entries:
+; CHECK-NEXT:      Type      {{.*}} Symbol's Name
+; CHECK-NEXT: R_AARCH64_NONE {{.*}} global
+
+ at global = global i32 1, sanitize_memtag
+
+define void @foo() {
+  ret void
+}
+
+define void @bar() #0 {
+  ret void
+}
+
+attributes #0 = { "target-features"="+execute-only" }