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    #[allow(dead_code)]
141    pub fn skip(&mut self, num_bytes: usize) -> Result<(), ValidationError> {
142        if num_bytes > self.full_wasm_binary.len() - self.pc {
143            return Err(ValidationError::Eof);
144        }
145        self.pc += num_bytes;
146        Ok(())
147    }
148
149    /// Consumes [Self], yielding back the internal reference to the WASM binary
150    pub fn into_inner(self) -> &'a [u8] {
151        self.full_wasm_binary
152    }
153
154    /// A wrapper function for reads with transaction-like behavior.
155    ///
156    /// The provided closure will be called with `&mut self` and its result will be returned.
157    /// However if the closure returns `Err(_)`, `self` will be reset as if the closure was never called.
158    #[allow(dead_code)]
159    pub fn handle_transaction<T, E>(
160        &mut self,
161        f: impl FnOnce(&mut WasmReader<'a>) -> Result<T, E>,
162    ) -> Result<T, E> {
163        let original = self.clone();
164        f(self).inspect_err(|_| {
165            *self = original;
166        })
167    }
168}
169
170pub trait WasmReadable: Sized {
171    /// Reads a new [`Self`] from given [`WasmReader`].
172    ///
173    /// Note that if this function returns `Err(_)`, the [`WasmReader`] may still have been advanced,
174    /// which may lead to unexpected behaviour.
175    /// To avoid this consider using the [`WasmReader::handle_transaction`] method to wrap this function call.
176    fn read(wasm: &mut WasmReader) -> Result<Self, ValidationError>;
177}
178
179pub mod span {
180    use core::ops::Index;
181
182    use crate::core::reader::WasmReader;
183
184    /// An index and offset to describe a (sub-) slice into WASM bytecode
185    ///
186    /// Can be used to index into a [WasmReader], yielding a byte slice. As it does not
187    /// actually own the indexed data, this struct is free of lifetimes. Caution is advised when
188    /// indexing unknown slices, as a [Span] does not validate the length of the indexed slice.
189    #[derive(Copy, Clone, Debug, Hash)]
190    pub struct Span {
191        pub from: usize,
192        pub len: usize,
193    }
194
195    impl Span {
196        /// Create a new [Span], starting from `from` and ranging `len` elements
197        pub const fn new(from: usize, len: usize) -> Self {
198            Self { from, len }
199        }
200
201        /// Returns the length of this [Span]
202        pub const fn len(&self) -> usize {
203            self.len
204        }
205
206        pub const fn from(&self) -> usize {
207            self.from
208        }
209    }
210
211    impl<'a> Index<Span> for WasmReader<'a> {
212        type Output = [u8];
213
214        fn index(&self, index: Span) -> &'a Self::Output {
215            &self.full_wasm_binary[index.from..(index.from + index.len)]
216        }
217    }
218}
219
220#[cfg(test)]
221mod test {
222    use crate::ValType;
223
224    use super::*;
225    use alloc::vec;
226
227    #[test]
228    fn move_start_to() {
229        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
230        let mut wasm_reader = WasmReader::new(&my_bytes);
231
232        let span = Span::new(0, 0);
233        wasm_reader.move_start_to(span).unwrap();
234        // this actually dangerous, we did not validate there to be more than 0 bytes using the Span
235        wasm_reader.peek_u8().unwrap();
236
237        let span = Span::new(0, my_bytes.len());
238        wasm_reader.move_start_to(span).unwrap();
239        wasm_reader.peek_u8().unwrap();
240        assert_eq!(wasm_reader[span], my_bytes);
241
242        let span = Span::new(my_bytes.len(), 0);
243        wasm_reader.move_start_to(span).unwrap();
244        // span had zero length, hence wasm_reader.peek_u8() would be allowed to fail
245
246        let span = Span::new(my_bytes.len() - 1, 1);
247        wasm_reader.move_start_to(span).unwrap();
248
249        assert_eq!(wasm_reader.peek_u8().unwrap(), *my_bytes.last().unwrap());
250    }
251
252    #[test]
253    fn move_start_to_out_of_bounds_1() {
254        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
255        let mut wasm_reader = WasmReader::new(&my_bytes);
256
257        let span = Span::new(my_bytes.len(), 1);
258        assert_eq!(wasm_reader.move_start_to(span), Err(ValidationError::Eof));
259    }
260
261    #[test]
262    fn move_start_to_out_of_bounds_2() {
263        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
264        let mut wasm_reader = WasmReader::new(&my_bytes);
265
266        let span = Span::new(0, my_bytes.len() + 1);
267        assert_eq!(wasm_reader.move_start_to(span), Err(ValidationError::Eof));
268    }
269
270    #[test]
271    fn remaining_bytes_1() {
272        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
273        let mut wasm_reader = WasmReader::new(&my_bytes);
274
275        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
276        wasm_reader.skip(4).unwrap();
277        assert_eq!(wasm_reader.peek_u8().unwrap(), 0x15);
278
279        assert_eq!(wasm_reader.remaining_bytes(), &my_bytes[4..]);
280    }
281
282    #[test]
283    fn remaining_bytes_2() {
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        wasm_reader.skip(5).unwrap();
289        assert_eq!(wasm_reader.remaining_bytes(), &my_bytes[5..]);
290        assert_eq!(wasm_reader.remaining_bytes(), &[]);
291    }
292
293    #[test]
294    fn strip_bytes_1() {
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        let stripped_bytes = wasm_reader.strip_bytes::<4>().unwrap();
300        assert_eq!(&stripped_bytes, &my_bytes[..4]);
301        assert_eq!(wasm_reader.remaining_bytes(), &[0x15]);
302    }
303
304    #[test]
305    fn strip_bytes_2() {
306        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
307        let mut wasm_reader = WasmReader::new(&my_bytes);
308
309        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
310        wasm_reader.skip(1).unwrap();
311        let stripped_bytes = wasm_reader.strip_bytes::<4>().unwrap();
312        assert_eq!(&stripped_bytes, &my_bytes[1..5]);
313        assert_eq!(wasm_reader.remaining_bytes(), &[]);
314    }
315
316    #[test]
317    fn strip_bytes_3() {
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(2).unwrap();
323        let stripped_bytes = wasm_reader.strip_bytes::<4>();
324        assert_eq!(stripped_bytes, Err(ValidationError::Eof));
325    }
326
327    #[test]
328    fn strip_bytes_4() {
329        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
330        let mut wasm_reader = WasmReader::new(&my_bytes);
331
332        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
333        wasm_reader.skip(5).unwrap();
334        let stripped_bytes = wasm_reader.strip_bytes::<0>().unwrap();
335        assert_eq!(stripped_bytes, [0u8; 0]);
336    }
337
338    #[test]
339    fn skip_1() {
340        let my_bytes = vec![0x11, 0x12, 0x13, 0x14, 0x15];
341        let mut wasm_reader = WasmReader::new(&my_bytes);
342        assert_eq!(wasm_reader.remaining_bytes(), my_bytes);
343        assert_eq!(wasm_reader.skip(6), Err(ValidationError::Eof));
344    }
345
346    #[test]
347    fn reader_transaction() {
348        let bytes = [0x1, 0x2, 0x3, 0x4, 0x5, 0x6];
349        let mut reader = WasmReader::new(&bytes);
350
351        assert_eq!(
352            reader.handle_transaction(|reader| { reader.strip_bytes::<2>() }),
353            Ok([0x1, 0x2]),
354        );
355
356        let transaction_result: Result<(), ValidationError> = reader.handle_transaction(|reader| {
357            assert_eq!(reader.strip_bytes::<2>(), Ok([0x3, 0x4]));
358
359            // The exact error type does not matter
360            Err(ValidationError::InvalidMagic)
361        });
362        assert_eq!(transaction_result, Err(ValidationError::InvalidMagic));
363
364        assert_eq!(reader.strip_bytes::<3>(), Ok([0x3, 0x4, 0x5]));
365    }
366
367    #[test]
368    fn reader_transaction_ergonomics() {
369        let bytes = [0x1, 0x2, 0x3, 0x4, 0x5, 0x6];
370        let mut reader = WasmReader::new(&bytes);
371
372        assert_eq!(reader.handle_transaction(WasmReader::read_u8), Ok(0x1));
373
374        assert_eq!(
375            reader.handle_transaction(ValType::read),
376            Err(ValidationError::InvalidValType)
377        );
378    }
379}