wasm/core/reader/
mod.rs

1use crate::core::reader::span::Span;
2use crate::ValidationError;
3
4pub mod section_header;
5pub mod types;
6
7/// A struct for managing and reading WASM bytecode
8///
9/// Its purpose is to abstract parsing basic WASM values from the bytecode.
10#[derive(Clone)]
11pub struct WasmReader<'a> {
12    /// Entire WASM binary as slice
13    pub full_wasm_binary: &'a [u8],
14
15    /// Current program counter, i. e. index of the next byte to be consumed from the WASM binary
16    ///
17    /// # Correctness Note
18    ///
19    /// The `pc` points to the next byte to be consumed from the WASM binary. Therefore, after
20    /// consuming last byte, this cursor will advance past the last byte; for a WASM binary that is
21    /// 100 bytes long (valid indexes start with 0 and end with 99), the `pc` therefore can become
22    /// 100. However, it can not advance further.
23    ///
24    /// The table below illustrates this with an example for a WASM binary that is 5 bytes long:
25    ///
26    /// |                     Index |   0  |   1  |   2  |   3  |   4  | 5 | 6 |
27    /// |--------------------------:|:----:|:----:|:----:|:----:|:----:|:-:|:-:|
28    /// | `full_wasm_binary[index]` | 0xaa | 0xbb | 0xcc | 0xee | 0xff | - | - |
29    /// |      Valid `pc` position? |   ✅  |   ✅  |   ✅  |   ✅  |   ✅  | ✅ | ❌ |
30    pub pc: usize,
31}
32
33impl<'a> WasmReader<'a> {
34    /// Initialize a new [WasmReader] from a WASM byte slice
35    pub const fn new(wasm: &'a [u8]) -> Self {
36        Self {
37            full_wasm_binary: wasm,
38            pc: 0,
39        }
40    }
41
42    /// Advance the cursor to the first byte of the provided [Span] and validates that entire [Span] fits the WASM binary
43    ///
44    /// # Note
45    ///
46    /// This allows setting the [`pc`](WasmReader::pc) to one byte *past* the end of
47    /// [full_wasm_binary](WasmReader::full_wasm_binary), **if** the [Span]'s length is 0. For
48    /// further information, refer to the [field documentation of `pc`](WasmReader::pc).
49    pub fn move_start_to(&mut self, span: Span) -> Result<(), ValidationError> {
50        if span.from + span.len > self.full_wasm_binary.len() {
51            return Err(ValidationError::Eof);
52        }
53
54        self.pc = span.from;
55
56        Ok(())
57    }
58
59    /// Byte slice to the remainder of the WASM binary, beginning from the current [`pc`](Self::pc)
60    pub fn remaining_bytes(&self) -> &[u8] {
61        &self.full_wasm_binary[self.pc..]
62    }
63
64    /// Create a [Span] starting from [`pc`](Self::pc) for the next `len` bytes
65    ///
66    /// Verifies the span to fit the WASM binary, i.e. using this span to index the WASM binary will
67    /// not yield an error.
68    pub fn make_span(&self, len: usize) -> Result<Span, ValidationError> {
69        if self.pc + len > self.full_wasm_binary.len() {
70            return Err(ValidationError::Eof);
71        }
72        Ok(Span::new(self.pc, len))
73    }
74
75    /// Take `N` bytes starting from [`pc`](Self::pc), then advance the [`pc`](Self::pc) by `N`
76    ///
77    /// This yields back an array of the correct length
78    ///
79    /// # Note
80    ///
81    /// This allows setting the [`pc`](WasmReader::pc) to one byte *past* the end of
82    /// [full_wasm_binary](WasmReader::full_wasm_binary), **if** `N` equals the remaining bytes
83    /// slice's length. For further information, refer to the [field documentation of `pc`]
84    /// (WasmReader::pc).
85    pub fn strip_bytes<const N: usize>(&mut self) -> Result<[u8; N], ValidationError> {
86        if N > self.full_wasm_binary.len() - self.pc {
87            return Err(ValidationError::Eof);
88        }
89
90        let bytes = &self.full_wasm_binary[self.pc..(self.pc + N)];
91        self.pc += N;
92
93        Ok(bytes.try_into().expect("the slice length to be exactly N"))
94    }
95
96    /// Read the current byte without advancing the [`pc`](Self::pc)
97    ///
98    /// May yield an error if the [`pc`](Self::pc) advanced past the end of the WASM binary slice
99    pub fn peek_u8(&self) -> Result<u8, ValidationError> {
100        self.full_wasm_binary
101            .get(self.pc)
102            .copied()
103            .ok_or(ValidationError::Eof)
104    }
105
106    /// Call a closure that may mutate the [WasmReader]
107    ///
108    /// Returns a tuple of the closure's return value and the number of bytes that the [`WasmReader`]
109    /// was advanced by.
110    ///
111    /// # Panics
112    ///
113    /// May panic if the closure moved the [`pc`](Self::pc) backwards, e.g. when
114    /// [move_start_to](Self::move_start_to) is called.
115    pub fn measure_num_read_bytes<T>(
116        &mut self,
117        f: impl FnOnce(&mut WasmReader) -> Result<T, ValidationError>,
118    ) -> Result<(T, usize), ValidationError> {
119        let before = self.pc;
120        let ret = f(self)?;
121
122        // TODO maybe use checked sub, that is slower but guarantees no surprises
123        debug_assert!(
124            self.pc >= before,
125            "pc was advanced backwards towards the start"
126        );
127
128        let num_read_bytes = self.pc - before;
129        Ok((ret, num_read_bytes))
130    }
131
132    /// Skip `num_bytes`, advancing the [`pc`](Self::pc) accordingly
133    ///
134    /// # Note
135    ///
136    /// This can move the [`pc`](Self::pc) past the last byte of the WASM binary, so that reading
137    /// more than 0 further bytes would panick. However, it can not move the [`pc`](Self::pc) any
138    /// further than that, instead an error is returned. For further information, refer to the
139    /// [field documentation of `pc`] (WasmReader::pc).
140    pub fn skip(&mut self, num_bytes: usize) -> Result<(), ValidationError> {
141        if num_bytes > self.full_wasm_binary.len() - self.pc {
142            return Err(ValidationError::Eof);
143        }
144        self.pc += num_bytes;
145        Ok(())
146    }
147
148    /// Consumes [Self], yielding back the internal reference to the WASM binary
149    pub fn into_inner(self) -> &'a [u8] {
150        self.full_wasm_binary
151    }
152
153    /// A wrapper function for reads with transaction-like behavior.
154    ///
155    /// The provided closure will be called with `&mut self` and its result will be returned.
156    /// However if the closure returns `Err(_)`, `self` will be reset as if the closure was never called.
157    pub fn handle_transaction<T, E>(
158        &mut self,
159        f: impl FnOnce(&mut WasmReader<'a>) -> Result<T, E>,
160    ) -> Result<T, E> {
161        let original = self.clone();
162        f(self).inspect_err(|_| {
163            *self = original;
164        })
165    }
166}
167
168pub mod span {
169    use core::ops::Index;
170
171    use crate::core::reader::WasmReader;
172
173    /// An index and offset to describe a (sub-) slice into WASM bytecode
174    ///
175    /// Can be used to index into a [WasmReader], yielding a byte slice. As it does not
176    /// actually own the indexed data, this struct is free of lifetimes. Caution is advised when
177    /// indexing unknown slices, as a [Span] does not validate the length of the indexed slice.
178    #[derive(Copy, Clone, Debug, Hash)]
179    pub struct Span {
180        pub from: usize,
181        pub len: usize,
182    }
183
184    impl Span {
185        /// Create a new [Span], starting from `from` and ranging `len` elements
186        pub const fn new(from: usize, len: usize) -> Self {
187            Self { from, len }
188        }
189
190        /// Returns the length of this [Span]
191        pub const fn len(&self) -> usize {
192            self.len
193        }
194
195        pub const fn from(&self) -> usize {
196            self.from
197        }
198    }
199
200    impl<'a> Index<Span> for WasmReader<'a> {
201        type Output = [u8];
202
203        fn index(&self, index: Span) -> &'a Self::Output {
204            &self.full_wasm_binary[index.from..(index.from + index.len)]
205        }
206    }
207}
208
209#[cfg(test)]
210mod test {
211    use crate::ValType;
212
213    use super::*;
214    use alloc::vec;
215
216    #[test]
217    fn move_start_to() {
218        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
219        let mut wasm_reader = WasmReader::new(&my_bytes);
220
221        let span = Span::new(0, 0);
222        wasm_reader.move_start_to(span).unwrap();
223        // this actually dangerous, we did not validate there to be more than 0 bytes using the Span
224        wasm_reader.peek_u8().unwrap();
225
226        let span = Span::new(0, my_bytes.len());
227        wasm_reader.move_start_to(span).unwrap();
228        wasm_reader.peek_u8().unwrap();
229        assert_eq!(wasm_reader[span], my_bytes);
230
231        let span = Span::new(my_bytes.len(), 0);
232        wasm_reader.move_start_to(span).unwrap();
233        // span had zero length, hence wasm_reader.peek_u8() would be allowed to fail
234
235        let span = Span::new(my_bytes.len() - 1, 1);
236        wasm_reader.move_start_to(span).unwrap();
237
238        assert_eq!(wasm_reader.peek_u8().unwrap(), *my_bytes.last().unwrap());
239    }
240
241    #[test]
242    fn move_start_to_out_of_bounds_1() {
243        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
244        let mut wasm_reader = WasmReader::new(&my_bytes);
245
246        let span = Span::new(my_bytes.len(), 1);
247        assert_eq!(wasm_reader.move_start_to(span), Err(ValidationError::Eof));
248    }
249
250    #[test]
251    fn move_start_to_out_of_bounds_2() {
252        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
253        let mut wasm_reader = WasmReader::new(&my_bytes);
254
255        let span = Span::new(0, my_bytes.len() + 1);
256        assert_eq!(wasm_reader.move_start_to(span), Err(ValidationError::Eof));
257    }
258
259    #[test]
260    fn remaining_bytes_1() {
261        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
262        let mut wasm_reader = WasmReader::new(&my_bytes);
263
264        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
265        wasm_reader.skip(4).unwrap();
266        assert_eq!(wasm_reader.peek_u8().unwrap(), 0x15);
267
268        assert_eq!(wasm_reader.remaining_bytes(), &my_bytes[4..]);
269    }
270
271    #[test]
272    fn remaining_bytes_2() {
273        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
274        let mut wasm_reader = WasmReader::new(&my_bytes);
275
276        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
277        wasm_reader.skip(5).unwrap();
278        assert_eq!(wasm_reader.remaining_bytes(), &my_bytes[5..]);
279        assert_eq!(wasm_reader.remaining_bytes(), &[]);
280    }
281
282    #[test]
283    fn strip_bytes_1() {
284        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
285        let mut wasm_reader = WasmReader::new(&my_bytes);
286
287        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
288        let stripped_bytes = wasm_reader.strip_bytes::<4>().unwrap();
289        assert_eq!(&stripped_bytes, &my_bytes[..4]);
290        assert_eq!(wasm_reader.remaining_bytes(), &[0x15]);
291    }
292
293    #[test]
294    fn strip_bytes_2() {
295        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
296        let mut wasm_reader = WasmReader::new(&my_bytes);
297
298        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
299        wasm_reader.skip(1).unwrap();
300        let stripped_bytes = wasm_reader.strip_bytes::<4>().unwrap();
301        assert_eq!(&stripped_bytes, &my_bytes[1..5]);
302        assert_eq!(wasm_reader.remaining_bytes(), &[]);
303    }
304
305    #[test]
306    fn strip_bytes_3() {
307        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
308        let mut wasm_reader = WasmReader::new(&my_bytes);
309
310        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
311        wasm_reader.skip(2).unwrap();
312        let stripped_bytes = wasm_reader.strip_bytes::<4>();
313        assert_eq!(stripped_bytes, Err(ValidationError::Eof));
314    }
315
316    #[test]
317    fn strip_bytes_4() {
318        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
319        let mut wasm_reader = WasmReader::new(&my_bytes);
320
321        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
322        wasm_reader.skip(5).unwrap();
323        let stripped_bytes = wasm_reader.strip_bytes::<0>().unwrap();
324        assert_eq!(stripped_bytes, [0u8; 0]);
325    }
326
327    #[test]
328    fn skip_1() {
329        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
330        let mut wasm_reader = WasmReader::new(&my_bytes);
331        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
332        assert_eq!(wasm_reader.skip(6), Err(ValidationError::Eof));
333    }
334
335    #[test]
336    fn reader_transaction() {
337        let bytes = [0x1, 0x2, 0x3, 0x4, 0x5, 0x6];
338        let mut reader = WasmReader::new(&bytes);
339
340        assert_eq!(
341            reader.handle_transaction(|reader| { reader.strip_bytes::<2>() }),
342            Ok([0x1, 0x2]),
343        );
344
345        let transaction_result: Result<(), ValidationError> = reader.handle_transaction(|reader| {
346            assert_eq!(reader.strip_bytes::<2>(), Ok([0x3, 0x4]));
347
348            // The exact error type does not matter
349            Err(ValidationError::InvalidMagic)
350        });
351        assert_eq!(transaction_result, Err(ValidationError::InvalidMagic));
352
353        assert_eq!(reader.strip_bytes::<3>(), Ok([0x3, 0x4, 0x5]));
354    }
355
356    #[test]
357    fn reader_transaction_ergonomics() {
358        let bytes = [0x1, 0x2, 0x3, 0x4, 0x5, 0x6];
359        let mut reader = WasmReader::new(&bytes);
360
361        assert_eq!(reader.handle_transaction(WasmReader::read_u8), Ok(0x1));
362
363        assert_eq!(
364            reader.handle_transaction(ValType::read),
365            Err(ValidationError::MalformedValType)
366        );
367    }
368}