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});