allow passive wasm segments for manganis (#5163)
* fix passive segment using tls calculation! * clean things up * clean diff * revert some pr changes for review * cleanups
This commit is contained in:
1 parent
ec8f31dece
commit
6b78816c1d
2 files changed
+116
-38
No files matched your search
@@ -367,6 +367,7 @@ fn find_native_symbol_offsets<'a, R: ReadRef<'a>>(
|
||||
Ok(offsets)
|
||||
}
|
||||
|
||||
/// Evaluate a walrus global expression to get its value.
|
||||
fn eval_walrus_global_expr(module: &walrus::Module, expr: &walrus::ConstExpr) -> Option<u64> {
|
||||
match expr {
|
||||
walrus::ConstExpr::Value(walrus::ir::Value::I32(value)) => Some(*value as u64),
|
||||
@@ -383,7 +384,26 @@ fn eval_walrus_global_expr(module: &walrus::Module, expr: &walrus::ConstExpr) ->
|
||||
}
|
||||
}
|
||||
|
||||
/// Find the value of a global export by name.
|
||||
fn find_global_export_value(module: &walrus::Module, name: &str) -> Option<u64> {
|
||||
for export in module.exports.iter() {
|
||||
if export.name == name {
|
||||
if let walrus::ExportItem::Global(g) = export.item {
|
||||
if let walrus::GlobalKind::Local(expr) = &module.globals.get(g).kind {
|
||||
return eval_walrus_global_expr(module, expr);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Find the offsets of any manganis symbols in the wasm file.
|
||||
///
|
||||
/// This handles both standard WASM builds and builds with advanced features like:
|
||||
/// - Bulk memory operations (passive data segments)
|
||||
/// - Thread Local Storage (TLS)
|
||||
/// - Atomics and shared memory
|
||||
fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
|
||||
file_contents: &[u8],
|
||||
file: &File<'a, R>,
|
||||
@@ -395,57 +415,85 @@ fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
|
||||
tracing::error!("Failed to find <data> section in WASM file");
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let Some((_, section_range_end)) = section.file_range() else {
|
||||
tracing::error!("Failed to find file range for <data> section in WASM file");
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
let section_size = section.data()?.len() as u64;
|
||||
let section_start = section_range_end - section_size;
|
||||
|
||||
// Translate the section_relative_address to the file offset
|
||||
// WASM files have a section address of 0 in object, reparse the data section with wasmparser
|
||||
// to get the correct address and section start
|
||||
// Note: We need to reparse just the data section with wasmparser to get the file offset because walrus does
|
||||
// not expose the file offset information
|
||||
// Parse data segments with wasmparser to get file offsets.
|
||||
// Walrus doesn't expose file offset information, so we need wasmparser for this.
|
||||
// With bulk memory operations, there may be multiple data segments.
|
||||
let reader = wasmparser::DataSectionReader::new(wasmparser::BinaryReader::new(
|
||||
&file_contents[section_start as usize..section_range_end as usize],
|
||||
0,
|
||||
))
|
||||
.context("Failed to create WASM data section reader")?;
|
||||
let main_memory = reader
|
||||
.into_iter()
|
||||
.next()
|
||||
.context("Failed find main memory from WASM data section")?
|
||||
.context("Failed to read main memory from WASM data section")?;
|
||||
// main_memory.data is a slice somewhere in file_contents. Find out the offset in the file
|
||||
let data_start_offset = (main_memory.data.as_ptr() as u64)
|
||||
.checked_sub(file_contents.as_ptr() as u64)
|
||||
.expect("Data section start offset should be within the file contents");
|
||||
|
||||
// Parse the wasm file to find the globals
|
||||
let module = walrus::Module::from_buffer(file_contents).unwrap();
|
||||
let mut offsets = Vec::new();
|
||||
// Collect all data segments with their file offsets and sizes
|
||||
let mut segment_file_info: Vec<(u64, u64)> = Vec::new();
|
||||
for segment in reader.into_iter() {
|
||||
let segment = segment.context("Failed to read data segment")?;
|
||||
segment_file_info.push((
|
||||
(segment.data.as_ptr() as u64)
|
||||
.checked_sub(file_contents.as_ptr() as u64)
|
||||
.expect("Data segment should be within file contents"),
|
||||
segment.data.len() as u64,
|
||||
));
|
||||
}
|
||||
|
||||
// Find the main memory offset
|
||||
let main_memory = module
|
||||
if segment_file_info.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
// Parse the wasm file with walrus to find globals and exports
|
||||
let module = walrus::Module::from_buffer(file_contents)
|
||||
.context("Failed to parse WASM module with walrus")?;
|
||||
|
||||
// Determine the memory base address for symbol lookup
|
||||
let main_memory_walrus = module
|
||||
.data
|
||||
.iter()
|
||||
.next()
|
||||
.context("Failed to find main memory in WASM module")?;
|
||||
|
||||
let walrus::DataKind::Active {
|
||||
offset: main_memory_offset,
|
||||
..
|
||||
} = main_memory.kind
|
||||
else {
|
||||
tracing::error!("Failed to find main memory offset in WASM module");
|
||||
return Ok(Vec::new());
|
||||
let main_memory_offset = match &main_memory_walrus.kind {
|
||||
walrus::DataKind::Active { offset, .. } => {
|
||||
// Active segments have an explicit offset expression
|
||||
eval_walrus_global_expr(&module, offset).unwrap_or_default()
|
||||
}
|
||||
walrus::DataKind::Passive => {
|
||||
// For passive segments (bulk memory operations), there's no static offset.
|
||||
// The memory.init instruction determines placement at runtime.
|
||||
//
|
||||
// Rust/LLVM places static data at 0x100000 (1MB) in linear memory.
|
||||
//
|
||||
// With TLS support, the linker calculates symbol addresses as if TLS data
|
||||
// is at 0x100000 followed by main data. But at runtime, TLS is stored
|
||||
// separately per-thread via __wasm_init_tls. The main data segment starts
|
||||
// at 0x100000 in the file, but symbol addresses include the TLS offset.
|
||||
//
|
||||
// We detect TLS by looking for __tls_size and remove the TLS segment from
|
||||
// our file info, then add the TLS size to the base address.
|
||||
let tls_size = find_global_export_value(&module, "__tls_size");
|
||||
|
||||
// If TLS is present and segment 0 matches TLS size, remove TLS segment
|
||||
if let Some(tls) = tls_size {
|
||||
if !segment_file_info.is_empty() && segment_file_info[0].1 == tls {
|
||||
segment_file_info.remove(0);
|
||||
}
|
||||
}
|
||||
|
||||
// Base = 0x100000 + TLS size (to account for TLS offset in symbol addresses)
|
||||
0x100000u64 + tls_size.unwrap_or(0)
|
||||
}
|
||||
};
|
||||
|
||||
// In the hot patch build, the main memory offset is a global from the main module and each global
|
||||
// is it's own global. Use an offset of 0 instead if we can't evaluate the global
|
||||
let main_memory_offset =
|
||||
eval_walrus_global_expr(&module, &main_memory_offset).unwrap_or_default();
|
||||
// Find all manganis symbols and calculate their file offsets
|
||||
let mut offsets = Vec::new();
|
||||
|
||||
for export in module.exports.iter() {
|
||||
let Some(version) = looks_like_manganis_symbol(&export.name) else {
|
||||
@@ -456,7 +504,8 @@ fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
|
||||
continue;
|
||||
};
|
||||
|
||||
let walrus::GlobalKind::Local(pointer) = module.globals.get(global).kind else {
|
||||
let global_data = module.globals.get(global);
|
||||
let walrus::GlobalKind::Local(pointer) = global_data.kind else {
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -468,11 +517,41 @@ fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
|
||||
continue;
|
||||
};
|
||||
|
||||
let section_relative_address: u64 = ((virtual_address as i128)
|
||||
- main_memory_offset as i128)
|
||||
.try_into()
|
||||
.expect("Virtual address should be greater than or equal to section address");
|
||||
let file_offset = data_start_offset + section_relative_address;
|
||||
// Calculate offset relative to the data base address
|
||||
let data_relative_offset =
|
||||
match (virtual_address as i128).checked_sub(main_memory_offset as i128) {
|
||||
Some(offset) if offset >= 0 => offset as u64,
|
||||
_ => {
|
||||
tracing::error!(
|
||||
"Virtual address 0x{:x} is below main memory offset 0x{:x}",
|
||||
virtual_address,
|
||||
main_memory_offset
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
// Find which segment this offset falls into.
|
||||
// Segments are laid out contiguously in memory.
|
||||
let mut cumulative_offset = 0u64;
|
||||
let mut file_offset = None;
|
||||
|
||||
for (seg_file_offset, seg_size) in segment_file_info.iter() {
|
||||
if data_relative_offset < cumulative_offset + seg_size {
|
||||
let offset_in_segment = data_relative_offset - cumulative_offset;
|
||||
file_offset = Some(seg_file_offset + offset_in_segment);
|
||||
break;
|
||||
}
|
||||
cumulative_offset += seg_size;
|
||||
}
|
||||
|
||||
let Some(file_offset) = file_offset else {
|
||||
tracing::error!(
|
||||
"Virtual address 0x{:x} is beyond all data segments",
|
||||
virtual_address
|
||||
);
|
||||
continue;
|
||||
};
|
||||
|
||||
offsets.push(ManganisSymbolOffset::new(version, file_offset));
|
||||
}
|
||||
@@ -530,7 +609,6 @@ pub(crate) async fn extract_assets_from_file(path: impl AsRef<Path>) -> Result<A
|
||||
let new_data = version.serialize(asset);
|
||||
|
||||
file.seek(std::io::SeekFrom::Start(offset))?;
|
||||
// Write the modified binary data back to the file
|
||||
file.write_all(new_data.as_ref())?;
|
||||
}
|
||||
|
||||
@@ -548,7 +626,6 @@ pub(crate) async fn extract_assets_from_file(path: impl AsRef<Path>) -> Result<A
|
||||
.arg(path)
|
||||
.output()
|
||||
.context("Failed to run codesign - is `codesign` in your path?")?;
|
||||
|
||||
if !output.status.success() {
|
||||
bail!(
|
||||
"Failed to re-sign the binary with codesign after finalizing the assets: {}",
|
||||
|
||||
@@ -58,6 +58,7 @@ impl WasmOpt {
|
||||
"--enable-bulk-memory",
|
||||
"--enable-mutable-globals",
|
||||
"--enable-nontrapping-float-to-int",
|
||||
"--enable-threads",
|
||||
];
|
||||
|
||||
if self.cfg.memory_packing {
|
||||
|
||||
Reference in new issue
Block a user