libsurfer/translation/
wasm_translator.rs

1use std::ffi::OsString;
2use std::fs::read_dir;
3use std::path::PathBuf;
4use std::sync::{Arc, Mutex};
5
6use camino::Utf8PathBuf;
7use extism::{Manifest, PTR, Plugin, PluginBuilder, Wasm, host_fn};
8use extism_manifest::MemoryOptions;
9use eyre::{Context, anyhow};
10use surfer_translation_types::plugin_types::TranslateParams;
11use surfer_translation_types::{
12    TranslationPreference, TranslationResult, Translator, VariableInfo, VariableMeta,
13    VariableNameInfo, VariableValue,
14};
15use tracing::{error, info, warn};
16
17use crate::config::{LOCAL_DIR, PROJECT_DIR};
18use crate::message::Message;
19use crate::wave_container::{ScopeId, VarId};
20
21pub static TRANSLATOR_DIR: &str = "translators";
22
23pub fn discover_wasm_translators() -> Vec<Message> {
24    let search_dirs = [
25        std::env::current_dir()
26            .ok()
27            .map(|dir| dir.join(LOCAL_DIR).join(TRANSLATOR_DIR)),
28        PROJECT_DIR
29            .as_ref()
30            .map(|dirs| dirs.data_dir().join(TRANSLATOR_DIR)),
31    ]
32    .into_iter()
33    .flatten()
34    .collect::<Vec<_>>();
35
36    let plugin_files = search_dirs
37        .into_iter()
38        .flat_map(|dir| {
39            info!("Looking for translators in {}", dir.display());
40            if !dir.exists() {
41                return vec![];
42            }
43            read_dir(&dir)
44                .map(|readdir| {
45                    readdir
46                        .filter_map(|entry| match entry {
47                            Ok(entry) => {
48                                let path = entry.path();
49                                if path.extension() == Some(&OsString::from("wasm")) {
50                                    info!("Found {}", path.display());
51                                    Some(path)
52                                } else {
53                                    None
54                                }
55                            }
56                            Err(e) => {
57                                warn!("Failed to read entry in {:?}. {e}", dir.to_string_lossy());
58                                None
59                            }
60                        })
61                        .collect::<Vec<_>>()
62                })
63                .map_err(|e| {
64                    warn!(
65                        "Failed to read dir entries in {}. {e}",
66                        dir.to_string_lossy()
67                    );
68                })
69                .unwrap_or_else(|()| vec![])
70        })
71        .filter_map(|file| {
72            file.clone()
73                .try_into()
74                .map_err(|_| {
75                    format!(
76                        "{} is not a valid UTF8 path, ignoring this translator",
77                        file.to_string_lossy()
78                    )
79                })
80                .ok()
81        });
82
83    plugin_files.map(Message::LoadWasmTranslator).collect()
84}
85
86pub struct PluginTranslator {
87    plugin: Arc<Mutex<Plugin>>,
88    file: PathBuf,
89}
90
91impl PluginTranslator {
92    pub fn new(file: PathBuf) -> eyre::Result<Self> {
93        let data = std::fs::read(&file)
94            .with_context(|| format!("Failed to read {}", file.to_string_lossy()))?;
95
96        let manifest = Manifest::new([Wasm::data(data)])
97            .with_memory_options(MemoryOptions::new().with_max_var_bytes(1024 * 1024 * 10));
98        let mut plugin = PluginBuilder::new(manifest)
99            .with_debug_info()
100            .with_function(
101                "read_file",
102                [PTR],
103                [PTR],
104                extism::UserData::new(()),
105                read_file,
106            )
107            .with_function(
108                "file_exists",
109                [PTR],
110                [PTR],
111                extism::UserData::new(()),
112                file_exists,
113            )
114            .build()
115            .map_err(|e| anyhow!("Failed to load plugin from {} {e}", file.to_string_lossy()))?;
116
117        if plugin.function_exists("new") {
118            plugin.call::<_, ()>("new", ()).map_err(|e| {
119                anyhow!(
120                    "Failed to call `new` on plugin from {}. {e}",
121                    file.to_string_lossy()
122                )
123            })?;
124        }
125
126        Ok(Self {
127            plugin: Arc::new(Mutex::new(plugin)),
128            file,
129        })
130    }
131}
132
133impl Translator<VarId, ScopeId, Message> for PluginTranslator {
134    fn name(&self) -> String {
135        self.plugin
136            .lock()
137            .unwrap()
138            .call::<_, &str>("name", ())
139            .map_err(|e| {
140                error!(
141                    "Failed to get translator name from {}. {e}",
142                    self.file.to_string_lossy()
143                );
144            })
145            .map(ToString::to_string)
146            .unwrap_or_default()
147    }
148
149    fn set_wave_source(&self, wave_source: Option<surfer_translation_types::WaveSource>) {
150        let mut plugin = self.plugin.lock().unwrap();
151        if plugin.function_exists("set_wave_source") {
152            plugin
153                .call::<_, ()>("set_wave_source", wave_source)
154                .map_err(|e| {
155                    error!(
156                        "Failed to set_wave_source on {}. {e}",
157                        self.file.to_string_lossy()
158                    );
159                })
160                .ok();
161        }
162    }
163
164    fn translate(
165        &self,
166        variable: &VariableMeta<VarId, ScopeId>,
167        value: &VariableValue,
168    ) -> eyre::Result<TranslationResult> {
169        let result = self
170            .plugin
171            .lock()
172            .unwrap()
173            .call(
174                "translate",
175                TranslateParams {
176                    variable: variable.clone().map_ids(|_| (), |_| ()),
177                    value: value.clone(),
178                },
179            )
180            .map_err(|e| {
181                anyhow!(
182                    "Failed to translate {} with {}. {e}",
183                    variable.var.name,
184                    self.file.to_string_lossy()
185                )
186            })?;
187        Ok(result)
188    }
189
190    fn variable_info(&self, variable: &VariableMeta<VarId, ScopeId>) -> eyre::Result<VariableInfo> {
191        let result = self
192            .plugin
193            .lock()
194            .unwrap()
195            .call("variable_info", variable.clone().map_ids(|_| (), |_| ()))
196            .map_err(|e| {
197                anyhow!(
198                    "Failed to get variable info for {} with {}. {e}",
199                    variable.var.name,
200                    self.file.to_string_lossy()
201                )
202            })?;
203        Ok(result)
204    }
205
206    fn translates(
207        &self,
208        variable: &VariableMeta<VarId, ScopeId>,
209    ) -> eyre::Result<TranslationPreference> {
210        match self
211            .plugin
212            .lock()
213            .unwrap()
214            .call("translates", variable.clone().map_ids(|_| (), |_| ()))
215        {
216            Ok(r) => Ok(r),
217            Err(e) => Err(anyhow!(e)),
218        }
219    }
220
221    fn reload(&self, _sender: std::sync::mpsc::Sender<Message>) {
222        let mut plugin = self.plugin.lock().unwrap();
223        if plugin.function_exists("reload") {
224            match plugin.call("reload", ()) {
225                Ok(()) => (),
226                Err(e) => error!("{e:#}"),
227            }
228        }
229    }
230
231    fn variable_name_info(
232        &self,
233        variable: &VariableMeta<VarId, ScopeId>,
234    ) -> Option<VariableNameInfo> {
235        let mut plugin = self.plugin.lock().unwrap();
236        if plugin.function_exists("variable_name_info") {
237            match plugin.call(
238                "variable_name_info",
239                variable.clone().map_ids(|_| (), |_| ()),
240            ) {
241                Ok(result) => result,
242                Err(e) => {
243                    error!("{e:#}");
244                    None
245                }
246            }
247        } else {
248            None
249        }
250    }
251}
252
253host_fn!(current_dir() -> String {
254    std::env::current_dir()
255        .with_context(|| "Failed to get current dir".to_string())
256        .and_then(|dir| {
257            dir.to_str().ok_or_else(|| {
258                anyhow!("{} is not valid utf8", dir.to_string_lossy())
259            }).map(ToString::to_string)
260        })
261        .map_err(|e| extism::Error::msg(format!("{e:#}")))
262});
263
264host_fn!(read_file(filename: String) -> Vec<u8> {
265    std::fs::read(Utf8PathBuf::from(&filename))
266        .with_context(|| format!("Failed to read {filename}"))
267        .map_err(|e| extism::Error::msg(format!("{e:#}")))
268});
269
270host_fn!(file_exists(filename: String) -> bool {
271    Ok(Utf8PathBuf::from(&filename).exists())
272});
libsurfer/translation/wasm_translator.rs

libsurfer/translation/
wasm_translator.rs
