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

libsurfer/translation/
wasm_translator.rs
