From 0a9f9fd6adb59c513783838297a5ef0087a536d3 Mon Sep 17 00:00:00 2001
From: Rerumu <25379555+Rerumu@users.noreply.github.com>
Date: Mon, 22 Aug 2022 14:10:37 -0400
Subject: [PATCH] Add function loading validity checks

---
 codegen/luajit/src/translator.rs |  2 +-
 codegen/luau/src/translator.rs   |  2 +-
 wasm-ast/src/factory.rs          | 23 +++++++++++------------
 3 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/codegen/luajit/src/translator.rs b/codegen/luajit/src/translator.rs
index 626b25a..0e68f99 100644
--- a/codegen/luajit/src/translator.rs
+++ b/codegen/luajit/src/translator.rs
@@ -211,7 +211,7 @@ fn build_func_list(wasm: &Module, type_info: &TypeInfo) -> Vec<FuncData> {
 	wasm.code_section()
 		.iter()
 		.enumerate()
-		.map(|f| builder.create_indexed(f.0 + offset, f.1))
+		.map(|f| builder.create_indexed(f.0 + offset, f.1).unwrap())
 		.collect()
 }
 
diff --git a/codegen/luau/src/translator.rs b/codegen/luau/src/translator.rs
index 9b9111a..fe2730b 100644
--- a/codegen/luau/src/translator.rs
+++ b/codegen/luau/src/translator.rs
@@ -211,7 +211,7 @@ fn build_func_list(wasm: &Module, type_info: &TypeInfo) -> Vec<FuncData> {
 	wasm.code_section()
 		.iter()
 		.enumerate()
-		.map(|f| builder.create_indexed(f.0 + offset, f.1))
+		.map(|f| builder.create_indexed(f.0 + offset, f.1).unwrap())
 		.collect()
 }
 
diff --git a/wasm-ast/src/factory.rs b/wasm-ast/src/factory.rs
index 07fd101..0742f7c 100644
--- a/wasm-ast/src/factory.rs
+++ b/wasm-ast/src/factory.rs
@@ -1,7 +1,7 @@
-use wasmparser::{BlockType, FunctionBody, MemoryImmediate, Operator};
+use wasmparser::{BlockType, FunctionBody, MemoryImmediate, Operator, Result};
 
 use crate::{
-	module::TypeInfo,
+	module::{read_checked, TypeInfo},
 	node::{
 		BinOp, BinOpType, Block, Br, BrIf, BrTable, Call, CallIndirect, CmpOp, CmpOpType,
 		Expression, FuncData, GetGlobal, GetLocal, If, LabelType, LoadAt, LoadType, MemoryGrow,
@@ -248,24 +248,23 @@ impl<'a> Factory<'a> {
 		}
 	}
 
-	/// # Panics
+	/// # Errors
 	///
-	/// Panics if the function is malformed.
-	#[must_use]
-	pub fn create_indexed(&mut self, index: usize, func: &FunctionBody) -> FuncData {
-		let code: Result<Vec<_>, _> = func.get_operators_reader().unwrap().into_iter().collect();
-		let local: Result<Vec<_>, _> = func.get_locals_reader().unwrap().into_iter().collect();
+	/// Returns an error if the function is malformed.
+	pub fn create_indexed(&mut self, index: usize, func: &FunctionBody) -> Result<FuncData> {
+		let code = read_checked(func.get_operators_reader()?)?;
+		let local = read_checked(func.get_locals_reader()?)?;
 
 		let (num_param, num_result) = self.type_info.by_func_index(index);
-		let data = self.build_stat_list(&code.unwrap(), num_result);
+		let data = self.build_stat_list(&code, num_result);
 
-		FuncData {
-			local_data: local.unwrap(),
+		Ok(FuncData {
+			local_data: local,
 			num_result,
 			num_param,
 			num_stack: data.stack.capacity,
 			code: data.into(),
-		}
+		})
 	}
 
 	fn start_block(&mut self, ty: BlockType, variant: BlockVariant) {