From 6747f665482388d28184a23949448ebf38902a55 Mon Sep 17 00:00:00 2001 From: Dmitry Savonin Date: Fri, 7 Aug 2026 20:33:15 +0400 Subject: [PATCH] fix(module): bound section allocations during rwasm decode --- src/isa/mod.rs | 13 +++----- src/module/mod.rs | 54 ++++++++++++++++++++++++++++--- src/types/codec.rs | 80 ++++++++++++++++++++++++++++++++++++++++++++++ src/types/mod.rs | 1 + 4 files changed, 135 insertions(+), 13 deletions(-) create mode 100644 src/types/codec.rs diff --git a/src/isa/mod.rs b/src/isa/mod.rs index ef7fed16..ee60de19 100644 --- a/src/isa/mod.rs +++ b/src/isa/mod.rs @@ -12,9 +12,9 @@ mod table; use crate::{ types::{ - AddressOffset, BlockFuel, BranchOffset, BranchTableTargets, CompiledFunc, DataSegmentIdx, - ElementSegmentIdx, GlobalIdx, LocalDepth, MaxStackHeight, Opcode, SignatureIdx, TableIdx, - UntypedValue, + codec::decode_section_vec, AddressOffset, BlockFuel, BranchOffset, BranchTableTargets, + CompiledFunc, DataSegmentIdx, ElementSegmentIdx, GlobalIdx, LocalDepth, MaxStackHeight, + Opcode, SignatureIdx, TableIdx, UntypedValue, }, CompilationError, NumLocals, SysFuncIdx, TrapCode, }; @@ -393,12 +393,7 @@ impl Encode for InstructionSet { impl Decode for InstructionSet { fn decode>(decoder: &mut D) -> Result { - let length: u64 = Decode::decode(decoder)?; - let mut instr: Vec = Vec::with_capacity(length as usize); - for _ in 0..length as usize { - let opcode: Opcode = Decode::decode(decoder)?; - instr.push(opcode); - } + let instr = decode_section_vec(decoder)?; Ok(Self { instr }) } } diff --git a/src/module/mod.rs b/src/module/mod.rs index 972caece..2e208c5b 100644 --- a/src/module/mod.rs +++ b/src/module/mod.rs @@ -1,4 +1,5 @@ use crate::{ + types::codec::{decode_section_bytes, decode_section_vec}, CompilationConfig, CompilationError, ConstructorParams, HintType, InstructionSet, ModuleParser, Opcode, }; @@ -182,9 +183,9 @@ impl Decode for RwasmModuleInner { return Err(DecodeError::Other("rwasm: not supported version")); } let code_section: InstructionSet = Decode::decode(decoder)?; - let data_section: Vec = Decode::decode(decoder)?; - let elem_section: Vec = Decode::decode(decoder)?; - let wasm_section: Vec = Decode::decode(decoder)?; + let data_section = decode_section_bytes(decoder)?; + let elem_section: Vec = decode_section_vec(decoder)?; + let wasm_section = decode_section_bytes(decoder)?; let source_pc: u32 = match Decode::decode(decoder) { Ok(source_pc) => source_pc, Err(DecodeError::UnexpectedEnd { additional }) => { @@ -289,7 +290,10 @@ impl From for RwasmModule { #[cfg(test)] mod tests { - use crate::{instruction_set, RwasmModule, RwasmModuleInner}; + use crate::{ + instruction_set, RwasmModule, RwasmModuleInner, RWASM_MAGIC_BYTE_0, RWASM_MAGIC_BYTE_1, + RWASM_VERSION_V1, + }; use bincode::error::DecodeError; use hex_literal::hex; @@ -365,6 +369,48 @@ mod tests { assert!(matches!(err, DecodeError::Other(_))); } + /// Every section length is attacker-controlled, so a truncated header must be rejected without + /// allocating the announced capacity. Before the fix these inputs panicked with + /// `capacity overflow`, or aborted through `handle_alloc_error` for large non-overflowing + /// lengths. + #[test] + fn test_decode_rejects_oversized_section_lengths() { + // Number of empty sections preceding the one under test, in encoding order: + // code, data, elem, hint. + for preceding_sections in 0..4 { + for bogus_length in [u64::MAX, 1u64 << 40] { + let mut sink = vec![RWASM_MAGIC_BYTE_0, RWASM_MAGIC_BYTE_1, RWASM_VERSION_V1]; + for _ in 0..preceding_sections { + sink.extend_from_slice(&0u64.to_le_bytes()); + } + sink.extend_from_slice(&bogus_length.to_le_bytes()); + + let err = RwasmModule::new_checked(&sink) + .expect_err("section length larger than the remaining input must be rejected"); + assert!( + matches!( + err, + DecodeError::UnexpectedEnd { .. } | DecodeError::OutsideUsizeRange(_) + ), + "unexpected error for section {preceding_sections} \ + with length {bogus_length}: {err:?}" + ); + } + } + } + + /// A section length that fits the input must still decode, so the bound above cannot be a + /// blanket size limit. + #[test] + fn test_decode_accepts_large_hint_section() { + let mut module = test_module(); + module.hint_section = vec![0xab; 512 * 1024]; + let encoded_module = bincode::encode_to_vec(&module, bincode::config::legacy()).unwrap(); + let (decoded, _): (RwasmModuleInner, usize) = + bincode::decode_from_slice(&encoded_module, bincode::config::legacy()).unwrap(); + assert_eq!(module, decoded); + } + #[test] fn test_endianness() { let module = vec![1, 2, 3]; diff --git a/src/types/codec.rs b/src/types/codec.rs new file mode 100644 index 00000000..9ee1ad56 --- /dev/null +++ b/src/types/codec.rs @@ -0,0 +1,80 @@ +//! Allocation-safe primitives for decoding length-prefixed rwasm sections. +//! +//! Every section in the rwasm binary format starts with a `u64` length taken straight from +//! untrusted input. Handing that length to `Vec::with_capacity` (which is what both the hand-written +//! and bincode-derived decoders used to do) lets an 11-byte binary request an arbitrary allocation +//! before a single element is read: `u64::MAX` panics with `capacity overflow`, and a large +//! non-overflowing value aborts the process through `handle_alloc_error`. +//! +//! The helpers below reserve a bounded amount up front and grow only as elements arrive, so the +//! peak allocation stays proportional to the input that is actually present. A truncated section +//! fails with `UnexpectedEnd` as soon as the reader runs dry, and any section backed by real input +//! still decodes exactly as it did before — this is not a size limit on the format. + +use alloc::vec::Vec; +use bincode::{ + de::{read::Reader, Decoder}, + error::DecodeError, + Decode, +}; + +/// Upper bound on how many elements a section decoder reserves before reading any input. +/// Anything beyond this grows on demand as elements are decoded. +const N_MAX_SECTION_PREALLOC: usize = 4096; + +/// How many bytes of a byte section are reserved and read at a time. +const N_MAX_BYTES_CHUNK: usize = 64 * 1024; + +/// Decodes a section length prefix and converts it to `usize` without trusting its magnitude. +fn decode_section_length>( + decoder: &mut D, +) -> Result { + let length: u64 = Decode::decode(decoder)?; + usize::try_from(length).map_err(|_| DecodeError::OutsideUsizeRange(length)) +} + +/// Reserves capacity for at most [`N_MAX_SECTION_PREALLOC`] elements of `length`. +fn reserve_capped(vec: &mut Vec, length: usize) -> Result<(), DecodeError> { + vec.try_reserve(length.min(N_MAX_SECTION_PREALLOC)) + .map_err(|_| DecodeError::Other("rwasm: failed to allocate section")) +} + +/// Decodes a length-prefixed `Vec` element by element. +pub(crate) fn decode_section_vec(decoder: &mut D) -> Result, DecodeError> +where + T: Decode, + D: Decoder, +{ + let length = decode_section_length(decoder)?; + decoder.claim_container_read::(length)?; + let mut items = Vec::new(); + reserve_capped(&mut items, length)?; + for _ in 0..length { + // See bincode's `unclaim_bytes_read` docs: the container read is claimed for the whole + // section, so every element must give its share back before decoding itself. + decoder.unclaim_bytes_read(size_of::()); + items.push(T::decode(decoder)?); + } + Ok(items) +} + +/// Decodes a length-prefixed `Vec` in bounded chunks, growing only as bytes arrive. +pub(crate) fn decode_section_bytes>( + decoder: &mut D, +) -> Result, DecodeError> { + let length = decode_section_length(decoder)?; + decoder.claim_container_read::(length)?; + let mut bytes = Vec::new(); + let mut filled = 0; + while filled < length { + let target = length.min(filled + N_MAX_BYTES_CHUNK); + bytes + .try_reserve(target - filled) + .map_err(|_| DecodeError::Other("rwasm: failed to allocate section"))?; + // The reservation above covers the whole chunk, so this cannot allocate again. + bytes.resize(target, 0); + decoder.reader().read(&mut bytes[filled..target])?; + filled = target; + } + Ok(bytes) +} diff --git a/src/types/mod.rs b/src/types/mod.rs index d83773ed..bfb6000c 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -1,4 +1,5 @@ mod branch_offset; +pub(crate) mod codec; mod constructor_params; mod error; mod func_ref;