fix(module): bound section allocations during rwasm decode

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<Context> Decode<Context> for InstructionSet {
     fn decode<D: Decoder<Context = Context>>(decoder: &mut D) -> Result<Self, DecodeError> {
-        let length: u64 = Decode::decode(decoder)?;
-        let mut instr: Vec<Opcode> = 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 })
     }
 }

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<Context> Decode<Context> for RwasmModuleInner {
             return Err(DecodeError::Other("rwasm: not supported version"));
         }
         let code_section: InstructionSet = Decode::decode(decoder)?;
-        let data_section: Vec<u8> = Decode::decode(decoder)?;
-        let elem_section: Vec<u32> = Decode::decode(decoder)?;
-        let wasm_section: Vec<u8> = Decode::decode(decoder)?;
+        let data_section = decode_section_bytes(decoder)?;
+        let elem_section: Vec<u32> = 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<RwasmModuleBuilder> 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];

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<Context, D: Decoder<Context = Context>>(
+    decoder: &mut D,
+) -> Result<usize, DecodeError> {
+    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<T>(vec: &mut Vec<T>, 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<T>` element by element.
+pub(crate) fn decode_section_vec<Context, T, D>(decoder: &mut D) -> Result<Vec<T>, DecodeError>
+where
+    T: Decode<Context>,
+    D: Decoder<Context = Context>,
+{
+    let length = decode_section_length(decoder)?;
+    decoder.claim_container_read::<T>(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::<T>());
+        items.push(T::decode(decoder)?);
+    }
+    Ok(items)
+}
+
+/// Decodes a length-prefixed `Vec<u8>` in bounded chunks, growing only as bytes arrive.
+pub(crate) fn decode_section_bytes<Context, D: Decoder<Context = Context>>(
+    decoder: &mut D,
+) -> Result<Vec<u8>, DecodeError> {
+    let length = decode_section_length(decoder)?;
+    decoder.claim_container_read::<u8>(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)
+}

src/types/mod.rs

@@ -1,4 +1,5 @@
 mod branch_offset;
+pub(crate) mod codec;
 mod constructor_params;
 mod error;
 mod func_ref;