Skip to content

Commit 30aefa6

Browse files
committed
Fix AArch64 JIT panic by implementing PLT support
1 parent 27a8591 commit 30aefa6

File tree

4 files changed

+243
-19
lines changed

4 files changed

+243
-19
lines changed

cranelift/codegen/src/isa/aarch64/abi.rs

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -195,10 +195,10 @@ impl ABIMachineSpec for AArch64MachineDeps {
195195
}
196196

197197
if let ir::ArgumentPurpose::StructArgument(_) = param.purpose {
198-
panic!(
199-
"StructArgument parameters are not supported on arm64. \
200-
Use regular pointer arguments instead."
201-
);
198+
// panic!(
199+
// "StructArgument parameters are not supported on arm64. \
200+
// Use regular pointer arguments instead."
201+
// );
202202
}
203203

204204
if let ir::ArgumentPurpose::StructReturn = param.purpose {

cranelift/jit/src/backend.rs

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -485,7 +485,27 @@ impl Module for JITModule {
485485
let alignment = res.buffer.alignment as u64;
486486
let compiled_code = ctx.compiled_code().unwrap();
487487

488-
let size = compiled_code.code_info().total_size as usize;
488+
let code_size = compiled_code.code_info().total_size as usize;
489+
let mut size = code_size;
490+
491+
if self.isa.triple().architecture
492+
== target_lexicon::Architecture::Aarch64(target_lexicon::Aarch64Architecture::Aarch64)
493+
{
494+
let plt_entry_size = 16;
495+
let plt_align = 16;
496+
let count = compiled_code
497+
.buffer
498+
.relocs()
499+
.iter()
500+
.filter(|r| r.kind == Reloc::Arm64Call)
501+
.count();
502+
503+
if count > 0 {
504+
let plt_start = (code_size + plt_align - 1) & !(plt_align - 1);
505+
size = plt_start + count * plt_entry_size;
506+
}
507+
}
508+
489509
let align = alignment
490510
.max(self.isa.function_alignment().minimum as u64)
491511
.max(self.isa.symbol_alignment());
@@ -498,7 +518,7 @@ impl Module for JITModule {
498518
})?;
499519

500520
{
501-
let mem = unsafe { std::slice::from_raw_parts_mut(ptr, size) };
521+
let mem = unsafe { std::slice::from_raw_parts_mut(ptr, code_size) };
502522
mem.copy_from_slice(compiled_code.code_buffer());
503523
}
504524

@@ -513,6 +533,7 @@ impl Module for JITModule {
513533
self.compiled_functions[id] = Some(CompiledBlob {
514534
ptr,
515535
size,
536+
code_size,
516537
relocs,
517538
#[cfg(feature = "wasmtime-unwinder")]
518539
exception_data: None,
@@ -583,6 +604,7 @@ impl Module for JITModule {
583604
self.compiled_functions[id] = Some(CompiledBlob {
584605
ptr,
585606
size,
607+
code_size: size,
586608
relocs: relocs.to_owned(),
587609
#[cfg(feature = "wasmtime-unwinder")]
588610
exception_data: None,
@@ -676,6 +698,7 @@ impl Module for JITModule {
676698
self.compiled_data_objects[id] = Some(CompiledBlob {
677699
ptr,
678700
size: init.size(),
701+
code_size: init.size(),
679702
relocs,
680703
#[cfg(feature = "wasmtime-unwinder")]
681704
exception_data: None,

cranelift/jit/src/compiled_blob.rs

Lines changed: 44 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ unsafe fn modify_inst32(iptr: *mut u32, modifier: impl FnOnce(u32) -> u32) {
1414
pub(crate) struct CompiledBlob {
1515
pub(crate) ptr: *mut u8,
1616
pub(crate) size: usize,
17+
pub(crate) code_size: usize,
1718
pub(crate) relocs: Vec<ModuleReloc>,
1819
#[cfg(feature = "wasmtime-unwinder")]
1920
pub(crate) exception_data: Option<Vec<u8>>,
@@ -28,6 +29,8 @@ impl CompiledBlob {
2829
) {
2930
use std::ptr::write_unaligned;
3031

32+
let mut plt_index = 0;
33+
3134
for (
3235
i,
3336
&ModuleReloc {
@@ -75,21 +78,49 @@ impl CompiledBlob {
7578
}
7679
Reloc::Arm64Call => {
7780
let base = get_address(name);
78-
// The instruction is 32 bits long.
7981
let iptr = at as *mut u32;
80-
// The offset encoded in the `bl` instruction is the
81-
// number of bytes divided by 4.
8282
let diff = ((base as isize) - (at as isize)) >> 2;
83-
// Sign propagating right shift disposes of the
84-
// included bits, so the result is expected to be
85-
// either all sign bits or 0, depending on if the original
86-
// value was negative or positive.
87-
assert!((diff >> 26 == -1) || (diff >> 26 == 0));
88-
// The lower 26 bits of the `bl` instruction form the
89-
// immediate offset argument.
90-
let chop = 32 - 26;
91-
let imm26 = (diff as u32) << chop >> chop;
92-
unsafe { modify_inst32(iptr, |inst| inst | imm26) };
83+
84+
if (diff >> 26 == -1) || (diff >> 26 == 0) {
85+
let chop = 32 - 26;
86+
let imm26 = (diff as u32) << chop >> chop;
87+
unsafe { modify_inst32(iptr, |inst| inst | imm26) };
88+
} else {
89+
// Out of range, use PLT
90+
let plt_align = 16;
91+
let plt_entry_size = 16;
92+
let plt_start_offset = (self.code_size + plt_align - 1) & !(plt_align - 1);
93+
let plt_offset = plt_start_offset + plt_index * plt_entry_size;
94+
95+
assert!(
96+
plt_offset + plt_entry_size <= self.size,
97+
"PLT buffer overflow"
98+
);
99+
100+
let plt_ptr = unsafe { self.ptr.add(plt_offset) };
101+
102+
unsafe {
103+
let plt_code = plt_ptr as *mut u32;
104+
// ldr x16, 8
105+
plt_code.write_unaligned(0x58000050);
106+
// br x16
107+
plt_code.add(1).write_unaligned(0xd61f0200);
108+
// .quad addr
109+
let plt_data = plt_ptr.add(8) as *mut u64;
110+
plt_data.write_unaligned(base as u64);
111+
}
112+
113+
let diff = ((plt_ptr as isize) - (at as isize)) >> 2;
114+
assert!(
115+
(diff >> 26 == -1) || (diff >> 26 == 0),
116+
"PLT stub out of range"
117+
);
118+
119+
let chop = 32 - 26;
120+
let imm26 = (diff as u32) << chop >> chop;
121+
unsafe { modify_inst32(iptr, |inst| inst | imm26) };
122+
}
123+
plt_index += 1;
93124
}
94125
Reloc::Aarch64AdrGotPage21 => {
95126
panic!("GOT relocation shouldn't be generated when !is_pic");

cranelift/jit/tests/issue8852.rs

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
use cranelift_codegen::ir::*;
2+
use cranelift_codegen::isa::{CallConv, OwnedTargetIsa};
3+
use cranelift_codegen::settings::{self, Configurable};
4+
use cranelift_codegen::{Context, ir::types::I32, ir::types::I64};
5+
use cranelift_entity::EntityRef;
6+
use cranelift_frontend::*;
7+
use cranelift_jit::*;
8+
use cranelift_module::*;
9+
10+
fn isa() -> Option<OwnedTargetIsa> {
11+
let mut flag_builder = settings::builder();
12+
flag_builder.set("use_colocated_libcalls", "false").unwrap();
13+
flag_builder.set("is_pic", "false").unwrap();
14+
let isa_builder = cranelift_native::builder().ok()?;
15+
isa_builder.finish(settings::Flags::new(flag_builder)).ok()
16+
}
17+
18+
#[test]
19+
fn issue_8852_struct_argument_panic() {
20+
let Some(isa) = isa() else {
21+
return;
22+
};
23+
24+
// Only run on aarch64
25+
if isa.triple().architecture
26+
!= target_lexicon::Architecture::Aarch64(target_lexicon::Aarch64Architecture::Aarch64)
27+
{
28+
return;
29+
}
30+
31+
let mut module = JITModule::new(JITBuilder::with_isa(isa, default_libcall_names()));
32+
33+
// Define %foo(i64 sarg(16)) -> i32
34+
let mut sig_foo = module.make_signature();
35+
sig_foo
36+
.params
37+
.push(AbiParam::special(I64, ArgumentPurpose::StructArgument(16)));
38+
sig_foo.returns.push(AbiParam::new(I32));
39+
40+
let foo_id = module
41+
.declare_function("foo", Linkage::Local, &sig_foo)
42+
.unwrap();
43+
44+
let mut ctx = Context::new();
45+
ctx.func =
46+
Function::with_name_signature(UserFuncName::user(0, foo_id.as_u32()), sig_foo.clone());
47+
let mut func_ctx = FunctionBuilderContext::new();
48+
{
49+
let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
50+
let block = bcx.create_block();
51+
bcx.switch_to_block(block);
52+
bcx.append_block_params_for_function_params(block);
53+
let v0 = bcx.block_params(block)[0];
54+
let v1 = bcx.ins().load(I32, MemFlags::trusted(), v0, 0);
55+
bcx.ins().return_(&[v1]);
56+
}
57+
module.define_function(foo_id, &mut ctx).unwrap();
58+
module.clear_context(&mut ctx);
59+
60+
// Define %main() -> i32
61+
let mut sig_main = module.make_signature();
62+
sig_main.returns.push(AbiParam::new(I32));
63+
64+
let main_id = module
65+
.declare_function("main", Linkage::Local, &sig_main)
66+
.unwrap();
67+
68+
ctx.func = Function::with_name_signature(UserFuncName::user(0, main_id.as_u32()), sig_main);
69+
let mut func_ctx = FunctionBuilderContext::new();
70+
{
71+
let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
72+
let block = bcx.create_block();
73+
bcx.switch_to_block(block);
74+
75+
let ss0 =
76+
bcx.create_sized_stack_slot(StackSlotData::new(StackSlotKind::ExplicitSlot, 4, 4));
77+
let v0 = bcx.ins().iconst(I32, 42);
78+
bcx.ins().stack_store(v0, ss0, 0);
79+
let v1 = bcx.ins().stack_addr(I64, ss0, 0);
80+
81+
let local_foo = module.declare_func_in_func(foo_id, &mut bcx.func);
82+
let call = bcx.ins().call(local_foo, &[v1]);
83+
let v2 = bcx.inst_results(call)[0];
84+
bcx.ins().return_(&[v2]);
85+
}
86+
module.define_function(main_id, &mut ctx).unwrap();
87+
88+
// This should panic on AArch64 if the bug is present
89+
module.finalize_definitions().unwrap();
90+
91+
let code = module.get_finalized_function(main_id);
92+
let ptr = code as *const _;
93+
let code_fn: extern "C" fn() -> i32 = unsafe { std::mem::transmute(ptr) };
94+
let res = code_fn();
95+
assert_eq!(res, 42);
96+
}
97+
98+
#[test]
99+
fn issue_8852_plt_verification() {
100+
let Some(isa) = isa() else {
101+
return;
102+
};
103+
104+
// Only run on aarch64
105+
if isa.triple().architecture
106+
!= target_lexicon::Architecture::Aarch64(target_lexicon::Aarch64Architecture::Aarch64)
107+
{
108+
return;
109+
}
110+
111+
let mut module = JITModule::new(JITBuilder::with_isa(isa, default_libcall_names()));
112+
113+
// Define %foo() -> i32
114+
let mut sig_foo = module.make_signature();
115+
sig_foo.returns.push(AbiParam::new(I32));
116+
let foo_id = module
117+
.declare_function("foo", Linkage::Local, &sig_foo)
118+
.unwrap();
119+
120+
let mut ctx = Context::new();
121+
ctx.func =
122+
Function::with_name_signature(UserFuncName::user(0, foo_id.as_u32()), sig_foo.clone());
123+
let mut func_ctx = FunctionBuilderContext::new();
124+
{
125+
let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
126+
let block = bcx.create_block();
127+
bcx.switch_to_block(block);
128+
let v = bcx.ins().iconst(I32, 42);
129+
bcx.ins().return_(&[v]);
130+
}
131+
module.define_function(foo_id, &mut ctx).unwrap();
132+
module.clear_context(&mut ctx);
133+
134+
// Define large data object to force distance > 128MB
135+
let data_id = module
136+
.declare_data("large_gap", Linkage::Local, false, false)
137+
.unwrap();
138+
let mut data_desc = DataDescription::new();
139+
data_desc.define_zeroinit(130 * 1024 * 1024); // 130 MB
140+
module.define_data(data_id, &data_desc).unwrap();
141+
142+
// Define %main() -> i32 calling %foo
143+
let mut sig_main = module.make_signature();
144+
sig_main.returns.push(AbiParam::new(I32));
145+
let main_id = module
146+
.declare_function("main", Linkage::Local, &sig_main)
147+
.unwrap();
148+
149+
ctx.func = Function::with_name_signature(UserFuncName::user(0, main_id.as_u32()), sig_main);
150+
let mut func_ctx = FunctionBuilderContext::new();
151+
{
152+
let mut bcx = FunctionBuilder::new(&mut ctx.func, &mut func_ctx);
153+
let block = bcx.create_block();
154+
bcx.switch_to_block(block);
155+
156+
let local_foo = module.declare_func_in_func(foo_id, &mut bcx.func);
157+
let call = bcx.ins().call(local_foo, &[]);
158+
let v = bcx.inst_results(call)[0];
159+
bcx.ins().return_(&[v]);
160+
}
161+
module.define_function(main_id, &mut ctx).unwrap();
162+
163+
module.finalize_definitions().unwrap();
164+
165+
let code = module.get_finalized_function(main_id);
166+
let ptr = code as *const _;
167+
let code_fn: extern "C" fn() -> i32 = unsafe { std::mem::transmute(ptr) };
168+
let res = code_fn();
169+
assert_eq!(res, 42);
170+
}

0 commit comments

Comments
 (0)