wasefire_scheduler/
lib.rs

1// Copyright 2022 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15#![no_std]
16#![feature(never_type)]
17#![feature(try_blocks)]
18
19extern crate alloc;
20#[cfg(feature = "std")]
21extern crate std;
22
23use alloc::collections::VecDeque;
24use alloc::vec::Vec;
25#[cfg(feature = "native")]
26use core::ffi::CStr;
27use core::marker::PhantomData;
28
29use derive_where::derive_where;
30use event::Key;
31use wasefire_applet_api::{self as api, Api, ArrayU32, Dispatch, Id, Signature, U32};
32#[cfg(feature = "board-api-storage")]
33use wasefire_board_api::Singleton;
34#[cfg(feature = "board-api-timer")]
35use wasefire_board_api::Support;
36use wasefire_board_api::{self as board, Api as Board, Failure, Trap};
37use wasefire_error::Error;
38#[cfg(feature = "wasm")]
39use wasefire_interpreter::{self as interpreter, Call, Module, RunAnswer, Val};
40use wasefire_logger as log;
41use wasefire_one_of::exactly_one_of;
42use wasefire_protocol::applet::ExitStatus;
43#[cfg(feature = "board-api-storage")]
44use wasefire_store as store;
45
46#[cfg(feature = "pulley")]
47use crate::applet::store::RunResult;
48use crate::applet::store::{Memory, Store, StoreApi};
49use crate::applet::{Applet, EventAction};
50use crate::event::InstId;
51
52mod applet;
53mod call;
54mod event;
55#[cfg(feature = "native")]
56mod native;
57#[cfg(feature = "internal-debug")]
58mod perf;
59mod protocol;
60
61#[cfg(all(feature = "native", not(target_pointer_width = "32")))]
62compile_error!("Only 32-bits architectures support native applets.");
63
64exactly_one_of!["native", "pulley", "wasm"];
65
66#[derive_where(Default)]
67pub struct Events<B: Board>(VecDeque<board::Event<B>>);
68
69impl<B: Board> Events<B> {
70    pub const fn new() -> Self {
71        Self(VecDeque::new())
72    }
73
74    pub fn is_empty(&self) -> bool {
75        self.0.is_empty()
76    }
77
78    pub fn push(&mut self, event: board::Event<B>) {
79        const MAX_EVENTS: usize = 10;
80        if self.0.contains(&event) {
81            log::trace!("Merging {:?}", event);
82        } else if self.0.len() < MAX_EVENTS {
83            log::debug!("Pushing {:?}", event);
84            self.0.push_back(event);
85        } else {
86            log::warn!("Dropping {:?}", event);
87        }
88    }
89
90    pub fn pop(&mut self) -> Option<board::Event<B>> {
91        self.0.pop_front().inspect(|event| log::debug!("Popping {:?}", event))
92    }
93}
94
95pub struct Scheduler<B: Board> {
96    #[cfg(feature = "board-api-storage")]
97    store: store::Store<B::Storage>,
98    host_funcs: Vec<Api<Id>>,
99    applet: applet::Slot<B>,
100    #[cfg(feature = "board-api-timer")]
101    timers: Vec<Option<Timer>>,
102    #[cfg(feature = "internal-debug")]
103    perf: perf::Perf<B>,
104    protocol: protocol::State,
105}
106
107#[cfg(feature = "board-api-timer")]
108#[derive(Clone)]
109struct Timer {
110    // TODO: Add AppletId.
111}
112
113impl<B: Board> core::fmt::Debug for Scheduler<B> {
114    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
115        f.debug_struct("Scheduler").finish()
116    }
117}
118
119struct SchedulerCallT<'a, B: Board> {
120    scheduler: &'a mut Scheduler<B>,
121    #[cfg(any(feature = "pulley", feature = "wasm"))]
122    args: Vec<u32>,
123    #[cfg(feature = "native")]
124    params: &'a [u32],
125    #[cfg(feature = "native")]
126    result: &'a mut i32,
127}
128
129impl<B: Board> core::fmt::Debug for SchedulerCallT<'_, B> {
130    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
131        f.debug_struct("SchedulerCallT").finish()
132    }
133}
134
135struct SchedulerCall<'a, B: Board, T> {
136    erased: SchedulerCallT<'a, B>,
137    phantom: PhantomData<T>,
138}
139
140impl<B: Board, T> core::fmt::Debug for SchedulerCall<'_, B, T> {
141    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
142        f.debug_struct("SchedulerCall").finish()
143    }
144}
145
146struct DispatchSchedulerCall<'a, B> {
147    phantom: PhantomData<&'a B>,
148}
149
150impl<'a, B: Board> Dispatch for DispatchSchedulerCall<'a, B> {
151    type Erased = SchedulerCallT<'a, B>;
152    type Merged<T: api::Signature> = SchedulerCall<'a, B, T>;
153
154    fn merge<T: api::Signature>(erased: Self::Erased) -> Self::Merged<T> {
155        SchedulerCall { erased, phantom: PhantomData }
156    }
157
158    fn erase<T: api::Signature>(merged: Self::Merged<T>) -> Self::Erased {
159        merged.erased
160    }
161}
162
163impl<'a, B: Board, T: Signature> SchedulerCall<'a, B, T> {
164    fn read(&self) -> T::Params {
165        #[cfg(any(feature = "pulley", feature = "wasm"))]
166        let params = &self.erased.args;
167        #[cfg(feature = "native")]
168        let params = self.erased.params;
169        *<T::Params as ArrayU32>::from(params)
170    }
171
172    #[cfg_attr(not(feature = "board-api-button"), allow(dead_code))]
173    fn inst(&mut self) -> InstId {
174        #[cfg(feature = "wasm")]
175        let id = self.call().inst();
176        #[cfg(any(feature = "pulley", feature = "native"))]
177        let id = InstId;
178        id
179    }
180
181    fn memory(&mut self) -> Memory<'_> {
182        self.store().memory()
183    }
184
185    fn scheduler(&mut self) -> &mut Scheduler<B> {
186        self.erased.scheduler
187    }
188
189    fn reply(self, result: Result<impl Into<U32<T::Result>>, Failure>) {
190        self.reply_(result.map(|x| x.into()))
191    }
192
193    fn reply_(mut self, result: Result<U32<T::Result>, Failure>) {
194        let result = Error::encode(match result {
195            Ok(x) => Ok(*x),
196            Err(e) => match e.split() {
197                Some(e) => Err(e),
198                None => return applet_trapped::<B>(self.scheduler(), Some(T::NAME)),
199            },
200        });
201        #[cfg(feature = "pulley")]
202        {
203            #[cfg(feature = "internal-debug")]
204            self.scheduler().perf.record(perf::Slot::Platform);
205            let result = self.store().resume(result as u32);
206            #[cfg(feature = "internal-debug")]
207            self.scheduler().perf.record(perf::Slot::Applets);
208            self.erased.scheduler.process_result(result);
209        }
210        #[cfg(feature = "wasm")]
211        {
212            let results = [Val::I32(result as u32)];
213            #[cfg(feature = "internal-debug")]
214            self.scheduler().perf.record(perf::Slot::Platform);
215            let answer = self.call().resume(&results).map(|x| x.forget());
216            #[cfg(feature = "internal-debug")]
217            self.scheduler().perf.record(perf::Slot::Applets);
218            self.erased.scheduler.process_answer(answer);
219        }
220        #[cfg(feature = "native")]
221        {
222            *self.erased.result = result;
223        }
224    }
225
226    fn applet(&mut self) -> &mut Applet<B> {
227        self.erased.scheduler.applet.get().unwrap()
228    }
229
230    fn store(&mut self) -> &mut Store {
231        self.applet().store_mut()
232    }
233
234    #[cfg(feature = "wasm")]
235    fn call(&mut self) -> Call<'_, 'static> {
236        self.store().last_call().unwrap()
237    }
238}
239
240impl<B: Board> Scheduler<B> {
241    #[cfg(any(feature = "pulley", feature = "wasm"))]
242    pub fn run() -> ! {
243        let mut scheduler = Self::new();
244        scheduler.start_applet();
245        loop {
246            log::trace!("Flushing events.");
247            scheduler.flush_events();
248            log::trace!("Processing applet.");
249            scheduler.process_applet();
250        }
251    }
252
253    #[cfg(feature = "native")]
254    pub fn run() -> ! {
255        native::set_scheduler(Self::new());
256        unsafe extern "C" {
257            unsafe fn applet_init();
258            unsafe fn applet_main();
259        }
260        #[cfg(feature = "internal-debug")]
261        native::with_scheduler(|x| x.perf_record(perf::Slot::Platform));
262        log::debug!("Execute init.");
263        unsafe { applet_init() };
264        log::debug!("Execute main.");
265        unsafe { applet_main() };
266        #[cfg(feature = "internal-debug")]
267        native::with_scheduler(|x| x.perf_record(perf::Slot::Applets));
268        log::debug!("Returned from main, executing callbacks only.");
269        loop {
270            native::with_scheduler(|scheduler| {
271                scheduler.flush_events();
272                scheduler.process_event();
273            });
274            native::execute_callback();
275        }
276    }
277
278    #[cfg(feature = "native")]
279    fn dispatch(&mut self, link: &CStr, params: *const u32) -> isize {
280        let name = link.to_str().unwrap();
281        let index = match self.host_funcs.binary_search_by_key(&name, |x| x.descriptor().name) {
282            Ok(x) => x,
283            Err(_) => {
284                let error = Error::world(wasefire_error::Code::NotImplemented);
285                return Error::encode(Err(error)) as isize;
286            }
287        };
288        let api_id = self.host_funcs[index].id();
289        let desc = api_id.descriptor();
290        let params = unsafe { core::slice::from_raw_parts(params, desc.params) };
291        let mut result = 0;
292        let erased = SchedulerCallT { scheduler: self, params, result: &mut result };
293        let call = api_id.merge(erased);
294        log::debug!("Calling {}", log::Debug2Format(&call.id()));
295        call::process(call);
296        result as isize
297    }
298
299    fn new() -> Self {
300        let mut host_funcs = Vec::new();
301        Api::<Id>::iter(&mut host_funcs, |x| x);
302        host_funcs.sort_by_key(|x| x.descriptor().name);
303        assert!(host_funcs.windows(2).all(|x| x[0].descriptor().name != x[1].descriptor().name));
304        protocol::enable::<B>();
305        Self {
306            #[cfg(feature = "board-api-storage")]
307            store: store::Store::new(board::Storage::<B>::take().unwrap()).ok().unwrap(),
308            host_funcs,
309            #[cfg(feature = "native")]
310            applet: applet::Slot::Running(Applet::default()),
311            #[cfg(any(feature = "pulley", feature = "wasm"))]
312            applet: applet::Slot::Empty,
313            #[cfg(feature = "board-api-timer")]
314            timers: alloc::vec![None; board::Timer::<B>::SUPPORT],
315            #[cfg(feature = "internal-debug")]
316            perf: perf::Perf::default(),
317            protocol: protocol::State::default(),
318        }
319    }
320
321    fn stop_applet(&mut self, status: ExitStatus) {
322        let applet = match self.applet.get() {
323            Some(x) => x,
324            None => return,
325        };
326        log::info!("Stopping applet.");
327        <board::Applet<B> as board::applet::Api>::notify_exit(status);
328        applet.free();
329        self.applet = applet::Slot::Exited(status);
330        #[cfg(feature = "native")]
331        {
332            log::debug!("Applet stopped, executing protocol events only.");
333            loop {
334                let event = B::wait_event();
335                if protocol::should_process_event(&event) {
336                    protocol::process_event(self, event);
337                }
338            }
339        }
340    }
341
342    #[cfg(any(feature = "pulley", feature = "wasm"))]
343    fn start_applet(&mut self) {
344        match self.start_applet_() {
345            Ok(()) => (),
346            Err(e) => log::warn!("Failed to start applet: {}", e),
347        }
348    }
349
350    #[cfg(feature = "pulley")]
351    fn start_applet_(&mut self) -> Result<(), Error> {
352        // SAFETY: We stop the applet before installing a new one.
353        let pulley = unsafe { <board::Applet<B> as board::applet::Api>::get()? };
354        if pulley.is_empty() {
355            log::info!("No applet to start.");
356            return Ok(());
357        }
358        // TODO: Check the wasmtime version. Or is it done by wasmtime already?
359        log::info!("Starting applet.");
360        <board::Applet<B> as board::applet::Api>::notify_start();
361        let mut store = applet::store::PreStore::default();
362        for (id, f) in self.host_funcs.iter().enumerate() {
363            let d = f.descriptor();
364            store.link_func(id, d.name, d.params)?;
365        }
366        // SAFETY: TODO: The module signature (attesting this check) is verified at install.
367        let store = unsafe { store.instantiate(pulley, self.host_funcs.len()) }?;
368        self.applet = applet::Slot::Running(Applet::new(store));
369        #[cfg(feature = "internal-debug")]
370        self.perf.record(perf::Slot::Platform);
371        self.call("init", &[]);
372        while let Some(call) = self.applet.get().unwrap().store_mut().last_call() {
373            match self.host_funcs[call.id].descriptor().name {
374                "dp" => (),
375                x => {
376                    log::warn!("init called {} into host", log::Debug2Format(&x));
377                    return Ok(applet_trapped(self, Some(x)));
378                }
379            }
380            self.process_applet();
381            if self.applet.get().is_none() {
382                return Ok(()); // applet trapped in process_applet()
383            }
384        }
385        assert!(matches!(self.applet.get().unwrap().pop(), EventAction::Reply));
386        #[cfg(feature = "internal-debug")]
387        self.perf.record(perf::Slot::Applets);
388        self.call("main", &[]);
389        Ok(())
390    }
391
392    #[cfg(feature = "wasm")]
393    fn start_applet_(&mut self) -> Result<(), Error> {
394        const MEMORY_SIZE: usize = memory_size();
395        #[repr(align(16))]
396        struct Memory([u8; MEMORY_SIZE]);
397        static mut MEMORY: Memory = Memory([0; MEMORY_SIZE]);
398
399        // SAFETY: We stop the applet before installing a new one.
400        let wasm = unsafe { <board::Applet<B> as board::applet::Api>::get()? };
401        if wasm.is_empty() {
402            log::info!("No applet to start.");
403            return Ok(());
404        }
405        log::info!("Starting applet.");
406        <board::Applet<B> as board::applet::Api>::notify_start();
407        self.applet = applet::Slot::Running(Applet::default());
408        #[cfg(not(feature = "unsafe-skip-validation"))]
409        let module = Module::new(wasm)?;
410        // SAFETY: The module is valid by the feature invariant.
411        #[cfg(feature = "unsafe-skip-validation")]
412        let module = unsafe { Module::new_unchecked(wasm) };
413        let applet = self.applet.get().unwrap();
414        let store = applet.store_mut();
415        for f in &self.host_funcs {
416            let d = f.descriptor();
417            store.link_func("env", d.name, d.params, 1)?;
418        }
419        store.link_func_default("env")?;
420        // SAFETY: We support only one applet at the moment, and the previous one is dropped when
421        // assigning to self.applet above.
422        #[allow(static_mut_refs)]
423        let inst = store.instantiate(module, unsafe { &mut MEMORY.0 })?;
424        #[cfg(feature = "internal-debug")]
425        self.perf.record(perf::Slot::Platform);
426        self.call(inst, "init", &[]);
427        while let Some(call) = self.applet.get().unwrap().store_mut().last_call() {
428            match self.host_funcs[call.index()].descriptor().name {
429                "dp" => (),
430                x => {
431                    log::warn!("init called {} into host", log::Debug2Format(&x));
432                    return Ok(applet_trapped(self, Some(x)));
433                }
434            }
435            self.process_applet();
436            if self.applet.get().is_none() {
437                return Ok(()); // applet trapped in process_applet()
438            }
439        }
440        assert!(matches!(self.applet.get().unwrap().pop(), EventAction::Reply));
441        #[cfg(feature = "internal-debug")]
442        self.perf.record(perf::Slot::Applets);
443        self.call(inst, "main", &[]);
444        Ok(())
445    }
446
447    fn flush_events(&mut self) {
448        while let Some(event) = B::try_event() {
449            self.triage_event(event);
450        }
451    }
452
453    fn triage_event(&mut self, event: board::Event<B>) {
454        if protocol::should_process_event(&event) {
455            return protocol::process_event(self, event);
456        }
457        match self.applet.get() {
458            None => log::warn!("{:?} matches no applet", event),
459            Some(applet) => applet.push(event),
460        }
461    }
462
463    /// Returns whether execution should resume.
464    fn process_event(&mut self) -> bool {
465        let Some(applet) = self.applet.get() else { return false };
466        match applet.pop() {
467            EventAction::Handle(event) => event::process(self, event),
468            EventAction::Wait => self.wait_event(),
469            #[cfg(any(feature = "pulley", feature = "wasm"))]
470            EventAction::Reply => return true,
471        }
472        false
473    }
474
475    fn wait_event(&mut self) {
476        #[cfg(feature = "internal-debug")]
477        self.perf.record(perf::Slot::Platform);
478        let event = B::wait_event();
479        #[cfg(feature = "internal-debug")]
480        self.perf.record(perf::Slot::Waiting);
481        self.triage_event(event);
482    }
483
484    #[cfg(feature = "pulley")]
485    fn process_applet(&mut self) {
486        let applet = match self.applet.get() {
487            Some(x) => x,
488            None => {
489                log::info!("There are no applets. Let's process events.");
490                while self.applet.get().is_none() {
491                    self.wait_event();
492                }
493                return;
494            }
495        };
496        let call = match applet.store_mut().last_call() {
497            Some(x) => x,
498            None => {
499                if applet.has_handlers() {
500                    // Continue processing events when main exits and at least one callback is
501                    // registered.
502                    let _ = self.process_event();
503                } else {
504                    self.stop_applet(ExitStatus::Exit);
505                }
506                return;
507            }
508        };
509        let api_id = match self.host_funcs.get(call.id) {
510            Some(x) => x.id(),
511            None => {
512                let error = Error::encode(Err(Error::world(wasefire_error::Code::NotImplemented)));
513                let result = applet.store_mut().resume(error as u32);
514                self.process_result(result);
515                return;
516            }
517        };
518        debug_assert_eq!(call.args.len(), api_id.descriptor().params);
519        let args = call.args.clone();
520        let erased = SchedulerCallT { scheduler: self, args };
521        let call = api_id.merge(erased);
522        log::debug!("Calling {}", log::Debug2Format(&call.id()));
523        call::process(call);
524    }
525
526    #[cfg(feature = "wasm")]
527    fn process_applet(&mut self) {
528        let applet = match self.applet.get() {
529            Some(x) => x,
530            None => {
531                log::info!("There are no applets. Let's process events.");
532                while self.applet.get().is_none() {
533                    self.wait_event();
534                }
535                return;
536            }
537        };
538        let call = match applet.store_mut().last_call() {
539            Some(x) => x,
540            None => {
541                if applet.has_handlers() {
542                    // Continue processing events when main exits and at least one callback is
543                    // registered.
544                    let _ = self.process_event();
545                } else {
546                    self.stop_applet(ExitStatus::Exit);
547                }
548                return;
549            }
550        };
551        let api_id = match self.host_funcs.get(call.index()) {
552            Some(x) => x.id(),
553            None => {
554                let error = Error::encode(Err(Error::world(wasefire_error::Code::NotImplemented)));
555                let answer = call.resume(&[Val::I32(error as u32)]).map(|x| x.forget());
556                self.process_answer(answer);
557                return;
558            }
559        };
560        let args = call.args();
561        debug_assert_eq!(args.len(), api_id.descriptor().params);
562        let args = args.iter().map(|x| x.unwrap_i32()).collect();
563        let erased = SchedulerCallT { scheduler: self, args };
564        let call = api_id.merge(erased);
565        log::debug!("Calling {}", log::Debug2Format(&call.id()));
566        call::process(call);
567    }
568
569    #[allow(dead_code)] // in case there are no applet-controlled events
570    fn disable_event(&mut self, key: Key<B>) -> Result<(), Trap> {
571        if let Some(applet) = self.applet.get() {
572            applet.disable(key)?;
573        }
574        self.flush_events();
575        Ok(())
576    }
577
578    #[cfg(feature = "pulley")]
579    fn call(&mut self, name: &str, args: &[u32]) {
580        log::debug!("Schedule thread {}{:?}.", name, args);
581        #[cfg(feature = "internal-debug")]
582        self.perf.record(perf::Slot::Platform);
583        let applet = self.applet.get().unwrap();
584        let result = applet.store_mut().invoke(name, args, 0);
585        #[cfg(feature = "internal-debug")]
586        self.perf.record(perf::Slot::Applets);
587        self.process_result(result);
588    }
589
590    #[cfg(feature = "wasm")]
591    fn call(&mut self, inst: InstId, name: &str, args: &[u32]) {
592        log::debug!("Schedule thread {}{:?}.", name, args);
593        let args = args.iter().map(|&x| Val::I32(x)).collect();
594        #[cfg(feature = "internal-debug")]
595        self.perf.record(perf::Slot::Platform);
596        let applet = self.applet.get().unwrap();
597        let answer = applet.store_mut().invoke(inst, name, args).map(|x| x.forget());
598        #[cfg(feature = "internal-debug")]
599        self.perf.record(perf::Slot::Applets);
600        self.process_answer(answer);
601    }
602
603    #[cfg(feature = "pulley")]
604    fn process_result(&mut self, result: Result<RunResult, Error>) {
605        match result {
606            Ok(RunResult::Done(x)) => {
607                log::debug!("Thread is done.");
608                debug_assert!(x.is_empty());
609                self.applet.get().unwrap().done();
610            }
611            Ok(RunResult::Host) => (),
612            Ok(RunResult::Trap) => applet_trapped::<B>(self, None),
613            Err(e) => log::panic!("{}", log::Debug2Format(&e)),
614        }
615    }
616
617    #[cfg(feature = "wasm")]
618    fn process_answer(&mut self, result: Result<RunAnswer, interpreter::Error>) {
619        match result {
620            Ok(RunAnswer::Done(x)) => {
621                log::debug!("Thread is done.");
622                debug_assert!(x.is_empty());
623                self.applet.get().unwrap().done();
624            }
625            Ok(RunAnswer::Host) => (),
626            Err(interpreter::Error::Trap) => applet_trapped::<B>(self, None),
627            Err(e) => log::panic!("{}", log::Debug2Format(&e)),
628        }
629    }
630}
631
632fn applet_trapped<B: Board>(scheduler: &mut Scheduler<B>, reason: Option<&'static str>) {
633    match reason {
634        None => log::warn!("Applet trapped in wasm (think segfault)."),
635        Some("sa") => log::warn!("Applet aborted (probably a panic)."),
636        Some("se") => log::info!("Applet exited."),
637        Some(name) => log::warn!("Applet trapped calling host {:?}.", name),
638    }
639    scheduler.stop_applet(match reason {
640        Some("se") => ExitStatus::Exit,
641        Some("sa") => ExitStatus::Abort,
642        _ => ExitStatus::Trap,
643    });
644}
645
646#[cfg(feature = "wasm")]
647const fn memory_size() -> usize {
648    let page = match option_env!("WASEFIRE_MEMORY_PAGE_COUNT") {
649        Some(x) => {
650            let x = x.as_bytes();
651            assert!(x.len() == 1, "not a single digit");
652            let x = x[0];
653            assert!(x.is_ascii_digit(), "not a single digit");
654            (x - b'0') as usize
655        }
656        None => 1,
657    };
658    page * 0x10000
659}