Skip to main content

crawlkit_engine/
plugin.rs

1use std::path::{Path, PathBuf};
2use std::sync::Arc;
3
4use parking_lot::RwLock;
5use serde::{Deserialize, Serialize};
6use thiserror::Error;
7
8/// Plugin errors.
9#[derive(Debug, Error)]
10pub enum PluginError {
11    #[error("plugin not found: {0}")]
12    NotFound(String),
13
14    #[error("plugin load failed: {0}")]
15    LoadFailed(String),
16
17    #[error("plugin init failed: {0}")]
18    InitFailed(String),
19
20    #[error("plugin analysis failed: {0}")]
21    AnalysisFailed(String),
22
23    #[error("incompatible API version: {0} (expected 1.0)")]
24    IncompatibleApiVersion(String),
25
26    #[error("manifest parse error: {0}")]
27    ManifestParse(String),
28
29    #[error("WASM execution error: {0}")]
30    WasmExecution(String),
31}
32
33/// Plugin manifest (crawlkit-plugin.toml).
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct PluginManifest {
36    pub plugin: PluginMetadata,
37}
38
39/// Plugin metadata from manifest.
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct PluginMetadata {
42    pub name: String,
43    pub version: String,
44    pub api_version: String,
45    pub author: String,
46    pub description: String,
47    pub license: String,
48    pub trust_level: Option<String>,
49    pub entry: PluginEntry,
50    pub permissions: Option<PluginPermissions>,
51    pub analyzer: Option<PluginAnalyzerInfo>,
52}
53
54/// Plugin entry point configuration.
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct PluginEntry {
57    pub wasm: Option<String>,
58    pub native: Option<String>,
59}
60
61/// WASM plugin permissions.
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct PluginPermissions {
64    pub network: Option<bool>,
65    pub filesystem: Option<bool>,
66    pub env_vars: Option<Vec<String>>,
67}
68
69/// Plugin analyzer metadata.
70#[derive(Debug, Clone, Serialize, Deserialize)]
71pub struct PluginAnalyzerInfo {
72    pub name: String,
73    pub category: Option<String>,
74    pub description: Option<String>,
75    pub severity: Option<String>,
76}
77
78/// Loaded WASM plugin instance.
79pub struct WasmPlugin {
80    pub manifest: PluginMetadata,
81    store: wasmtime::Store<()>,
82    instance: wasmtime::Instance,
83    memory: wasmtime::Memory,
84}
85
86impl WasmPlugin {
87    /// Load a WASM plugin from a directory.
88    pub fn load(plugin_dir: &Path) -> Result<Self, PluginError> {
89        let manifest_path = plugin_dir.join("crawlkit-plugin.toml");
90        let manifest_str = std::fs::read_to_string(&manifest_path)
91            .map_err(|e| PluginError::ManifestParse(format!("Failed to read manifest: {}", e)))?;
92        let manifest: PluginManifest = toml::from_str(&manifest_str)
93            .map_err(|e| PluginError::ManifestParse(format!("Invalid manifest: {}", e)))?;
94
95        if !manifest.plugin.api_version.starts_with("1.") {
96            return Err(PluginError::IncompatibleApiVersion(
97                manifest.plugin.api_version,
98            ));
99        }
100
101        let wasm_file =
102            manifest.plugin.entry.wasm.as_ref().ok_or_else(|| {
103                PluginError::LoadFailed("No WASM entry point specified".to_string())
104            })?;
105        let wasm_path = plugin_dir.join(wasm_file);
106
107        let engine = wasmtime::Engine::default();
108        let module = wasmtime::Module::from_file(&engine, &wasm_path)
109            .map_err(|e| PluginError::LoadFailed(format!("Failed to compile WASM: {}", e)))?;
110
111        let mut store = wasmtime::Store::new(&engine, ());
112        let linker = wasmtime::Linker::new(&engine);
113
114        let instance = linker
115            .instantiate(&mut store, &module)
116            .map_err(|e| PluginError::LoadFailed(format!("Failed to instantiate WASM: {}", e)))?;
117
118        let memory = instance
119            .get_memory(&mut store, "memory")
120            .ok_or_else(|| PluginError::LoadFailed("No memory export found".to_string()))?;
121
122        let init_func = instance
123            .get_typed_func::<i32, i32>(&mut store, "crawlkit_plugin_init")
124            .map_err(|e| PluginError::InitFailed(format!("Init function not found: {}", e)))?;
125
126        let result = init_func
127            .call(&mut store, 0)
128            .map_err(|e| PluginError::InitFailed(format!("Init failed: {}", e)))?;
129
130        if result != 0 {
131            return Err(PluginError::InitFailed(format!(
132                "Init returned error code: {}",
133                result
134            )));
135        }
136
137        Ok(Self {
138            manifest: manifest.plugin,
139            store,
140            instance,
141            memory,
142        })
143    }
144
145    /// Analyze HTML content using the plugin.
146    pub fn analyze(&mut self, html: &str, url: &str) -> Result<String, PluginError> {
147        let analyze_func = self
148            .instance
149            .get_typed_func::<(i32, i32, i32, i32), i32>(&mut self.store, "crawlkit_plugin_analyze")
150            .map_err(|e| {
151                PluginError::AnalysisFailed(format!("Analyze function not found: {}", e))
152            })?;
153
154        let html_bytes = html.as_bytes();
155        let url_bytes = url.as_bytes();
156
157        let alloc_func = self
158            .instance
159            .get_typed_func::<i32, i32>(&mut self.store, "crawlkit_plugin_alloc")
160            .map_err(|e| PluginError::AnalysisFailed(format!("Alloc function not found: {}", e)))?;
161
162        let html_ptr = alloc_func
163            .call(&mut self.store, html_bytes.len() as i32)
164            .map_err(|e| {
165                PluginError::AnalysisFailed(format!("Failed to allocate for HTML: {}", e))
166            })?;
167
168        let url_ptr = alloc_func
169            .call(&mut self.store, url_bytes.len() as i32)
170            .map_err(|e| {
171                PluginError::AnalysisFailed(format!("Failed to allocate for URL: {}", e))
172            })?;
173
174        self.memory.data_mut(&mut self.store)
175            [html_ptr as usize..(html_ptr as usize + html_bytes.len())]
176            .copy_from_slice(html_bytes);
177
178        self.memory.data_mut(&mut self.store)
179            [url_ptr as usize..(url_ptr as usize + url_bytes.len())]
180            .copy_from_slice(url_bytes);
181
182        let result_ptr = analyze_func
183            .call(
184                &mut self.store,
185                (
186                    html_ptr,
187                    html_bytes.len() as i32,
188                    url_ptr,
189                    url_bytes.len() as i32,
190                ),
191            )
192            .map_err(|e| PluginError::AnalysisFailed(format!("Analyze failed: {}", e)))?;
193
194        let result = self.read_string(result_ptr as usize)?;
195
196        let free_func = self
197            .instance
198            .get_typed_func::<i32, ()>(&mut self.store, "crawlkit_plugin_free")
199            .map_err(|e| PluginError::AnalysisFailed(format!("Free function not found: {}", e)))?;
200        let _ = free_func.call(&mut self.store, html_ptr);
201        let _ = free_func.call(&mut self.store, url_ptr);
202        let _ = free_func.call(&mut self.store, result_ptr);
203
204        Ok(result)
205    }
206
207    /// Read a null-terminated string from WASM memory.
208    fn read_string(&self, ptr: usize) -> Result<String, PluginError> {
209        let data = self.memory.data(&self.store);
210        let end = data[ptr..]
211            .iter()
212            .position(|&b| b == 0)
213            .unwrap_or(data.len() - ptr);
214        let bytes = &data[ptr..ptr + end];
215        String::from_utf8(bytes.to_vec())
216            .map_err(|e| PluginError::WasmExecution(format!("Invalid UTF-8: {}", e)))
217    }
218
219    /// Get plugin metadata.
220    pub fn metadata(&self) -> &PluginMetadata {
221        &self.manifest
222    }
223}
224
225/// Plugin registry managing all loaded plugins.
226pub struct PluginRegistry {
227    plugins: Arc<RwLock<Vec<WasmPlugin>>>,
228    search_paths: Vec<PathBuf>,
229}
230
231impl PluginRegistry {
232    /// Create empty plugin registry.
233    pub fn new() -> Self {
234        Self {
235            plugins: Arc::new(RwLock::new(Vec::new())),
236            search_paths: Vec::new(),
237        }
238    }
239
240    /// Add a plugin search path.
241    pub fn add_search_path(&mut self, path: PathBuf) {
242        self.search_paths.push(path);
243    }
244
245    /// Scan search paths and load all valid plugins.
246    pub fn load_all(&mut self) -> Vec<PluginError> {
247        let mut errors = Vec::new();
248
249        for search_path in &self.search_paths {
250            if !search_path.exists() {
251                continue;
252            }
253
254            if let Ok(entries) = std::fs::read_dir(search_path) {
255                for entry in entries.flatten() {
256                    let plugin_dir = entry.path();
257                    if plugin_dir.is_dir() {
258                        match WasmPlugin::load(&plugin_dir) {
259                            Ok(plugin) => {
260                                tracing::info!(
261                                    "Loaded plugin: {} v{}",
262                                    plugin.metadata().name,
263                                    plugin.metadata().version
264                                );
265                                self.plugins.write().push(plugin);
266                            }
267                            Err(e) => {
268                                tracing::warn!(
269                                    "Failed to load plugin from {}: {}",
270                                    plugin_dir.display(),
271                                    e
272                                );
273                                errors.push(e);
274                            }
275                        }
276                    }
277                }
278            }
279        }
280
281        errors
282    }
283
284    /// Get list of loaded plugin names.
285    pub fn list(&self) -> Vec<String> {
286        self.plugins
287            .read()
288            .iter()
289            .map(|p| p.metadata().name.clone())
290            .collect()
291    }
292
293    /// Get plugin count.
294    pub fn count(&self) -> usize {
295        self.plugins.read().len()
296    }
297
298    /// Run analysis through all loaded plugins.
299    pub fn analyze_all(&self, html: &str, url: &str) -> Vec<Result<String, PluginError>> {
300        let mut results = Vec::new();
301        let mut plugins = self.plugins.write();
302
303        for plugin in plugins.iter_mut() {
304            results.push(plugin.analyze(html, url));
305        }
306
307        results
308    }
309}
310
311impl Default for PluginRegistry {
312    fn default() -> Self {
313        Self::new()
314    }
315}
316
317// ---------------------------------------------------------------------------
318// Tests
319// ---------------------------------------------------------------------------
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[test]
326    fn test_plugin_registry_default() {
327        let registry = PluginRegistry::new();
328        assert_eq!(registry.count(), 0);
329        assert!(registry.list().is_empty());
330    }
331
332    #[test]
333    fn test_plugin_loader_add_search_path() {
334        let mut registry = PluginRegistry::new();
335        registry.add_search_path(PathBuf::from("/tmp/plugins"));
336        assert_eq!(registry.search_paths.len(), 1);
337    }
338
339    #[test]
340    fn test_plugin_loader_nonexistent_path() {
341        let mut registry = PluginRegistry::new();
342        registry.add_search_path(PathBuf::from("/nonexistent/path"));
343        let errors = registry.load_all();
344        assert!(errors.is_empty());
345    }
346}