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                "translators_config_dir",
118                [],
119                [PTR],
120                extism::UserData::new(()),
121                translators_config_dir,
122            )
123            .build()
124            .map_err(|e| anyhow!("Failed to load plugin from {} {e}", file.to_string_lossy()))?;
125
126        if plugin.function_exists("new") {
127            plugin
128                .call::<_, ()>("new", ())
129                .map_err(|e| Self::enrich_error(e, "new", &file, max_memory_mib))?;
130        }
131
132        Ok(Self {
133            plugin: Arc::new(Mutex::new(plugin)),
134            file,
135            max_memory_mib,
136        })
137    }
138
139    fn enrich_error(
140        e: extism::Error,
141        func: &str,
142        file: &Path,
143        max_memory_mib: u64,
144    ) -> eyre::Report {
145        let mut msg = format!(
146            "Failed to call `{func}` on plugin from {}. {e}",
147            file.to_string_lossy()
148        );
149        if e.to_string().contains("oom") {
150            msg.push_str(&format!(
151                "\nPlugin ran out of memory ({max_memory_mib} MiB). \
152                Increase `plugin.max_memory_mib` in your surfer config."
153            ));
154        }
155        anyhow!("{msg}")
156    }
157}
158
159impl Translator<VarId, ScopeId, Message> for PluginTranslator {
160    fn name(&self) -> String {
161        self.plugin
162            .lock()
163            .unwrap()
164            .call::<_, &str>("name", ())
165            .map_err(|e| {
166                error!(
167                    "{:#}",
168                    Self::enrich_error(e, "name", &self.file, self.max_memory_mib)
169                );
170            })
171            .map(ToString::to_string)
172            .unwrap_or_default()
173    }
174
175    fn set_wave_source(&self, wave_source: Option<surfer_translation_types::WaveSource>) {
176        let mut plugin = self.plugin.lock().unwrap();
177        if plugin.function_exists("set_wave_source") {
178            plugin
179                .call::<_, ()>("set_wave_source", extism_convert::Json(wave_source))
180                .map_err(|e| {
181                    error!(
182                        "{:#}",
183                        Self::enrich_error(e, "set_wave_source", &self.file, self.max_memory_mib)
184                    );
185                })
186                .ok();
187        }
188    }
189
190    fn translate(
191        &self,
192        variable: &VariableMeta<VarId, ScopeId>,
193        value: &VariableValue,
194    ) -> eyre::Result<TranslationResult> {
195        let result = self
196            .plugin
197            .lock()
198            .unwrap()
199            .call(
200                "translate",
201                TranslateParams {
202                    variable: variable.clone().map_ids(|_| (), |_| ()),
203                    value: value.clone(),
204                },
205            )
206            .map_err(|e| Self::enrich_error(e, "translate", &self.file, self.max_memory_mib))?;
207        Ok(result)
208    }
209
210    fn variable_info(&self, variable: &VariableMeta<VarId, ScopeId>) -> eyre::Result<VariableInfo> {
211        let result = self
212            .plugin
213            .lock()
214            .unwrap()
215            .call("variable_info", variable.clone().map_ids(|_| (), |_| ()))
216            .map_err(|e| Self::enrich_error(e, "variable_info", &self.file, self.max_memory_mib))?;
217        Ok(result)
218    }
219
220    fn translates(
221        &self,
222        variable: &VariableMeta<VarId, ScopeId>,
223    ) -> eyre::Result<TranslationPreference> {
224        self.plugin
225            .lock()
226            .unwrap()
227            .call("translates", variable.clone().map_ids(|_| (), |_| ()))
228            .map_err(|e| Self::enrich_error(e, "translates", &self.file, self.max_memory_mib))
229    }
230
231    fn reload(&self, _sender: std::sync::mpsc::Sender<Message>) {
232        let mut plugin = self.plugin.lock().unwrap();
233        if plugin.function_exists("reload")
234            && let Err(e) = plugin.call::<_, ()>("reload", ())
235        {
236            error!(
237                "{:#}",
238                Self::enrich_error(e, "reload", &self.file, self.max_memory_mib)
239            );
240        }
241    }
242
243    fn variable_name_info(
244        &self,
245        variable: &VariableMeta<VarId, ScopeId>,
246    ) -> Option<VariableNameInfo> {
247        let mut plugin = self.plugin.lock().unwrap();
248        if plugin.function_exists("variable_name_info") {
249            match plugin.call(
250                "variable_name_info",
251                variable.clone().map_ids(|_| (), |_| ()),
252            ) {
253                Ok(result) => result,
254                Err(e) => {
255                    error!(
256                        "{:#}",
257                        Self::enrich_error(
258                            e,
259                            "variable_name_info",
260                            &self.file,
261                            self.max_memory_mib
262                        )
263                    );
264                    None
265                }
266            }
267        } else {
268            None
269        }
270    }
271}
272
273host_fn!(current_dir() -> String {
274    std::env::current_dir()
275        .with_context(|| "Failed to get current dir".to_string())
276        .and_then(|dir| {
277            dir.to_str().ok_or_else(|| {
278                anyhow!("{} is not valid utf8", dir.to_string_lossy())
279            }).map(ToString::to_string)
280        })
281        .map_err(|e| extism::Error::msg(format!("{e:#}")))
282});
283
284host_fn!(translators_config_dir() -> extism_convert::Json(Option<String>) {
285    // Check local .surfer/translators/ first, then fall back to global config dir
286    let local = std::env::current_dir()
287        .ok()
288        .map(|dir| dir.join(LOCAL_DIR).join(TRANSLATOR_DIR))
289        .filter(|dir| dir.exists());
290
291    let global = PROJECT_DIR.as_ref()
292        .map(|dirs| dirs.config_dir().join("translators"));
293
294    Ok(extism_convert::Json(local.or(global)
295        .and_then(|dir| {
296            dir.to_str().ok_or_else(|| {
297                anyhow!("{} is not valid utf8", dir.to_string_lossy())
298            }).map(std::string::ToString::to_string).ok()
299        })))
300});
301
302host_fn!(read_file(filename: String) -> Vec<u8> {
303    std::fs::read(Utf8PathBuf::from(&filename))
304        .with_context(|| format!("Failed to read {filename}"))
305        .map_err(|e| extism::Error::msg(format!("{e:#}")))
306});
307
308host_fn!(file_exists(filename: String) -> bool {
309    Ok(Utf8PathBuf::from(&filename).exists())
310});
libsurfer/translation/wasm_translator.rs

libsurfer/translation/
wasm_translator.rs
