Skip to content

Commit 7c2f3ec

Browse files
fffonionfffonion
committed
fix(jit): validate inherited continuation targets
1 parent e492fb1 commit 7c2f3ec

7 files changed

Lines changed: 141 additions & 17 deletions

File tree

‎src/vm/jit/native/lower.rs‎

Lines changed: 60 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,9 @@ const INHERITED_STATE_FRAME_KEY_OFFSET: i32 = 8;
6262
const INHERITED_STATE_STACK_BASE_OFFSET: i32 = 16;
6363
const INHERITED_STATE_LOCAL_BASE_OFFSET: i32 = 24;
6464
const INHERITED_STATE_DYNAMIC_TARGET_OFFSET: i32 = 32;
65-
const INHERITED_STATE_VALUES_OFFSET: i32 = 40;
65+
const INHERITED_STATE_TARGET_IP_OFFSET: i32 = 40;
66+
const INHERITED_STATE_VALUE_COUNT_OFFSET: i32 = 48;
67+
const INHERITED_STATE_VALUES_OFFSET: i32 = 56;
6668

6769
type TaggedConstants = (Box<[Value]>, HashMap<SsaValueId, usize>);
6870

@@ -260,7 +262,7 @@ pub(crate) fn compile_system_inherited_tail_wrapper(
260262
let vm_ptr = builder.block_params(entry)[0];
261263
let pointer_bytes = pointer_type.bits() / 8;
262264
let packet_bytes = pointer_bytes
263-
.checked_mul((MAX_INHERITED_ENTRY_VALUES + 5) as u32)
265+
.checked_mul((MAX_INHERITED_ENTRY_VALUES + 7) as u32)
264266
.ok_or_else(|| {
265267
VmError::JitNative("native inherited-state packet is too large".to_string())
266268
})?;
@@ -359,7 +361,7 @@ pub(crate) fn compile_tail_trace_dispatcher(
359361
let dynamic_target_is_null = builder.ins().icmp_imm(IntCC::Equal, dynamic_target, 0);
360362
builder.ins().brif(
361363
dynamic_target_is_null,
362-
static_dispatch,
364+
return_status,
363365
&[status.into()],
364366
dynamic_linked,
365367
&[],
@@ -858,8 +860,40 @@ fn try_compile_ssa_trace(
858860
let frame_matches = b
859861
.ins()
860862
.icmp(IntCC::Equal, inherited_frame_key, expected_frame_key);
863+
let inherited_target_ip = b.ins().load(
864+
pointer_type,
865+
MemFlags::new(),
866+
inherited_state_ptr,
867+
INHERITED_STATE_TARGET_IP_OFFSET,
868+
);
869+
let expected_target_ip = b.ins().iconst(
870+
pointer_type,
871+
i64::try_from(trace.root_ip).map_err(|_| {
872+
VmError::JitNative("SSA inherited target ip exceeds i64".to_string())
873+
})?,
874+
);
875+
let target_matches =
876+
b.ins()
877+
.icmp(IntCC::Equal, inherited_target_ip, expected_target_ip);
878+
let inherited_value_count = b.ins().load(
879+
pointer_type,
880+
MemFlags::new(),
881+
inherited_state_ptr,
882+
INHERITED_STATE_VALUE_COUNT_OFFSET,
883+
);
884+
let expected_value_count = b.ins().iconst(
885+
pointer_type,
886+
i64::try_from(entry_ssa_block.params.len()).map_err(|_| {
887+
VmError::JitNative("SSA inherited value count exceeds i64".to_string())
888+
})?,
889+
);
890+
let value_count_matches =
891+
b.ins()
892+
.icmp(IntCC::Equal, inherited_value_count, expected_value_count);
893+
let metadata_matches = b.ins().band(frame_matches, target_matches);
894+
let metadata_matches = b.ins().band(metadata_matches, value_count_matches);
861895
b.ins()
862-
.brif(frame_matches, inherited_entry, &[], normal_entry, &[]);
896+
.brif(metadata_matches, inherited_entry, &[], normal_entry, &[]);
863897

864898
b.switch_to_block(inherited_entry);
865899
let active_stack_base = b.ins().load(
@@ -4472,6 +4506,28 @@ fn write_inherited_state_packet(
44724506
ctx.inherited_state_ptr,
44734507
INHERITED_STATE_LOCAL_BASE_OFFSET,
44744508
);
4509+
let target_ip = b.ins().iconst(
4510+
ctx.pointer_type,
4511+
i64::try_from(exit.exit_ip)
4512+
.map_err(|_| VmError::JitNative("SSA inherited target ip exceeds i64".to_string()))?,
4513+
);
4514+
b.ins().store(
4515+
MemFlags::new(),
4516+
target_ip,
4517+
ctx.inherited_state_ptr,
4518+
INHERITED_STATE_TARGET_IP_OFFSET,
4519+
);
4520+
let value_count = b.ins().iconst(
4521+
ctx.pointer_type,
4522+
i64::try_from(entry_value_count)
4523+
.map_err(|_| VmError::JitNative("SSA inherited value count exceeds i64".to_string()))?,
4524+
);
4525+
b.ins().store(
4526+
MemFlags::new(),
4527+
value_count,
4528+
ctx.inherited_state_ptr,
4529+
INHERITED_STATE_VALUE_COUNT_OFFSET,
4530+
);
44754531
let stack_ptr = b.ins().load(
44764532
ctx.pointer_type,
44774533
MemFlags::new(),

‎src/vm/jit/native/mod.rs‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,6 @@ impl Default for NativeSideLinkSlot {
116116
}
117117
}
118118

119-
pub(crate) const LINKED_CONTINUE_SLOT_ID: u32 = u32::MAX;
120119
pub(crate) const CONTINUE_SLOT_ID: u32 = u32::MAX - 1;
121120

122121
pub(crate) struct CompiledTraceDispatcher {
@@ -145,9 +144,6 @@ pub(crate) fn compile_native_trace_dispatcher(
145144
slots.insert(exit.id.raw(), slot);
146145
}
147146
let slot = Arc::new(NativeSideLinkSlot::new());
148-
descriptors.push((STATUS_LINKED_CONTINUE, slot.address() as usize));
149-
slots.insert(LINKED_CONTINUE_SLOT_ID, slot);
150-
let slot = Arc::new(NativeSideLinkSlot::new());
151147
descriptors.push((STATUS_CONTINUE, slot.address() as usize));
152148
slots.insert(CONTINUE_SLOT_ID, slot);
153149
let dispatcher = lower::compile_tail_trace_dispatcher(trace_entry, trace_id, &descriptors)?;

‎src/vm/jit/runtime.rs‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1203,11 +1203,6 @@ impl Vm {
12031203
if let Some(next_trace_id) = next_trace_id
12041204
&& next_trace_id != current_trace_id
12051205
{
1206-
self.publish_native_direct_slot(
1207-
current_trace_id,
1208-
native::LINKED_CONTINUE_SLOT_ID,
1209-
next_trace_id,
1210-
)?;
12111206
self.record_jit_link_handoff();
12121207
current_trace_id = next_trace_id;
12131208
if let Some(state) = self.cached_native_trace_state(

‎src/vm/native/bridge.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -821,7 +821,8 @@ fn write_inherited_state_packet(vm: &Vm, packet: *mut u8) -> VmResult<()> {
821821
use super::{
822822
INHERITED_STATE_ACTIVE_OFFSET, INHERITED_STATE_DYNAMIC_TARGET_OFFSET,
823823
INHERITED_STATE_FRAME_KEY_OFFSET, INHERITED_STATE_LOCAL_BASE_OFFSET,
824-
INHERITED_STATE_STACK_BASE_OFFSET, INHERITED_STATE_VALUES_OFFSET,
824+
INHERITED_STATE_STACK_BASE_OFFSET, INHERITED_STATE_TARGET_IP_OFFSET,
825+
INHERITED_STATE_VALUE_COUNT_OFFSET, INHERITED_STATE_VALUES_OFFSET,
825826
MAX_INHERITED_ENTRY_VALUES,
826827
};
827828

@@ -861,6 +862,14 @@ fn write_inherited_state_packet(vm: &Vm, packet: *mut u8) -> VmResult<()> {
861862
.add(INHERITED_STATE_LOCAL_BASE_OFFSET as usize)
862863
.cast::<usize>()
863864
.write(state.local_base);
865+
packet
866+
.add(INHERITED_STATE_TARGET_IP_OFFSET as usize)
867+
.cast::<usize>()
868+
.write(vm.ip);
869+
packet
870+
.add(INHERITED_STATE_VALUE_COUNT_OFFSET as usize)
871+
.cast::<usize>()
872+
.write(value_count);
864873
let values = packet
865874
.add(INHERITED_STATE_VALUES_OFFSET as usize)
866875
.cast::<*const Value>();

‎src/vm/native/mod.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,9 @@ pub(crate) const INHERITED_STATE_FRAME_KEY_OFFSET: i32 = 8;
6060
pub(crate) const INHERITED_STATE_STACK_BASE_OFFSET: i32 = 16;
6161
pub(crate) const INHERITED_STATE_LOCAL_BASE_OFFSET: i32 = 24;
6262
pub(crate) const INHERITED_STATE_DYNAMIC_TARGET_OFFSET: i32 = 32;
63-
pub(crate) const INHERITED_STATE_VALUES_OFFSET: i32 = 40;
63+
pub(crate) const INHERITED_STATE_TARGET_IP_OFFSET: i32 = 40;
64+
pub(crate) const INHERITED_STATE_VALUE_COUNT_OFFSET: i32 = 48;
65+
pub(crate) const INHERITED_STATE_VALUES_OFFSET: i32 = 56;
6466

6567
#[cfg(feature = "cranelift-jit")]
6668
pub(crate) fn selected_codegen_backend() -> &'static str {

‎tests/jit/jit_tests.rs‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6190,6 +6190,41 @@ fn trace_jit_direct_links_cross_frame_call_and_return_edges() {
61906190
assert_eq!(vm.jit_native_region_count(), 0);
61916191
}
61926192

6193+
#[test]
6194+
fn trace_jit_missing_dynamic_return_target_never_uses_stale_static_continuation() {
6195+
if !native_jit_supported() {
6196+
return;
6197+
}
6198+
let source = r#"
6199+
let add: fn(int) -> int = |value| value + 1;
6200+
let mut i = 0;
6201+
let mut total = 0;
6202+
while i < 256 {
6203+
total = add(total);
6204+
total = total + 10;
6205+
i = i + 1;
6206+
}
6207+
total = add(total);
6208+
total = total + 100;
6209+
total;
6210+
"#;
6211+
let compiled = compile_source(source).expect("multiple continuation sites should compile");
6212+
let mut vm = Vm::new(compiled.program.with_local_count(compiled.locals));
6213+
vm.set_jit_native_direct_links_enabled(true);
6214+
vm.set_jit_config(JitConfig {
6215+
enabled: true,
6216+
hot_loop_threshold: 1,
6217+
max_trace_len: 512,
6218+
});
6219+
6220+
assert_eq!(vm.run().unwrap(), VmStatus::Halted);
6221+
assert_eq!(vm.stack(), &[Value::Int(2_917)]);
6222+
6223+
vm.reset_for_reuse();
6224+
assert_eq!(vm.run().unwrap(), VmStatus::Halted);
6225+
assert_eq!(vm.stack(), &[Value::Int(2_917)]);
6226+
}
6227+
61936228
#[test]
61946229
fn trace_jit_direct_link_slots_clear_and_republish_after_mode_toggle() {
61956230
if !native_jit_supported() {
@@ -6385,6 +6420,37 @@ fn trace_jit_call_value_yields_and_resumes_host_callable_without_losing_frame_st
63856420
assert!(vm.dump_jit_info().contains("interpreter fallbacks: 0"));
63866421
}
63876422

6423+
#[test]
6424+
fn trace_jit_missing_dynamic_return_target_does_not_use_stale_static_slot() {
6425+
if !native_jit_supported() {
6426+
return;
6427+
}
6428+
let source = r#"
6429+
fn inc(x: int) -> int { x + 1 }
6430+
let mut i = 0;
6431+
let mut value = 0;
6432+
while i < 32 {
6433+
value = inc(value);
6434+
i = i + 1;
6435+
}
6436+
value = inc(value) + 100;
6437+
value;
6438+
"#;
6439+
let compiled = compile_source(source).expect("stale return target fixture should compile");
6440+
let mut vm = Vm::new(compiled.program.with_local_count(compiled.locals));
6441+
vm.set_jit_native_direct_links_enabled(true);
6442+
vm.set_jit_config(JitConfig {
6443+
enabled: true,
6444+
hot_loop_threshold: 1,
6445+
max_trace_len: 256,
6446+
});
6447+
assert_eq!(
6448+
vm.run().expect("stale return target fixture should run"),
6449+
VmStatus::Halted
6450+
);
6451+
assert_eq!(vm.stack(), &[Value::Int(133)]);
6452+
}
6453+
63886454
#[test]
63896455
fn trace_jit_links_nested_dynamic_script_callables_without_interpreter_handoff() {
63906456
if !native_jit_supported() {

‎tests/jit/perf_tests.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1447,7 +1447,7 @@ fn build_host_call_latency_source(name: &str, arity: u8, calls: i64) -> String {
14471447
match arity {
14481448
0 => format!(
14491449
r#"
1450-
fn {name}();
1450+
fn {name}() -> int;
14511451
let mut i = 0;
14521452
let mut sum = 0;
14531453
while i < {calls} {{
@@ -1459,7 +1459,7 @@ fn build_host_call_latency_source(name: &str, arity: u8, calls: i64) -> String {
14591459
),
14601460
1 => format!(
14611461
r#"
1462-
fn {name}(x);
1462+
fn {name}(x) -> int;
14631463
let mut i = 0;
14641464
let mut sum = 0;
14651465
while i < {calls} {{

0 commit comments

Comments
 (0)