Validate the effective type size of modules (#215)

* Validate the effective type size of modules

This commit implements a fix for #214 where the effective size of the
type of any item in a module (or a module itself) is now required to be
bounded. The goal here is to accept reasonable modules but at the same
time protecting implementations like wasmtime from having to worry about
deeply recursive and nested module types. The general fix here is to
account for the size of all types as we parse them, and at all times the
size is bounded.

Additionally wasm-smith is updated with similar accounting for the size
of types and such to prevent generating modules with massively recursive
types.

Closes #214

* Clarify doc comments

crates/wasm-smith/src/config.rs

@@ -254,6 +254,19 @@ pub trait Config: Arbitrary + Default + Clone {
     fn max_nesting_depth(&self) -> usize {
         10
     }
+
+    /// Returns the maximal effective size of any type generated by wasm-smith.
+    ///
+    /// Note that this number is roughly in units of "how many types would be
+    /// needed to represent the recursive type". A function with 8 parameters
+    /// and 2 results would take 11 types (one for the type, 10 for
+    /// params/results). A module type with 2 imports and 3 exports would
+    /// take 6 (module + imports + exports) plus the size of each import/export
+    /// type. This is a somewhat rough measurement that is not intended to be
+    /// very precise.
+    fn max_type_size(&self) -> u32 {
+        1_000
+    }
 }
 
 /// The default configuration.

crates/wasm-smith/src/lib.rs

@@ -191,6 +191,10 @@ where
     elems: Vec<ElementSegment>,
     code: Vec<Code>,
     data: Vec<DataSegment>,
+
+    /// The predicted size of the effective type of this module, based on this
+    /// module's size of the types of imports/exports.
+    type_size: u32,
 }
 
 impl<C: Config> ConfiguredModule<C> {
@@ -256,11 +260,13 @@ struct FuncType {
 
 #[derive(Clone, Debug, Default)]
 struct InstanceType {
+    type_size: u32,
     exports: indexmap::IndexMap<String, EntityType>,
 }
 
 #[derive(Clone, Debug)]
 struct ModuleType {
+    type_size: u32,
     imports: Vec<(String, Option<String>, EntityType)>,
     import_types: indexmap::IndexMap<String, EntityType>,
     /// The list of exports can be found in the `InstanceType` indirection here,
@@ -822,10 +828,15 @@ where
         let mut imports = Vec::new();
         let mut import_types = indexmap::IndexMap::new();
         let mut names = HashMap::new();
+        let mut type_size = exports.type_size;
         if !entities.max_reached(&self.config) {
             arbitrary_loop(u, 0, self.config.max_imports(), |u| {
                 let (module, name) = unique_import_strings(1_000, &mut names, true, u)?;
                 let ty = self.arbitrary_entity_type(u, entities)?;
+                match type_size.checked_add(ty.size()) {
+                    Some(s) if s < self.config.max_type_size() => type_size = s,
+                    _ => return Ok(false),
+                }
                 if let Some(name) = &name {
                     let ity = import_types.entry(module.clone()).or_insert_with(|| {
                         EntityType::Instance(u32::max_value(), Default::default())
@@ -843,6 +854,7 @@ where
             })?;
         }
         Ok(Rc::new(ModuleType {
+            type_size,
             imports,
             import_types,
             exports,
@@ -856,15 +868,20 @@ where
     ) -> Result<Rc<InstanceType>> {
         let mut export_names = HashSet::new();
         let mut exports = indexmap::IndexMap::new();
+        let mut type_size = 0u32;
         if !entities.max_reached(&self.config) {
             arbitrary_loop(u, 0, self.config.max_exports(), |u| {
                 let name = unique_string(1_000, &mut export_names, u)?;
                 let ty = self.arbitrary_entity_type(u, entities)?;
+                match type_size.checked_add(ty.size()) {
+                    Some(s) if s < self.config.max_type_size() => type_size = s,
+                    _ => return Ok(false),
+                }
                 exports.insert(name, ty);
                 Ok(!entities.max_reached(&self.config))
             })?;
         }
-        Ok(Rc::new(InstanceType { exports }))
+        Ok(Rc::new(InstanceType { type_size, exports }))
     }
 
     fn arbitrary_entity_type(
@@ -1038,8 +1055,9 @@ where
             }
 
             self.num_imports += 1;
+            self.type_size += ty.size();
             imports.push((module, name, ty));
-            Ok(true)
+            Ok(self.type_size < self.config.max_type_size())
         })?;
         if !imports.is_empty() || u.arbitrary()? {
             self.initial_sections.push(InitialSection::Import(imports));
@@ -1230,9 +1248,16 @@ where
                 exports.insert(name.clone(), ty);
             }
             let ty = Rc::new(ModuleType {
+                type_size: module.type_size,
                 imports,
                 import_types,
-                exports: Rc::new(InstanceType { exports }),
+                exports: Rc::new(InstanceType {
+                    // This type size isn't quite right since it takes into
+                    // account imports, but that's ok for now since it's just a
+                    // predictor about how big things should get.
+                    type_size: module.type_size,
+                    exports,
+                }),
             });
 
             // And then given the type of the module we copy it over to
@@ -1447,47 +1472,58 @@ where
     }
 
     fn arbitrary_exports(&mut self, u: &mut Unstructured) -> Result<()> {
+        if self.type_size < self.config.max_type_size() {
+            return Ok(());
+        }
+
         let mut choices: Vec<fn(&mut Unstructured, &mut ConfiguredModule<C>) -> Result<_>> =
             Vec::with_capacity(4);
 
         if self.funcs.len() > 0 {
             choices.push(|u, m| {
                 let idx = u.int_in_range(0..=m.funcs.len() - 1)?;
+                let sig = &m.funcs[idx].1;
+                m.type_size += (sig.params.len() + sig.results.len()) as u32;
                 Ok(Export::Func(idx as u32))
             });
         }
 
         if self.tables.len() > 0 {
             choices.push(|u, m| {
                 let idx = u.int_in_range(0..=m.tables.len() - 1)?;
+                m.type_size += 1;
                 Ok(Export::Table(idx as u32))
             });
         }
 
         if self.memories.len() > 0 {
             choices.push(|u, m| {
                 let idx = u.int_in_range(0..=m.memories.len() - 1)?;
+                m.type_size += 1;
                 Ok(Export::Memory(idx as u32))
             });
         }
 
         if self.globals.len() > 0 {
             choices.push(|u, m| {
                 let idx = u.int_in_range(0..=m.globals.len() - 1)?;
+                m.type_size += 1;
                 Ok(Export::Global(idx as u32))
             });
         }
 
         if self.instances.len() > 0 {
             choices.push(|u, m| {
                 let idx = u.int_in_range(0..=m.instances.len() - 1)?;
+                m.type_size += m.instances[idx].type_size;
                 Ok(Export::Instance(idx as u32))
             });
         }
 
         if self.modules.len() > 0 {
             choices.push(|u, m| {
                 let idx = u.int_in_range(0..=m.modules.len() - 1)?;
+                m.type_size += m.modules[idx].type_size;
                 Ok(Export::Module(idx as u32))
             });
         }
@@ -1506,7 +1542,7 @@ where
                 let f = u.choose(&choices)?;
                 let export = f(u, self)?;
                 self.exports.push((name, export));
-                Ok(true)
+                Ok(self.type_size < self.config.max_type_size())
             },
         )
     }
@@ -1996,6 +2032,18 @@ fn arbitrary_vec_u8(u: &mut Unstructured) -> Result<Vec<u8>> {
     Ok(u.get_bytes(size)?.to_vec())
 }
 
+impl EntityType {
+    fn size(&self) -> u32 {
+        let base = match self {
+            EntityType::Global(_) | EntityType::Table(_) | EntityType::Memory(_) => 1,
+            EntityType::Func(_, ty) => (ty.params.len() + ty.results.len()) as u32,
+            EntityType::Instance(_, ty) => ty.type_size,
+            EntityType::Module(_, ty) => ty.type_size,
+        };
+        base + 1
+    }
+}
+
 /// This is a helper structure used during the `arbitrary_initial_sections`
 /// phase of generating a module to generate the alias section.
 ///

crates/wasmparser/src/limits.rs

@@ -36,3 +36,4 @@ pub const MAX_WASM_MEMORIES: usize = 100;
 pub const MAX_WASM_MODULES: usize = 1_000;
 pub const MAX_WASM_INSTANCES: usize = 1_000;
 pub const MAX_WASM_EVENTS: usize = 1_000_000;
+pub const MAX_TYPE_SIZE: u32 = 100_000;

crates/wasmparser/src/validator.rs

@@ -146,7 +146,7 @@ struct ModuleState {
     function_references: HashSet<u32>,
 
     // This is populated when we hit the export section
-    exports: HashMap<String, EntityType>,
+    exports: NameSet,
 
     // This is populated as we visit import sections, which might be
     // incrementally in the face of a module-linking-using module.
@@ -257,12 +257,14 @@ impl TypeDef {
 }
 
 struct ModuleType {
+    type_size: u32,
     imports: HashMap<String, EntityType>,
     exports: HashMap<String, EntityType>,
 }
 
 #[derive(Default)]
 struct InstanceType {
+    type_size: u32,
     exports: HashMap<String, EntityType>,
 }
 
@@ -594,6 +596,11 @@ impl Validator {
                     exports.push(self.offset, e.name, None, ty, &mut self.types, "export")?;
                 }
                 TypeDef::Module(ModuleType {
+                    type_size: combine_type_sizes(
+                        self.offset,
+                        imports.type_size,
+                        exports.type_size,
+                    )?,
                     imports: imports.set,
                     exports: exports.set,
                 })
@@ -608,6 +615,7 @@ impl Validator {
                     exports.push(self.offset, e.name, None, ty, &mut self.types, "export")?;
                 }
                 TypeDef::Instance(InstanceType {
+                    type_size: exports.type_size,
                     exports: exports.set,
                 })
             }
@@ -1002,7 +1010,12 @@ impl Validator {
         // Create a synthetic type declaration for this instance's type and
         // record its type in the global type list. We might not have another
         // `TypeDef::Instance` to point to if the module was locally declared.
+        //
+        // Note that the predicted size of this type is inflated due to
+        // accounting for the imports on the original module, but that should be
+        // ok for now since it's only used to limit the size of larger types.
         let instance_ty = InstanceType {
+            type_size: ty.type_size,
             exports: ty.exports.clone(),
         };
         self.cur.state.assert_mut().instances.push(self.types.len());
@@ -1297,9 +1310,9 @@ impl Validator {
             }
             let ty = me.check_external_kind("exported", e.kind, e.index)?;
             let state = me.cur.state.assert_mut();
-            if state.exports.insert(e.field.to_string(), ty).is_some() {
-                return me.create_error("duplicate export name");
-            }
+            state
+                .exports
+                .push(me.offset, e.field, None, ty, &mut me.types, "export")?;
             Ok(())
         })
     }
@@ -1532,15 +1545,27 @@ impl Validator {
                 return self.create_error("function and code sections have inconsistent lengths");
             }
         }
+
+        // Ensure that the effective type size of this module is of a bounded
+        // size. This is primarily here for the module linking proposal, and
+        // we'll record this in the module type below if we're part of a nested
+        // module.
+        let type_size = combine_type_sizes(
+            self.offset,
+            self.cur.state.imports.type_size,
+            self.cur.state.exports.type_size,
+        )?;
+
         // If we have a parent then we're going to exit this module's context
         // and resume where we left off in the parent. We inject a new type for
         // our module we just validated in the parent's module index space, and
         // then we reset our current state to the parent.
         if let Some(mut parent) = self.parents.pop() {
             let module_type = self.types.len();
             self.types.push(TypeDef::Module(ModuleType {
+                type_size,
                 imports: self.cur.state.imports.set.clone(),
-                exports: self.cur.state.exports.clone(),
+                exports: self.cur.state.exports.set.clone(),
             }));
             parent.state.assert_mut().submodules.push(module_type);
             self.cur = parent;
@@ -1549,6 +1574,16 @@ impl Validator {
     }
 }
 
+fn combine_type_sizes(offset: usize, a: u32, b: u32) -> Result<u32> {
+    match a.checked_add(b) {
+        Some(sum) if sum < MAX_TYPE_SIZE => Ok(sum),
+        _ => Err(BinaryReaderError::new(
+            "effective type size too large".to_string(),
+            offset,
+        )),
+    }
+}
+
 impl WasmFeatures {
     pub(crate) fn check_value_type(&self, ty: Type) -> Result<(), &'static str> {
         match ty {
@@ -1696,6 +1731,7 @@ mod arc {
 struct NameSet {
     set: HashMap<String, EntityType>,
     implicit: HashSet<String>,
+    type_size: u32,
 }
 
 impl NameSet {
@@ -1723,6 +1759,7 @@ impl NameSet {
         types: &mut SnapshotList<TypeDef>,
         desc: &str,
     ) -> Result<Option<usize>> {
+        self.type_size = combine_type_sizes(offset, self.type_size, ty.size(types))?;
         let name = match name {
             Some(name) => name,
             // If the `name` is not provided then this is a module-linking style
@@ -1797,6 +1834,28 @@ impl NameSet {
     }
 }
 
+impl EntityType {
+    fn size(&self, list: &SnapshotList<TypeDef>) -> u32 {
+        let recursive_size = match self {
+            // Note that this function computes the size of the *type*, not the
+            // size of the value, so these "leaves" all count as 1
+            EntityType::Global(_) | EntityType::Memory(_) | EntityType::Table(_) => 1,
+
+            // These types have recursive sizes so we look up the size in the
+            // type tables.
+            EntityType::Func(i)
+            | EntityType::Module(i)
+            | EntityType::Instance(i)
+            | EntityType::Event(i) => match &list[*i] {
+                TypeDef::Func(f) => (f.params.len() + f.returns.len()) as u32,
+                TypeDef::Module(m) => m.type_size,
+                TypeDef::Instance(i) => i.type_size,
+            },
+        };
+        recursive_size.saturating_add(1)
+    }
+}
+
 /// This is a type which mirrors a subset of the `Vec<T>` API, but is intended
 /// to be able to be cheaply snapshotted and cloned.
 ///