From d054b94c85be3b40658ee68a83ddd00196a19406 Mon Sep 17 00:00:00 2001 From: cezarbbb Date: Mon, 4 Aug 2025 14:54:38 +0800 Subject: [PATCH] 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 --- compiler/rustc_codegen_llvm/src/attributes.rs | 15 +++- compiler/rustc_middle/src/ty/context.rs | 5 +- compiler/rustc_mir_transform/src/lib.rs | 13 ++++ .../src/stack_protector.rs | 70 +++++++++++++++++++ compiler/rustc_target/src/spec/mod.rs | 8 +++ ...otector-heuristics-effect-windows-64bit.rs | 34 ++++++++- .../stack-protector-heuristics-effect.rs | 35 +++++++++- 7 files changed, 174 insertions(+), 6 deletions(-) create mode 100644 compiler/rustc_mir_transform/src/stack_protector.rs diff --git a/compiler/rustc_codegen_llvm/src/attributes.rs b/compiler/rustc_codegen_llvm/src/attributes.rs index c548f4675834f..7fa9123354b16 100644 --- a/compiler/rustc_codegen_llvm/src/attributes.rs +++ b/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")); diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs index db56082c71aaa..5216969d2b87c 100644 --- a/compiler/rustc_middle/src/ty/context.rs +++ b/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>, + 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, }); diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs index 08f25276cecc1..6c7e5f8d4c43f 100644 --- a/compiler/rustc_mir_transform/src/lib.rs +++ b/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)) } diff --git a/compiler/rustc_mir_transform/src/stack_protector.rs b/compiler/rustc_mir_transform/src/stack_protector.rs new file mode 100644 index 0000000000000..b5f0263004eba --- /dev/null +++ b/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 + } +} diff --git a/compiler/rustc_target/src/spec/mod.rs b/compiler/rustc_target/src/spec/mod.rs index 033590e01a67d..4a0269167b57d 100644 --- a/compiler/rustc_target/src/spec/mod.rs +++ b/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(()), }) diff --git a/tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect-windows-64bit.rs b/tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect-windows-64bit.rs index 9a3dabc74dded..d09d078d82591 100644 --- a/tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect-windows-64bit.rs +++ b/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 diff --git a/tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect.rs b/tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect.rs index ae281cb95da5f..0cc8ebc85f329 100644 --- a/tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect.rs +++ b/tests/assembly-llvm/stack-protector/stack-protector-heuristics-effect.rs @@ -1,10 +1,11 @@ -//@ revisions: all strong basic none missing +//@ revisions: all rusty strong basic none missing //@ assembly-output: emit-asm //@ ignore-apple slightly different policy on stack protection of arrays //@ ignore-msvc stack check code uses different function names //@ ignore-nvptx64 stack protector is not supported //@ ignore-wasm32-bare //@ [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 @@ -23,6 +24,7 @@ #[no_mangle] pub fn emptyfn() { // all: __stack_chk_fail + // rusty-NOT: __stack_chk_fail // strong-NOT: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -41,6 +43,7 @@ pub fn array_char(f: fn(*const char)) { f(&c as *const _); // all: __stack_chk_fail + // rusty: __stack_chk_fail // strong: __stack_chk_fail // basic: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -57,6 +60,7 @@ pub fn array_u8_1(f: fn(*const u8)) { // array variables regardless of their size. // all: __stack_chk_fail + // rusty: __stack_chk_fail // strong: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -74,6 +78,7 @@ pub fn array_u8_small(f: fn(*const u8)) { // Small arrays do not lead to stack protection by the 'basic' heuristic. // all: __stack_chk_fail + // rusty: __stack_chk_fail // strong: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -90,6 +95,7 @@ pub fn array_u8_large(f: fn(*const u8)) { // will also protect this function. // all: __stack_chk_fail + // rusty: __stack_chk_fail // strong: __stack_chk_fail // basic: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -109,6 +115,7 @@ pub fn array_bytesizednewtype_9(f: fn(*const ByteSizedNewtype)) { // also protect this function. // all: __stack_chk_fail + // rusty: __stack_chk_fail // strong: __stack_chk_fail // basic: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -136,6 +143,7 @@ pub fn local_var_addr_used_indirectly(f: fn(bool)) { // ``` // all: __stack_chk_fail + // rusty: __stack_chk_fail // strong: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -152,6 +160,7 @@ pub fn local_string_addr_taken(f: fn(&String)) { // protection. It does not matter that the reference is not mut. // all: __stack_chk_fail + // rusty: __stack_chk_fail // strong: __stack_chk_fail // basic: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -181,6 +190,10 @@ pub fn local_var_addr_taken_used_locally_only(factory: fn() -> i32, sink: fn(i32 // the `strong` heuristic. // all: __stack_chk_fail + + // FIXME: rusty stack smash protection needs to support inline scenario detection + // rusty: __stack_chk_fail + // strong-NOT: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -218,6 +231,10 @@ pub fn local_large_var_moved(f: fn(Gigastruct)) { // ``` // all: __stack_chk_fail + + // FIXME: How does the rust compiler handle moves of large structures? + // rusty-NOT: __stack_chk_fail + // strong: __stack_chk_fail // basic: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -247,6 +264,10 @@ pub fn local_large_var_cloned(f: fn(Gigastruct)) { // ``` // all: __stack_chk_fail + + // FIXME: How does the rust compiler handle moves of large structures? + // rusty-NOT: __stack_chk_fail + // strong: __stack_chk_fail // basic: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -287,6 +308,12 @@ pub fn alloca_small_compile_time_constant_arg(f: fn(*mut ())) { f(unsafe { alloca(8) }); // all: __stack_chk_fail + + // 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 + // strong-NOT: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -299,6 +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 // strong-NOT: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -311,6 +339,7 @@ pub fn alloca_dynamic_arg(f: fn(*mut ()), n: usize) { f(unsafe { alloca(n) }); // all: __stack_chk_fail + // rusty: __stack_chk_fail // strong-NOT: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail @@ -338,6 +367,10 @@ pub fn unsized_fn_param(s: [u8], l: bool, f: fn([u8])) { // heuristics. // all: __stack_chk_fail + + // FIXME: Does Rusty need to handle this type of optimization? + // rusty: __stack_chk_fail + // strong-NOT: __stack_chk_fail // basic-NOT: __stack_chk_fail // none-NOT: __stack_chk_fail