1 #![cfg(not(miri))]
2
3 use super::REALLOC_AND_FREE;
4 use crate::call_hook::{Context, State, sync_call_hook};
5 use std::future::Future;
6 use std::pin::Pin;
7 use std::task::{self, Poll};
8 use wasmtime::component::*;
9 use wasmtime::{CallHook, CallHookHandler, Engine, Result, Store, StoreContextMut, bail};
10
11 // Crate a synchronous Func, call it directly:
12 #[test]
call_wrapped_func() -> Result<()>13 fn call_wrapped_func() -> Result<()> {
14 let wat = r#"
15 (component
16 (import "f" (func $f))
17
18 (core func $f_lower
19 (canon lower (func $f))
20 )
21 (core module $m
22 (import "" "" (func $f))
23
24 (func $export
25 (call $f)
26 )
27
28 (export "export" (func $export))
29 )
30 (core instance $i (instantiate $m
31 (with "" (instance
32 (export "" (func $f_lower))
33 ))
34 ))
35 (func (export "export")
36 (canon lift
37 (core func $i "export")
38 )
39 )
40 )
41 "#;
42
43 let engine = Engine::default();
44 let component = Component::new(&engine, wat)?;
45
46 let mut linker = Linker::<State>::new(&engine);
47 linker
48 .root()
49 .func_wrap("f", |_, _: ()| -> Result<()> { Ok(()) })?;
50
51 let mut store = Store::new(&engine, State::default());
52 store.call_hook(sync_call_hook);
53 let inst = linker
54 .instantiate(&mut store, &component)
55 .expect("instantiate");
56
57 let export = inst
58 .get_typed_func::<(), ()>(&mut store, "export")
59 .expect("looking up `export`");
60
61 export.call(&mut store, ())?;
62
63 let s = store.into_data();
64 assert_eq!(s.calls_into_host, 1);
65 assert_eq!(s.returns_from_host, 1);
66 assert_eq!(s.calls_into_wasm, 1);
67 assert_eq!(s.returns_from_wasm, 1);
68
69 Ok(())
70 }
71
72 // Call a func that turns a `list<u8>` into a `string`, to ensure that `realloc` calls are counted.
73 #[test]
call_func_with_realloc() -> Result<()>74 fn call_func_with_realloc() -> Result<()> {
75 let wat = format!(
76 r#"(component
77 (core module $m
78 (memory (export "memory") 1)
79 (func (export "roundtrip") (param i32 i32) (result i32)
80 (local $base i32)
81 (local.set $base
82 (call $realloc
83 (i32.const 0)
84 (i32.const 0)
85 (i32.const 4)
86 (i32.const 8)))
87 (i32.store offset=0
88 (local.get $base)
89 (local.get 0))
90 (i32.store offset=4
91 (local.get $base)
92 (local.get 1))
93 (local.get $base)
94 )
95
96 {REALLOC_AND_FREE}
97 )
98 (core instance $i (instantiate $m))
99
100 (func (export "list8-to-str") (param "a" (list u8)) (result string)
101 (canon lift
102 (core func $i "roundtrip")
103 (memory $i "memory")
104 (realloc (func $i "realloc"))
105 )
106 )
107 )"#
108 );
109
110 let engine = Engine::default();
111 let component = Component::new(&engine, wat)?;
112 let linker = Linker::<State>::new(&engine);
113 let mut store = Store::new(&engine, State::default());
114 store.call_hook(sync_call_hook);
115 let inst = linker
116 .instantiate(&mut store, &component)
117 .expect("instantiate");
118
119 let export = inst
120 .get_typed_func::<(&[u8],), (WasmStr,)>(&mut store, "list8-to-str")
121 .expect("looking up `list8-to-str`");
122
123 let message = String::from("hello, world!");
124 let res = export.call(&mut store, (message.as_bytes(),))?.0;
125 let result = res.to_str(&store)?;
126 assert_eq!(&message, &result);
127
128 // There are two wasm calls for the `list8-to-str` call and the guest realloc call for the list
129 // argument.
130 let s = store.into_data();
131 assert_eq!(s.calls_into_host, 0);
132 assert_eq!(s.returns_from_host, 0);
133 assert_eq!(s.calls_into_wasm, 2);
134 assert_eq!(s.returns_from_wasm, 2);
135
136 Ok(())
137 }
138
139 // Call a guest function that also defines a post-return.
140 #[test]
call_func_with_post_return() -> Result<()>141 fn call_func_with_post_return() -> Result<()> {
142 let wat = r#"
143 (component
144 (core module $m
145 (func (export "roundtrip"))
146 (func (export "post-return"))
147 )
148 (core instance $i (instantiate $m))
149
150 (func (export "export")
151 (canon lift
152 (core func $i "roundtrip")
153 (post-return (func $i "post-return"))
154 )
155 )
156 )"#;
157
158 let engine = Engine::default();
159 let component = Component::new(&engine, wat)?;
160 let linker = Linker::<State>::new(&engine);
161 let mut store = Store::new(&engine, State::default());
162 store.call_hook(sync_call_hook);
163 let inst = linker
164 .instantiate(&mut store, &component)
165 .expect("instantiate");
166
167 let export = inst
168 .get_typed_func::<(), ()>(&mut store, "export")
169 .expect("looking up `export`");
170
171 export.call(&mut store, ())?;
172
173 // There are no host calls in this example, but the post-return does increment the count of
174 // wasm calls by 1, putting the total number of wasm calls at 2.
175 let s = store.into_data();
176 assert_eq!(s.calls_into_host, 0);
177 assert_eq!(s.returns_from_host, 0);
178 assert_eq!(s.calls_into_wasm, 2);
179 assert_eq!(s.returns_from_wasm, 2);
180
181 Ok(())
182 }
183
184 // Create an async Func, call it directly:
185 #[tokio::test]
call_wrapped_async_func() -> Result<()>186 async fn call_wrapped_async_func() -> Result<()> {
187 let wat = r#"
188 (component
189 (import "f" (func $f))
190
191 (core func $f_lower
192 (canon lower (func $f))
193 )
194 (core module $m
195 (import "" "" (func $f))
196
197 (func $export
198 (call $f)
199 )
200
201 (export "export" (func $export))
202 )
203 (core instance $i (instantiate $m
204 (with "" (instance
205 (export "" (func $f_lower))
206 ))
207 ))
208 (func (export "export")
209 (canon lift
210 (core func $i "export")
211 )
212 )
213 )
214 "#;
215
216 let engine = Engine::default();
217
218 let component = Component::new(&engine, wat)?;
219
220 let mut linker = Linker::<State>::new(&engine);
221 linker
222 .root()
223 .func_wrap_async("f", |_, _: ()| Box::new(async { Ok(()) }))?;
224
225 let mut store = Store::new(&engine, State::default());
226 store.call_hook(sync_call_hook);
227
228 let inst = linker
229 .instantiate_async(&mut store, &component)
230 .await
231 .expect("instantiate");
232
233 let export = inst
234 .get_typed_func::<(), ()>(&mut store, "export")
235 .expect("looking up `export`");
236
237 export.call_async(&mut store, ()).await?;
238
239 let s = store.into_data();
240 assert_eq!(s.calls_into_host, 1);
241 assert_eq!(s.returns_from_host, 1);
242 assert_eq!(s.calls_into_wasm, 1);
243 assert_eq!(s.returns_from_wasm, 1);
244
245 Ok(())
246 }
247
248 #[test]
trapping() -> Result<()>249 fn trapping() -> Result<()> {
250 const TRAP_IN_F: i32 = 0;
251 const TRAP_NEXT_CALL_HOST: i32 = 1;
252 const TRAP_NEXT_RETURN_HOST: i32 = 2;
253 const TRAP_NEXT_CALL_WASM: i32 = 3;
254 const TRAP_NEXT_RETURN_WASM: i32 = 4;
255 const DO_NOTHING: i32 = 5;
256
257 let engine = Engine::default();
258
259 let mut linker = Linker::<State>::new(&engine);
260
261 linker
262 .root()
263 .func_wrap("f", |mut store: _, (action,): (i32,)| -> Result<()> {
264 assert_eq!(store.data().context.last(), Some(&Context::Host));
265 assert_eq!(store.data().calls_into_host, store.data().calls_into_wasm);
266
267 match action {
268 TRAP_IN_F => bail!("trapping in f"),
269 TRAP_NEXT_CALL_HOST => store.data_mut().trap_next_call_host = true,
270 TRAP_NEXT_RETURN_HOST => store.data_mut().trap_next_return_host = true,
271 TRAP_NEXT_CALL_WASM => store.data_mut().trap_next_call_wasm = true,
272 TRAP_NEXT_RETURN_WASM => store.data_mut().trap_next_return_wasm = true,
273 _ => {} // Do nothing
274 }
275
276 Ok(())
277 })?;
278
279 let wat = r#"
280 (component
281 (import "f" (func $f (param "action" s32)))
282
283 (core func $f_lower
284 (canon lower (func $f))
285 )
286 (core module $m
287 (import "" "" (func $f (param i32)))
288
289 (func $export (param i32)
290 (call $f (local.get 0))
291 )
292
293 (export "export" (func $export))
294 )
295 (core instance $i (instantiate $m
296 (with "" (instance
297 (export "" (func $f_lower))
298 ))
299 ))
300 (func (export "export") (param "action" s32)
301 (canon lift
302 (core func $i "export")
303 )
304 )
305 )
306 "#;
307
308 let component = Component::new(&engine, wat)?;
309
310 let run = |action: i32, again: bool| -> (State, Option<wasmtime::Error>) {
311 let mut store = Store::new(&engine, State::default());
312 store.call_hook(sync_call_hook);
313 let inst = linker
314 .instantiate(&mut store, &component)
315 .expect("instantiate");
316
317 let export = inst
318 .get_typed_func::<(i32,), ()>(&mut store, "export")
319 .expect("looking up `export`");
320
321 let mut r = export.call(&mut store, (action,));
322 if r.is_ok() && again {
323 r = export.call(&mut store, (action,));
324 }
325 (store.into_data(), r.err())
326 };
327
328 let (s, e) = run(DO_NOTHING, false);
329 assert!(e.is_none());
330 assert_eq!(s.calls_into_host, 1);
331 assert_eq!(s.returns_from_host, 1);
332 assert_eq!(s.calls_into_wasm, 1);
333 assert_eq!(s.returns_from_wasm, 1);
334
335 let (s, e) = run(DO_NOTHING, true);
336 assert!(e.is_none());
337 assert_eq!(s.calls_into_host, 2);
338 assert_eq!(s.returns_from_host, 2);
339 assert_eq!(s.calls_into_wasm, 2);
340 assert_eq!(s.returns_from_wasm, 2);
341
342 let (s, e) = run(TRAP_IN_F, false);
343 assert!(format!("{:?}", e.unwrap()).contains("trapping in f"));
344 assert_eq!(s.calls_into_host, 1);
345 assert_eq!(s.returns_from_host, 1);
346 assert_eq!(s.calls_into_wasm, 1);
347 assert_eq!(s.returns_from_wasm, 1);
348
349 // // trap in next call to host. No calls after the bit is set, so this trap shouldn't happen
350 let (s, e) = run(TRAP_NEXT_CALL_HOST, false);
351 assert!(e.is_none());
352 assert_eq!(s.calls_into_host, 1);
353 assert_eq!(s.returns_from_host, 1);
354 assert_eq!(s.calls_into_wasm, 1);
355 assert_eq!(s.returns_from_wasm, 1);
356
357 // trap in next call to host. call again, so the second call into host traps:
358 let (s, e) = run(TRAP_NEXT_CALL_HOST, true);
359 println!("{:?}", e.as_ref().unwrap());
360 assert!(format!("{:?}", e.unwrap()).contains("call_hook: trapping on CallingHost"));
361 assert_eq!(s.calls_into_host, 2);
362 assert_eq!(s.returns_from_host, 1);
363 assert_eq!(s.calls_into_wasm, 2);
364 assert_eq!(s.returns_from_wasm, 2);
365
366 // trap in the return from host. should trap right away, without a second call
367 let (s, e) = run(TRAP_NEXT_RETURN_HOST, false);
368 assert!(format!("{:?}", e.unwrap()).contains("call_hook: trapping on ReturningFromHost"));
369 assert_eq!(s.calls_into_host, 1);
370 assert_eq!(s.returns_from_host, 1);
371 assert_eq!(s.calls_into_wasm, 1);
372 assert_eq!(s.returns_from_wasm, 1);
373
374 // trap in next call to wasm. No calls after the bit is set, so this trap shouldn't happen:
375 let (s, e) = run(TRAP_NEXT_CALL_WASM, false);
376 assert!(e.is_none());
377 assert_eq!(s.calls_into_host, 1);
378 assert_eq!(s.returns_from_host, 1);
379 assert_eq!(s.calls_into_wasm, 1);
380 assert_eq!(s.returns_from_wasm, 1);
381
382 // trap in next call to wasm. call again, so the second call into wasm traps:
383 let (s, e) = run(TRAP_NEXT_CALL_WASM, true);
384 assert!(format!("{:?}", e.unwrap()).contains("call_hook: trapping on CallingWasm"));
385 assert_eq!(s.calls_into_host, 1);
386 assert_eq!(s.returns_from_host, 1);
387 assert_eq!(s.calls_into_wasm, 2);
388 assert_eq!(s.returns_from_wasm, 1);
389
390 // trap in the return from wasm. should trap right away, without a second call
391 let (s, e) = run(TRAP_NEXT_RETURN_WASM, false);
392 assert!(format!("{:?}", e.unwrap()).contains("call_hook: trapping on ReturningFromWasm"));
393 assert_eq!(s.calls_into_host, 1);
394 assert_eq!(s.returns_from_host, 1);
395 assert_eq!(s.calls_into_wasm, 1);
396 assert_eq!(s.returns_from_wasm, 1);
397
398 Ok(())
399 }
400
401 #[tokio::test]
timeout_async_hook() -> Result<()>402 async fn timeout_async_hook() -> Result<()> {
403 struct HandlerR;
404
405 #[async_trait::async_trait]
406 impl CallHookHandler<State> for HandlerR {
407 async fn handle_call_event(
408 &self,
409 mut ctx: StoreContextMut<'_, State>,
410 ch: CallHook,
411 ) -> Result<()> {
412 let obj = ctx.data_mut();
413 if obj.calls_into_host > 200 {
414 bail!("timeout");
415 }
416
417 match ch {
418 CallHook::CallingHost => obj.calls_into_host += 1,
419 CallHook::CallingWasm => obj.calls_into_wasm += 1,
420 CallHook::ReturningFromHost => obj.returns_from_host += 1,
421 CallHook::ReturningFromWasm => obj.returns_from_wasm += 1,
422 }
423
424 Ok(())
425 }
426 }
427
428 let wat = r#"
429 (component
430 (import "f" (func $f))
431
432 (core func $f_lower
433 (canon lower (func $f))
434 )
435 (core module $m
436 (import "" "" (func $f))
437
438 (func $export
439 (loop $start
440 (call $f)
441 (br $start))
442 )
443
444 (export "export" (func $export))
445 )
446 (core instance $i (instantiate $m
447 (with "" (instance
448 (export "" (func $f_lower))
449 ))
450 ))
451 (func (export "export")
452 (canon lift
453 (core func $i "export")
454 )
455 )
456 )
457 "#;
458
459 let engine = Engine::default();
460
461 let component = Component::new(&engine, wat)?;
462
463 let mut linker = Linker::<State>::new(&engine);
464 linker
465 .root()
466 .func_wrap_async("f", |_, _: ()| Box::new(async { Ok(()) }))?;
467
468 let mut store = Store::new(&engine, State::default());
469 store.call_hook_async(HandlerR {});
470
471 let inst = linker
472 .instantiate_async(&mut store, &component)
473 .await
474 .expect("instantiate");
475
476 let export = inst
477 .get_typed_func::<(), ()>(&mut store, "export")
478 .expect("looking up `export`");
479
480 let r = export.call_async(&mut store, ()).await;
481 assert!(format!("{:?}", r.unwrap_err()).contains("timeout"));
482
483 let s = store.into_data();
484 assert!(s.calls_into_host > 1);
485 assert!(s.returns_from_host > 1);
486 assert_eq!(s.calls_into_wasm, 1);
487 assert_eq!(s.returns_from_wasm, 0);
488
489 Ok(())
490 }
491
492 #[tokio::test]
drop_suspended_async_hook() -> Result<()>493 async fn drop_suspended_async_hook() -> Result<()> {
494 struct Handler;
495
496 #[async_trait::async_trait]
497 impl CallHookHandler<u32> for Handler {
498 async fn handle_call_event(
499 &self,
500 mut ctx: StoreContextMut<'_, u32>,
501 _ch: CallHook,
502 ) -> Result<()> {
503 let state = ctx.data_mut();
504 assert_eq!(*state, 0);
505 *state += 1;
506 let _dec = Decrement(state);
507
508 // Simulate some sort of event which takes a number of yields
509 for _ in 0..500 {
510 tokio::task::yield_now().await;
511 }
512 Ok(())
513 }
514 }
515
516 let wat = r#"
517 (component
518 (import "f" (func $f))
519
520 (core func $f_lower
521 (canon lower (func $f))
522 )
523 (core module $m
524 (import "" "" (func $f))
525
526 (func $export
527 (call $f)
528 )
529
530 (export "export" (func $export))
531 )
532 (core instance $i (instantiate $m
533 (with "" (instance
534 (export "" (func $f_lower))
535 ))
536 ))
537 (func (export "export")
538 (canon lift
539 (core func $i "export")
540 )
541 )
542 )
543 "#;
544
545 let engine = Engine::default();
546
547 let component = Component::new(&engine, wat)?;
548
549 let mut linker = Linker::<u32>::new(&engine);
550 linker.root().func_wrap_async("f", |mut store, _: ()| {
551 Box::new(async move {
552 let state = store.data_mut();
553 assert_eq!(*state, 0);
554 *state += 1;
555 let _dec = Decrement(state);
556 for _ in 0.. {
557 tokio::task::yield_now().await;
558 }
559 Ok(())
560 })
561 })?;
562
563 let mut store = Store::new(&engine, 0);
564 store.call_hook_async(Handler);
565
566 let inst = linker
567 .instantiate_async(&mut store, &component)
568 .await
569 .expect("instantiate");
570
571 let export = inst
572 .get_typed_func::<(), ()>(&mut store, "export")
573 .expect("looking up `export`");
574
575 // Test that if we drop in the middle of an async hook that everything
576 // is alright.
577 PollNTimes {
578 future: Box::pin(export.call_async(&mut store, ())),
579 times: 200,
580 }
581 .await;
582 assert_eq!(*store.data(), 0); // double-check user dtors ran
583
584 return Ok(());
585
586 // A helper struct to poll an inner `future` N `times` and then resolve.
587 // This is used above to test that when futures are dropped while they're
588 // pending everything works and is cleaned up on the Wasmtime side of
589 // things.
590 struct PollNTimes<F> {
591 future: F,
592 times: u32,
593 }
594
595 impl<F: Future + Unpin> Future for PollNTimes<F>
596 where
597 F::Output: std::fmt::Debug,
598 {
599 type Output = ();
600 fn poll(mut self: Pin<&mut Self>, task: &mut task::Context<'_>) -> Poll<()> {
601 for i in 0..self.times {
602 match Pin::new(&mut self.future).poll(task) {
603 Poll::Ready(v) => panic!("future should not be ready at {i}; result is {v:?}"),
604 Poll::Pending => {}
605 }
606 }
607
608 Poll::Ready(())
609 }
610 }
611
612 // helper struct to decrement a counter on drop
613 struct Decrement<'a>(&'a mut u32);
614
615 impl Drop for Decrement<'_> {
616 fn drop(&mut self) {
617 *self.0 -= 1;
618 }
619 }
620 }
621