kms_secp256k1_api/
wasm_loader.rs

1use std::{error::Error, sync::Arc};
2use tokio::sync::OnceCell;
3use tracing::info;
4use wasmtime::{Engine, Func, Instance, Memory, Module, Store};
5
6pub struct WasmLoader {
7    engine: Engine,
8    module: Module,
9}
10
11pub struct WasmInstance {
12    pub memory: Memory,
13    pub alloc: Func,
14    pub free: Func,
15    pub public_key: Func,
16    pub convert: Func,
17    pub unconvert: Func,
18    pub verify: Func,
19    pub verify_eip155: Func,
20    pub recover_v: Func,
21    pub address_eth: Func,
22    pub address_cosmos: Func,
23    pub store: Store<()>,
24    pub instance: Instance,
25}
26
27static WASM_INSTANCE: OnceCell<Arc<WasmLoader>> = OnceCell::const_new();
28
29impl WasmLoader {
30    /// Loads a WASM module from the specified file path and initializes a singleton instance.
31    ///
32    /// Returns an `Arc` pointing to the initialized `Self`.
33    ///
34    /// # Errors
35    ///
36    /// Returns an error if the WASM module fails to load from the given path.
37    pub async fn new(wasm_path: &str) -> Result<Arc<Self>, Box<dyn Error>> {
38        WASM_INSTANCE
39            .get_or_try_init(|| async move {
40                let engine = Engine::default();
41                let module = Module::from_file(&engine, wasm_path)?;
42                info!("WASM module loaded from {wasm_path}");
43                Ok(Arc::new(Self { engine, module }))
44            })
45            .await
46            .map(Arc::clone)
47    }
48
49    /// Instantiates the WASM module, initializing memory and required exported functions.
50    ///
51    /// Returns a `WasmInstance` containing references to the module’s memory and exported functions.
52    ///
53    /// # Errors
54    ///
55    /// Returns an error if the instance creation fails or if any of the required exports
56    /// (memory, alloc, free, `public_key`, verify, convert, unconvert) cannot be found.
57    pub fn instantiate(&self) -> Result<WasmInstance, Box<dyn Error>> {
58        let mut store = Store::new(&self.engine, ());
59        let instance = Instance::new(&mut store, &self.module, &[])?;
60
61        let memory = instance
62            .get_memory(&mut store, "memory")
63            .ok_or("failed to find memory export")?;
64
65        let alloc = instance
66            .get_func(&mut store, "alloc")
67            .ok_or("failed to find alloc export")?;
68
69        let free = instance
70            .get_func(&mut store, "free")
71            .ok_or("failed to find free export")?;
72
73        let public_key = instance
74            .get_func(&mut store, "public_key")
75            .ok_or("failed to find public_key export")?;
76
77        let verify = instance
78            .get_func(&mut store, "verify")
79            .ok_or("failed to find verify export")?;
80
81        let verify_eip155 = instance
82            .get_func(&mut store, "verify_eip155")
83            .ok_or("failed to find verify_eip155 export")?;
84
85        let convert = instance
86            .get_func(&mut store, "convert")
87            .ok_or("failed to find convert export")?;
88
89        let unconvert = instance
90            .get_func(&mut store, "unconvert")
91            .ok_or("failed to find unconvert export")?;
92
93        let recover_v = instance
94            .get_func(&mut store, "recover_v")
95            .ok_or("failed to find recover_v export")?;
96
97        let address_eth = instance
98            .get_func(&mut store, "address_eth")
99            .ok_or("failed to find address_eth export")?;
100
101        let address_cosmos = instance
102            .get_func(&mut store, "address_cosmos")
103            .ok_or("failed to find address_cosmos export")?;
104
105        Ok(WasmInstance {
106            memory,
107            alloc,
108            free,
109            public_key,
110            convert,
111            unconvert,
112            verify,
113            verify_eip155,
114            recover_v,
115            address_eth,
116            address_cosmos,
117            store,
118            instance,
119        })
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use crate::constants::WASM_PATH;
126
127    use super::*;
128
129    #[tokio::test]
130    async fn test_wasm_loader_singleton_loads_module() {
131        // First load - should succeed and initialize the singleton
132        let loader1 = WasmLoader::new(WASM_PATH).await.unwrap();
133
134        // Second load - should return the same instance
135        let loader2 = WasmLoader::new(WASM_PATH).await.unwrap();
136
137        assert!(Arc::ptr_eq(&loader1, &loader2));
138
139        let module_bytes = loader1.module.serialize().unwrap();
140        assert!(!module_bytes.is_empty(), "Module bytes should not be empty");
141    }
142
143    #[tokio::test]
144    async fn test_instantiate_wasm_instance() {
145        let loader = WasmLoader::new(WASM_PATH)
146            .await
147            .expect("Failed to load wasm module");
148
149        let wasm_instance = loader
150            .instantiate()
151            .expect("Failed to instantiate wasm instance");
152
153        // Check that memory and all functions are present
154        assert!(
155            wasm_instance.memory.size(&wasm_instance.store) > 0,
156            "Memory size should be > 0"
157        );
158
159        // Check alloc function params/results count
160        let alloc_ty = wasm_instance.alloc.ty(&wasm_instance.store);
161        assert_eq!(alloc_ty.params().len(), 1, "alloc should take 1 parameter");
162        assert_eq!(alloc_ty.results().len(), 1, "alloc should return 1 result");
163
164        // Check free function params/results count
165        let free_ty = wasm_instance.free.ty(&wasm_instance.store);
166        assert_eq!(free_ty.params().len(), 2, "free should take 2 parameters");
167        assert_eq!(free_ty.results().len(), 0, "free should return no results");
168
169        // Check public_key function is present and has some params/results
170        let public_key_ty = wasm_instance.public_key.ty(&wasm_instance.store);
171        assert!(
172            public_key_ty.params().len() > 0,
173            "public_key should have at least one parameter"
174        );
175
176        // Check verify function is present
177        let verify_ty = wasm_instance.verify.ty(&wasm_instance.store);
178        assert!(
179            verify_ty.params().len() > 0,
180            "verify should have at least one parameter"
181        );
182
183        // Check convert function is present
184        let convert_ty = wasm_instance.convert.ty(&wasm_instance.store);
185        assert!(
186            convert_ty.params().len() > 0,
187            "convert should have at least one parameter"
188        );
189
190        // Check unconvert function is present
191        let unconvert_ty = wasm_instance.unconvert.ty(&wasm_instance.store);
192        assert!(
193            unconvert_ty.params().len() > 0,
194            "unconvert should have at least one parameter"
195        );
196
197        // Check recover function is present
198        let recover_v_ty = wasm_instance.recover_v.ty(&wasm_instance.store);
199        assert!(
200            recover_v_ty.params().len() > 0,
201            "recover_v should have at least one parameter"
202        );
203
204        let address_eth_ty = wasm_instance.address_eth.ty(&wasm_instance.store);
205        assert!(
206            address_eth_ty.params().len() > 0,
207            "address_eth should have at least one parameter"
208        );
209
210        let address_cosmos_ty = wasm_instance.address_cosmos.ty(&wasm_instance.store);
211        assert!(
212            address_cosmos_ty.params().len() > 0,
213            "address_cosmos should have at least one parameter"
214        );
215    }
216}
kms_secp256k1_api/wasm_loader.rs

kms_secp256k1_api/
wasm_loader.rs
