1 use clap::Parser;
2 use std::collections::HashMap;
3 use std::ffi::OsStr;
4 use std::fs;
5 use std::path::{Path, PathBuf};
6 use wasmtime::component::{Component, Linker, bindgen};
7 use wasmtime::{Engine, Store, format_err};
8 
9 // Generates bindings for the plugin world defined in the wit/calculator.wit file.
10 bindgen!("plugin");
11 
12 /// A CLI that implements a plugin-based calculator whose
13 /// operations are implemented by independent WebAssembly components.
14 
15 #[derive(Parser)]
16 struct BinaryOperation {
17     /// The name of the operation
18     op: String,
19     /// The first operand
20     x: i32,
21     /// The second operand
22     y: i32,
23 }
24 
25 fn lookup_plugin<'a>(
26     state: &'a mut CalculatorState,
27     op: &str,
28 ) -> wasmtime::Result<&'a mut PluginDesc> {
29     state
30         .plugin_descs
31         .get_mut(op)
32         .ok_or(format_err!("unknown operation {}", op))
33 }
34 
35 impl BinaryOperation {
36     fn run(self, state: &mut CalculatorState) -> wasmtime::Result<()> {
37         let desc = lookup_plugin(state, self.op.as_ref())?;
38         let result = desc.plugin.call_evaluate(&mut desc.store, self.x, self.y)?;
39         println!("{}({}, {}) = {}", self.op, self.x, self.y, result);
40         Ok(())
41     }
42 }
43 
44 pub struct CalculatorState {
45     // Mapping from operation names to descriptors for the plugins
46     pub plugin_descs: HashMap<String, PluginDesc>,
47 }
48 
49 impl Default for CalculatorState {
50     fn default() -> Self {
51         Self::new()
52     }
53 }
54 
55 impl CalculatorState {
56     pub fn new() -> Self {
57         CalculatorState {
58             plugin_descs: HashMap::new(),
59         }
60     }
61 }
62 
63 pub struct PluginDesc {
64     pub plugin: Plugin,
65     pub store: wasmtime::Store<()>,
66 }
67 
68 fn load_plugin(
69     state: &mut CalculatorState,
70     engine: &Engine,
71     linker: &Linker<()>,
72     path: PathBuf,
73 ) -> wasmtime::Result<()> {
74     println!("Loading plugin from file {:?}", path);
75 
76     // Creates a component from a .wasm file
77     let component = Component::from_file(engine, &path)?;
78 
79     let (plugin_name, plugin, store) = {
80         // Creates a Store for each plugin.
81         // A Store represents dynamic state, so each plugin
82         // has its own Store.
83         let mut store = Store::new(engine, ());
84         // Instantiate the plugin to the store we just created.
85         let plugin = Plugin::instantiate(&mut store, &component, linker)?;
86         (plugin.call_get_plugin_name(&mut store)?, plugin, store)
87     };
88 
89     state
90         .plugin_descs
91         .insert(plugin_name, PluginDesc { plugin, store });
92 
93     Ok(())
94 }
95 
96 fn load_plugins(state: &mut CalculatorState, plugins_dir: &Path) -> wasmtime::Result<()> {
97     let engine = Engine::default();
98     let linker: Linker<()> = Linker::new(&engine);
99 
100     if !plugins_dir.is_dir() {
101         wasmtime::bail!("plugins directory does not exist");
102     }
103 
104     for entry in fs::read_dir(plugins_dir)? {
105         let path = entry?.path();
106         if path.is_file() && path.extension().and_then(OsStr::to_str) == Some("wasm") {
107             load_plugin(state, &engine, &linker, path)?;
108         }
109     }
110     Ok(())
111 }
112 
113 #[derive(Parser)]
114 #[command(name = "calculator-host", version = env!("CARGO_PKG_VERSION"))]
115 #[command(about = "A calculator with plugin support")]
116 struct Args {
117     #[command(flatten)]
118     op: BinaryOperation,
119 
120     #[arg(long, help = "Plugin directory")]
121     plugins: PathBuf,
122 }
123 
124 fn main() -> wasmtime::Result<()> {
125     // Get plugins directory
126     let args = Args::parse();
127 
128     // Initialize mapping from plugin names to plugins
129     let mut state = CalculatorState::new();
130 
131     // Load plugins from plugins directory
132     load_plugins(&mut state, args.plugins.as_path())?;
133 
134     // Evaluate the expression given on the command line
135     args.op.run(&mut state)
136 }
137