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:
Jonathan Kelley authored and GitHub committed 2026-01-07 02:35:08 -08:00
1 parent ec8f31dece
commit 6b78816c1d
2 files changed
+115 -37

No files matched your search

+114 -37
View File
@@ -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)
// 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 section start offset should be within the file contents");
.expect("Data segment should be within file contents"),
segment.data.len() as u64,
));
}
// Parse the wasm file to find the globals
let module = walrus::Module::from_buffer(file_contents).unwrap();
let mut offsets = Vec::new();
if segment_file_info.is_empty() {
return Ok(Vec::new());
}
// Find the main memory offset
let main_memory = module
// 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: {}",
+1
View File
@@ -58,6 +58,7 @@ impl WasmOpt {
"--enable-bulk-memory",
"--enable-mutable-globals",
"--enable-nontrapping-float-to-int",
"--enable-threads",
];
if self.cfg.memory_packing {