use anyhow::{Context, Result, bail};
use wasm_encoder::{BlockType, Function, MemArg, ValType};
use wasmparser::{Parser, Payload};
pub const ARENA_IMPORT_MODULE: &str = "env";
pub const ARENA_IMPORT_FIELD: &str = "__cabi_arena_realloc";
#[derive(Debug)]
pub enum ArenaBind {
NoArenaImport,
KeptHostSeam(&'static str),
Bound(BoundArena),
}
#[derive(Debug)]
pub struct BoundArena {
pub bytes: Vec<u8>,
pub arena_base: u32,
pub arena_end: u32,
}
struct Scan {
arena_type_idx: Option<u32>,
arena_sig_ok: bool,
total_imports: u32,
memory_pages: Option<u64>,
memory64: bool,
defined_globals: u32,
data_end: u32,
non_const_data_offset: bool,
global_inits: Vec<u32>,
has_function_section: bool,
has_code_section: bool,
}
fn scan(wasm: &[u8]) -> Result<Scan> {
let mut s = Scan {
arena_type_idx: None,
arena_sig_ok: false,
total_imports: 0,
memory_pages: None,
memory64: false,
defined_globals: 0,
data_end: 0,
non_const_data_offset: false,
global_inits: Vec::new(),
has_function_section: false,
has_code_section: false,
};
let mut func_types: Vec<bool> = Vec::new(); for payload in Parser::new(0).parse_all(wasm) {
match payload.context("parse wasm (#418 arena-bind scan)")? {
Payload::TypeSection(reader) => {
for rec_group in reader {
for sub_ty in rec_group.context("parse type section (#418)")?.types() {
let ok = match &sub_ty.composite_type.inner {
wasmparser::CompositeInnerType::Func(f) => {
f.params().len() == 4
&& f.params().iter().all(|t| *t == wasmparser::ValType::I32)
&& f.results() == [wasmparser::ValType::I32]
}
_ => false,
};
func_types.push(ok);
}
}
}
Payload::ImportSection(reader) => {
for import in reader.into_imports() {
let import = import.context("parse import (#418)")?;
s.total_imports += 1;
if import.module == ARENA_IMPORT_MODULE
&& import.name == ARENA_IMPORT_FIELD
&& let wasmparser::TypeRef::Func(type_idx) = import.ty
{
s.arena_type_idx = Some(type_idx);
s.arena_sig_ok =
func_types.get(type_idx as usize).copied().unwrap_or(false);
}
}
}
Payload::MemorySection(reader) => {
for (i, mem) in reader.into_iter().enumerate() {
let mem = mem.context("parse memory (#418)")?;
if i == 0 {
s.memory_pages = Some(mem.initial);
s.memory64 = mem.memory64;
}
}
}
Payload::GlobalSection(reader) => {
for global in reader {
let global = global.context("parse global (#418)")?;
s.defined_globals += 1;
let mut ops = global.init_expr.get_operators_reader();
if let Ok(wasmparser::Operator::I32Const { value }) = ops.read()
&& value > 0
{
s.global_inits.push(value as u32);
}
}
}
Payload::DataSection(reader) => {
for seg in reader {
let seg = seg.context("parse data segment (#418)")?;
if let wasmparser::DataKind::Active {
memory_index,
offset_expr,
} = seg.kind
{
if memory_index != 0 {
continue; }
let mut ops = offset_expr.get_operators_reader();
match ops.read() {
Ok(wasmparser::Operator::I32Const { value }) => {
let end = (value as u32).saturating_add(seg.data.len() as u32);
s.data_end = s.data_end.max(end);
}
_ => s.non_const_data_offset = true,
}
}
}
}
Payload::FunctionSection(_) => s.has_function_section = true,
Payload::CodeSectionStart { .. } => s.has_code_section = true,
_ => {}
}
}
Ok(s)
}
fn allocator_body(cursor_global: u32, arena_end: u32) -> Function {
let mem = MemArg {
offset: 0,
align: 0,
memory_index: 0,
};
let mut f = Function::new([(4, ValType::I32)]);
f.instructions()
.local_get(1)
.i32_eqz()
.local_get(3)
.i32_eqz()
.i32_and()
.if_(BlockType::Empty)
.local_get(2)
.return_()
.end()
.local_get(2)
.i32_eqz()
.if_(BlockType::Empty)
.unreachable()
.end()
.global_get(cursor_global)
.local_get(2)
.i32_add()
.i32_const(1)
.i32_sub()
.local_get(2)
.i32_const(1)
.i32_sub()
.i32_const(-1)
.i32_xor()
.i32_and()
.local_set(4)
.local_get(4)
.global_get(cursor_global)
.i32_lt_u()
.if_(BlockType::Empty)
.unreachable()
.end()
.local_get(4)
.local_get(3)
.i32_add()
.local_tee(5)
.local_get(4)
.i32_lt_u()
.if_(BlockType::Empty)
.unreachable()
.end()
.local_get(5)
.i32_const(arena_end as i32)
.i32_gt_u()
.if_(BlockType::Empty)
.unreachable()
.end()
.local_get(5)
.global_set(cursor_global)
.local_get(1)
.local_get(3)
.local_get(1)
.local_get(3)
.i32_lt_u()
.select()
.local_set(6)
.block(BlockType::Empty)
.loop_(BlockType::Empty)
.local_get(7)
.local_get(6)
.i32_ge_u()
.br_if(1)
.local_get(4)
.local_get(7)
.i32_add()
.local_get(0)
.local_get(7)
.i32_add()
.i32_load8_u(mem)
.i32_store8(mem)
.local_get(7)
.i32_const(1)
.i32_add()
.local_set(7)
.br(0)
.end()
.end()
.local_get(4)
.end();
f
}
fn read_uleb(bytes: &[u8]) -> Result<(u32, usize)> {
let mut value: u32 = 0;
let mut shift = 0;
for (i, &b) in bytes.iter().enumerate().take(5) {
value |= u32::from(b & 0x7F) << shift;
if b & 0x80 == 0 {
return Ok((value, i + 1));
}
shift += 7;
}
bail!("malformed LEB128 count in section (#418)");
}
fn write_uleb(mut value: u32, out: &mut Vec<u8>) {
loop {
let mut b = (value & 0x7F) as u8;
value >>= 7;
if value != 0 {
b |= 0x80;
}
out.push(b);
if value == 0 {
return;
}
}
}
fn write_sleb(mut value: i32, out: &mut Vec<u8>) {
loop {
let b = (value & 0x7F) as u8;
value >>= 7;
let sign = b & 0x40;
if (value == 0 && sign == 0) || (value == -1 && sign != 0) {
out.push(b);
return;
}
out.push(b | 0x80);
}
}
fn prepend_entry(contents: &[u8], entry: &[u8]) -> Result<Vec<u8>> {
let (count, len) = read_uleb(contents)?;
let mut out = Vec::with_capacity(contents.len() + entry.len() + 1);
write_uleb(count + 1, &mut out);
out.extend_from_slice(entry);
out.extend_from_slice(&contents[len..]);
Ok(out)
}
fn append_entry(contents: &[u8], entry: &[u8]) -> Result<Vec<u8>> {
let (count, len) = read_uleb(contents)?;
let mut out = Vec::with_capacity(contents.len() + entry.len() + 1);
write_uleb(count + 1, &mut out);
out.extend_from_slice(&contents[len..]);
out.extend_from_slice(entry);
Ok(out)
}
fn cursor_global_entry(arena_base: u32) -> Vec<u8> {
let mut e = vec![0x7F, 0x01, 0x41]; write_sleb(arena_base as i32, &mut e);
e.push(0x0B); e
}
pub fn bind_cabi_arena_realloc(wasm: &[u8]) -> Result<ArenaBind> {
let s = scan(wasm)?;
let Some(arena_type_idx) = s.arena_type_idx else {
return Ok(ArenaBind::NoArenaImport);
};
if s.total_imports > 1 {
return Ok(ArenaBind::KeptHostSeam(
"module has other imports — keeping the host-linked seam",
));
}
if !s.arena_sig_ok {
bail!(
"#418: env::{ARENA_IMPORT_FIELD} is imported with a signature \
other than (i32, i32, i32, i32) -> i32 — not the canonical-ABI \
arena realloc contract; refusing to bind (compile with \
--no-bind-cabi-arena to keep it an external symbol)"
);
}
let Some(pages) = s.memory_pages else {
bail!(
"#418: cannot bind env::{ARENA_IMPORT_FIELD}: the module declares \
no linear memory to allocate from"
);
};
if s.memory64 {
bail!("#418: cannot bind env::{ARENA_IMPORT_FIELD}: memory64 module");
}
if s.non_const_data_offset {
bail!(
"#418: cannot bind env::{ARENA_IMPORT_FIELD}: a data segment has \
a non-constant offset, so the static-data extent (the arena \
floor) cannot be derived soundly"
);
}
if !s.has_function_section || !s.has_code_section {
bail!(
"#418: cannot bind env::{ARENA_IMPORT_FIELD}: the module defines \
no functions (nothing synth could route the binding through)"
);
}
let arena_end: u32 = u32::try_from(pages.saturating_mul(64 * 1024))
.unwrap_or(u32::MAX)
.min(0xFFFF_0000);
let global_top = s
.global_inits
.iter()
.copied()
.filter(|&v| v <= arena_end)
.max()
.unwrap_or(0);
let arena_base = s.data_end.max(global_top).max(16).next_multiple_of(16);
if arena_base >= arena_end {
bail!(
"#418: cannot bind env::{ARENA_IMPORT_FIELD}: the static layout \
(data + stack + wasm-ld layout globals) extends to {arena_base} \
bytes but linear memory is only {arena_end} bytes — no arena \
region left; every allocation would trap"
);
}
let cursor_global = s.defined_globals;
let mut body = Vec::new();
wasm_encoder::Encode::encode(&allocator_body(cursor_global, arena_end), &mut body);
let mut module = wasm_encoder::Module::new();
let mut global_emitted = false;
let mut function_emitted = false;
let ensure_globals = |module: &mut wasm_encoder::Module, emitted: &mut bool| {
if !*emitted {
let mut out = Vec::new();
write_uleb(1, &mut out);
out.extend_from_slice(&cursor_global_entry(arena_base));
module.section(&wasm_encoder::RawSection {
id: wasm_encoder::SectionId::Global as u8,
data: &out,
});
*emitted = true;
}
};
for payload in Parser::new(0).parse_all(wasm) {
let payload = payload.context("parse wasm (#418 arena-bind rewrite)")?;
match &payload {
Payload::Version { .. } | Payload::End(_) => {}
Payload::ImportSection(_) => {
}
Payload::FunctionSection(reader) => {
let mut entry = Vec::new();
write_uleb(arena_type_idx, &mut entry);
let contents = &wasm[reader.range()];
module.section(&wasm_encoder::RawSection {
id: wasm_encoder::SectionId::Function as u8,
data: &prepend_entry(contents, &entry)?,
});
function_emitted = true;
}
Payload::GlobalSection(reader) => {
let contents = &wasm[reader.range()];
module.section(&wasm_encoder::RawSection {
id: wasm_encoder::SectionId::Global as u8,
data: &append_entry(contents, &cursor_global_entry(arena_base))?,
});
global_emitted = true;
}
Payload::ExportSection(_)
| Payload::StartSection { .. }
| Payload::ElementSection(_)
| Payload::DataCountSection { .. }
| Payload::DataSection(_) => {
ensure_globals(&mut module, &mut global_emitted);
copy_raw(&mut module, &payload, wasm)?;
}
Payload::CodeSectionStart { range, .. } => {
ensure_globals(&mut module, &mut global_emitted);
let contents = &wasm[range.clone()];
module.section(&wasm_encoder::RawSection {
id: wasm_encoder::SectionId::Code as u8,
data: &prepend_entry(contents, &body)?,
});
}
Payload::CodeSectionEntry(_) => {} other => copy_raw(&mut module, other, wasm)?,
}
}
if !function_emitted {
bail!("#418 internal: function section not re-emitted"); }
let bytes = module.finish();
wasmparser::Validator::new()
.validate_all(&bytes)
.context("#418 internal: arena-bind rewrite produced an invalid module (bug)")?;
Ok(ArenaBind::Bound(BoundArena {
bytes,
arena_base,
arena_end,
}))
}
fn copy_raw(module: &mut wasm_encoder::Module, payload: &Payload<'_>, wasm: &[u8]) -> Result<()> {
let Some((id, range)) = payload.as_section() else {
bail!("#418 internal: unhandled non-section payload {payload:?}");
};
module.section(&wasm_encoder::RawSection {
id,
data: &wasm[range],
});
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
fn fixture() -> Vec<u8> {
wat::parse_str(
r#"(module
(import "env" "__cabi_arena_realloc"
(func $arena (param i32 i32 i32 i32) (result i32)))
(memory (export "memory") 1)
(global $sp (mut i32) (i32.const 4096))
(global (export "__heap_base") i32 (i32.const 6144))
(data (i32.const 5120) "0123456789abcdef")
(func (export "cabi_realloc") (param i32 i32 i32 i32) (result i32)
local.get 0 local.get 1 local.get 2 local.get 3 call $arena))"#,
)
.unwrap()
}
#[test]
fn binds_sole_arena_import() {
let ArenaBind::Bound(b) = bind_cabi_arena_realloc(&fixture()).unwrap() else {
panic!("expected Bound");
};
assert_eq!(b.arena_base, 6144);
assert_eq!(b.arena_end, 65536);
let mut num_imports = 0;
let mut num_funcs = 0;
let mut num_globals = 0;
for p in Parser::new(0).parse_all(&b.bytes) {
match p.unwrap() {
Payload::ImportSection(r) => num_imports += r.count(),
Payload::FunctionSection(r) => num_funcs = r.count(),
Payload::GlobalSection(r) => num_globals = r.count(),
_ => {}
}
}
assert_eq!(num_imports, 0);
assert_eq!(num_funcs, 2); assert_eq!(num_globals, 3); }
#[test]
fn bound_module_executes_contract() {
let ArenaBind::Bound(b) = bind_cabi_arena_realloc(&fixture()).unwrap() else {
panic!("expected Bound");
};
wasmparser::Validator::new().validate_all(&b.bytes).unwrap();
}
#[test]
fn no_arena_import_passes_through() {
let wasm =
wat::parse_str(r#"(module (memory 1) (func (export "f") (result i32) i32.const 7))"#)
.unwrap();
assert!(matches!(
bind_cabi_arena_realloc(&wasm).unwrap(),
ArenaBind::NoArenaImport
));
}
#[test]
fn other_imports_keep_host_seam() {
let wasm = wat::parse_str(
r#"(module
(import "env" "k_spin_lock" (func (param i32)))
(import "env" "__cabi_arena_realloc"
(func $arena (param i32 i32 i32 i32) (result i32)))
(memory 1)
(func (export "f") (param i32 i32 i32 i32) (result i32)
local.get 0 local.get 1 local.get 2 local.get 3 call $arena))"#,
)
.unwrap();
assert!(matches!(
bind_cabi_arena_realloc(&wasm).unwrap(),
ArenaBind::KeptHostSeam(_)
));
}
#[test]
fn wrong_signature_declines_loudly() {
let wasm = wat::parse_str(
r#"(module
(import "env" "__cabi_arena_realloc"
(func $arena (param i32 i32) (result i32)))
(memory 1)
(func (export "f") (param i32 i32) (result i32)
local.get 0 local.get 1 call $arena))"#,
)
.unwrap();
let err = bind_cabi_arena_realloc(&wasm).unwrap_err().to_string();
assert!(err.contains("#418"), "{err}");
assert!(err.contains("signature"), "{err}");
}
#[test]
fn no_memory_declines_loudly() {
let wasm = wat::parse_str(
r#"(module
(import "env" "__cabi_arena_realloc"
(func $arena (param i32 i32 i32 i32) (result i32)))
(func (export "f") (result i32)
i32.const 0 i32.const 0 i32.const 8 i32.const 4 call $arena))"#,
)
.unwrap();
let err = bind_cabi_arena_realloc(&wasm).unwrap_err().to_string();
assert!(err.contains("no linear memory"), "{err}");
}
#[test]
fn full_static_layout_declines_loudly() {
let wasm = wat::parse_str(
r#"(module
(import "env" "__cabi_arena_realloc"
(func $arena (param i32 i32 i32 i32) (result i32)))
(memory 1)
(global (export "__heap_base") i32 (i32.const 65536))
(func (export "f") (result i32)
i32.const 0 i32.const 0 i32.const 8 i32.const 4 call $arena))"#,
)
.unwrap();
let err = bind_cabi_arena_realloc(&wasm).unwrap_err().to_string();
assert!(err.contains("no arena region left"), "{err}");
}
#[test]
fn module_without_globals_gets_global_section() {
let wasm = wat::parse_str(
r#"(module
(import "env" "__cabi_arena_realloc"
(func $arena (param i32 i32 i32 i32) (result i32)))
(memory 1)
(data (i32.const 64) "xyzw")
(func (export "f") (result i32)
i32.const 0 i32.const 0 i32.const 8 i32.const 4 call $arena))"#,
)
.unwrap();
let ArenaBind::Bound(b) = bind_cabi_arena_realloc(&wasm).unwrap() else {
panic!("expected Bound");
};
assert_eq!(b.arena_base, 80); let mut num_globals = 0;
for p in Parser::new(0).parse_all(&b.bytes) {
if let Payload::GlobalSection(r) = p.unwrap() {
num_globals = r.count();
}
}
assert_eq!(num_globals, 1);
}
}