modified: compiler/rustc_codegen_llvm/src/attributes.rs

	modified:   compiler/rustc_mir_transform/src/stack_protector.rs
	modified:   tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect.rs

Author: cezarbbb

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -6,6 +6,7 @@ use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, PatchableFuncti
 use rustc_middle::ty::{self, TyCtxt};
 use rustc_session::config::{BranchProtection, FunctionReturn, OptLevel, PAuthKey, PacRet};
 use rustc_symbol_mangling::mangle_internal_symbol;
+use rustc_target::callconv::PassMode;
 use rustc_target::spec::{FramePointer, SanitizerSet, StackProbeType, StackProtector};
 use smallvec::SmallVec;
 
@@ -270,13 +271,22 @@ fn probestack_attr<'ll>(cx: &CodegenCx<'ll, '_>) -> Option<&'ll Attribute> {
     Some(llvm::CreateAttrStringValue(cx.llcx, "probe-stack", attr_value))
 }
 
-fn stackprotector_attr<'ll>(cx: &CodegenCx<'ll, '_>, def_id: DefId) -> Option<&'ll Attribute> {
+fn stackprotector_attr<'ll, 'tcx>(
+    cx: &CodegenCx<'ll, 'tcx>,
+    instance: ty::Instance<'tcx>,
+) -> Option<&'ll Attribute> {
+    let fn_abi = cx.fn_abi_of_instance(instance, ty::List::empty());
     let sspattr = match cx.sess().stack_protector() {
         StackProtector::None => return None,
         StackProtector::All => AttributeKind::StackProtectReq,
 
         StackProtector::Rusty => {
-            if cx.tcx.stack_protector.borrow().contains(&def_id) {
+            if cx.tcx.stack_protector.borrow().contains(&instance.def_id())
+                || matches!(
+                    &fn_abi.ret.mode,
+                    PassMode::Indirect { attrs: _, meta_attrs: _, on_stack: false }
+                )
+            {
                 AttributeKind::StackProtectStrong
             } else {
                 return None;
@@ -395,7 +405,7 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
     to_add.extend(probestack_attr(cx));
 
     // stack protector
-    to_add.extend(stackprotector_attr(cx, instance.def_id()));
+    to_add.extend(stackprotector_attr(cx, instance));
 
     if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_BUILTINS) {
         to_add.push(llvm::CreateAttrString(cx.llcx, "no-builtins"));

compiler/rustc_mir_transform/src/stack_protector.rs

@@ -1,11 +1,9 @@
 //! Validates the MIR to ensure that invariants are upheld.
-
 use std::ops::Deref;
 
 use rustc_middle::mir::*;
 use rustc_middle::ty;
-use rustc_middle::ty::{Instance, TyCtxt};
-use rustc_target::callconv::PassMode;
+use rustc_middle::ty::TyCtxt;
 
 pub(super) struct StackProtectorFinder;
 
@@ -38,29 +36,6 @@ impl<'tcx> crate::MirPass<'tcx> for StackProtectorFinder {
                     }
                 }
             }
-
-            let instance = Instance::mono(tcx, def_id);
-            let Ok(fn_abi) = tcx.fn_abi_of_instance(
-                ty::TypingEnv::fully_monomorphized().as_query_input((instance, ty::List::empty())),
-            ) else {
-                // FIXME: Find when an Err() message is returned
-
-                return;
-            };
-
-            // for arg in fn_abi.args.iter() {
-            //     if matches!(&arg.mode, PassMode::Indirect { attrs: _, meta_attrs: _, on_stack: false }) {
-            //         tcx.stack_protector.borrow_mut().insert(def_id);
-            //         return;
-            //     }
-            // }
-
-            let ret = &fn_abi.ret;
-            if matches!(&ret.mode, PassMode::Indirect { attrs: _, meta_attrs: _, on_stack: false })
-            {
-                tcx.stack_protector.borrow_mut().insert(def_id);
-                return;
-            }
         }
     }
 

tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect.rs

@@ -192,7 +192,7 @@ pub fn local_var_addr_taken_used_locally_only(factory: fn() -> i32, sink: fn(i32
     // all: __stack_chk_fail
 
     // FIXME: rusty stack smash protection needs to support inline scenario detection
-    // rusty: __stack_chk_fail
+    // rusty-NOT: __stack_chk_fail
 
     // strong-NOT: __stack_chk_fail
     // basic-NOT: __stack_chk_fail
@@ -312,7 +312,7 @@ pub fn alloca_small_compile_time_constant_arg(f: fn(*mut ())) {
     // FIXME: Rusty thinks a function that returns a mutable raw pointer may
     // be a stack memory allocation function, so it performs stack smash protection.
     // Is it possible to optimize the heuristics?
-    // rusty: __stack_chk_fail
+    // rusty-NOT: __stack_chk_fail
 
     // strong-NOT: __stack_chk_fail
     // basic-NOT: __stack_chk_fail
@@ -326,7 +326,7 @@ pub fn alloca_large_compile_time_constant_arg(f: fn(*mut ())) {
     f(unsafe { alloca(9) });
 
     // all: __stack_chk_fail
-    // rusty: __stack_chk_fail
+    // rusty-NOT: __stack_chk_fail
     // strong-NOT: __stack_chk_fail
     // basic-NOT: __stack_chk_fail
     // none-NOT: __stack_chk_fail
@@ -339,7 +339,7 @@ pub fn alloca_dynamic_arg(f: fn(*mut ()), n: usize) {
     f(unsafe { alloca(n) });
 
     // all: __stack_chk_fail
-    // rusty: __stack_chk_fail
+    // rusty-NOT: __stack_chk_fail
     // strong-NOT: __stack_chk_fail
     // basic-NOT: __stack_chk_fail
     // none-NOT: __stack_chk_fail
@@ -369,7 +369,7 @@ pub fn unsized_fn_param(s: [u8], l: bool, f: fn([u8])) {
     // all: __stack_chk_fail
 
     // FIXME: Does Rusty need to handle this type of optimization?
-    // rusty: __stack_chk_fail
+    // rusty-NOT: __stack_chk_fail
 
     // strong-NOT: __stack_chk_fail
     // basic-NOT: __stack_chk_fail