Skip to main content

hyperlight_wasm/sandbox/
wasm_sandbox.rs

1/*
2Copyright 2024 The Hyperlight Authors.
3
4Licensed under the Apache License, Version 2.0 (the "License");
5you may not use this file except in compliance with the License.
6You may obtain a copy of the License at
7
8    http://www.apache.org/licenses/LICENSE-2.0
9
10Unless required by applicable law or agreed to in writing, software
11distributed under the License is distributed on an "AS IS" BASIS,
12WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13See the License for the specific language governing permissions and
14limitations under the License.
15*/
16
17use std::path::Path;
18use std::sync::Arc;
19
20#[cfg(target_os = "linux")]
21use hyperlight_host::mem::memory_region::{MemoryRegion, MemoryRegionFlags, MemoryRegionType};
22use hyperlight_host::sandbox::snapshot::Snapshot;
23use hyperlight_host::{MultiUseSandbox, Result, new_error};
24
25use super::loaded_wasm_sandbox::LoadedWasmSandbox;
26use crate::sandbox::metrics::{
27    METRIC_ACTIVE_WASM_SANDBOXES, METRIC_SANDBOX_LOADS, METRIC_TOTAL_WASM_SANDBOXES,
28};
29
30// All the logic around when to restore is nicely encapsulated here,
31// so that it would be harder for a `WasmSandbox` to end up in an
32// un-restored state.
33mod backing_sandbox {
34    use super::*;
35    #[derive(Debug)]
36    pub(super) enum BackingSandbox {
37        /// A sandbox which has a clean copy of the runtime in it
38        Clean(MultiUseSandbox),
39        /// A sandbox which has had a wasm component/module loaded into
40        /// it, but has not yet run any code from that
41        Loaded(MultiUseSandbox),
42        /// A sandbox which came from a `LoadedWasmSandbox`, and
43        /// therefore presumably has run user code
44        Dirty(MultiUseSandbox),
45        /// A non-existent sandbox, used as an internal implementation
46        /// detail of a few methods.
47        Missing,
48    }
49    impl BackingSandbox {
50        pub(super) fn clean(&mut self, snapshot: Arc<Snapshot>) -> Result<()> {
51            *self = match std::mem::replace(self, BackingSandbox::Missing) {
52                BackingSandbox::Clean(x) => BackingSandbox::Clean(x),
53                BackingSandbox::Loaded(_) => {
54                    return Err(new_error!(
55                        "internal invariant violation: cleaning loaded backing sandbox"
56                    ));
57                }
58                BackingSandbox::Dirty(mut x) => {
59                    x.restore(snapshot)?;
60                    BackingSandbox::Clean(x)
61                }
62                BackingSandbox::Missing => {
63                    return Err(new_error!(
64                        "internal invariant violation: cleaning missing backing sandbox"
65                    ));
66                }
67            };
68            Ok(())
69        }
70        pub(super) fn load_via_restore(&mut self, snapshot: Arc<Snapshot>) -> Result<()> {
71            *self = match std::mem::replace(self, BackingSandbox::Missing) {
72                BackingSandbox::Clean(mut x) | BackingSandbox::Dirty(mut x) => {
73                    x.restore(snapshot)?;
74                    BackingSandbox::Loaded(x)
75                }
76                BackingSandbox::Loaded(_) => {
77                    return Err(new_error!(
78                        "internal invariant violation: loading loaded backing sandbox"
79                    ));
80                }
81                BackingSandbox::Missing => {
82                    return Err(new_error!(
83                        "internal invariant violation: loading missing backing sandbox"
84                    ));
85                }
86            };
87            Ok(())
88        }
89        pub(super) fn load_via_fn(
90            &mut self,
91            load: impl FnOnce(&mut MultiUseSandbox) -> Result<()>,
92        ) -> Result<()> {
93            *self = match std::mem::replace(self, BackingSandbox::Missing) {
94                BackingSandbox::Clean(mut x) => {
95                    load(&mut x)?;
96                    BackingSandbox::Loaded(x)
97                }
98                _ => {
99                    return Err(new_error!(
100                        "internal invariant violation: loading non-clean backing sandbox"
101                    ));
102                }
103            };
104            Ok(())
105        }
106        pub(super) fn get_loaded(&mut self) -> Result<MultiUseSandbox> {
107            match std::mem::replace(self, BackingSandbox::Missing) {
108                BackingSandbox::Loaded(x) => Ok(x),
109                _ => Err(new_error!(
110                    "internal invariant violation: encountered non-loaded backing sandbox"
111                )),
112            }
113        }
114    }
115
116    #[cfg(test)]
117    mod tests {
118        use super::super::tests::*;
119        use super::*;
120        #[test]
121        fn test_backing_sandbox_use_marks_dirty() -> Result<()> {
122            let mut sb = SandboxBuilder::new().build()?;
123            sb.register(
124                "GetTimeSinceBootMicrosecond",
125                get_time_since_boot_microsecond,
126            )?;
127            let sb = sb.load_runtime()?;
128            let lb = sb.load_module(get_test_file_path("RunWasm.aot")?)?;
129            let sb = lb.unload_module()?;
130            assert!(matches!(sb.inner, super::BackingSandbox::Dirty(_)));
131            Ok(())
132        }
133
134        #[test]
135        fn test_dirty_backing_sandbox_cannot_be_loaded_via_fn() -> Result<()> {
136            let mut sb = SandboxBuilder::new().build()?;
137            sb.register(
138                "GetTimeSinceBootMicrosecond",
139                get_time_since_boot_microsecond,
140            )?;
141            let sb = sb.load_runtime()?;
142            let lb = sb.load_module(get_test_file_path("RunWasm.aot")?)?;
143            let mut sb = lb.unload_module()?;
144            assert!(sb.inner.load_via_fn(|_| Ok(())).is_err());
145            Ok(())
146        }
147
148        #[test]
149        fn test_dirty_backing_sandbox_cannot_be_gotten_as_loaded() -> Result<()> {
150            let mut sb = SandboxBuilder::new().build()?;
151            sb.register(
152                "GetTimeSinceBootMicrosecond",
153                get_time_since_boot_microsecond,
154            )?;
155            let sb = sb.load_runtime()?;
156            let lb = sb.load_module(get_test_file_path("RunWasm.aot")?)?;
157            let mut sb = lb.unload_module()?;
158            assert!(sb.inner.get_loaded().is_err());
159            Ok(())
160        }
161    }
162}
163use backing_sandbox::*;
164
165/// A sandbox with just the Wasm engine loaded into memory. `WasmSandbox`es
166/// are not yet ready to execute guest functions.
167///
168/// Before you can call guest functions, you must call the `load_module`
169/// function to load a Wasm module into memory. That function will return a
170/// `LoadedWasmSandbox` able to execute code in the loaded Wasm Module.
171pub struct WasmSandbox {
172    // inner is an Option<MultiUseSandbox> as we need to take ownership of it
173    // We implement drop on the WasmSandbox to decrement the count of Sandboxes when it is dropped
174    // because of this we cannot implement drop without making inner an Option (alternatively we could make MultiUseSandbox Copy but that would introduce other issues)
175    inner: BackingSandbox,
176    // Snapshot of state of an initial WasmSandbox (runtime loaded, but no guest module code loaded).
177    // Used for LoadedWasmSandbox to be able restore state back to WasmSandbox
178    snapshot: Option<Arc<Snapshot>>,
179}
180
181const MAPPED_BINARY_VA: u64 = 0x1_0000_0000u64;
182impl WasmSandbox {
183    /// Create a new WasmSandBox from a `MultiUseSandbox`.
184    /// This function should be used to create a new `WasmSandbox` from a ProtoWasmSandbox.
185    /// The difference between this function and creating  a `WasmSandbox` directly is that
186    /// this function will increment the metrics for the number of `WasmSandbox`es in the system.
187    pub(super) fn new(mut inner: MultiUseSandbox) -> Result<Self> {
188        let snapshot = inner.snapshot()?;
189        metrics::gauge!(METRIC_ACTIVE_WASM_SANDBOXES).increment(1);
190        metrics::counter!(METRIC_TOTAL_WASM_SANDBOXES).increment(1);
191        Ok(WasmSandbox {
192            inner: BackingSandbox::Clean(inner),
193            snapshot: Some(snapshot),
194        })
195    }
196
197    /// Same as new, but doesn't take a new snapshot. Useful if `new` has already been called,
198    /// for example when creating a `WasmSandbox` from a `LoadedWasmSandbox`, since
199    /// the snapshot has already been created in that case.
200    /// Expects a snapshot of the state where wasm runtime is loaded, but no guest module code is loaded.
201    pub(super) fn new_from_loaded(
202        loaded: MultiUseSandbox,
203        snapshot: Arc<Snapshot>,
204    ) -> Result<Self> {
205        metrics::gauge!(METRIC_ACTIVE_WASM_SANDBOXES).increment(1);
206        metrics::counter!(METRIC_TOTAL_WASM_SANDBOXES).increment(1);
207        Ok(WasmSandbox {
208            inner: BackingSandbox::Dirty(loaded),
209            snapshot: Some(snapshot),
210        })
211    }
212
213    fn clean_inner(&mut self) -> Result<()> {
214        let snapshot = self.snapshot.as_ref().ok_or(new_error!(
215            "internal invariant violation: Snapshot is missing"
216        ))?;
217        self.inner.clean(snapshot.clone())
218    }
219
220    /// Load a Wasm module at the given path into the sandbox and return a `LoadedWasmSandbox`
221    /// able to execute code in the loaded Wasm Module.
222    ///
223    /// Before you can call guest functions in the sandbox, you must call
224    /// this function and use the returned value to call guest functions.
225    pub fn load_module(mut self, file: impl AsRef<Path>) -> Result<LoadedWasmSandbox> {
226        self.clean_inner()?;
227
228        self.inner.load_via_fn(|inner| {
229            if let Ok(len) = inner.map_file_cow(file.as_ref(), MAPPED_BINARY_VA) {
230                inner.call::<()>("LoadWasmModulePhys", (MAPPED_BINARY_VA, len))?;
231            } else {
232                let wasm_bytes = std::fs::read(file)?;
233                load_wasm_module_from_bytes(inner, wasm_bytes)?;
234            }
235            Ok(())
236        })?;
237
238        self.finalize_module_load()
239    }
240
241    /// Load a Wasm module by restoring a Hyperlight snapshot taken
242    /// from a `LoadedWasmSandbox`.
243    pub fn load_from_snapshot(mut self, snapshot: Arc<Snapshot>) -> Result<LoadedWasmSandbox> {
244        self.inner.load_via_restore(snapshot)?;
245
246        self.finalize_module_load()
247    }
248
249    /// Load a Wasm module that is currently present in a buffer in
250    /// host memory, by mapping the host memory directly into the
251    /// sandbox.
252    ///
253    /// Depending on the host platform, there are likely alignment
254    /// requirements of at least one page for base and len
255    ///
256    /// # Safety
257    /// It is the caller's responsibility to ensure that the host side
258    /// of the region remains intact and is not written to until the
259    /// produced LoadedWasmSandbox is discarded or devolved.
260    #[cfg(target_os = "linux")]
261    pub unsafe fn load_module_by_mapping(
262        mut self,
263        base: *mut libc::c_void,
264        len: usize,
265    ) -> Result<LoadedWasmSandbox> {
266        self.clean_inner()?;
267
268        self.inner.load_via_fn(|inner| {
269            let guest_base: usize = MAPPED_BINARY_VA as usize;
270            let rgn = MemoryRegion {
271                host_region: base as usize..base.wrapping_add(len) as usize,
272                guest_region: guest_base..guest_base + len,
273                flags: MemoryRegionFlags::READ | MemoryRegionFlags::EXECUTE,
274                region_type: MemoryRegionType::Heap,
275            };
276            if let Ok(()) = unsafe { inner.map_region(&rgn) } {
277                inner.call::<()>("LoadWasmModulePhys", (MAPPED_BINARY_VA, len as u64))?;
278            } else {
279                let wasm_bytes =
280                    unsafe { std::slice::from_raw_parts(base as *const u8, len).to_vec() };
281                load_wasm_module_from_bytes(inner, wasm_bytes)?;
282            }
283            Ok(())
284        })?;
285
286        self.finalize_module_load()
287    }
288
289    /// Load a Wasm module from a buffer of bytes into the sandbox and return a `LoadedWasmSandbox`
290    /// able to execute code in the loaded Wasm Module.
291    ///
292    /// Before you can call guest functions in the sandbox, you must call
293    /// this function and use the returned value to call guest functions.
294    pub fn load_module_from_buffer(mut self, buffer: &[u8]) -> Result<LoadedWasmSandbox> {
295        self.clean_inner()?;
296
297        // TODO: get rid of this clone
298        self.inner
299            .load_via_fn(|inner| load_wasm_module_from_bytes(inner, buffer.to_vec()))?;
300
301        self.finalize_module_load()
302    }
303
304    /// Helper function to finalize module loading and create LoadedWasmSandbox
305    fn finalize_module_load(mut self) -> Result<LoadedWasmSandbox> {
306        metrics::counter!(METRIC_SANDBOX_LOADS).increment(1);
307
308        let sandbox = self.inner.get_loaded()?;
309
310        let snapshot = self.snapshot.take().ok_or(new_error!(
311            "internal invariant violation: Snapshot is missing"
312        ))?;
313
314        LoadedWasmSandbox::new(sandbox, snapshot)
315    }
316}
317
318fn load_wasm_module_from_bytes(inner: &mut MultiUseSandbox, wasm_bytes: Vec<u8>) -> Result<()> {
319    let res: i32 = inner.call(
320        "LoadWasmModule",
321        (wasm_bytes.clone(), wasm_bytes.len() as i32),
322    )?;
323    if res != 0 {
324        return Err(new_error!(
325            "LoadWasmModule Failed with error code {:?}",
326            res
327        ));
328    }
329    Ok(())
330}
331
332impl std::fmt::Debug for WasmSandbox {
333    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
334        f.debug_struct("WasmSandbox").finish()
335    }
336}
337
338impl Drop for WasmSandbox {
339    fn drop(&mut self) {
340        metrics::gauge!(METRIC_ACTIVE_WASM_SANDBOXES).decrement(1);
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use std::env;
347    use std::path::Path;
348
349    use hyperlight_host::{HyperlightError, is_hypervisor_present};
350
351    use super::*;
352    pub(super) use crate::sandbox::sandbox_builder::SandboxBuilder;
353
354    #[test]
355    fn test_new_sandbox() -> Result<()> {
356        let _sandbox = SandboxBuilder::new().build()?;
357        Ok(())
358    }
359
360    pub(super) fn get_time_since_boot_microsecond() -> Result<i64> {
361        let res = std::time::SystemTime::now()
362            .duration_since(std::time::SystemTime::UNIX_EPOCH)?
363            .as_micros();
364        i64::try_from(res).map_err(HyperlightError::IntConversionFailure)
365    }
366
367    #[test]
368    fn test_termination() -> Result<()> {
369        let mut sandbox = SandboxBuilder::new().build()?;
370
371        sandbox.register(
372            "GetTimeSinceBootMicrosecond",
373            get_time_since_boot_microsecond,
374        )?;
375
376        let loaded = sandbox.load_runtime()?;
377
378        let run_wasm = get_test_file_path("RunWasm.aot")?;
379
380        let mut loaded = loaded.load_module(run_wasm)?;
381
382        let interrupt = loaded.interrupt_handle()?;
383
384        std::thread::spawn(move || {
385            std::thread::sleep(std::time::Duration::from_millis(1000));
386            interrupt.kill();
387        });
388
389        let result = loaded.call_guest_function::<i32>("KeepCPUBusy", 10000i32);
390
391        match result {
392            Ok(_) => panic!("Expected error"),
393            Err(e) => match e {
394                HyperlightError::ExecutionCanceledByHost() => {}
395                _ => panic!("Unexpected error: {:?}", e),
396            },
397        }
398
399        // Verify sandbox is poisoned after interruption
400        assert!(
401            loaded.status()?.is_poisoned(),
402            "Sandbox should be poisoned after interruption"
403        );
404
405        Ok(())
406    }
407
408    #[test]
409    fn test_sandbox_is_poisoned_after_interruption() -> Result<()> {
410        let mut sandbox = SandboxBuilder::new().build()?;
411
412        sandbox.register(
413            "GetTimeSinceBootMicrosecond",
414            get_time_since_boot_microsecond,
415        )?;
416
417        let loaded = sandbox.load_runtime()?;
418        let run_wasm = get_test_file_path("RunWasm.aot")?;
419        let mut loaded = loaded.load_module(run_wasm)?;
420
421        // Verify sandbox is not poisoned initially
422        assert!(
423            !loaded.status()?.is_poisoned(),
424            "Sandbox should not be poisoned initially"
425        );
426
427        let interrupt = loaded.interrupt_handle()?;
428
429        std::thread::spawn(move || {
430            std::thread::sleep(std::time::Duration::from_millis(500));
431            interrupt.kill();
432        });
433
434        // This call will be interrupted
435        let _ = loaded.call_guest_function::<i32>("KeepCPUBusy", 100000i32);
436
437        // Verify sandbox is now poisoned
438        assert!(
439            loaded.status()?.is_poisoned(),
440            "Sandbox should be poisoned after interruption"
441        );
442
443        Ok(())
444    }
445
446    #[test]
447    fn test_call_guest_function_fails_when_poisoned() -> Result<()> {
448        let mut sandbox = SandboxBuilder::new().build()?;
449
450        sandbox.register(
451            "GetTimeSinceBootMicrosecond",
452            get_time_since_boot_microsecond,
453        )?;
454
455        let loaded = sandbox.load_runtime()?;
456        let run_wasm = get_test_file_path("RunWasm.aot")?;
457        let mut loaded = loaded.load_module(run_wasm)?;
458
459        let interrupt = loaded.interrupt_handle()?;
460
461        std::thread::spawn(move || {
462            std::thread::sleep(std::time::Duration::from_millis(500));
463            interrupt.kill();
464        });
465
466        // First call will be interrupted
467        let _ = loaded.call_guest_function::<i32>("KeepCPUBusy", 100000i32);
468
469        // Second call should fail with PoisonedSandbox
470        let result = loaded.call_guest_function::<i32>("PrintOutput", 42i32);
471
472        match result {
473            Ok(_) => panic!("Expected PoisonedSandbox error"),
474            Err(HyperlightError::PoisonedSandbox) => {
475                // Expected error
476            }
477            Err(e) => panic!("Unexpected error: {:?}", e),
478        }
479
480        Ok(())
481    }
482
483    #[test]
484    fn test_snapshot_fails_when_poisoned() -> Result<()> {
485        let mut sandbox = SandboxBuilder::new().build()?;
486
487        sandbox.register(
488            "GetTimeSinceBootMicrosecond",
489            get_time_since_boot_microsecond,
490        )?;
491
492        let loaded = sandbox.load_runtime()?;
493        let run_wasm = get_test_file_path("RunWasm.aot")?;
494        let mut loaded = loaded.load_module(run_wasm)?;
495
496        let interrupt = loaded.interrupt_handle()?;
497
498        std::thread::spawn(move || {
499            std::thread::sleep(std::time::Duration::from_millis(500));
500            interrupt.kill();
501        });
502
503        // Call will be interrupted, poisoning the sandbox
504        let _ = loaded.call_guest_function::<i32>("KeepCPUBusy", 100000i32);
505
506        // Snapshot should fail on poisoned sandbox
507        let result = loaded.snapshot();
508
509        match result {
510            Ok(_) => panic!("Expected PoisonedSandbox error"),
511            Err(HyperlightError::PoisonedSandbox) => {
512                // Expected error
513            }
514            Err(e) => panic!("Unexpected error: {:?}", e),
515        }
516
517        Ok(())
518    }
519
520    #[test]
521    fn test_restore_recovers_poisoned_sandbox() -> Result<()> {
522        let mut sandbox = SandboxBuilder::new().build()?;
523
524        sandbox.register(
525            "GetTimeSinceBootMicrosecond",
526            get_time_since_boot_microsecond,
527        )?;
528
529        let loaded = sandbox.load_runtime()?;
530        let run_wasm = get_test_file_path("RunWasm.aot")?;
531        let mut loaded = loaded.load_module(run_wasm)?;
532
533        // Take a snapshot before poisoning
534        let snapshot = loaded.snapshot()?;
535
536        let interrupt = loaded.interrupt_handle()?;
537
538        std::thread::spawn(move || {
539            std::thread::sleep(std::time::Duration::from_millis(500));
540            interrupt.kill();
541        });
542
543        // Call will be interrupted, poisoning the sandbox
544        let _ = loaded.call_guest_function::<i32>("KeepCPUBusy", 100000i32);
545
546        assert!(loaded.status()?.is_poisoned(), "Sandbox should be poisoned");
547
548        // Restore should recover the sandbox
549        loaded.restore(snapshot)?;
550
551        assert!(
552            !loaded.status()?.is_poisoned(),
553            "Sandbox should not be poisoned after restore"
554        );
555
556        // Should be able to call guest functions again
557        let result: i32 = loaded.call_guest_function("CalcFib", 10i32)?;
558        assert_eq!(result, 55);
559
560        Ok(())
561    }
562
563    #[test]
564    fn test_unload_module_recovers_poisoned_sandbox() -> Result<()> {
565        let mut sandbox = SandboxBuilder::new().build()?;
566
567        sandbox.register(
568            "GetTimeSinceBootMicrosecond",
569            get_time_since_boot_microsecond,
570        )?;
571
572        let loaded = sandbox.load_runtime()?;
573        let run_wasm = get_test_file_path("RunWasm.aot")?;
574        let mut loaded = loaded.load_module(run_wasm)?;
575
576        let interrupt = loaded.interrupt_handle()?;
577
578        std::thread::spawn(move || {
579            std::thread::sleep(std::time::Duration::from_millis(500));
580            interrupt.kill();
581        });
582
583        // Call will be interrupted, poisoning the sandbox
584        let _ = loaded.call_guest_function::<i32>("KeepCPUBusy", 100000i32);
585
586        assert!(loaded.status()?.is_poisoned(), "Sandbox should be poisoned");
587
588        // unload_module should recover the sandbox (it calls restore internally)
589        let wasm_sandbox = loaded.unload_module()?;
590
591        // Should be able to load a new module and call functions
592        let helloworld_wasm = get_test_file_path("HelloWorld.aot")?;
593        let mut new_loaded = wasm_sandbox.load_module(helloworld_wasm)?;
594
595        assert!(
596            !new_loaded.status()?.is_poisoned(),
597            "New sandbox should not be poisoned"
598        );
599
600        let result: i32 = new_loaded.call_guest_function("HelloWorld", "Test".to_string())?;
601        assert_eq!(result, 0);
602
603        Ok(())
604    }
605
606    #[test]
607    fn test_load_module_file() {
608        let sandboxes = get_test_wasm_sandboxes().unwrap();
609
610        for sbox_test in sandboxes {
611            let name = sbox_test.name;
612            println!("test_load_module: {name}");
613            let wasm_sandbox = sbox_test.sbox;
614
615            let helloworld_wasm = get_test_file_path("HelloWorld.aot").unwrap();
616            let mut loaded_wasm_sandbox = wasm_sandbox.load_module(helloworld_wasm).unwrap();
617            let result: i32 = loaded_wasm_sandbox
618                .call_guest_function("HelloWorld", "Message from Rust Test".to_string())
619                .unwrap();
620
621            // TODO: Validate the output from the Wasm Modules.
622            println!("({name}) Result {:?}", result);
623        }
624    }
625
626    #[test]
627    fn test_load_from_snapshot() {
628        let mut sandbox = SandboxBuilder::new().build().unwrap();
629        sandbox
630            .register(
631                "GetTimeSinceBootMicrosecond",
632                get_time_since_boot_microsecond,
633            )
634            .unwrap();
635        let sb = sandbox.load_runtime().unwrap();
636
637        let helloworld_wasm = get_test_file_path("HelloWorld.aot").unwrap();
638        let runwasm_wasm = get_test_file_path("RunWasm.aot").unwrap();
639
640        // load one module, and make sure that a function in it
641        // can be called
642        let mut lb1 = sb.load_module(helloworld_wasm).unwrap();
643        let result: i32 = lb1
644            .call_guest_function("HelloWorld", "Message from Rust Test".to_string())
645            .unwrap();
646        assert_eq!(result, 0);
647        let snapshot = lb1.snapshot().unwrap();
648
649        // load another module, and make sure that a function in
650        // it can be called
651        let sb = lb1.unload_module().unwrap();
652        let mut lb2 = sb.load_module(runwasm_wasm).unwrap();
653        let result: i32 = lb2.call_guest_function("CalcFib", 10i32).unwrap();
654        assert_eq!(result, 55);
655
656        // reload the first module via snapshot, and make sure the
657        // original function can be called again
658        let sb = lb2.unload_module().unwrap();
659        let mut lb3 = sb.load_from_snapshot(snapshot).unwrap();
660        let result: i32 = lb3
661            .call_guest_function("HelloWorld", "Message from Rust Test".to_string())
662            .unwrap();
663        assert_eq!(result, 0);
664    }
665
666    #[test]
667    fn test_load_module_buffer() {
668        let sandboxes = get_test_wasm_sandboxes().unwrap();
669
670        for sbox_test in sandboxes {
671            let name = sbox_test.name;
672            println!("test_load_module: {name}");
673            let wasm_sandbox = sbox_test.sbox;
674
675            let wasm_module_buffer: Vec<u8> =
676                std::fs::read(get_test_file_path("HelloWorld.aot").unwrap()).unwrap();
677            let mut loaded_wasm_sandbox = wasm_sandbox
678                .load_module_from_buffer(&wasm_module_buffer)
679                .unwrap();
680            let result: i32 = loaded_wasm_sandbox
681                .call_guest_function("HelloWorld", "Message from Rust Test".to_string())
682                .unwrap();
683
684            // TODO: Validate the output from the Wasm Modules.
685            println!("({name}) Result {:?}", result);
686        }
687    }
688
689    pub(super) fn get_test_file_path(filename: &str) -> Result<String> {
690        #[cfg(debug_assertions)]
691        let config = "debug";
692        #[cfg(not(debug_assertions))]
693        let config = "release";
694        let proj_dir = env::var_os("CARGO_MANIFEST_DIR").unwrap_or_else(|| {
695            env::var_os("RUST_DIR_FOR_DEBUGGING_TESTS")
696                .expect("Failed to get CARGO_MANIFEST_DIR  or RUST_DIR_FOR_DEBUGGING_TESTS env var")
697        });
698
699        let relative_path = "../../x64";
700
701        let filename_path = Path::new(&proj_dir)
702            .join(relative_path)
703            .join(config)
704            .join(filename);
705
706        let full_path = filename_path
707            .canonicalize()
708            .unwrap()
709            .to_str()
710            .unwrap()
711            .to_string();
712
713        Ok(full_path)
714    }
715
716    struct SandboxTest {
717        sbox: WasmSandbox,
718        name: String,
719    }
720
721    fn get_test_wasm_sandboxes() -> Result<Vec<SandboxTest>> {
722        let builder = SandboxBuilder::new()
723            .with_guest_input_buffer_size(0x8000)
724            .with_guest_output_buffer_size(0x8000)
725            .with_guest_scratch_size(0x2000)
726            .with_guest_heap_size(0x100000);
727
728        let mut sandboxes: Vec<SandboxTest> = Vec::new();
729        if is_hypervisor_present() {
730            sandboxes.push(SandboxTest {
731                sbox: builder.clone().build()?.load_runtime()?,
732                name: "regular in-hypervisor".to_string(),
733            });
734        }
735
736        Ok(sandboxes)
737    }
738}