Skip to main content

tinywasm_types/
archive.rs

1use core::fmt::{Display, Formatter};
2
3use alloc::vec::Vec;
4
5use crate::Module;
6
7#[rustfmt::skip]
8const 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];
9const TWASM_MAGIC_PREFIX: &[u8; 4] = b"TWAS";
10const TWASM_VERSION: &[u8; 2] = b"05";
11
12fn validate_magic(wasm: &[u8]) -> Result<usize, TwasmError> {
13    if wasm.len() < TWASM_MAGIC.len() || &wasm[..TWASM_MAGIC_PREFIX.len()] != TWASM_MAGIC_PREFIX {
14        return Err(TwasmError::InvalidMagic);
15    }
16    if &wasm[TWASM_MAGIC_PREFIX.len()..TWASM_MAGIC_PREFIX.len() + TWASM_VERSION.len()] != TWASM_VERSION {
17        return Err(TwasmError::InvalidVersion);
18    }
19    if wasm[TWASM_MAGIC_PREFIX.len() + TWASM_VERSION.len()..TWASM_MAGIC.len()] != [0; 10] {
20        return Err(TwasmError::InvalidPadding);
21    }
22
23    Ok(TWASM_MAGIC.len())
24}
25
26#[derive(Debug, PartialEq, Eq)]
27pub enum TwasmError {
28    InvalidMagic,
29    InvalidVersion,
30    InvalidPadding,
31    InvalidArchive(postcard::Error),
32}
33
34impl Display for TwasmError {
35    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
36        match self {
37            Self::InvalidMagic => write!(f, "Invalid twasm: invalid magic number"),
38            Self::InvalidVersion => write!(f, "Invalid twasm: invalid version"),
39            Self::InvalidPadding => write!(f, "Invalid twasm: invalid padding"),
40            Self::InvalidArchive(e) => write!(f, "Invalid twasm: {e}"),
41        }
42    }
43}
44
45impl core::error::Error for TwasmError {}
46
47impl Module {
48    /// Creates a [`Module`] from internal `twasm` archive bytes.
49    ///
50    /// Archives are version-specific and are not validated as untrusted input.
51    /// Only load archives from a trusted source.
52    pub fn try_from_twasm(wasm: &[u8]) -> Result<Self, TwasmError> {
53        let len = validate_magic(wasm)?;
54        postcard::from_bytes(&wasm[len..]).map_err(TwasmError::InvalidArchive)
55    }
56
57    /// Serializes the [`Module`] into a vector of bytes.
58    pub fn serialize_twasm(&self) -> Result<Vec<u8>, TwasmError> {
59        let buf = Vec::from(TWASM_MAGIC);
60        postcard::to_extend(self, buf).map_err(TwasmError::InvalidArchive)
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67    use crate::Operand128Idx;
68    use crate::{
69        AbstractHeapType, ConstInstruction, Global, GlobalType, Instruction, ModuleFuncIdx, ModuleInner, Operand128,
70        RefType, Shared, WasmFunction, WasmType,
71    };
72    use alloc::boxed::Box;
73
74    #[test]
75    fn test_invalid_magic() {
76        let wasm = Module::default();
77        let mut twasm = wasm.serialize_twasm().expect("should serialize");
78        twasm[0] = 0;
79        assert!(matches!(Module::try_from_twasm(&twasm), Err(TwasmError::InvalidMagic)));
80    }
81
82    #[test]
83    fn test_invalid_version() {
84        let wasm = Module::default();
85        let mut twasm = wasm.serialize_twasm().expect("should serialize");
86        twasm[4] = 0;
87        assert!(matches!(Module::try_from_twasm(&twasm), Err(TwasmError::InvalidVersion)));
88    }
89
90    #[test]
91    fn v128_operands_round_trip_archive() {
92        let bytes = [0x00, 0x01, 0x02, 0x03, 0x7f, 0x80, 0xfe, 0xff, 0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x90];
93        let mut function = WasmFunction::default();
94        let constant = Operand128Idx::new(0);
95        function.data.operands128 = Box::new([Operand128::<[u8; 16]>::new(bytes).cast()]);
96        function.instructions = Box::new([Instruction::Const128(constant), Instruction::I8x16Shuffle(constant)]);
97        let module = Module::from(ModuleInner { funcs: Box::new([Shared::new(function)]), ..ModuleInner::default() });
98
99        let archive = module.serialize_twasm().expect("serialize archive");
100        assert_eq!(&archive[..6], b"TWAS05");
101        let decoded = Module::try_from_twasm(&archive).expect("deserialize archive");
102        let function = &decoded.funcs[0];
103
104        for instruction in function.instructions.iter() {
105            let index = match instruction {
106                Instruction::Const128(index) | Instruction::I8x16Shuffle(index) => *index,
107                _ => panic!("unexpected instruction"),
108            };
109            assert_eq!(index.resolve(&function.data).value(), bytes);
110        }
111    }
112
113    #[test]
114    fn const_reference_expressions_round_trip_archive() {
115        let null_type = RefType::new_abstract(true, AbstractHeapType::Extern);
116        let module = Module::from(ModuleInner {
117            globals: Box::new([
118                Global {
119                    ty: GlobalType::new(WasmType::Ref(null_type), false),
120                    init: Box::new([ConstInstruction::RefNull(null_type)]),
121                },
122                Global {
123                    ty: GlobalType::new(WasmType::Ref(RefType::FUNCREF), false),
124                    init: Box::new([ConstInstruction::RefFunc(ModuleFuncIdx::new(7))]),
125                },
126                Global {
127                    ty: GlobalType::new(WasmType::Ref(RefType::new_abstract(false, AbstractHeapType::I31)), false),
128                    init: Box::new([ConstInstruction::I32Const(42), ConstInstruction::RefI31]),
129                },
130            ]),
131            ..ModuleInner::default()
132        });
133
134        let archive = module.serialize_twasm().expect("serialize archive");
135        let decoded = Module::try_from_twasm(&archive).expect("deserialize archive");
136
137        assert!(matches!(decoded.globals[0].init.as_ref(), [ConstInstruction::RefNull(ty)] if *ty == null_type));
138        assert!(matches!(
139            decoded.globals[1].init.as_ref(),
140            [ConstInstruction::RefFunc(index)] if index.index() == 7
141        ));
142        assert!(matches!(decoded.globals[2].init.as_ref(), [ConstInstruction::I32Const(42), ConstInstruction::RefI31]));
143    }
144}