[llvm] [SafeStack] Add 'unsafealloc' metadata to identify unsafe AllocaInsts (PR #78941)

[Nicholas Mosier via llvm-commits]

https://github.com/nmosier updated https://github.com/llvm/llvm-project/pull/78941

>From b86527e6fa98a2ad0790963626cb950108ed946d Mon Sep 17 00:00:00 2001
From: Nicholas Mosier <nmosier at stanford.edu>
Date: Sat, 20 Jan 2024 22:07:32 +0000
Subject: [PATCH] [SafeStack] Add 'unsafealloc' metadata to identify unsafe
 AllocaInsts

This patch adds a new metadata type for identifying alloca's that SafeStack
should consider unsafe. Currently, SafeStack offers no means for users or
other passes to identify stack allocations that it should consider safe
(instead, it relies solely on its own internal, unconfigurable analysis).
---
 llvm/docs/LangRef.rst                         | 17 ++++++++++++++++
 llvm/include/llvm/IR/FixedMetadataKinds.def   |  1 +
 llvm/lib/CodeGen/SafeStack.cpp                |  6 ++++++
 .../SafeStack/X86/unsafealloc-metadata.ll     | 20 +++++++++++++++++++
 4 files changed, 44 insertions(+)
 create mode 100644 llvm/test/Transforms/SafeStack/X86/unsafealloc-metadata.ll

diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index d881deb30049a20..5a01ac8de73fd17 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -7694,6 +7694,23 @@ it will contain a list of ids, including the ids of the callsites in the
 full inline sequence, in order from the leaf-most call's id to the outermost
 inlined call.
 
+'``unsafealloc``' Metadata
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The ``unsafealloc`` metadata is used to indicate that this alloca
+should be moved to the unsafe stack. This metadata only makes sense if
+the containing function has the ``safestack`` attribute (and thus has
+`SafeStack <https://clang.llvm.org/docs/SafeStack.html>`_
+enabled). The metadata is only used as a flag, so the associated node
+must be empty.
+
+Example:
+
+.. code-block:: text
+
+    %x = alloca [16 x i8], !unsafealloc !0
+    !0 = !{}
+
 Module Flags Metadata
 =====================
 
diff --git a/llvm/include/llvm/IR/FixedMetadataKinds.def b/llvm/include/llvm/IR/FixedMetadataKinds.def
index b375d0f0912060f..c13c82254fc590b 100644
--- a/llvm/include/llvm/IR/FixedMetadataKinds.def
+++ b/llvm/include/llvm/IR/FixedMetadataKinds.def
@@ -51,3 +51,4 @@ LLVM_FIXED_MD_KIND(MD_kcfi_type, "kcfi_type", 36)
 LLVM_FIXED_MD_KIND(MD_pcsections, "pcsections", 37)
 LLVM_FIXED_MD_KIND(MD_DIAssignID, "DIAssignID", 38)
 LLVM_FIXED_MD_KIND(MD_coro_outside_frame, "coro.outside.frame", 39)
+LLVM_FIXED_MD_KIND(MD_unsafealloc, "unsafealloc", 40)
diff --git a/llvm/lib/CodeGen/SafeStack.cpp b/llvm/lib/CodeGen/SafeStack.cpp
index 0a26247a4d16590..303f9560e830b57 100644
--- a/llvm/lib/CodeGen/SafeStack.cpp
+++ b/llvm/lib/CodeGen/SafeStack.cpp
@@ -273,6 +273,12 @@ bool SafeStack::IsMemIntrinsicSafe(const MemIntrinsic *MI, const Use &U,
 /// stack or not. The function analyzes all uses of AI and checks whether it is
 /// only accessed in a memory safe way (as decided statically).
 bool SafeStack::IsSafeStackAlloca(const Value *AllocaPtr, uint64_t AllocaSize) {
+  // Consider the allocation unsafe if it is an AllocaInst with 'unsafealloca'
+  // metadata attached.
+  if (const auto *AI = dyn_cast<AllocaInst>(AllocaPtr);
+      AI && AI->hasMetadata(LLVMContext::MD_unsafealloc))
+    return false;
+
   // Go through all uses of this alloca and check whether all accesses to the
   // allocated object are statically known to be memory safe and, hence, the
   // object can be placed on the safe stack.
diff --git a/llvm/test/Transforms/SafeStack/X86/unsafealloc-metadata.ll b/llvm/test/Transforms/SafeStack/X86/unsafealloc-metadata.ll
new file mode 100644
index 000000000000000..96871053e7feb44
--- /dev/null
+++ b/llvm/test/Transforms/SafeStack/X86/unsafealloc-metadata.ll
@@ -0,0 +1,20 @@
+; RUN: opt -safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -passes=safe-stack -S -mtriple=i386-pc-linux-gnu < %s -o - | FileCheck %s
+; RUN: opt -passes=safe-stack -S -mtriple=x86_64-pc-linux-gnu < %s -o - | FileCheck %s
+;
+; Check whether SafeStack recognizes 'unsafealloc' metadata.
+
+define void @test() safestack {
+  ; CHECK-LABEL: @test(
+  ; CHECK: %[[SP:.*]] = load ptr, ptr @__safestack_unsafe_stack_ptr
+  ; CHECK: %[[STATICTOP:.*]] = getelementptr i8, ptr %[[SP]], i32 -16
+  ; CHECK: store ptr %[[STATICTOP]], ptr @__safestack_unsafe_stack_ptr
+  ; CHECK: %safe = alloca i32, align 4
+  ; CHECK: store ptr %[[SP]], ptr @__safestack_unsafe_stack_ptr
+  ; CHECK: ret void
+  %safe = alloca i32, align 4
+  %unsafe = alloca i32, align 4, !unsafealloc !{}
+  ret void
+}
+