#[derive(Debug, Clone, Copy)]
pub struct WasmSection<'a> {
pub id: u8,
pub name: Option<&'a str>,
pub payload_offset: usize,
pub payload_size: usize,
}
pub fn walk(data: &[u8]) -> WasmSectionIter<'_> {
if !data.starts_with(b"\x00asm\x01\x00\x00\x00") {
return WasmSectionIter {
data,
pos: data.len(),
};
}
WasmSectionIter { data, pos: 8 }
}
pub struct WasmSectionIter<'a> {
data: &'a [u8],
pos: usize,
}
impl<'a> Iterator for WasmSectionIter<'a> {
type Item = WasmSection<'a>;
fn next(&mut self) -> Option<WasmSection<'a>> {
if self.pos >= self.data.len() {
return None;
}
let id = *self.data.get(self.pos)?;
let after_id = self.pos.checked_add(1)?;
let (size, after_len) = read_uleb128(self.data, after_id)?;
let size = usize::try_from(size).ok()?;
let payload_end = after_len.checked_add(size)?;
if payload_end > self.data.len() {
self.pos = self.data.len();
return None;
}
let (payload_offset, payload_size, name) = if id == 0 {
let (name_len, after_name_len) = read_uleb128(self.data, after_len)?;
let name_len = usize::try_from(name_len).ok()?;
let name_end = after_name_len.checked_add(name_len)?;
if name_end > payload_end {
self.pos = self.data.len();
return None;
}
let name_bytes = self.data.get(after_name_len..name_end)?;
let name = std::str::from_utf8(name_bytes).ok();
let body_size = payload_end.checked_sub(name_end)?;
(name_end, body_size, name)
} else {
(after_len, size, None)
};
self.pos = payload_end;
Some(WasmSection {
id,
name,
payload_offset,
payload_size,
})
}
}
#[derive(Debug, Clone, Copy)]
pub struct WasmDataSegment {
pub mem_offset: u64,
pub file_offset: usize,
pub size: usize,
}
pub fn data_segments(data: &[u8]) -> Vec<WasmDataSegment> {
let mut out = Vec::new();
for sec in walk(data) {
if sec.id != 11 {
continue;
}
let section_end = match sec.payload_offset.checked_add(sec.payload_size) {
Some(e) => e,
None => return out,
};
let mut pos = sec.payload_offset;
let count = match read_uleb128(data, pos) {
Some((c, p)) => {
pos = p;
c
}
None => return out,
};
for _ in 0..count {
if pos >= section_end {
break;
}
let flags = match read_uleb128(data, pos) {
Some((f, p)) => {
pos = p;
f
}
None => return out,
};
let mem_offset = if flags == 0 {
match data.get(pos) {
Some(0x41) => {
pos = match pos.checked_add(1) {
Some(p) => p,
None => return out,
};
let (val, np) = match read_uleb128_signed(data, pos) {
Some(v) => v,
None => return out,
};
pos = np;
if data.get(pos) != Some(&0x0b) {
return out;
}
pos = match pos.checked_add(1) {
Some(p) => p,
None => return out,
};
u64::from(val as u32)
}
_ => {
return out;
}
}
} else {
let (_, np) = match read_uleb128(data, pos) {
Some(v) => v,
None => return out,
};
pos = np;
continue;
};
let payload_size = match read_uleb128(data, pos) {
Some((s, p)) => {
pos = p;
s as usize
}
None => return out,
};
let payload_end = match pos.checked_add(payload_size) {
Some(e) => e,
None => return out,
};
if payload_end > section_end {
return out;
}
out.push(WasmDataSegment {
mem_offset,
file_offset: pos,
size: payload_size,
});
pos = payload_end;
}
}
out
}
pub fn build_linear_memory_image(data: &[u8], max_size_bytes: usize) -> Option<Vec<u8>> {
let segs = data_segments(data);
if segs.is_empty() {
return None;
}
let mut max_end: u64 = 0;
for s in &segs {
let end = s.mem_offset.checked_add(s.size as u64)?;
if end > max_end {
max_end = end;
}
}
let total = usize::try_from(max_end).ok()?;
if total > max_size_bytes {
return None;
}
let mut image = vec![0u8; total];
for s in &segs {
let dst_start = usize::try_from(s.mem_offset).ok()?;
let dst_end = dst_start.checked_add(s.size)?;
let src = data.get(s.file_offset..s.file_offset.checked_add(s.size)?)?;
let dst = image.get_mut(dst_start..dst_end)?;
dst.copy_from_slice(src);
}
Some(image)
}
fn read_uleb128(data: &[u8], offset: usize) -> Option<(u32, usize)> {
let mut result: u32 = 0;
let mut shift: u32 = 0;
let mut pos = offset;
for _ in 0..5 {
let byte = *data.get(pos)?;
pos = pos.checked_add(1)?;
let chunk = u32::from(byte & 0x7f);
let shifted = chunk.checked_shl(shift)?;
result = result.checked_add(shifted)?;
if byte & 0x80 == 0 {
return Some((result, pos));
}
shift = shift.checked_add(7)?;
}
None
}
fn read_uleb128_signed(data: &[u8], offset: usize) -> Option<(i32, usize)> {
let mut result: i64 = 0;
let mut shift: u32 = 0;
let mut pos = offset;
let mut byte: u8;
loop {
byte = *data.get(pos)?;
pos = pos.checked_add(1)?;
let chunk = i64::from(byte & 0x7f);
let shifted = chunk.checked_shl(shift)?;
result = result.checked_add(shifted)?;
shift = shift.checked_add(7)?;
if byte & 0x80 == 0 {
break;
}
if shift >= 35 {
return None;
}
}
if shift < 32 && (byte & 0x40) != 0 {
result |= !0i64 << shift;
}
let v = i32::try_from(result).ok()?;
Some((v, pos))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn walks_minimal_module() {
let mut bytes = Vec::new();
bytes.extend_from_slice(b"\x00asm\x01\x00\x00\x00");
bytes.push(0);
bytes.push(13);
bytes.push(10); bytes.extend_from_slice(b"go:buildid");
bytes.extend_from_slice(&[0xff, 0x42]);
let s = walk(&bytes).next().unwrap();
assert_eq!(s.id, 0);
assert_eq!(s.name, Some("go:buildid"));
assert_eq!(
&bytes[s.payload_offset..s.payload_offset + s.payload_size],
&[0xff, 0x42]
);
}
#[test]
fn rejects_non_wasm() {
assert!(walk(b"\x7fELF\x02\x01\x01\x00").next().is_none());
}
#[test]
fn truncated_length_aborts() {
let mut bytes = Vec::new();
bytes.extend_from_slice(b"\x00asm\x01\x00\x00\x00");
bytes.push(0);
assert!(walk(&bytes).next().is_none());
}
#[test]
fn malformed_length_aborts() {
let mut bytes = Vec::new();
bytes.extend_from_slice(b"\x00asm\x01\x00\x00\x00");
bytes.push(0); bytes.extend_from_slice(&[0xe8, 0x07, 0x00]);
assert!(walk(&bytes).next().is_none());
}
}