#[macro_use]
extern crate failure;
extern crate parity_wasm;
use std::collections::HashMap;
use failure::{Error, ResultExt};
use parity_wasm::elements::*;
const PAGE_SIZE: u32 = 1 << 16;
pub struct Config {
maximum_memory: u32,
thread_stack_size: u32,
}
impl Config {
pub fn new() -> Config {
Config {
maximum_memory: 1 << 30,
thread_stack_size: 1 << 20,
}
}
pub fn maximum_memory(&mut self, max: u32) -> &mut Config {
self.maximum_memory = max;
self
}
pub fn thread_stack_size(&mut self, size: u32) -> &mut Config {
self.thread_stack_size = size;
self
}
pub fn run(&self, module: &mut Module) -> Result<(), Error> {
let segments = switch_data_segments_to_passive(module)?;
import_memory_zero(module)?;
share_imported_memory_zero(module, self.maximum_memory)?;
let stack_pointer_idx = find_stack_pointer(module)?;
let globals = inject_thread_globals(module);
let addr = inject_thread_id_counter(module)?;
start_with_init_memory(
module,
&segments,
&globals,
addr,
stack_pointer_idx,
self.thread_stack_size,
);
implement_thread_intrinsics(module, &globals)?;
Ok(())
}
}
struct PassiveSegment {
idx: u32,
offset: u32,
len: u32,
}
fn switch_data_segments_to_passive(module: &mut Module) -> Result<Vec<PassiveSegment>, Error> {
let section = match module.data_section_mut() {
Some(section) => section,
None => return Ok(Vec::new()),
};
let mut ret = Vec::new();
for (i, segment) in section.entries_mut().iter_mut().enumerate() {
let mut offset = match segment.offset_mut().take() {
Some(offset) => offset,
None => continue,
};
assert!(!segment.passive());
let offset = *get_offset(&mut offset)
.with_context(|_| format!("failed to read data segment {}", i))?;
*segment.passive_mut() = true;
ret.push(PassiveSegment {
idx: i as u32,
offset: offset as u32,
len: segment.value().len() as u32,
});
}
Ok(ret)
}
fn get_offset(offset: &mut InitExpr) -> Result<&mut i32, Error> {
if offset.code().len() != 2 || offset.code()[1] != Instruction::End {
bail!("unrecognized offset")
}
match &mut offset.code_mut()[0] {
Instruction::I32Const(n) => Ok(n),
_ => bail!("unrecognized offset"),
}
}
fn import_memory_zero(module: &mut Module) -> Result<(), Error> {
let limits = {
let section = match module.memory_section_mut() {
Some(section) => section,
None => return Ok(()),
};
let limits = match section.entries_mut().pop() {
Some(limits) => limits,
None => return Ok(()),
};
if section.entries().len() > 0 {
bail!("too many memories in wasm module for this tool to work");
}
limits
};
module.sections_mut().retain(|s| match s {
Section::Memory(_) => false,
_ => true,
});
if let Some(s) = module.export_section_mut() {
s.entries_mut().retain(|s| match s.internal() {
Internal::Memory(_) => false,
_ => true,
});
}
let pos = maybe_add_import_section(module);
let imports = match &mut module.sections_mut()[pos] {
Section::Import(s) => s,
_ => unreachable!(),
};
let entry = ImportEntry::new(
"env".to_string(),
"memory".to_string(),
External::Memory(limits),
);
imports.entries_mut().push(entry);
Ok(())
}
fn maybe_add_import_section(module: &mut Module) -> usize {
let mut pos = None;
for i in 0..module.sections().len() {
match &mut module.sections_mut()[i] {
Section::Type(_) => continue,
Section::Import(_) => return i,
_ => {}
}
pos = Some(i);
break;
}
let empty = ImportSection::with_entries(Vec::new());
let section = Section::Import(empty);
let len = module.sections().len();
let pos = pos.unwrap_or_else(|| len - 1);
module.sections_mut().insert(pos, section);
return pos;
}
fn share_imported_memory_zero(module: &mut Module, memory_max: u32) -> Result<(), Error> {
assert!(memory_max % PAGE_SIZE == 0);
let imports = match module.import_section_mut() {
Some(s) => s,
None => panic!("failed to find an import section"),
};
for entry in imports.entries_mut() {
let mem = match entry.external_mut() {
External::Memory(m) => m,
_ => continue,
};
*mem = MemoryType::new(
mem.limits().initial(),
Some(mem.limits().maximum().unwrap_or(memory_max / PAGE_SIZE)),
true,
);
return Ok(());
}
panic!("failed to find an imported memory")
}
struct Globals {
thread_id: u32,
thread_tcb: u32,
}
fn inject_thread_globals(module: &mut Module) -> Globals {
let pos = maybe_add_global_section(module);
let globals = match &mut module.sections_mut()[pos] {
Section::Global(s) => s,
_ => unreachable!(),
};
globals.entries_mut().push(GlobalEntry::new(
GlobalType::new(ValueType::I32, true),
InitExpr::new(vec![Instruction::I32Const(0), Instruction::End]),
));
globals.entries_mut().push(GlobalEntry::new(
GlobalType::new(ValueType::I32, true),
InitExpr::new(vec![Instruction::I32Const(0), Instruction::End]),
));
let len = globals.entries().len() as u32;
Globals {
thread_id: len - 2,
thread_tcb: len - 1,
}
}
fn maybe_add_global_section(module: &mut Module) -> usize {
let mut pos = None;
for i in 0..module.sections().len() {
match &mut module.sections_mut()[i] {
Section::Type(_)
| Section::Import(_)
| Section::Function(_)
| Section::Table(_)
| Section::Memory(_) => continue,
Section::Global(_) => return i,
_ => {}
}
pos = Some(i);
break;
}
let empty = GlobalSection::with_entries(Vec::new());
let section = Section::Global(empty);
let len = module.sections().len();
let pos = pos.unwrap_or_else(|| len - 1);
module.sections_mut().insert(pos, section);
return pos;
}
fn inject_thread_id_counter(module: &mut Module) -> Result<u32, Error> {
let heap_base = {
let exports = match module.export_section() {
Some(s) => s,
None => bail!("failed to find `__heap_base` for injecting thread id"),
};
exports
.entries()
.iter()
.filter(|e| e.field() == "__heap_base")
.filter_map(|e| match e.internal() {
Internal::Global(idx) => Some(*idx),
_ => None,
})
.next()
};
let heap_base = match heap_base {
Some(idx) => idx,
None => bail!("failed to find `__heap_base` for injecting thread id"),
};
let (address, add_a_page) = {
let globals = match module.global_section_mut() {
Some(s) => s,
None => bail!("failed to find globals section"),
};
let entry = match globals.entries_mut().get_mut(heap_base as usize) {
Some(i) => i,
None => bail!("the `__heap_base` export index is out of bounds"),
};
if entry.global_type().content_type() != ValueType::I32 {
bail!("the `__heap_base` global doesn't have the type `i32`");
}
if entry.global_type().is_mutable() {
bail!("the `__heap_base` global is unexpectedly mutable");
}
let offset = get_offset(entry.init_expr_mut())?;
let address = (*offset as u32 + 3) & !3;
let add_a_page = (address + 4) / PAGE_SIZE != address / PAGE_SIZE;
*offset = (address + 4) as i32;
(address, add_a_page)
};
if add_a_page {
add_one_to_imported_memory_limits_minimum(module);
}
Ok(address)
}
fn add_one_to_imported_memory_limits_minimum(module: &mut Module) {
let imports = match module.import_section_mut() {
Some(s) => s,
None => panic!("failed to find import section"),
};
for entry in imports.entries_mut() {
let mem = match entry.external_mut() {
External::Memory(m) => m,
_ => continue,
};
*mem = MemoryType::new(
mem.limits().initial() + 1,
mem.limits().maximum().map(|m| {
if m == mem.limits().initial() {
m + 1
} else {
m
}
}),
mem.limits().shared(),
);
return;
}
panic!("failed to find an imported memory")
}
fn find_stack_pointer(module: &mut Module) -> Result<Option<u32>, Error> {
let globals = match module.global_section() {
Some(s) => s,
None => bail!("failed to find the stack pointer"),
};
let candidates = globals
.entries()
.iter()
.enumerate()
.filter(|(_, g)| g.global_type().content_type() == ValueType::I32)
.filter(|(_, g)| g.global_type().is_mutable())
.collect::<Vec<_>>();
if candidates.len() == 0 {
return Ok(None);
}
if candidates[0].0 == 0 {
return Ok(Some(0));
}
bail!(
"the first global wasn't a mutable i32, has LLD changed or was \
this wasm file not produced by LLD?"
)
}
fn start_with_init_memory(
module: &mut Module,
segments: &[PassiveSegment],
globals: &Globals,
addr: u32,
stack_pointer_idx: Option<u32>,
stack_size: u32,
) {
assert!(stack_size % PAGE_SIZE == 0);
let mut instrs = Vec::new();
instrs.push(Instruction::I32Const(addr as i32));
instrs.push(Instruction::I32Const(1));
let mem = parity_wasm::elements::MemArg {
align: 2,
offset: 0,
};
instrs.push(Instruction::I32AtomicRmwAdd(mem));
instrs.push(Instruction::TeeLocal(0));
instrs.push(Instruction::SetGlobal(globals.thread_id));
instrs.push(Instruction::GetLocal(0));
instrs.push(Instruction::If(BlockType::NoResult));
if let Some(stack_pointer_idx) = stack_pointer_idx {
instrs.push(Instruction::I32Const((stack_size / PAGE_SIZE) as i32));
instrs.push(Instruction::GrowMemory(0));
instrs.push(Instruction::SetLocal(0));
instrs.push(Instruction::Block(BlockType::NoResult));
instrs.push(Instruction::GetLocal(0));
instrs.push(Instruction::I32Const(-1));
instrs.push(Instruction::I32Ne);
instrs.push(Instruction::BrIf(0));
instrs.push(Instruction::Unreachable);
instrs.push(Instruction::End);
instrs.push(Instruction::GetLocal(0));
instrs.push(Instruction::I32Const(PAGE_SIZE as i32));
instrs.push(Instruction::I32Mul);
instrs.push(Instruction::I32Const(stack_size as i32));
instrs.push(Instruction::I32Add);
instrs.push(Instruction::SetGlobal(stack_pointer_idx));
}
instrs.push(Instruction::Else);
for segment in segments {
instrs.push(Instruction::I32Const(segment.offset as i32));
instrs.push(Instruction::I32Const(0));
instrs.push(Instruction::I32Const(segment.len as i32));
instrs.push(Instruction::MemoryInit(segment.idx));
}
instrs.push(Instruction::End);
for segment in segments {
instrs.push(Instruction::MemoryDrop(segment.idx));
}
if let Some(idx) = module.start_section() {
instrs.push(Instruction::Call(idx));
}
instrs.push(Instruction::End);
let instrs = Instructions::new(instrs);
let local = Local::new(1, ValueType::I32);
let body = FuncBody::new(vec![local], instrs);
let code_idx = {
let s = module.code_section_mut().expect("module had no code");
s.bodies_mut().push(body);
(s.bodies().len() - 1) as u32
};
let type_idx = {
let section = module
.type_section_mut()
.expect("module has no type section");
let pos = section
.types()
.iter()
.map(|t| match t {
Type::Function(t) => t,
})
.position(|t| t.params().is_empty() && t.return_type().is_none());
match pos {
Some(i) => i as u32,
None => {
let f = FunctionType::new(Vec::new(), None);
section.types_mut().push(Type::Function(f));
(section.types().len() - 1) as u32
}
}
};
module
.function_section_mut()
.expect("module has no function section")
.entries_mut()
.push(Func::new(type_idx));
let idx = code_idx + (module.import_count(ImportCountType::Function) as u32);
update_start_section(module, idx);
}
fn update_start_section(module: &mut Module, start: u32) {
let mut pos = None;
for i in 0..module.sections().len() {
match &mut module.sections_mut()[i] {
Section::Type(_)
| Section::Import(_)
| Section::Function(_)
| Section::Table(_)
| Section::Memory(_)
| Section::Global(_)
| Section::Export(_) => continue,
Section::Start(start_idx) => {
*start_idx = start;
return;
}
_ => {}
}
pos = Some(i);
break;
}
let section = Section::Start(start);
let len = module.sections().len();
let pos = pos.unwrap_or_else(|| len - 1);
module.sections_mut().insert(pos, section);
}
fn implement_thread_intrinsics(module: &mut Module, globals: &Globals) -> Result<(), Error> {
let mut map = HashMap::new();
{
let imports = match module.import_section() {
Some(i) => i,
None => return Ok(()),
};
let entries = imports
.entries()
.iter()
.filter(|i| match i.external() {
External::Function(_) => true,
_ => false,
})
.enumerate()
.filter(|(_, entry)| entry.module() == "__wbindgen_thread_xform__");
for (idx, entry) in entries {
let type_idx = match entry.external() {
External::Function(i) => *i,
_ => unreachable!(),
};
let types = module.type_section().unwrap();
let fty = match &types.types()[type_idx as usize] {
Type::Function(f) => f,
};
match entry.field() {
"__wbindgen_thread_id" => {
if !fty.params().is_empty() || fty.return_type() != Some(ValueType::I32) {
bail!("__wbindgen_thread_id intrinsic has the wrong signature");
}
map.insert(idx as u32, Instruction::GetGlobal(globals.thread_id));
}
"__wbindgen_tcb_get" => {
if !fty.params().is_empty() || fty.return_type() != Some(ValueType::I32) {
bail!("__wbindgen_tcb_get intrinsic has the wrong signature");
}
map.insert(idx as u32, Instruction::GetGlobal(globals.thread_tcb));
}
"__wbindgen_tcb_set" => {
if fty.params().len() != 1 || fty.return_type().is_some() {
bail!("__wbindgen_tcb_set intrinsic has the wrong signature");
}
map.insert(idx as u32, Instruction::SetGlobal(globals.thread_tcb));
}
other => bail!("unknown thread intrinsic: {}", other),
}
}
};
for body in module.code_section_mut().unwrap().bodies_mut() {
for instr in body.code_mut().elements_mut() {
let other = match instr {
Instruction::Call(idx) => match map.get(idx) {
Some(other) => other,
None => continue,
},
_ => continue,
};
*instr = other.clone();
}
}
Ok(())
}