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 color_eyre::eyre::{anyhow, Context};
8use directories::ProjectDirs;
9use extism::{host_fn, Manifest, Plugin, PluginBuilder, Wasm, PTR};
10use extism_manifest::MemoryOptions;
11use log::{error, warn};
12use surfer_translation_types::plugin_types::TranslateParams;
13use surfer_translation_types::{
14    TranslationPreference, TranslationResult, Translator, VariableInfo, VariableMeta, VariableValue,
15};
16
17use crate::message::Message;
18use crate::wave_container::{ScopeId, VarId};
19
20pub fn discover_wasm_translators() -> Vec<Message> {
21    let search_dirs = [
22        std::env::current_dir()
23            .ok()
24            .map(|dir| dir.join(".surfer").join("translators")),
25        ProjectDirs::from("org", "surfer-project", "surfer")
26            .map(|dirs| dirs.data_dir().join("translators")),
27    ]
28    .into_iter()
29    .filter_map(|x| x)
30    .collect::<Vec<_>>();
31
32    let plugin_files = search_dirs
33        .into_iter()
34        .map(|dir| {
35            if !dir.exists() {
36                return vec![];
37            }
38            read_dir(&dir)
39                .map(|readdir| {
40                    readdir
41                        .filter_map(|entry| match entry {
42                            Ok(entry) => {
43                                let path = entry.path();
44                                if path.extension() == Some(&OsString::from("wasm")) {
45                                    Some(path)
46                                } else {
47                                    None
48                                }
49                            }
50                            Err(e) => {
51                                warn!("Failed to read entry in {:?}. {e}", dir.to_string_lossy());
52                                None
53                            }
54                        })
55                        .collect::<Vec<_>>()
56                })
57                .map_err(|e| {
58                    warn!(
59                        "Failed to read dir entries in {}. {e}",
60                        dir.to_string_lossy()
61                    )
62                })
63                .unwrap_or_else(|_| vec![])
64        })
65        .flatten()
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
79        .map(|file| Message::LoadWasmTranslator(file))
80        .collect()
81}
82
83pub struct PluginTranslator {
84    plugin: Arc<Mutex<Plugin>>,
85    file: PathBuf,
86}
87
88impl PluginTranslator {
89    pub fn new(file: PathBuf) -> color_eyre::Result<Self> {
90        let data = std::fs::read(&file)
91            .with_context(|| format!("Failed to read {}", file.to_string_lossy()))?;
92
93        let manifest = Manifest::new([Wasm::data(data)])
94            .with_memory_options(MemoryOptions::new().with_max_var_bytes(1024 * 1024 * 10));
95        let mut plugin = PluginBuilder::new(manifest)
96            .with_function(
97                "read_file",
98                [PTR],
99                [PTR],
100                extism::UserData::new(()),
101                read_file,
102            )
103            .with_function(
104                "file_exists",
105                [PTR],
106                [PTR],
107                extism::UserData::new(()),
108                file_exists,
109            )
110            .build()
111            .map_err(|e| anyhow!("Failed to load plugin from {} {e}", file.to_string_lossy()))?;
112
113        plugin.call::<_, ()>("new", ()).map_err(|e| {
114            anyhow!(
115                "Failed to call `new` on plugin from {}. {e}",
116                file.to_string_lossy()
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    ) -> color_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(
185        &self,
186        variable: &VariableMeta<VarId, ScopeId>,
187    ) -> color_eyre::Result<VariableInfo> {
188        let result = self
189            .plugin
190            .lock()
191            .unwrap()
192            .call("variable_info", variable.clone().map_ids(|_| (), |_| ()))
193            .map_err(|e| {
194                anyhow!(
195                    "Failed to get variable info for {} with {}. {e}",
196                    variable.var.name,
197                    self.file.to_string_lossy()
198                )
199            })?;
200        Ok(result)
201    }
202
203    fn translates(
204        &self,
205        variable: &VariableMeta<VarId, ScopeId>,
206    ) -> color_eyre::Result<TranslationPreference> {
207        match self
208            .plugin
209            .lock()
210            .unwrap()
211            .call("translates", variable.clone().map_ids(|_| (), |_| ()))
212        {
213            Ok(r) => Ok(r),
214            Err(e) => Err(anyhow!(e)),
215        }
216    }
217
218    fn reload(&self, _sender: std::sync::mpsc::Sender<Message>) {
219        let mut plugin = self.plugin.lock().unwrap();
220        if plugin.function_exists("reload") {
221            match plugin.call("reload", ()) {
222                Ok(()) => (),
223                Err(e) => error!("{e:#}"),
224            }
225        }
226    }
227}
228
229host_fn!(current_dir() -> String {
230    std::env::current_dir()
231        .with_context(|| format!("Failed to get current dir"))
232        .and_then(|dir| {
233            dir.to_str().ok_or_else(|| {
234                anyhow!("{} is not valid utf8", dir.to_string_lossy())
235            }).map(|s| s.to_string())
236        })
237        .map_err(|e| extism::Error::msg(format!("{e:#}")))
238});
239
240host_fn!(read_file(filename: String) -> Vec<u8> {
241    std::fs::read(Utf8PathBuf::from(&filename))
242        .with_context(|| format!("Failed to read {filename}"))
243        .map_err(|e| extism::Error::msg(format!("{e:#}")))
244});
245
246host_fn!(file_exists(filename: String) -> bool {
247    Ok(Utf8PathBuf::from(&filename).exists())
248});
libsurfer/translation/wasm_translator.rs

libsurfer/translation/
wasm_translator.rs
