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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-ir

@llvm/pr-subscribers-llvm-transforms

Author: Nicholas Mosier (nmosier)

<details>
<summary>Changes</summary>

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).

---
Full diff: https://github.com/llvm/llvm-project/pull/78941.diff


3 Files Affected:

- (modified) llvm/include/llvm/IR/FixedMetadataKinds.def (+1) 
- (modified) llvm/lib/CodeGen/SafeStack.cpp (+6) 
- (added) llvm/test/Transforms/SafeStack/X86/unsafealloc-metadata.ll (+20) 


``````````diff
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
+}
+

``````````

</details>


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