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