wasm/core/decoding/
reader.rs

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

wasm/core/decoding/
reader.rs
