1use anyhow::{Context, Result, bail};
49use wasm_encoder::{BlockType, Function, MemArg, ValType};
50use wasmparser::{Parser, Payload};
51
52pub const ARENA_IMPORT_MODULE: &str = "env";
55pub const ARENA_IMPORT_FIELD: &str = "__cabi_arena_realloc";
56
57#[derive(Debug)]
59pub enum ArenaBind {
60 NoArenaImport,
63 KeptHostSeam(&'static str),
66 Bound(BoundArena),
68}
69
70#[derive(Debug)]
72pub struct BoundArena {
73 pub bytes: Vec<u8>,
75 pub arena_base: u32,
77 pub arena_end: u32,
79}
80
81struct Scan {
83 arena_type_idx: Option<u32>,
85 arena_sig_ok: bool,
87 total_imports: u32,
88 memory_pages: Option<u64>,
90 memory64: bool,
91 defined_globals: u32,
92 data_end: u32,
94 non_const_data_offset: bool,
96 global_inits: Vec<u32>,
99 has_function_section: bool,
100 has_code_section: bool,
101}
102
103fn scan(wasm: &[u8]) -> Result<Scan> {
104 let mut s = Scan {
105 arena_type_idx: None,
106 arena_sig_ok: false,
107 total_imports: 0,
108 memory_pages: None,
109 memory64: false,
110 defined_globals: 0,
111 data_end: 0,
112 non_const_data_offset: false,
113 global_inits: Vec::new(),
114 has_function_section: false,
115 has_code_section: false,
116 };
117 let mut func_types: Vec<bool> = Vec::new(); for payload in Parser::new(0).parse_all(wasm) {
119 match payload.context("parse wasm (#418 arena-bind scan)")? {
120 Payload::TypeSection(reader) => {
121 for rec_group in reader {
122 for sub_ty in rec_group.context("parse type section (#418)")?.types() {
123 let ok = match &sub_ty.composite_type.inner {
124 wasmparser::CompositeInnerType::Func(f) => {
125 f.params().len() == 4
126 && f.params().iter().all(|t| *t == wasmparser::ValType::I32)
127 && f.results() == [wasmparser::ValType::I32]
128 }
129 _ => false,
130 };
131 func_types.push(ok);
132 }
133 }
134 }
135 Payload::ImportSection(reader) => {
136 for import in reader.into_imports() {
139 let import = import.context("parse import (#418)")?;
140 s.total_imports += 1;
141 if import.module == ARENA_IMPORT_MODULE
142 && import.name == ARENA_IMPORT_FIELD
143 && let wasmparser::TypeRef::Func(type_idx) = import.ty
144 {
145 s.arena_type_idx = Some(type_idx);
146 s.arena_sig_ok =
147 func_types.get(type_idx as usize).copied().unwrap_or(false);
148 }
149 }
150 }
151 Payload::MemorySection(reader) => {
152 for (i, mem) in reader.into_iter().enumerate() {
153 let mem = mem.context("parse memory (#418)")?;
154 if i == 0 {
155 s.memory_pages = Some(mem.initial);
156 s.memory64 = mem.memory64;
157 }
158 }
159 }
160 Payload::GlobalSection(reader) => {
161 for global in reader {
162 let global = global.context("parse global (#418)")?;
163 s.defined_globals += 1;
164 let mut ops = global.init_expr.get_operators_reader();
167 if let Ok(wasmparser::Operator::I32Const { value }) = ops.read()
168 && value > 0
169 {
170 s.global_inits.push(value as u32);
171 }
172 }
173 }
174 Payload::DataSection(reader) => {
175 for seg in reader {
176 let seg = seg.context("parse data segment (#418)")?;
177 if let wasmparser::DataKind::Active {
178 memory_index,
179 offset_expr,
180 } = seg.kind
181 {
182 if memory_index != 0 {
183 continue; }
185 let mut ops = offset_expr.get_operators_reader();
186 match ops.read() {
187 Ok(wasmparser::Operator::I32Const { value }) => {
188 let end = (value as u32).saturating_add(seg.data.len() as u32);
189 s.data_end = s.data_end.max(end);
190 }
191 _ => s.non_const_data_offset = true,
192 }
193 }
194 }
195 }
196 Payload::FunctionSection(_) => s.has_function_section = true,
197 Payload::CodeSectionStart { .. } => s.has_code_section = true,
198 _ => {}
199 }
200 }
201 Ok(s)
202}
203
204fn allocator_body(cursor_global: u32, arena_end: u32) -> Function {
209 let mem = MemArg {
210 offset: 0,
211 align: 0,
212 memory_index: 0,
213 };
214 let mut f = Function::new([(4, ValType::I32)]);
215 f.instructions()
217 .local_get(1)
218 .i32_eqz()
219 .local_get(3)
220 .i32_eqz()
221 .i32_and()
222 .if_(BlockType::Empty)
223 .local_get(2)
224 .return_()
225 .end()
226 .local_get(2)
229 .i32_eqz()
230 .if_(BlockType::Empty)
231 .unreachable()
232 .end()
233 .global_get(cursor_global)
235 .local_get(2)
236 .i32_add()
237 .i32_const(1)
238 .i32_sub()
239 .local_get(2)
240 .i32_const(1)
241 .i32_sub()
242 .i32_const(-1)
243 .i32_xor()
244 .i32_and()
245 .local_set(4)
246 .local_get(4)
248 .global_get(cursor_global)
249 .i32_lt_u()
250 .if_(BlockType::Empty)
251 .unreachable()
252 .end()
253 .local_get(4)
255 .local_get(3)
256 .i32_add()
257 .local_tee(5)
258 .local_get(4)
259 .i32_lt_u()
260 .if_(BlockType::Empty)
261 .unreachable()
262 .end()
263 .local_get(5)
265 .i32_const(arena_end as i32)
266 .i32_gt_u()
267 .if_(BlockType::Empty)
268 .unreachable()
269 .end()
270 .local_get(5)
272 .global_set(cursor_global)
273 .local_get(1)
275 .local_get(3)
276 .local_get(1)
277 .local_get(3)
278 .i32_lt_u()
279 .select()
280 .local_set(6)
281 .block(BlockType::Empty)
283 .loop_(BlockType::Empty)
284 .local_get(7)
285 .local_get(6)
286 .i32_ge_u()
287 .br_if(1)
288 .local_get(4)
289 .local_get(7)
290 .i32_add()
291 .local_get(0)
292 .local_get(7)
293 .i32_add()
294 .i32_load8_u(mem)
295 .i32_store8(mem)
296 .local_get(7)
297 .i32_const(1)
298 .i32_add()
299 .local_set(7)
300 .br(0)
301 .end()
302 .end()
303 .local_get(4)
304 .end();
305 f
306}
307
308fn read_uleb(bytes: &[u8]) -> Result<(u32, usize)> {
310 let mut value: u32 = 0;
311 let mut shift = 0;
312 for (i, &b) in bytes.iter().enumerate().take(5) {
313 value |= u32::from(b & 0x7F) << shift;
314 if b & 0x80 == 0 {
315 return Ok((value, i + 1));
316 }
317 shift += 7;
318 }
319 bail!("malformed LEB128 count in section (#418)");
320}
321
322fn write_uleb(mut value: u32, out: &mut Vec<u8>) {
323 loop {
324 let mut b = (value & 0x7F) as u8;
325 value >>= 7;
326 if value != 0 {
327 b |= 0x80;
328 }
329 out.push(b);
330 if value == 0 {
331 return;
332 }
333 }
334}
335
336fn write_sleb(mut value: i32, out: &mut Vec<u8>) {
337 loop {
338 let b = (value & 0x7F) as u8;
339 value >>= 7;
340 let sign = b & 0x40;
341 if (value == 0 && sign == 0) || (value == -1 && sign != 0) {
342 out.push(b);
343 return;
344 }
345 out.push(b | 0x80);
346 }
347}
348
349fn prepend_entry(contents: &[u8], entry: &[u8]) -> Result<Vec<u8>> {
352 let (count, len) = read_uleb(contents)?;
353 let mut out = Vec::with_capacity(contents.len() + entry.len() + 1);
354 write_uleb(count + 1, &mut out);
355 out.extend_from_slice(entry);
356 out.extend_from_slice(&contents[len..]);
357 Ok(out)
358}
359
360fn append_entry(contents: &[u8], entry: &[u8]) -> Result<Vec<u8>> {
363 let (count, len) = read_uleb(contents)?;
364 let mut out = Vec::with_capacity(contents.len() + entry.len() + 1);
365 write_uleb(count + 1, &mut out);
366 out.extend_from_slice(&contents[len..]);
367 out.extend_from_slice(entry);
368 Ok(out)
369}
370
371fn cursor_global_entry(arena_base: u32) -> Vec<u8> {
373 let mut e = vec![0x7F, 0x01, 0x41]; write_sleb(arena_base as i32, &mut e);
375 e.push(0x0B); e
377}
378
379pub fn bind_cabi_arena_realloc(wasm: &[u8]) -> Result<ArenaBind> {
390 let s = scan(wasm)?;
391 let Some(arena_type_idx) = s.arena_type_idx else {
392 return Ok(ArenaBind::NoArenaImport);
393 };
394 if s.total_imports > 1 {
395 return Ok(ArenaBind::KeptHostSeam(
399 "module has other imports — keeping the host-linked seam",
400 ));
401 }
402 if !s.arena_sig_ok {
403 bail!(
404 "#418: env::{ARENA_IMPORT_FIELD} is imported with a signature \
405 other than (i32, i32, i32, i32) -> i32 — not the canonical-ABI \
406 arena realloc contract; refusing to bind (compile with \
407 --no-bind-cabi-arena to keep it an external symbol)"
408 );
409 }
410 let Some(pages) = s.memory_pages else {
411 bail!(
412 "#418: cannot bind env::{ARENA_IMPORT_FIELD}: the module declares \
413 no linear memory to allocate from"
414 );
415 };
416 if s.memory64 {
417 bail!("#418: cannot bind env::{ARENA_IMPORT_FIELD}: memory64 module");
418 }
419 if s.non_const_data_offset {
420 bail!(
421 "#418: cannot bind env::{ARENA_IMPORT_FIELD}: a data segment has \
422 a non-constant offset, so the static-data extent (the arena \
423 floor) cannot be derived soundly"
424 );
425 }
426 if !s.has_function_section || !s.has_code_section {
427 bail!(
428 "#418: cannot bind env::{ARENA_IMPORT_FIELD}: the module defines \
429 no functions (nothing synth could route the binding through)"
430 );
431 }
432
433 let arena_end: u32 = u32::try_from(pages.saturating_mul(64 * 1024))
434 .unwrap_or(u32::MAX)
435 .min(0xFFFF_0000);
436 let global_top = s
441 .global_inits
442 .iter()
443 .copied()
444 .filter(|&v| v <= arena_end)
445 .max()
446 .unwrap_or(0);
447 let arena_base = s.data_end.max(global_top).max(16).next_multiple_of(16);
448 if arena_base >= arena_end {
449 bail!(
450 "#418: cannot bind env::{ARENA_IMPORT_FIELD}: the static layout \
451 (data + stack + wasm-ld layout globals) extends to {arena_base} \
452 bytes but linear memory is only {arena_end} bytes — no arena \
453 region left; every allocation would trap"
454 );
455 }
456
457 let cursor_global = s.defined_globals;
461 let mut body = Vec::new();
462 wasm_encoder::Encode::encode(&allocator_body(cursor_global, arena_end), &mut body);
463
464 let mut module = wasm_encoder::Module::new();
466 let mut global_emitted = false;
467 let mut function_emitted = false;
468 let ensure_globals = |module: &mut wasm_encoder::Module, emitted: &mut bool| {
471 if !*emitted {
472 let mut out = Vec::new();
473 write_uleb(1, &mut out);
474 out.extend_from_slice(&cursor_global_entry(arena_base));
475 module.section(&wasm_encoder::RawSection {
476 id: wasm_encoder::SectionId::Global as u8,
477 data: &out,
478 });
479 *emitted = true;
480 }
481 };
482
483 for payload in Parser::new(0).parse_all(wasm) {
484 let payload = payload.context("parse wasm (#418 arena-bind rewrite)")?;
485 match &payload {
486 Payload::Version { .. } | Payload::End(_) => {}
487 Payload::ImportSection(_) => {
488 }
490 Payload::FunctionSection(reader) => {
491 let mut entry = Vec::new();
492 write_uleb(arena_type_idx, &mut entry);
493 let contents = &wasm[reader.range()];
494 module.section(&wasm_encoder::RawSection {
495 id: wasm_encoder::SectionId::Function as u8,
496 data: &prepend_entry(contents, &entry)?,
497 });
498 function_emitted = true;
499 }
500 Payload::GlobalSection(reader) => {
501 let contents = &wasm[reader.range()];
502 module.section(&wasm_encoder::RawSection {
503 id: wasm_encoder::SectionId::Global as u8,
504 data: &append_entry(contents, &cursor_global_entry(arena_base))?,
505 });
506 global_emitted = true;
507 }
508 Payload::ExportSection(_)
509 | Payload::StartSection { .. }
510 | Payload::ElementSection(_)
511 | Payload::DataCountSection { .. }
512 | Payload::DataSection(_) => {
513 ensure_globals(&mut module, &mut global_emitted);
514 copy_raw(&mut module, &payload, wasm)?;
515 }
516 Payload::CodeSectionStart { range, .. } => {
517 ensure_globals(&mut module, &mut global_emitted);
518 let contents = &wasm[range.clone()];
521 module.section(&wasm_encoder::RawSection {
522 id: wasm_encoder::SectionId::Code as u8,
523 data: &prepend_entry(contents, &body)?,
524 });
525 }
526 Payload::CodeSectionEntry(_) => {} other => copy_raw(&mut module, other, wasm)?,
528 }
529 }
530 if !function_emitted {
531 bail!("#418 internal: function section not re-emitted"); }
533
534 let bytes = module.finish();
535 wasmparser::Validator::new()
538 .validate_all(&bytes)
539 .context("#418 internal: arena-bind rewrite produced an invalid module (bug)")?;
540 Ok(ArenaBind::Bound(BoundArena {
541 bytes,
542 arena_base,
543 arena_end,
544 }))
545}
546
547fn copy_raw(module: &mut wasm_encoder::Module, payload: &Payload<'_>, wasm: &[u8]) -> Result<()> {
549 let Some((id, range)) = payload.as_section() else {
550 bail!("#418 internal: unhandled non-section payload {payload:?}");
551 };
552 module.section(&wasm_encoder::RawSection {
553 id,
554 data: &wasm[range],
555 });
556 Ok(())
557}
558
559#[cfg(test)]
560mod tests {
561 use super::*;
562
563 fn fixture() -> Vec<u8> {
564 wat::parse_str(
565 r#"(module
566 (import "env" "__cabi_arena_realloc"
567 (func $arena (param i32 i32 i32 i32) (result i32)))
568 (memory (export "memory") 1)
569 (global $sp (mut i32) (i32.const 4096))
570 (global (export "__heap_base") i32 (i32.const 6144))
571 (data (i32.const 5120) "0123456789abcdef")
572 (func (export "cabi_realloc") (param i32 i32 i32 i32) (result i32)
573 local.get 0 local.get 1 local.get 2 local.get 3 call $arena))"#,
574 )
575 .unwrap()
576 }
577
578 #[test]
579 fn binds_sole_arena_import() {
580 let ArenaBind::Bound(b) = bind_cabi_arena_realloc(&fixture()).unwrap() else {
581 panic!("expected Bound");
582 };
583 assert_eq!(b.arena_base, 6144);
585 assert_eq!(b.arena_end, 65536);
586 let mut num_imports = 0;
589 let mut num_funcs = 0;
590 let mut num_globals = 0;
591 for p in Parser::new(0).parse_all(&b.bytes) {
592 match p.unwrap() {
593 Payload::ImportSection(r) => num_imports += r.count(),
594 Payload::FunctionSection(r) => num_funcs = r.count(),
595 Payload::GlobalSection(r) => num_globals = r.count(),
596 _ => {}
597 }
598 }
599 assert_eq!(num_imports, 0);
600 assert_eq!(num_funcs, 2); assert_eq!(num_globals, 3); }
603
604 #[test]
605 fn bound_module_executes_contract() {
606 let ArenaBind::Bound(b) = bind_cabi_arena_realloc(&fixture()).unwrap() else {
610 panic!("expected Bound");
611 };
612 wasmparser::Validator::new().validate_all(&b.bytes).unwrap();
613 }
614
615 #[test]
616 fn no_arena_import_passes_through() {
617 let wasm =
618 wat::parse_str(r#"(module (memory 1) (func (export "f") (result i32) i32.const 7))"#)
619 .unwrap();
620 assert!(matches!(
621 bind_cabi_arena_realloc(&wasm).unwrap(),
622 ArenaBind::NoArenaImport
623 ));
624 }
625
626 #[test]
627 fn other_imports_keep_host_seam() {
628 let wasm = wat::parse_str(
629 r#"(module
630 (import "env" "k_spin_lock" (func (param i32)))
631 (import "env" "__cabi_arena_realloc"
632 (func $arena (param i32 i32 i32 i32) (result i32)))
633 (memory 1)
634 (func (export "f") (param i32 i32 i32 i32) (result i32)
635 local.get 0 local.get 1 local.get 2 local.get 3 call $arena))"#,
636 )
637 .unwrap();
638 assert!(matches!(
639 bind_cabi_arena_realloc(&wasm).unwrap(),
640 ArenaBind::KeptHostSeam(_)
641 ));
642 }
643
644 #[test]
645 fn wrong_signature_declines_loudly() {
646 let wasm = wat::parse_str(
647 r#"(module
648 (import "env" "__cabi_arena_realloc"
649 (func $arena (param i32 i32) (result i32)))
650 (memory 1)
651 (func (export "f") (param i32 i32) (result i32)
652 local.get 0 local.get 1 call $arena))"#,
653 )
654 .unwrap();
655 let err = bind_cabi_arena_realloc(&wasm).unwrap_err().to_string();
656 assert!(err.contains("#418"), "{err}");
657 assert!(err.contains("signature"), "{err}");
658 }
659
660 #[test]
661 fn no_memory_declines_loudly() {
662 let wasm = wat::parse_str(
663 r#"(module
664 (import "env" "__cabi_arena_realloc"
665 (func $arena (param i32 i32 i32 i32) (result i32)))
666 (func (export "f") (result i32)
667 i32.const 0 i32.const 0 i32.const 8 i32.const 4 call $arena))"#,
668 )
669 .unwrap();
670 let err = bind_cabi_arena_realloc(&wasm).unwrap_err().to_string();
671 assert!(err.contains("no linear memory"), "{err}");
672 }
673
674 #[test]
675 fn full_static_layout_declines_loudly() {
676 let wasm = wat::parse_str(
678 r#"(module
679 (import "env" "__cabi_arena_realloc"
680 (func $arena (param i32 i32 i32 i32) (result i32)))
681 (memory 1)
682 (global (export "__heap_base") i32 (i32.const 65536))
683 (func (export "f") (result i32)
684 i32.const 0 i32.const 0 i32.const 8 i32.const 4 call $arena))"#,
685 )
686 .unwrap();
687 let err = bind_cabi_arena_realloc(&wasm).unwrap_err().to_string();
688 assert!(err.contains("no arena region left"), "{err}");
689 }
690
691 #[test]
692 fn module_without_globals_gets_global_section() {
693 let wasm = wat::parse_str(
694 r#"(module
695 (import "env" "__cabi_arena_realloc"
696 (func $arena (param i32 i32 i32 i32) (result i32)))
697 (memory 1)
698 (data (i32.const 64) "xyzw")
699 (func (export "f") (result i32)
700 i32.const 0 i32.const 0 i32.const 8 i32.const 4 call $arena))"#,
701 )
702 .unwrap();
703 let ArenaBind::Bound(b) = bind_cabi_arena_realloc(&wasm).unwrap() else {
704 panic!("expected Bound");
705 };
706 assert_eq!(b.arena_base, 80); let mut num_globals = 0;
708 for p in Parser::new(0).parse_all(&b.bytes) {
709 if let Payload::GlobalSection(r) = p.unwrap() {
710 num_globals = r.count();
711 }
712 }
713 assert_eq!(num_globals, 1);
714 }
715}