Skip to content

Commit c6ea71a

Browse files
cezarbbbcezarbbb
committed
add new rusty option
1 parent cc0a5b7 commit c6ea71a

File tree

7 files changed

+171
-8
lines changed

7 files changed

+171
-8
lines changed

‎compiler/rustc_codegen_llvm/src/attributes.rs‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -270,10 +270,19 @@ fn probestack_attr<'ll>(cx: &CodegenCx<'ll, '_>) -> Option<&'ll Attribute> {
270270
Some(llvm::CreateAttrStringValue(cx.llcx, "probe-stack", attr_value))
271271
}
272272

273-
fn stackprotector_attr<'ll>(cx: &CodegenCx<'ll, '_>) -> Option<&'ll Attribute> {
273+
fn stackprotector_attr<'ll>(cx: &CodegenCx<'ll, '_>, def_id: DefId) -> Option<&'ll Attribute> {
274274
let sspattr = match cx.sess().stack_protector() {
275275
StackProtector::None => return None,
276276
StackProtector::All => AttributeKind::StackProtectReq,
277+
278+
StackProtector::Rusty => {
279+
if cx.tcx.stack_protector.borrow().contains(&def_id) {
280+
AttributeKind::StackProtectStrong
281+
} else {
282+
return None;
283+
}
284+
}
285+
277286
StackProtector::Strong => AttributeKind::StackProtectStrong,
278287
StackProtector::Basic => AttributeKind::StackProtect,
279288
};
@@ -384,7 +393,9 @@ pub(crate) fn llfn_attrs_from_instance<'ll, 'tcx>(
384393
to_add.extend(instrument_function_attr(cx));
385394
to_add.extend(nojumptables_attr(cx));
386395
to_add.extend(probestack_attr(cx));
387-
to_add.extend(stackprotector_attr(cx));
396+
397+
// stack protector
398+
to_add.extend(stackprotector_attr(cx, instance.def_id()));
388399

389400
if codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::NO_BUILTINS) {
390401
to_add.push(llvm::CreateAttrString(cx.llcx, "no-builtins"));

‎compiler/rustc_middle/src/ty/context.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use rustc_ast as ast;
2020
use rustc_attr_data_structures::{AttributeKind, find_attr};
2121
use rustc_data_structures::defer;
2222
use rustc_data_structures::fingerprint::Fingerprint;
23-
use rustc_data_structures::fx::FxHashMap;
23+
use rustc_data_structures::fx::{FxHashMap, FxHashSet};
2424
use rustc_data_structures::intern::Interned;
2525
use rustc_data_structures::jobserver::Proxy;
2626
use rustc_data_structures::profiling::SelfProfilerRef;
@@ -1519,6 +1519,8 @@ pub struct GlobalCtxt<'tcx> {
15191519
/// Stores memory for globals (statics/consts).
15201520
pub(crate) alloc_map: interpret::AllocMap<'tcx>,
15211521

1522+
pub stack_protector: Lock<FxHashSet<DefId>>,
1523+
15221524
current_gcx: CurrentGcx,
15231525

15241526
/// A jobserver reference used to release then acquire a token while waiting on a query.
@@ -1746,6 +1748,7 @@ impl<'tcx> TyCtxt<'tcx> {
17461748
clauses_cache: Default::default(),
17471749
data_layout,
17481750
alloc_map: interpret::AllocMap::new(),
1751+
stack_protector: Default::default(),
17491752
current_gcx,
17501753
jobserver_proxy,
17511754
});

‎compiler/rustc_mir_transform/src/lib.rs‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ mod pass_manager;
4040
use std::sync::LazyLock;
4141

4242
use pass_manager::{self as pm, Lint, MirLint, MirPass, WithMinOptLevel};
43+
use rustc_target::spec::StackProtector;
4344

4445
mod check_pointers;
4546
mod cost_checker;
@@ -193,6 +194,7 @@ declare_passes! {
193194
mod single_use_consts : SingleUseConsts;
194195
mod sroa : ScalarReplacementOfAggregates;
195196
mod strip_debuginfo : StripDebugInfo;
197+
mod stack_protector: StackProtectorFinder;
196198
mod unreachable_enum_branching : UnreachableEnumBranching;
197199
mod unreachable_prop : UnreachablePropagation;
198200
mod validate : Validator;
@@ -450,6 +452,17 @@ fn mir_promoted(
450452
lint_tail_expr_drop_order::run_lint(tcx, def, &body);
451453

452454
let promoted = promote_pass.promoted_fragments.into_inner();
455+
456+
if tcx.sess.stack_protector() == StackProtector::Rusty {
457+
pm::run_passes(
458+
tcx,
459+
&mut body,
460+
&[&stack_protector::StackProtectorFinder],
461+
None,
462+
pm::Optimizations::Allowed,
463+
)
464+
}
465+
453466
(tcx.alloc_steal_mir(body), tcx.alloc_steal_promoted(promoted))
454467
}
455468

‎compiler/rustc_mir_transform/src/stack_protector.rs‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
//! Validates the MIR to ensure that invariants are upheld.
2+
3+
use std::ops::Deref;
4+
5+
use rustc_middle::mir::*;
6+
use rustc_middle::ty;
7+
use rustc_middle::ty::Instance;
8+
use rustc_middle::ty::TyCtxt;
9+
use rustc_target::callconv::PassMode;
10+
11+
pub(super) struct StackProtectorFinder;
12+
13+
impl<'tcx> crate::MirPass<'tcx> for StackProtectorFinder {
14+
fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
15+
use Rvalue::*;
16+
let def_id = body.source.def_id();
17+
18+
for block in body.basic_blocks.iter() {
19+
for stmt in block.statements.iter() {
20+
if let StatementKind::Assign(assign) = &stmt.kind {
21+
let (_, rvalue) = assign.deref();
22+
match rvalue {
23+
// Get a reference/pointer to a variable
24+
Ref(..) | ThreadLocalRef(_) | RawPtr(..) => {
25+
tcx.stack_protector.borrow_mut().insert(def_id);
26+
return;
27+
}
28+
_ => continue,
29+
}
30+
}
31+
}
32+
33+
if let Some(terminator) = block.terminator.as_ref() {
34+
if let TerminatorKind::Call { destination: place, .. } = &terminator.kind {
35+
// Returns a mutable raw pointer, possibly a memory allocation function
36+
if let ty::RawPtr(_, Mutability::Mut) = place.ty(body, tcx).ty.kind() {
37+
tcx.stack_protector.borrow_mut().insert(def_id);
38+
return;
39+
}
40+
}
41+
}
42+
43+
let instance = Instance::mono(tcx, def_id);
44+
let Ok(fn_abi) = tcx.fn_abi_of_instance(ty::TypingEnv::fully_monomorphized().as_query_input((instance, ty::List::empty())),) else { return;};
45+
46+
// for arg in fn_abi.args.iter() {
47+
// if matches!(&arg.mode, PassMode::Indirect { attrs: _, meta_attrs: _, on_stack: false }) {
48+
// tcx.stack_protector.borrow_mut().insert(def_id);
49+
// return;
50+
// }
51+
// }
52+
53+
let ret = &fn_abi.ret;
54+
if matches!(&ret.mode, PassMode::Indirect { attrs: _, meta_attrs: _, on_stack: false }) {
55+
tcx.stack_protector.borrow_mut().insert(def_id);
56+
return;
57+
}
58+
59+
}
60+
}
61+
62+
fn is_required(&self) -> bool {
63+
true
64+
}
65+
}

‎compiler/rustc_target/src/spec/mod.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1760,6 +1760,12 @@ pub enum StackProtector {
17601760
/// the address of a local variable.
17611761
Strong,
17621762

1763+
/// Stack protection for Rust code, the following are function check rules
1764+
/// that require stack protection in Rust:
1765+
/// - calls to stack memory allocation
1766+
/// - obtaining reference/pointer of local variables
1767+
Rusty,
1768+
17631769
/// Generate stack canaries in all functions.
17641770
All,
17651771
}
@@ -1770,6 +1776,7 @@ impl StackProtector {
17701776
StackProtector::None => "none",
17711777
StackProtector::Basic => "basic",
17721778
StackProtector::Strong => "strong",
1779+
StackProtector::Rusty => "rusty",
17731780
StackProtector::All => "all",
17741781
}
17751782
}
@@ -1783,6 +1790,7 @@ impl FromStr for StackProtector {
17831790
"none" => StackProtector::None,
17841791
"basic" => StackProtector::Basic,
17851792
"strong" => StackProtector::Strong,
1793+
"rusty" => StackProtector::Rusty,
17861794
"all" => StackProtector::All,
17871795
_ => return Err(()),
17881796
})

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

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
//@ revisions: all strong basic none missing
1+
//@ revisions: all rusty strong basic none missing
22
//@ assembly-output: emit-asm
33
//@ only-windows
44
//@ only-msvc
55
//@ ignore-32bit 64-bit table based SEH has slightly different behaviors than classic SEH
66
//@ [all] compile-flags: -Z stack-protector=all
7+
//@ [rusty] compile-flags: -Z stack-protector=rusty
78
//@ [strong] compile-flags: -Z stack-protector=strong
89
//@ [basic] compile-flags: -Z stack-protector=basic
910
//@ [none] compile-flags: -Z stack-protector=none
@@ -16,6 +17,7 @@
1617
#[no_mangle]
1718
pub fn emptyfn() {
1819
// all: __security_check_cookie
20+
// rusty-NOT: __security_check_cookie
1921
// strong-NOT: __security_check_cookie
2022
// basic-NOT: __security_check_cookie
2123
// none-NOT: __security_check_cookie
@@ -34,6 +36,7 @@ pub fn array_char(f: fn(*const char)) {
3436
f(&c as *const _);
3537

3638
// all: __security_check_cookie
39+
// rusty: __security_check_cookie
3740
// strong: __security_check_cookie
3841
// basic: __security_check_cookie
3942
// none-NOT: __security_check_cookie
@@ -50,6 +53,7 @@ pub fn array_u8_1(f: fn(*const u8)) {
5053
// array variables regardless of their size.
5154

5255
// all: __security_check_cookie
56+
// rusty: __security_check_cookie
5357
// strong: __security_check_cookie
5458
// basic-NOT: __security_check_cookie
5559
// none-NOT: __security_check_cookie
@@ -67,6 +71,7 @@ pub fn array_u8_small(f: fn(*const u8)) {
6771
// Small arrays do not lead to stack protection by the 'basic' heuristic.
6872

6973
// all: __security_check_cookie
74+
// rusty: __security_check_cookie
7075
// strong: __security_check_cookie
7176
// basic-NOT: __security_check_cookie
7277
// none-NOT: __security_check_cookie
@@ -83,6 +88,7 @@ pub fn array_u8_large(f: fn(*const u8)) {
8388
// will also protect this function.
8489

8590
// all: __security_check_cookie
91+
// rusty: __security_check_cookie
8692
// strong: __security_check_cookie
8793
// basic: __security_check_cookie
8894
// none-NOT: __security_check_cookie
@@ -102,6 +108,7 @@ pub fn array_bytesizednewtype_9(f: fn(*const ByteSizedNewtype)) {
102108
// also protect this function.
103109

104110
// all: __security_check_cookie
111+
// rusty: __security_check_cookie
105112
// strong: __security_check_cookie
106113
// basic: __security_check_cookie
107114
// none-NOT: __security_check_cookie
@@ -129,6 +136,7 @@ pub fn local_var_addr_used_indirectly(f: fn(bool)) {
129136
// ```
130137

131138
// all: __security_check_cookie
139+
// rusty: __security_check_cookie
132140
// strong: __security_check_cookie
133141
// basic-NOT: __security_check_cookie
134142
// none-NOT: __security_check_cookie
@@ -164,9 +172,10 @@ pub fn local_string_addr_taken(f: fn(&String)) {
164172
// LLVM does not support generating stack protectors in functions with funclet
165173
// based EH personalities.
166174
// https://github.com/llvm/llvm-project/blob/37fd3c96b917096d8a550038f6e61cdf0fc4174f/llvm/lib/CodeGen/StackProtector.cpp#L103C1-L109C4
175+
167176
// all-NOT: __security_check_cookie
177+
// rusty-NOT: __security_check_cookie
168178
// strong-NOT: __security_check_cookie
169-
170179
// basic-NOT: __security_check_cookie
171180
// none-NOT: __security_check_cookie
172181
// missing-NOT: __security_check_cookie
@@ -197,6 +206,10 @@ pub fn local_var_addr_taken_used_locally_only(factory: fn() -> i32, sink: fn(i32
197206
// the `strong` heuristic.
198207

199208
// all: __security_check_cookie
209+
210+
// FIXME: rusty stack smash protection needs to support inline scenario detection
211+
// rusty: __security_check_cookie
212+
200213
// strong-NOT: __security_check_cookie
201214
// basic-NOT: __security_check_cookie
202215
// none-NOT: __security_check_cookie
@@ -234,6 +247,10 @@ pub fn local_large_var_moved(f: fn(Gigastruct)) {
234247
// ```
235248

236249
// all: __security_check_cookie
250+
251+
// FIXME: How does the rust compiler handle moves of large structures?
252+
// rusty-NOT: __security_check_cookie
253+
237254
// strong: __security_check_cookie
238255
// basic: __security_check_cookie
239256
// none-NOT: __security_check_cookie
@@ -263,6 +280,10 @@ pub fn local_large_var_cloned(f: fn(Gigastruct)) {
263280
// ```
264281

265282
// all: __security_check_cookie
283+
284+
// FIXME: How does the rust compiler handle moves of large structures?
285+
// rusty-NOT: __security_check_cookie
286+
266287
// strong: __security_check_cookie
267288
// basic: __security_check_cookie
268289
// none-NOT: __security_check_cookie
@@ -301,8 +322,14 @@ extern "C" {
301322
#[no_mangle]
302323
pub fn alloca_small_compile_time_constant_arg(f: fn(*mut ())) {
303324
f(unsafe { alloca(8) });
304-
325+
305326
// all: __security_check_cookie
327+
328+
// FIXME: Rusty thinks a function that returns a mutable raw pointer may
329+
// be a stack memory allocation function, so it performs stack smash protection.
330+
// Is it possible to optimize the heuristics?
331+
// rusty: __security_check_cookie
332+
306333
// strong-NOT: __security_check_cookie
307334
// basic-NOT: __security_check_cookie
308335
// none-NOT: __security_check_cookie
@@ -315,6 +342,7 @@ pub fn alloca_large_compile_time_constant_arg(f: fn(*mut ())) {
315342
f(unsafe { alloca(9) });
316343

317344
// all: __security_check_cookie
345+
// rusty: __security_check_cookie
318346
// strong-NOT: __security_check_cookie
319347
// basic-NOT: __security_check_cookie
320348
// none-NOT: __security_check_cookie
@@ -327,6 +355,7 @@ pub fn alloca_dynamic_arg(f: fn(*mut ()), n: usize) {
327355
f(unsafe { alloca(n) });
328356

329357
// all: __security_check_cookie
358+
// rusty: __security_check_cookie
330359
// strong-NOT: __security_check_cookie
331360
// basic-NOT: __security_check_cookie
332361
// none-NOT: __security_check_cookie
@@ -357,9 +386,10 @@ pub fn unsized_fn_param(s: [u8], l: bool, f: fn([u8])) {
357386
// LLVM does not support generating stack protectors in functions with funclet
358387
// based EH personalities.
359388
// https://github.com/llvm/llvm-project/blob/37fd3c96b917096d8a550038f6e61cdf0fc4174f/llvm/lib/CodeGen/StackProtector.cpp#L103C1-L109C4
389+
360390
// all-NOT: __security_check_cookie
391+
// rusty-NOT: __security_check_cookie
361392
// strong-NOT: __security_check_cookie
362-
363393
// basic-NOT: __security_check_cookie
364394
// none-NOT: __security_check_cookie
365395
// missing-NOT: __security_check_cookie

0 commit comments

Comments
 (0)