1use alloc::vec::Vec;
11
12use crate::binary::leb128::{self, Cursor};
13use crate::binary::section::RawSection;
14use crate::binary::typeparser::parse_val_type as parse_binary_val_type;
15use crate::error::{ByteOffset, DecodeContext, DecodeError, DecodeErrorKind};
16use crate::types::{CodeBody, LocalDecl, ValType};
17
18const MAX_TOTAL_LOCALS: u64 = 50_000;
22
23pub fn parse_code_section<'a>(section: &RawSection<'a>) -> Result<Vec<CodeBody<'a>>, DecodeError> {
24 let mut cursor = Cursor::new(section.data);
25 let count = decode_u32_in_code_section(&mut cursor, section.offset)?;
26
27 let mut codes = Vec::with_capacity(cursor.capacity_hint(count));
28 for _ in 0..count {
29 codes.push(parse_code_body(&mut cursor, section.offset)?);
30 }
31
32 if !cursor.is_empty() {
33 return Err(DecodeError {
34 offset: ByteOffset(section.offset + cursor.position()),
35 context: DecodeContext::CodeSection,
36 kind: DecodeErrorKind::SectionSizeMismatch {
37 expected: section.data.len() as u32,
38 consumed: cursor.position() as u32,
39 },
40 });
41 }
42
43 Ok(codes)
44}
45
46fn parse_code_body<'a>(
47 cursor: &mut Cursor<'a>,
48 base_offset: usize,
49) -> Result<CodeBody<'a>, DecodeError> {
50 let body_size = decode_u32_in_code_section(cursor, base_offset)? as usize;
51 let body_offset = cursor.position();
52 let body_bytes = cursor.read_bytes(body_size).map_err(|_| DecodeError {
53 offset: ByteOffset(base_offset + body_offset),
54 context: DecodeContext::CodeSection,
55 kind: DecodeErrorKind::UnexpectedEof,
56 })?;
57
58 let mut body_cursor = Cursor::new(body_bytes);
59 let local_count = decode_u32_in_code_section(&mut body_cursor, base_offset + body_offset)?;
60
61 let mut locals = Vec::with_capacity(body_cursor.capacity_hint(local_count));
62 let mut total_locals: u64 = 0;
63 for _ in 0..local_count {
64 let count = decode_u32_in_code_section(&mut body_cursor, base_offset + body_offset)?;
65 let val_type = parse_val_type(&mut body_cursor, base_offset + body_offset)?;
66 total_locals += u64::from(count);
67 if total_locals > MAX_TOTAL_LOCALS {
72 return Err(DecodeError {
73 offset: ByteOffset(base_offset + body_offset),
74 context: DecodeContext::CodeSection,
75 kind: DecodeErrorKind::TooManyLocals,
76 });
77 }
78 locals.push(LocalDecl { count, val_type });
79 }
80
81 let instr_offset = body_cursor.position();
82 let body = &body_bytes[instr_offset..];
83
84 Ok(CodeBody {
85 locals,
86 body,
87 body_offset: base_offset + body_offset + instr_offset,
88 })
89}
90
91fn parse_val_type(cursor: &mut Cursor<'_>, base_offset: usize) -> Result<ValType, DecodeError> {
92 parse_binary_val_type(cursor, base_offset, DecodeContext::CodeSection)
93}
94
95fn decode_u32_in_code_section(
96 cursor: &mut Cursor<'_>,
97 base_offset: usize,
98) -> Result<u32, DecodeError> {
99 leb128::decode_u32(cursor).map_err(|mut e| {
100 e.context = DecodeContext::CodeSection;
101 e.offset = ByteOffset(base_offset + e.offset.0);
102 e
103 })
104}
105
106#[cfg(test)]
107mod tests {
108 use super::*;
109 use crate::binary::section::SectionId;
110 use crate::types::NumType;
111
112 fn raw_code_section(data: &[u8]) -> RawSection<'_> {
113 RawSection {
114 id: SectionId::Code,
115 offset: 40,
116 data,
117 }
118 }
119
120 #[test]
121 fn parse_empty_code_section() {
122 let section = raw_code_section(&[0x00]);
123 let codes = parse_code_section(§ion).unwrap();
124 assert!(codes.is_empty());
125 }
126
127 #[test]
128 fn parse_single_body_without_locals() {
129 let section = raw_code_section(&[
130 0x01, 0x02, 0x00, 0x0B, ]);
135 let codes = parse_code_section(§ion).unwrap();
136 assert_eq!(codes.len(), 1);
137 assert!(codes[0].locals.is_empty());
138 assert_eq!(codes[0].body, &[0x0B]);
139 assert_eq!(codes[0].body_offset, 43);
140 }
141
142 #[test]
146 fn reject_huge_local_decl_count_without_oom() {
147 let bytes = [
148 0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00, 0x01, 0x04, 0x01, 0x60, 0x00, 0x00, 0x03, 0x02, 0x01, 0x00, 0x0a, 0x07, 0x01, 0x05, 0xff, 0xff, 0xff, 0xff, 0x0b, ];
153 if let Ok(module) = crate::binary::module::Module::decode(&bytes) {
156 assert!(module.lower().is_err());
157 }
158 }
159
160 #[test]
164 fn reject_local_count_above_engine_cap() {
165 let section = raw_code_section(&[
166 0x01, 0x07, 0x01, 0xff, 0xff, 0xff, 0xff, 0x00, 0x7f, 0x0B, ]);
173 let error = parse_code_section(§ion).unwrap_err();
174 assert_eq!(error.kind, DecodeErrorKind::TooManyLocals);
175 }
176
177 #[test]
178 fn parse_body_with_locals() {
179 let section = raw_code_section(&[
180 0x01, 0x06, 0x02, 0x01, 0x7F, 0x02, 0x7E, 0x0B, ]);
187 let codes = parse_code_section(§ion).unwrap();
188 assert_eq!(codes[0].locals.len(), 2);
189 assert_eq!(
190 codes[0].locals[0],
191 LocalDecl {
192 count: 1,
193 val_type: ValType::Num(NumType::I32),
194 }
195 );
196 assert_eq!(
197 codes[0].locals[1],
198 LocalDecl {
199 count: 2,
200 val_type: ValType::Num(NumType::I64),
201 }
202 );
203 assert_eq!(codes[0].body, &[0x0B]);
204 }
205
206 #[test]
207 fn reject_invalid_local_type() {
208 let section = raw_code_section(&[
209 0x01, 0x03, 0x01, 0x01, 0x01, ]);
214 let err = parse_code_section(§ion).unwrap_err();
215 assert_eq!(err.kind, DecodeErrorKind::UnknownValType { byte: 0x01 });
216 }
217
218 #[test]
219 fn reject_trailing_bytes() {
220 let section = raw_code_section(&[0x00, 0xFF]);
221 let err = parse_code_section(§ion).unwrap_err();
222 assert_eq!(
223 err.kind,
224 DecodeErrorKind::SectionSizeMismatch {
225 expected: 2,
226 consumed: 1,
227 }
228 );
229 }
230}