libsurfer/translation/
wasm_translator.rs

1use std::ffi::OsString;
2use std::fs::read_dir;
3use std::path::{Path, PathBuf};
4use std::sync::{Arc, Mutex};
5
6use camino::Utf8PathBuf;
7use extism::{Manifest, PTR, Plugin, PluginBuilder, Wasm, host_fn};
8use extism_convert;
9use extism_manifest::MemoryOptions;
10use eyre::{WrapErr as _, anyhow};
11use surfer_translation_types::plugin_types::TranslateParams;
12use surfer_translation_types::{
13    TranslationPreference, TranslationResult, Translator, VariableInfo, VariableMeta,
14    VariableNameInfo, VariableValue,
15};
16use tracing::{error, info, warn};
17
18use crate::config::{LOCAL_DIR, PROJECT_DIR};
19use crate::message::Message;
20use crate::wave_container::{ScopeId, VarId};
21
22pub static TRANSLATOR_DIR: &str = "translators";
23
24pub fn discover_wasm_translators() -> Vec<Message> {
25    let search_dirs = [
26        std::env::current_dir()
27            .ok()
28            .map(|dir| dir.join(LOCAL_DIR).join(TRANSLATOR_DIR)),
29        PROJECT_DIR
30            .as_ref()
31            .map(|dirs| dirs.data_dir().join(TRANSLATOR_DIR)),
32    ]
33    .into_iter()
34    .flatten();
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    max_memory_mib: u64,
90}
91
92impl PluginTranslator {
93    pub fn new(file: PathBuf, max_memory_mib: u64) -> eyre::Result<Self> {
94        let data = std::fs::read(&file)
95            .with_context(|| format!("Failed to read {}", file.to_string_lossy()))?;
96
97        let manifest = Manifest::new([Wasm::data(data)]).with_memory_options(
98            MemoryOptions::new().with_max_var_bytes(max_memory_mib * 1024 * 1024),
99        );
100        let mut plugin = PluginBuilder::new(manifest)
101            .with_debug_info()
102            .with_function(
103                "read_file",
104                [PTR],
105                [PTR],
106                extism::UserData::new(()),
107                read_file,
108            )
109            .with_function(
110                "file_exists",
111                [PTR],
112                [PTR],
113                extism::UserData::new(()),
114                file_exists,
115            )
116            .with_function(
117                "absolute_path",
118                [PTR],
119                [PTR],
120                extism::UserData::new(()),
121                absolute_path,
122            )
123            .with_function(
124                "translators_config_dir",
125                [],
126                [PTR],
127                extism::UserData::new(()),
128                translators_config_dir,
129            )
130            .build()
131            .map_err(|e| anyhow!("Failed to load plugin from {} {e}", file.to_string_lossy()))?;
132
133        if plugin.function_exists("new") {
134            plugin
135                .call::<_, ()>("new", ())
136                .map_err(|e| Self::enrich_error(e, "new", &file, max_memory_mib))?;
137        }
138
139        Ok(Self {
140            plugin: Arc::new(Mutex::new(plugin)),
141            file,
142            max_memory_mib,
143        })
144    }
145
146    fn enrich_error(
147        e: extism::Error,
148        func: &str,
149        file: &Path,
150        max_memory_mib: u64,
151    ) -> eyre::Report {
152        let mut msg = format!(
153            "Failed to call `{func}` on plugin from {}. {e}",
154            file.to_string_lossy()
155        );
156        if e.to_string().contains("oom") {
157            msg.push_str(&format!(
158                "\nPlugin ran out of memory ({max_memory_mib} MiB). \
159                Increase `plugin.max_memory_mib` in your surfer config."
160            ));
161        }
162        anyhow!("{msg}")
163    }
164}
165
166impl Translator<VarId, ScopeId, Message> for PluginTranslator {
167    fn name(&self) -> String {
168        self.plugin
169            .lock()
170            .unwrap()
171            .call::<_, &str>("name", ())
172            .map_err(|e| {
173                error!(
174                    "{:#}",
175                    Self::enrich_error(e, "name", &self.file, self.max_memory_mib)
176                );
177            })
178            .map(ToString::to_string)
179            .unwrap_or_default()
180    }
181
182    fn set_wave_source(&self, wave_source: Option<surfer_translation_types::WaveSource>) {
183        let mut plugin = self.plugin.lock().unwrap();
184        if plugin.function_exists("set_wave_source") {
185            plugin
186                .call::<_, ()>("set_wave_source", extism_convert::Json(wave_source))
187                .map_err(|e| {
188                    error!(
189                        "{:#}",
190                        Self::enrich_error(e, "set_wave_source", &self.file, self.max_memory_mib)
191                    );
192                })
193                .ok();
194        }
195    }
196
197    fn translate(
198        &self,
199        variable: &VariableMeta<VarId, ScopeId>,
200        value: &VariableValue,
201    ) -> eyre::Result<TranslationResult> {
202        let result = self
203            .plugin
204            .lock()
205            .unwrap()
206            .call(
207                "translate",
208                TranslateParams {
209                    variable: variable.clone().map_ids(|_| (), |_| ()),
210                    value: value.clone(),
211                },
212            )
213            .map_err(|e| Self::enrich_error(e, "translate", &self.file, self.max_memory_mib))?;
214        Ok(result)
215    }
216
217    fn variable_info(&self, variable: &VariableMeta<VarId, ScopeId>) -> eyre::Result<VariableInfo> {
218        let result = self
219            .plugin
220            .lock()
221            .unwrap()
222            .call("variable_info", variable.clone().map_ids(|_| (), |_| ()))
223            .map_err(|e| Self::enrich_error(e, "variable_info", &self.file, self.max_memory_mib))?;
224        Ok(result)
225    }
226
227    fn translates(
228        &self,
229        variable: &VariableMeta<VarId, ScopeId>,
230    ) -> eyre::Result<TranslationPreference> {
231        self.plugin
232            .lock()
233            .unwrap()
234            .call("translates", variable.clone().map_ids(|_| (), |_| ()))
235            .map_err(|e| Self::enrich_error(e, "translates", &self.file, self.max_memory_mib))
236    }
237
238    fn reload(&self, _sender: std::sync::mpsc::Sender<Message>) {
239        let mut plugin = self.plugin.lock().unwrap();
240        if plugin.function_exists("reload")
241            && let Err(e) = plugin.call::<_, ()>("reload", ())
242        {
243            error!(
244                "{:#}",
245                Self::enrich_error(e, "reload", &self.file, self.max_memory_mib)
246            );
247        }
248    }
249
250    fn variable_name_info(
251        &self,
252        variable: &VariableMeta<VarId, ScopeId>,
253    ) -> Option<VariableNameInfo> {
254        let mut plugin = self.plugin.lock().unwrap();
255        if plugin.function_exists("variable_name_info") {
256            match plugin.call(
257                "variable_name_info",
258                variable.clone().map_ids(|_| (), |_| ()),
259            ) {
260                Ok(result) => result,
261                Err(e) => {
262                    error!(
263                        "{:#}",
264                        Self::enrich_error(
265                            e,
266                            "variable_name_info",
267                            &self.file,
268                            self.max_memory_mib
269                        )
270                    );
271                    None
272                }
273            }
274        } else {
275            None
276        }
277    }
278}
279
280host_fn!(current_dir() -> String {
281    std::env::current_dir()
282        .with_context(|| "Failed to get current dir".to_string())
283        .and_then(|dir| {
284            dir.to_str().ok_or_else(|| {
285                anyhow!("{} is not valid utf8", dir.to_string_lossy())
286            }).map(ToString::to_string)
287        })
288        .map_err(|e| extism::Error::msg(format!("{e:#}")))
289});
290
291host_fn!(translators_config_dir() -> extism_convert::Json(Option<String>) {
292    // Check local .surfer/translators/ first, then fall back to global config dir
293    let local = std::env::current_dir()
294        .ok()
295        .map(|dir| dir.join(LOCAL_DIR).join(TRANSLATOR_DIR))
296        .filter(|dir| dir.exists());
297
298    let global = PROJECT_DIR.as_ref()
299        .map(|dirs| dirs.config_dir().join("translators"));
300
301    Ok(extism_convert::Json(local.or(global)
302        .and_then(|dir| {
303            dir.to_str().ok_or_else(|| {
304                anyhow!("{} is not valid utf8", dir.to_string_lossy())
305            }).map(std::string::ToString::to_string).ok()
306        })))
307});
308
309host_fn!(read_file(filename: String) -> Vec<u8> {
310    std::fs::read(Utf8PathBuf::from(&filename))
311        .with_context(|| format!("Failed to read {filename}"))
312        .map_err(|e| extism::Error::msg(format!("{e:#}")))
313});
314
315host_fn!(file_exists(filename: String) -> bool {
316    Ok(Utf8PathBuf::from(&filename).exists())
317});
318
319host_fn!(absolute_path(path: String) -> String {
320    Ok(camino::absolute_utf8(Utf8PathBuf::from(&path))?.into_string())
321});
libsurfer/translation/wasm_translator.rs

libsurfer/translation/
wasm_translator.rs
