modified: compiler/rustc_codegen_llvm/src/attributes.rs

	modified:   compiler/rustc_middle/src/ty/context.rs
	modified:   compiler/rustc_mir_transform/src/lib.rs
	new file:   compiler/rustc_mir_transform/src/stack_protector.rs
	modified:   compiler/rustc_target/src/spec/mod.rs
	modified:   tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect-windows-64bit.rs
	modified:   tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect.rs

Author: cezarbbb

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -270,10 +270,19 @@ 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, '_>) -> Option<&'ll Attribute> {
+fn stackprotector_attr<'ll>(cx: &CodegenCx<'ll, '_>, def_id: DefId) -> Option<&'ll Attribute> {
     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) {
+                AttributeKind::StackProtectStrong
+            } else {
+                return None;
+            }
+        }
+
         StackProtector::Strong => AttributeKind::StackProtectStrong,
         StackProtector::Basic => AttributeKind::StackProtect,
     };
@@ -384,7 +393,9 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
     to_add.extend(instrument_function_attr(cx));
     to_add.extend(nojumptables_attr(cx));
     to_add.extend(probestack_attr(cx));
-    to_add.extend(stackprotector_attr(cx));
+
+    // stack protector
+    to_add.extend(stackprotector_attr(cx, instance.def_id()));
 
     if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_BUILTINS) {
         to_add.push(llvm::CreateAttrString(cx.llcx, "no-builtins"));

compiler/rustc_middle/src/ty/context.rs

@@ -19,7 +19,7 @@ use rustc_abi::{ExternAbi, FieldIdx, Layout, LayoutData, TargetDataLayout, Varia
 use rustc_ast as ast;
 use rustc_data_structures::defer;
 use rustc_data_structures::fingerprint::Fingerprint;
-use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::intern::Interned;
 use rustc_data_structures::jobserver::Proxy;
 use rustc_data_structures::profiling::SelfProfilerRef;
@@ -1519,6 +1519,8 @@ pub struct GlobalCtxt<'tcx> {
     /// Stores memory for globals (statics/consts).
     pub(crate) alloc_map: interpret::AllocMap<'tcx>,
 
+    pub stack_protector: Lock<FxHashSet<DefId>>,
+
     current_gcx: CurrentGcx,
 
     /// A jobserver reference used to release then acquire a token while waiting on a query.
@@ -1746,6 +1748,7 @@ impl<'tcx> TyCtxt<'tcx> {
             clauses_cache: Default::default(),
             data_layout,
             alloc_map: interpret::AllocMap::new(),
+            stack_protector: Default::default(),
             current_gcx,
             jobserver_proxy,
         });

compiler/rustc_mir_transform/src/lib.rs

@@ -40,6 +40,7 @@ mod pass_manager;
 use std::sync::LazyLock;
 
 use pass_manager::{self as pm, Lint, MirLint, MirPass, WithMinOptLevel};
+use rustc_target::spec::StackProtector;
 
 mod check_pointers;
 mod cost_checker;
@@ -193,6 +194,7 @@ declare_passes! {
     mod single_use_consts : SingleUseConsts;
     mod sroa : ScalarReplacementOfAggregates;
     mod strip_debuginfo : StripDebugInfo;
+    mod stack_protector: StackProtectorFinder;
     mod unreachable_enum_branching : UnreachableEnumBranching;
     mod unreachable_prop : UnreachablePropagation;
     mod validate : Validator;
@@ -450,6 +452,17 @@ fn mir_promoted(
     lint_tail_expr_drop_order::run_lint(tcx, def, &body);
 
     let promoted = promote_pass.promoted_fragments.into_inner();
+
+    if tcx.sess.stack_protector() == StackProtector::Rusty {
+        pm::run_passes(
+            tcx,
+            &mut body,
+            &[&stack_protector::StackProtectorFinder],
+            None,
+            pm::Optimizations::Allowed,
+        )
+    }
+
     (tcx.alloc_steal_mir(body), tcx.alloc_steal_promoted(promoted))
 }
 

compiler/rustc_mir_transform/src/stack_protector.rs

@@ -0,0 +1,70 @@
+//! 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;
+
+pub(super) struct StackProtectorFinder;
+
+impl<'tcx> crate::MirPass<'tcx> for StackProtectorFinder {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        use Rvalue::*;
+        let def_id = body.source.def_id();
+
+        for block in body.basic_blocks.iter() {
+            for stmt in block.statements.iter() {
+                if let StatementKind::Assign(assign) = &stmt.kind {
+                    let (_, rvalue) = assign.deref();
+                    match rvalue {
+                        // Get a reference/pointer to a variable
+                        Ref(..) | ThreadLocalRef(_) | RawPtr(..) => {
+                            tcx.stack_protector.borrow_mut().insert(def_id);
+                            return;
+                        }
+                        _ => continue,
+                    }
+                }
+            }
+
+            if let Some(terminator) = block.terminator.as_ref() {
+                if let TerminatorKind::Call { destination: place, .. } = &terminator.kind {
+                    // Returns a mutable raw pointer, possibly a memory allocation function
+                    if let ty::RawPtr(_, Mutability::Mut) = place.ty(body, tcx).ty.kind() {
+                        tcx.stack_protector.borrow_mut().insert(def_id);
+                        return;
+                    }
+                }
+            }
+
+            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;
+            }
+        }
+    }
+
+    fn is_required(&self) -> bool {
+        true
+    }
+}

compiler/rustc_target/src/spec/mod.rs

@@ -1760,6 +1760,12 @@ pub enum StackProtector {
     /// the address of a local variable.
     Strong,
 
+    /// Stack protection for Rust code, the following are function check rules
+    /// that require stack protection in Rust:
+    /// - calls to stack memory allocation
+    /// - obtaining reference/pointer of local variables
+    Rusty,
+
     /// Generate stack canaries in all functions.
     All,
 }
@@ -1770,6 +1776,7 @@ impl StackProtector {
             StackProtector::None => "none",
             StackProtector::Basic => "basic",
             StackProtector::Strong => "strong",
+            StackProtector::Rusty => "rusty",
             StackProtector::All => "all",
         }
     }
@@ -1783,6 +1790,7 @@ impl FromStr for StackProtector {
             "none" => StackProtector::None,
             "basic" => StackProtector::Basic,
             "strong" => StackProtector::Strong,
+            "rusty" => StackProtector::Rusty,
             "all" => StackProtector::All,
             _ => return Err(()),
         })

tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect-windows-64bit.rs

@@ -1,9 +1,10 @@
-//@ revisions: all strong basic none missing
+//@ revisions: all rusty strong basic none missing
 //@ assembly-output: emit-asm
 //@ only-windows
 //@ only-msvc
 //@ ignore-32bit 64-bit table based SEH has slightly different behaviors than classic SEH
 //@ [all] compile-flags: -Z stack-protector=all
+//@ [rusty] compile-flags: -Z stack-protector=rusty
 //@ [strong] compile-flags: -Z stack-protector=strong
 //@ [basic] compile-flags: -Z stack-protector=basic
 //@ [none] compile-flags: -Z stack-protector=none
@@ -16,6 +17,7 @@
 #[no_mangle]
 pub fn emptyfn() {
     // all: __security_check_cookie
+    // rusty-NOT: __security_check_cookie
     // strong-NOT: __security_check_cookie
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -34,6 +36,7 @@ pub fn array_char(f: fn(*const char)) {
     f(&c as *const _);
 
     // all: __security_check_cookie
+    // rusty: __security_check_cookie
     // strong: __security_check_cookie
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -50,6 +53,7 @@ pub fn array_u8_1(f: fn(*const u8)) {
     // array variables regardless of their size.
 
     // all: __security_check_cookie
+    // rusty: __security_check_cookie
     // strong: __security_check_cookie
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -67,6 +71,7 @@ pub fn array_u8_small(f: fn(*const u8)) {
     // Small arrays do not lead to stack protection by the 'basic' heuristic.
 
     // all: __security_check_cookie
+    // rusty: __security_check_cookie
     // strong: __security_check_cookie
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -83,6 +88,7 @@ pub fn array_u8_large(f: fn(*const u8)) {
     // will also protect this function.
 
     // all: __security_check_cookie
+    // rusty: __security_check_cookie
     // strong: __security_check_cookie
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -102,6 +108,7 @@ pub fn array_bytesizednewtype_9(f: fn(*const ByteSizedNewtype)) {
     // also protect this function.
 
     // all: __security_check_cookie
+    // rusty: __security_check_cookie
     // strong: __security_check_cookie
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -129,6 +136,7 @@ pub fn local_var_addr_used_indirectly(f: fn(bool)) {
     // ```
 
     // all: __security_check_cookie
+    // rusty: __security_check_cookie
     // strong: __security_check_cookie
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -164,9 +172,10 @@ pub fn local_string_addr_taken(f: fn(&String)) {
     // LLVM does not support generating stack protectors in functions with funclet
     // based EH personalities.
     // https://github.com/llvm/llvm-project/blob/37fd3c96b917096d8a550038f6e61cdf0fc4174f/llvm/lib/CodeGen/StackProtector.cpp#L103C1-L109C4
+
     // all-NOT: __security_check_cookie
+    // rusty-NOT: __security_check_cookie
     // strong-NOT: __security_check_cookie
-
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
@@ -197,6 +206,10 @@ pub fn local_var_addr_taken_used_locally_only(factory: fn() -> i32, sink: fn(i32
     // the `strong` heuristic.
 
     // all: __security_check_cookie
+
+    // FIXME: rusty stack smash protection needs to support inline scenario detection
+    // rusty: __security_check_cookie
+
     // strong-NOT: __security_check_cookie
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -234,6 +247,10 @@ pub fn local_large_var_moved(f: fn(Gigastruct)) {
     // ```
 
     // all: __security_check_cookie
+
+    // FIXME: How does the rust compiler handle moves of large structures?
+    // rusty-NOT: __security_check_cookie
+
     // strong: __security_check_cookie
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -263,6 +280,10 @@ pub fn local_large_var_cloned(f: fn(Gigastruct)) {
     // ```
 
     // all: __security_check_cookie
+
+    // FIXME: How does the rust compiler handle moves of large structures?
+    // rusty-NOT: __security_check_cookie
+
     // strong: __security_check_cookie
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -303,6 +324,12 @@ pub fn alloca_small_compile_time_constant_arg(f: fn(*mut ())) {
     f(unsafe { alloca(8) });
 
     // all: __security_check_cookie
+
+    // 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: __security_check_cookie
+
     // strong-NOT: __security_check_cookie
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -315,6 +342,7 @@ pub fn alloca_large_compile_time_constant_arg(f: fn(*mut ())) {
     f(unsafe { alloca(9) });
 
     // all: __security_check_cookie
+    // rusty: __security_check_cookie
     // strong-NOT: __security_check_cookie
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -327,6 +355,7 @@ pub fn alloca_dynamic_arg(f: fn(*mut ()), n: usize) {
     f(unsafe { alloca(n) });
 
     // all: __security_check_cookie
+    // rusty: __security_check_cookie
     // strong-NOT: __security_check_cookie
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
@@ -358,6 +387,7 @@ pub fn unsized_fn_param(s: [u8], l: bool, f: fn([u8])) {
     // based EH personalities.
     // https://github.com/llvm/llvm-project/blob/37fd3c96b917096d8a550038f6e61cdf0fc4174f/llvm/lib/CodeGen/StackProtector.cpp#L103C1-L109C4
     // all-NOT: __security_check_cookie
+    // rusty-NOT: __security_check_cookie
     // strong-NOT: __security_check_cookie
 
     // basic-NOT: __security_check_cookie