rusty stack protector

Author: SparrowLii

compiler/rustc_codegen_llvm/src/attributes.rs

@@ -254,12 +254,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::Strong => AttributeKind::StackProtectStrong,
         StackProtector::Basic => AttributeKind::StackProtect,
+        StackProtector::Rusty => {
+            if cx.tcx.stack_protector.borrow().contains(&def_id) {
+                AttributeKind::StackProtectReq
+            } else {
+                return None;
+            }
+        }
     };
 
     Some(sspattr.create_attr(cx.llcx))
@@ -384,7 +391,10 @@ 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

@@ -17,7 +17,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::profiling::SelfProfilerRef;
 use rustc_data_structures::sharded::{IntoPointer, ShardedHashMap};
@@ -1392,6 +1392,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,
 }
 
@@ -1608,6 +1610,7 @@ impl<'tcx> TyCtxt<'tcx> {
             canonical_param_env_cache: Default::default(),
             data_layout,
             alloc_map: interpret::AllocMap::new(),
+            stack_protector: Default::default(),
             current_gcx,
         });
 

compiler/rustc_mir_transform/src/lib.rs

@@ -44,6 +44,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;
@@ -57,6 +58,7 @@ mod lint_tail_expr_drop_order;
 mod patch;
 mod shim;
 mod ssa;
+mod stack_protector;
 
 /// We import passes via this macro so that we can have a static list of pass names
 /// (used to verify CLI arguments). It takes a list of modules, followed by the passes
@@ -451,6 +453,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,46 @@
+//! 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::TyCtxt;
+
+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;
+                    }
+                }
+            }
+        }
+    }
+
+    fn is_required(&self) -> bool {
+        true
+    }
+}

compiler/rustc_target/src/spec/mod.rs

@@ -1605,6 +1605,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,
 }
@@ -1615,6 +1621,7 @@ impl StackProtector {
             StackProtector::None => "none",
             StackProtector::Basic => "basic",
             StackProtector::Strong => "strong",
+            StackProtector::Rusty => "rusty",
             StackProtector::All => "all",
         }
     }
@@ -1628,6 +1635,7 @@ impl FromStr for StackProtector {
             "none" => StackProtector::None,
             "basic" => StackProtector::Basic,
             "strong" => StackProtector::Strong,
+            "rusty" => StackProtector::Rusty,
             "all" => StackProtector::All,
             _ => return Err(()),
         })

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

@@ -1,9 +1,10 @@
-//@ revisions: all strong basic none missing
+//@ revisions: all strong basic none missing rusty
 //@ 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
@@ -21,6 +22,8 @@ pub fn emptyfn() {
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty-NOT: __security_check_cookie
 }
 
 // CHECK-LABEL: array_char
@@ -39,6 +42,8 @@ pub fn array_char(f: fn(*const char)) {
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty: __security_check_cookie
 }
 
 // CHECK-LABEL: array_u8_1
@@ -55,6 +60,8 @@ pub fn array_u8_1(f: fn(*const u8)) {
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty: __security_check_cookie
 }
 
 // CHECK-LABEL: array_u8_small:
@@ -72,6 +79,8 @@ pub fn array_u8_small(f: fn(*const u8)) {
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty: __security_check_cookie
 }
 
 // CHECK-LABEL: array_u8_large:
@@ -88,6 +97,8 @@ pub fn array_u8_large(f: fn(*const u8)) {
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty: __security_check_cookie
 }
 
 #[derive(Copy, Clone)]
@@ -107,6 +118,8 @@ pub fn array_bytesizednewtype_9(f: fn(*const ByteSizedNewtype)) {
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty: __security_check_cookie
 }
 
 // CHECK-LABEL: local_var_addr_used_indirectly
@@ -134,6 +147,8 @@ pub fn local_var_addr_used_indirectly(f: fn(bool)) {
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty: __security_check_cookie
 }
 
 // CHECK-LABEL: local_string_addr_taken
@@ -172,6 +187,8 @@ pub fn local_string_addr_taken(f: fn(&String)) {
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
 
+    // rusty-NOT: __security_check_cookie
+
     // CHECK-DAG: .seh_endproc
 }
 
@@ -202,6 +219,9 @@ pub fn local_var_addr_taken_used_locally_only(factory: fn() -> i32, sink: fn(i32
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // FIXME: rusty stack smash protection needs to support inline scenario detection
+    // rusty: __security_check_cookie
 }
 
 pub struct Gigastruct {
@@ -239,6 +259,9 @@ pub fn local_large_var_moved(f: fn(Gigastruct)) {
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // TODO: How does the rust compiler handle moves of large structures?
+    // rusty-NOT: __security_check_cookie
 }
 
 // CHECK-LABEL: local_large_var_cloned
@@ -268,6 +291,9 @@ pub fn local_large_var_cloned(f: fn(Gigastruct)) {
     // basic: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // TODO: How does the rust compiler handle moves of large structures?
+    // rusty-NOT: __security_check_cookie
 }
 
 extern "C" {
@@ -308,6 +334,11 @@ pub fn alloca_small_compile_time_constant_arg(f: fn(*mut ())) {
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // TODO: 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
 }
 
 // CHECK-LABEL: alloca_large_compile_time_constant_arg
@@ -320,6 +351,11 @@ pub fn alloca_large_compile_time_constant_arg(f: fn(*mut ())) {
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // TODO: 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
 }
 
 // CHECK-LABEL: alloca_dynamic_arg
@@ -332,6 +368,11 @@ pub fn alloca_dynamic_arg(f: fn(*mut ()), n: usize) {
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // TODO: 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
 }
 
 // The question then is: in what ways can Rust code generate array-`alloca`
@@ -364,6 +405,8 @@ pub fn unsized_fn_param(s: [u8], l: bool, f: fn([u8])) {
     // basic-NOT: __security_check_cookie
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty-NOT: __security_check_cookie
 }
 
 // CHECK-LABEL: unsized_local
@@ -388,4 +431,6 @@ pub fn unsized_local(s: &[u8], l: bool, f: fn(&mut [u8])) {
 
     // none-NOT: __security_check_cookie
     // missing-NOT: __security_check_cookie
+
+    // rusty-NOT: __security_check_cookie
 }