tinywasm_types/
archive.rs

1use core::fmt::{Display, Formatter};
2
3use crate::TinyWasmModule;
4use rkyv::{
5    check_archived_root,
6    ser::{serializers::AllocSerializer, Serializer},
7    Deserialize,
8};
9
10const TWASM_MAGIC_PREFIX: &[u8; 4] = b"TWAS";
11const TWASM_VERSION: &[u8; 2] = b"01";
12#[rustfmt::skip]
13const TWASM_MAGIC: [u8; 16] = [ TWASM_MAGIC_PREFIX[0], TWASM_MAGIC_PREFIX[1], TWASM_MAGIC_PREFIX[2], TWASM_MAGIC_PREFIX[3], TWASM_VERSION[0], TWASM_VERSION[1], 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
14
15pub use rkyv::AlignedVec;
16
17fn validate_magic(wasm: &[u8]) -> Result<usize, TwasmError> {
18    if wasm.len() < TWASM_MAGIC.len() || &wasm[..TWASM_MAGIC_PREFIX.len()] != TWASM_MAGIC_PREFIX {
19        return Err(TwasmError::InvalidMagic);
20    }
21    if &wasm[TWASM_MAGIC_PREFIX.len()..TWASM_MAGIC_PREFIX.len() + TWASM_VERSION.len()] != TWASM_VERSION {
22        return Err(TwasmError::InvalidVersion);
23    }
24    if wasm[TWASM_MAGIC_PREFIX.len() + TWASM_VERSION.len()..TWASM_MAGIC.len()] != [0; 10] {
25        return Err(TwasmError::InvalidPadding);
26    }
27
28    Ok(TWASM_MAGIC.len())
29}
30
31#[derive(Debug)]
32pub enum TwasmError {
33    InvalidMagic,
34    InvalidVersion,
35    InvalidPadding,
36    InvalidArchive,
37}
38
39impl Display for TwasmError {
40    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
41        match self {
42            TwasmError::InvalidMagic => write!(f, "Invalid twasm: invalid magic number"),
43            TwasmError::InvalidVersion => write!(f, "Invalid twasm: invalid version"),
44            TwasmError::InvalidPadding => write!(f, "Invalid twasm: invalid padding"),
45            TwasmError::InvalidArchive => write!(f, "Invalid twasm: invalid archive"),
46        }
47    }
48}
49
50#[cfg(feature = "std")]
51extern crate std;
52
53#[cfg(feature = "std")]
54impl std::error::Error for TwasmError {}
55
56impl TinyWasmModule {
57    /// Creates a `TinyWasmModule` from a slice of bytes.
58    pub fn from_twasm(wasm: &[u8]) -> Result<TinyWasmModule, TwasmError> {
59        let len = validate_magic(wasm)?;
60        let root = check_archived_root::<Self>(&wasm[len..]).map_err(|_e| TwasmError::InvalidArchive)?;
61        Ok(root.deserialize(&mut rkyv::Infallible).unwrap())
62    }
63
64    /// Serializes the `TinyWasmModule` into a vector of bytes.
65    /// `AlignedVec` can be deferenced as a slice of bytes and
66    /// implements `io::Write` when the `std` feature is enabled.
67    pub fn serialize_twasm(&self) -> rkyv::AlignedVec {
68        let mut serializer = AllocSerializer::<0>::default();
69        serializer.pad(TWASM_MAGIC.len()).unwrap();
70        serializer.serialize_value(self).unwrap();
71        let mut out = serializer.into_serializer().into_inner();
72        out[..TWASM_MAGIC.len()].copy_from_slice(&TWASM_MAGIC);
73        out
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[test]
82    fn test_serialize() {
83        let wasm = TinyWasmModule::default();
84        let twasm = wasm.serialize_twasm();
85        let wasm2 = TinyWasmModule::from_twasm(&twasm).unwrap();
86        assert_eq!(wasm, wasm2);
87    }
88}