[IR] Do not inline functions carrying "zeroize-stack"

llvm/docs/LangRef.md

@@ -2678,8 +2678,28 @@ fn -> other_fn -> other_fn ; fn is norecurse
     provenance.
 
     A transform may not leave a function carrying this attribute with a frame
-    that is not cleared on every return from it. How the attribute constrains
-    inlining in particular is specified separately.
+    that is not cleared on every return from it.
+
+    Inlining is constrained accordingly. A function that carries this attribute
+    is inlined only into a caller that carries it as well, and only where the
+    caller's mode clears at least as much of the frame as the callee's mode does.
+    Inlining dissolves the callee's frame into the caller's, so the bytes the
+    callee promised to clear before returning become bytes of the caller's frame:
+    a caller clearing at least as much still clears them, while an unannotated
+    caller, or one clearing less, leaves them with nothing to record the
+    obligation. Because `"used"` is the widest mode, a `"used"` caller may inline
+    a callee of either mode, and a `"sensitive"` caller may inline a
+    `"sensitive"` callee but not a `"used"` one.
+
+    A request to inline does not override this, however it is spelled. An
+    implementation may report a warning when it declines to inline for this
+    reason, and should not report an error: inlining is advisory, and refusing it
+    changes no semantics other than the frame the clear covers.
+
+    Inlining a function that does not carry the attribute into one that does is
+    not constrained by this, and is desirable: the callee's frame lies below the
+    stack pointer once it returns and no clear reaches it, whereas inlining
+    turns those bytes into frame bytes of the caller, which are cleared.
 
 `returns_twice`
 :   This attribute indicates that this function can return twice. The C

llvm/include/llvm/IR/Attributes.td

@@ -469,6 +469,7 @@ def : CompatRule<"isEqual<UseSampleProfileAttr>">;
 def : CompatRule<"isEqual<NoProfileAttr>">;
 def : CompatRule<"checkDenormMode">;
 def : CompatRule<"checkStrictFP">;
+def : CompatRule<"checkZeroizeStack">;
 def : CompatRuleStrAttr<"isEqual", "sign-return-address">;
 def : CompatRuleStrAttr<"isEqual", "sign-return-address-key">;
 def : CompatRuleStrAttr<"isEqual", "branch-protection-pauth-lr">;

llvm/lib/IR/Attributes.cpp

@@ -2603,6 +2603,26 @@ static bool checkStrictFP(const Function &Caller, const Function &Callee) {
          Caller.getAttributes().hasFnAttr(Attribute::StrictFP);
 }
 
+static bool checkZeroizeStack(const Function &Caller, const Function &Callee) {
+  // Do not inline a function that carries "zeroize-stack" into any caller. The
+  // attribute is a promise about the callee's own frame, and inlining dissolves
+  // that frame into the caller's: the bytes the callee promised to clear before
+  // returning become bytes of a frame that outlives the point where the clear
+  // was due, and nothing is left in the IR to record the obligation.
+  //
+  // This is a callee-side rule with no exemption for a caller that carries the
+  // attribute itself. The caller's attribute constrains the caller's own frame
+  // and its own returns; it does not reproduce the clear the callee owed at the
+  // point the callee would have returned, and the two functions may in any case
+  // ask for different amounts of the frame to be cleared.
+  //
+  // Inlining an unprotected callee into a protected caller is unaffected, and
+  // is worth encouraging: it moves the callee's frame, which sits below the
+  // stack pointer at the caller's return and which no clear reaches, into the
+  // frame bytes the protected caller does clear.
+  return !Callee.hasFnAttribute("zeroize-stack");
+}
+
 template<typename AttrClass>
 static bool isEqual(const Function &Caller, const Function &Callee) {
   return Caller.getFnAttribute(AttrClass::getKind()) ==

llvm/test/Transforms/Inline/zeroize-stack-lto.ll

@@ -0,0 +1,76 @@
+; The rule that a function carrying "zeroize-stack" is not inlined into any
+; caller has to hold across module boundaries too, where the callee reaches the
+; inliner through the LTO link rather than through the module it was compiled
+; in. Both regular LTO, which merges the modules before optimizing, and ThinLTO,
+; which imports the definition into the caller's module, are checked. Inlining
+; an unannotated callee into an annotated caller must still happen, so that a
+; missing inline is evidence of the rule rather than of nothing being inlined.
+
+; REQUIRES: x86-registered-target
+
+; RUN: split-file %s %t
+
+;; Regular LTO: one combined module, task 0.
+; RUN: llvm-as %t/caller.ll -o %t/caller.bc
+; RUN: llvm-as %t/callee.ll -o %t/callee.bc
+; RUN: llvm-lto2 run %t/caller.bc %t/callee.bc -save-temps -o %t/lto \
+; RUN:   -r %t/caller.bc,call_zeroize,plx \
+; RUN:   -r %t/caller.bc,call_plain,plx \
+; RUN:   -r %t/caller.bc,zeroize_callee,l \
+; RUN:   -r %t/caller.bc,plain_callee,l \
+; RUN:   -r %t/callee.bc,zeroize_callee,plx \
+; RUN:   -r %t/callee.bc,plain_callee,plx
+; RUN: llvm-dis %t/lto.0.4.opt.bc -o - | FileCheck %s
+
+;; ThinLTO: the caller's module is task 1, and the callee is imported into it.
+; RUN: opt -module-summary %t/caller.ll -o %t/caller.thin.bc
+; RUN: opt -module-summary %t/callee.ll -o %t/callee.thin.bc
+; RUN: llvm-lto2 run %t/caller.thin.bc %t/callee.thin.bc -save-temps -o %t/thin \
+; RUN:   -r %t/caller.thin.bc,call_zeroize,plx \
+; RUN:   -r %t/caller.thin.bc,call_plain,plx \
+; RUN:   -r %t/caller.thin.bc,zeroize_callee,l \
+; RUN:   -r %t/caller.thin.bc,plain_callee,l \
+; RUN:   -r %t/callee.thin.bc,zeroize_callee,plx \
+; RUN:   -r %t/callee.thin.bc,plain_callee,plx
+; RUN: llvm-dis %t/thin.1.4.opt.bc -o - | FileCheck %s
+
+;; The annotated callee is still called, from an annotated caller at that.
+; CHECK-LABEL: define {{.*}}@call_zeroize(
+; CHECK: call i32 @zeroize_callee(
+; CHECK-NOT: mul i32 %{{.*}}, 7
+
+;; The unannotated one is gone, its body folded into the annotated caller.
+; CHECK-LABEL: define {{.*}}@call_plain(
+; CHECK-NOT: call i32 @plain_callee(
+; CHECK: mul i32 %{{.*}}, 11
+
+;--- caller.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare i32 @zeroize_callee(i32)
+declare i32 @plain_callee(i32)
+
+define i32 @call_zeroize(i32 %x) "zeroize-stack"="used" {
+  %r = call i32 @zeroize_callee(i32 %x)
+  ret i32 %r
+}
+
+define i32 @call_plain(i32 %x) "zeroize-stack"="used" {
+  %r = call i32 @plain_callee(i32 %x)
+  ret i32 %r
+}
+
+;--- callee.ll
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @zeroize_callee(i32 %x) "zeroize-stack"="used" {
+  %r = mul i32 %x, 7
+  ret i32 %r
+}
+
+define i32 @plain_callee(i32 %x) {
+  %r = mul i32 %x, 11
+  ret i32 %r
+}

llvm/test/Transforms/Inline/zeroize-stack.ll

@@ -0,0 +1,102 @@
+; A function carrying "zeroize-stack" clears its own frame before returning.
+; Inlining it would dissolve that frame into the caller's and lose the
+; obligation, so it is never inlined, whatever the caller is annotated with.
+; Inlining an unannotated callee into an annotated caller is still allowed.
+
+; RUN: opt < %s -passes=inline -S | FileCheck %s
+; RUN: opt < %s -passes='default<O2>' -S | FileCheck %s --check-prefix=O2
+
+declare void @sink(i32)
+
+define void @zeroize_used_callee(i32 %x) "zeroize-stack"="used" {
+  call void @sink(i32 %x)
+  ret void
+}
+
+define void @zeroize_sensitive_callee(i32 %x) "zeroize-stack"="sensitive" {
+  call void @sink(i32 %x)
+  ret void
+}
+
+define void @plain_callee(i32 %x) {
+  call void @sink(i32 %x)
+  ret void
+}
+
+;; Not inlined into a caller that has no attribute: the caller clears nothing,
+;; so the callee's frame bytes would be left behind.
+
+define void @unannotated_caller(i32 %x) {
+; CHECK-LABEL: define void @unannotated_caller(
+; CHECK: call void @zeroize_used_callee(
+; CHECK: call void @zeroize_sensitive_callee(
+; CHECK-NOT: call void @sink(
+;
+; O2-LABEL: define void @unannotated_caller(
+; O2: call void @zeroize_used_callee(
+; O2: call void @zeroize_sensitive_callee(
+; O2-NOT: call void @sink(
+  call void @zeroize_used_callee(i32 %x)
+  call void @zeroize_sensitive_callee(i32 %x)
+  ret void
+}
+
+;; Not inlined into a caller that has the attribute either. The rule has no
+;; same-attribute exemption: the caller's clear covers the caller's own frame at
+;; the caller's returns, not the clear the callee owed where it would have
+;; returned, and the two need not ask for the same amount of the frame.
+
+define void @zeroize_used_caller(i32 %x) "zeroize-stack"="used" {
+; CHECK-LABEL: define void @zeroize_used_caller(
+; CHECK: call void @zeroize_used_callee(
+; CHECK: call void @zeroize_sensitive_callee(
+;
+; O2-LABEL: define void @zeroize_used_caller(
+; O2: call void @zeroize_used_callee(
+; O2: call void @zeroize_sensitive_callee(
+  call void @zeroize_used_callee(i32 %x)
+  call void @zeroize_sensitive_callee(i32 %x)
+  ret void
+}
+
+define void @zeroize_sensitive_caller(i32 %x) "zeroize-stack"="sensitive" {
+; CHECK-LABEL: define void @zeroize_sensitive_caller(
+; CHECK: call void @zeroize_used_callee(
+; CHECK: call void @zeroize_sensitive_callee(
+;
+; O2-LABEL: define void @zeroize_sensitive_caller(
+; O2: call void @zeroize_used_callee(
+; O2: call void @zeroize_sensitive_callee(
+  call void @zeroize_used_callee(i32 %x)
+  call void @zeroize_sensitive_callee(i32 %x)
+  ret void
+}
+
+;; An unannotated callee is still inlined into an annotated caller, and that is
+;; the direction worth encouraging: the callee's frame sits below the stack
+;; pointer at the caller's return and no clear reaches it, while inlining turns
+;; those bytes into frame bytes the caller does clear.
+
+define void @plain_into_zeroize_caller(i32 %x) "zeroize-stack"="used" {
+; CHECK-LABEL: define void @plain_into_zeroize_caller(
+; CHECK-NOT: call void @plain_callee(
+; CHECK: call void @sink(
+;
+; O2-LABEL: define void @plain_into_zeroize_caller(
+; O2-NOT: call void @plain_callee(
+; O2: call void @sink(
+  call void @plain_callee(i32 %x)
+  ret void
+}
+
+define void @plain_into_sensitive_caller(i32 %x) "zeroize-stack"="sensitive" {
+; CHECK-LABEL: define void @plain_into_sensitive_caller(
+; CHECK-NOT: call void @plain_callee(
+; CHECK: call void @sink(
+;
+; O2-LABEL: define void @plain_into_sensitive_caller(
+; O2-NOT: call void @plain_callee(
+; O2: call void @sink(
+  call void @plain_callee(i32 %x)
+  ret void
+}