1use std::collections::HashMap;
26use std::path::{Path, PathBuf};
27use std::sync::Arc;
28use std::time::SystemTime;
29
30use anyhow::{Context, Result};
31use wasmtime::{Config, Engine, Linker, Module, Store};
32
33use crate::types::{PluginInput, PluginOutput};
34
35#[derive(Debug, Clone)]
37pub struct RuntimeConfig {
38 pub max_memory: usize,
40 pub max_time_secs: u64,
42}
43
44impl Default for RuntimeConfig {
45 fn default() -> Self {
46 Self {
47 max_memory: 256 * 1024 * 1024, max_time_secs: 30,
49 }
50 }
51}
52
53pub fn validate_plugin_module(bytes: &[u8]) -> Result<()> {
64 let engine = Engine::default();
65 let module = Module::new(&engine, bytes)?;
66
67 if let Some(import) = module.imports().next() {
69 anyhow::bail!(
70 "plugin has forbidden import: {}::{}",
71 import.module(),
72 import.name()
73 );
74 }
75
76 let exports: Vec<_> = module.exports().map(|e| e.name()).collect();
78
79 if !exports.contains(&"memory") {
80 anyhow::bail!("plugin must export 'memory'");
81 }
82 if !exports.contains(&"alloc") {
83 anyhow::bail!("plugin must export 'alloc' function");
84 }
85 if !exports.contains(&"process") {
86 anyhow::bail!("plugin must export 'process' function");
87 }
88
89 Ok(())
90}
91
92pub struct Plugin {
94 name: String,
96 module: Module,
98 engine: Arc<Engine>,
100}
101
102impl Plugin {
103 pub fn load(path: &Path, _config: &RuntimeConfig) -> Result<Self> {
105 let name = path
106 .file_stem()
107 .and_then(|s| s.to_str())
108 .unwrap_or("unknown")
109 .to_string();
110
111 let mut engine_config = Config::new();
113 engine_config.consume_fuel(true); let engine = Arc::new(Engine::new(&engine_config)?);
116
117 let wasm_bytes =
119 std::fs::read(path).with_context(|| format!("failed to read {}", path.display()))?;
120
121 let module = Module::new(&engine, &wasm_bytes)
122 .map_err(anyhow::Error::from)
123 .with_context(|| format!("failed to compile {}", path.display()))?;
124
125 Ok(Self {
126 name,
127 module,
128 engine,
129 })
130 }
131
132 pub fn load_bytes(
134 name: impl Into<String>,
135 bytes: &[u8],
136 _config: &RuntimeConfig,
137 ) -> Result<Self> {
138 let name = name.into();
139
140 let mut engine_config = Config::new();
141 engine_config.consume_fuel(true);
142
143 let engine = Arc::new(Engine::new(&engine_config)?);
144 let module = Module::new(&engine, bytes)?;
145
146 Ok(Self {
147 name,
148 module,
149 engine,
150 })
151 }
152
153 pub fn name(&self) -> &str {
155 &self.name
156 }
157
158 pub fn execute(&self, input: &PluginInput, config: &RuntimeConfig) -> Result<PluginOutput> {
160 let mut store = Store::new(&self.engine, ());
162
163 let fuel = config.max_time_secs * 1_000_000;
165 store.set_fuel(fuel)?;
166
167 let linker = Linker::new(&self.engine);
170
171 let instance = linker.instantiate(&mut store, &self.module)?;
173
174 let input_bytes = rmp_serde::to_vec(input)?;
176
177 let memory = instance
179 .get_memory(&mut store, "memory")
180 .ok_or_else(|| anyhow::anyhow!("plugin must export 'memory'"))?;
181
182 let alloc = instance
184 .get_typed_func::<u32, u32>(&mut store, "alloc")
185 .map_err(anyhow::Error::from)
186 .context("plugin must export 'alloc' function")?;
187
188 let input_ptr = alloc.call(&mut store, input_bytes.len() as u32)?;
190
191 memory.write(&mut store, input_ptr as usize, &input_bytes)?;
193
194 let process = instance
196 .get_typed_func::<(u32, u32), u64>(&mut store, "process")
197 .map_err(anyhow::Error::from)
198 .context("plugin must export 'process' function")?;
199
200 let result = process.call(&mut store, (input_ptr, input_bytes.len() as u32))?;
201
202 let output_ptr = (result >> 32) as u32;
204 let output_len = (result & 0xFFFF_FFFF) as u32;
205
206 let mut output_bytes = vec![0u8; output_len as usize];
208 memory.read(&store, output_ptr as usize, &mut output_bytes)?;
209
210 let output: PluginOutput = rmp_serde::from_slice(&output_bytes)?;
212
213 Ok(output)
214 }
215}
216
217pub struct PluginManager {
219 config: RuntimeConfig,
221 plugins: Vec<Plugin>,
223}
224
225impl PluginManager {
226 pub fn new() -> Self {
228 Self::with_config(RuntimeConfig::default())
229 }
230
231 pub const fn with_config(config: RuntimeConfig) -> Self {
233 Self {
234 config,
235 plugins: Vec::new(),
236 }
237 }
238
239 pub fn load(&mut self, path: &Path) -> Result<usize> {
241 let plugin = Plugin::load(path, &self.config)?;
242 let index = self.plugins.len();
243 self.plugins.push(plugin);
244 Ok(index)
245 }
246
247 pub fn load_bytes(&mut self, name: impl Into<String>, bytes: &[u8]) -> Result<usize> {
249 let plugin = Plugin::load_bytes(name, bytes, &self.config)?;
250 let index = self.plugins.len();
251 self.plugins.push(plugin);
252 Ok(index)
253 }
254
255 pub fn execute(&self, index: usize, input: &PluginInput) -> Result<PluginOutput> {
257 let plugin = self
258 .plugins
259 .get(index)
260 .context("plugin index out of bounds")?;
261 plugin.execute(input, &self.config)
262 }
263
264 pub fn execute_all(&self, mut input: PluginInput) -> Result<PluginOutput> {
266 let mut all_errors = Vec::new();
272
273 for plugin in &self.plugins {
274 let output = plugin.execute(&input, &self.config)?;
275 all_errors.extend(output.errors);
276 input.directives = output.directives;
277 }
278
279 Ok(PluginOutput {
280 directives: input.directives,
281 errors: all_errors,
282 })
283 }
284
285 pub const fn len(&self) -> usize {
287 self.plugins.len()
288 }
289
290 pub const fn is_empty(&self) -> bool {
292 self.plugins.is_empty()
293 }
294}
295
296impl Default for PluginManager {
297 fn default() -> Self {
298 Self::new()
299 }
300}
301
302struct TrackedPlugin {
304 plugin: Plugin,
306 path: PathBuf,
308 modified: SystemTime,
310}
311
312pub struct WatchingPluginManager {
332 config: RuntimeConfig,
334 plugins: Vec<TrackedPlugin>,
336 name_index: HashMap<String, usize>,
338 on_reload: Option<Box<dyn Fn(&str) + Send + Sync>>,
340}
341
342impl WatchingPluginManager {
343 pub fn new() -> Self {
345 Self::with_config(RuntimeConfig::default())
346 }
347
348 pub fn with_config(config: RuntimeConfig) -> Self {
350 Self {
351 config,
352 plugins: Vec::new(),
353 name_index: HashMap::new(),
354 on_reload: None,
355 }
356 }
357
358 pub fn on_reload<F>(&mut self, callback: F)
360 where
361 F: Fn(&str) + Send + Sync + 'static,
362 {
363 self.on_reload = Some(Box::new(callback));
364 }
365
366 pub fn load(&mut self, path: impl AsRef<Path>) -> Result<usize> {
368 let path = path.as_ref();
369 let abs_path = path.canonicalize().unwrap_or_else(|_| path.to_path_buf());
371
372 let metadata = std::fs::metadata(&abs_path)
374 .with_context(|| format!("failed to stat {}", abs_path.display()))?;
375 let modified = metadata.modified()?;
376
377 let plugin = Plugin::load(&abs_path, &self.config)?;
379 let name = plugin.name().to_string();
380 let index = self.plugins.len();
381
382 self.plugins.push(TrackedPlugin {
384 plugin,
385 path: abs_path,
386 modified,
387 });
388 self.name_index.insert(name, index);
389
390 Ok(index)
391 }
392
393 pub fn check_and_reload(&mut self) -> Result<bool> {
397 let mut reloaded = false;
398
399 for tracked in &mut self.plugins {
400 let metadata = match std::fs::metadata(&tracked.path) {
402 Ok(m) => m,
403 Err(_) => continue, };
405
406 let current_modified = match metadata.modified() {
407 Ok(m) => m,
408 Err(_) => continue,
409 };
410
411 if current_modified > tracked.modified {
413 match Plugin::load(&tracked.path, &self.config) {
415 Ok(new_plugin) => {
416 let name = tracked.plugin.name().to_string();
417 tracked.plugin = new_plugin;
418 tracked.modified = current_modified;
419 reloaded = true;
420
421 if let Some(ref callback) = self.on_reload {
423 callback(&name);
424 }
425 }
426 Err(e) => {
427 eprintln!(
429 "warning: failed to reload plugin {}: {}",
430 tracked.path.display(),
431 e
432 );
433 }
434 }
435 }
436 }
437
438 Ok(reloaded)
439 }
440
441 pub fn reload_all(&mut self) -> Result<()> {
443 for tracked in &mut self.plugins {
444 let new_plugin = Plugin::load(&tracked.path, &self.config)?;
445 let metadata = std::fs::metadata(&tracked.path)?;
446 tracked.plugin = new_plugin;
447 tracked.modified = metadata.modified()?;
448 }
449 Ok(())
450 }
451
452 pub fn get(&self, name: &str) -> Option<&Plugin> {
454 self.name_index.get(name).map(|&i| &self.plugins[i].plugin)
455 }
456
457 pub fn execute(&self, index: usize, input: &PluginInput) -> Result<PluginOutput> {
459 let tracked = self
460 .plugins
461 .get(index)
462 .context("plugin index out of bounds")?;
463 tracked.plugin.execute(input, &self.config)
464 }
465
466 pub fn execute_by_name(&self, name: &str, input: &PluginInput) -> Result<PluginOutput> {
468 let index = self
469 .name_index
470 .get(name)
471 .with_context(|| format!("plugin '{name}' not found"))?;
472 self.execute(*index, input)
473 }
474
475 pub fn execute_all(&self, mut input: PluginInput) -> Result<PluginOutput> {
477 let mut all_errors = Vec::new();
478
479 for tracked in &self.plugins {
480 let output = tracked.plugin.execute(&input, &self.config)?;
481 all_errors.extend(output.errors);
482 input.directives = output.directives;
483 }
484
485 Ok(PluginOutput {
486 directives: input.directives,
487 errors: all_errors,
488 })
489 }
490
491 pub const fn len(&self) -> usize {
493 self.plugins.len()
494 }
495
496 pub const fn is_empty(&self) -> bool {
498 self.plugins.is_empty()
499 }
500
501 pub fn plugin_info(&self) -> Vec<(&Path, SystemTime)> {
503 self.plugins
504 .iter()
505 .map(|t| (t.path.as_path(), t.modified))
506 .collect()
507 }
508}
509
510impl Default for WatchingPluginManager {
511 fn default() -> Self {
512 Self::new()
513 }
514}
515
516#[cfg(test)]
517mod tests {
518 use super::*;
519
520 #[test]
524 fn test_valid_plugin_validation() {
525 let wasm = wat::parse_str(
531 r#"
532 (module
533 (memory (export "memory") 1)
534 (func (export "alloc") (param i32) (result i32)
535 i32.const 0
536 )
537 (func (export "process") (param i32 i32) (result i64)
538 i64.const 0
539 )
540 )
541 "#,
542 )
543 .expect("valid wat");
544
545 let result = validate_plugin_module(&wasm);
546 assert!(
547 result.is_ok(),
548 "valid plugin should pass validation: {:?}",
549 result.err()
550 );
551 }
552
553 #[test]
555 fn test_wasi_import_rejected() {
556 let wasm = wat::parse_str(
558 r#"
559 (module
560 (import "wasi_snapshot_preview1" "fd_write"
561 (func $fd_write (param i32 i32 i32 i32) (result i32))
562 )
563 (memory (export "memory") 1)
564 (func (export "alloc") (param i32) (result i32)
565 i32.const 0
566 )
567 (func (export "process") (param i32 i32) (result i64)
568 i64.const 0
569 )
570 )
571 "#,
572 )
573 .expect("valid wat");
574
575 let result = validate_plugin_module(&wasm);
576 assert!(
577 result.is_err(),
578 "module with WASI import should be rejected"
579 );
580 let err = result.unwrap_err().to_string();
581 assert!(
582 err.contains("forbidden import"),
583 "error should mention forbidden import: {err}"
584 );
585 assert!(
586 err.contains("wasi_snapshot_preview1"),
587 "error should mention WASI: {err}"
588 );
589 }
590
591 #[test]
593 fn test_env_import_rejected() {
594 let wasm = wat::parse_str(
596 r#"
597 (module
598 (import "env" "some_func" (func $some_func))
599 (memory (export "memory") 1)
600 (func (export "alloc") (param i32) (result i32)
601 i32.const 0
602 )
603 (func (export "process") (param i32 i32) (result i64)
604 i64.const 0
605 )
606 )
607 "#,
608 )
609 .expect("valid wat");
610
611 let result = validate_plugin_module(&wasm);
612 assert!(result.is_err(), "module with env import should be rejected");
613 }
614
615 #[test]
617 fn test_missing_exports_rejected() {
618 let wasm = wat::parse_str(
620 r#"
621 (module
622 (memory (export "memory") 1)
623 (func (export "process") (param i32 i32) (result i64)
624 i64.const 0
625 )
626 )
627 "#,
628 )
629 .expect("valid wat");
630
631 let result = validate_plugin_module(&wasm);
632 assert!(result.is_err(), "module missing alloc should be rejected");
633 assert!(result.unwrap_err().to_string().contains("alloc"));
634 }
635
636 #[test]
638 fn test_runtime_config_defaults() {
639 let config = RuntimeConfig::default();
640 assert_eq!(config.max_memory, 256 * 1024 * 1024); assert_eq!(config.max_time_secs, 30);
642 }
643
644 #[test]
646 fn test_missing_memory_rejected() {
647 let wasm = wat::parse_str(
648 r#"
649 (module
650 (func (export "alloc") (param i32) (result i32)
651 i32.const 0
652 )
653 (func (export "process") (param i32 i32) (result i64)
654 i64.const 0
655 )
656 )
657 "#,
658 )
659 .expect("valid wat");
660
661 let result = validate_plugin_module(&wasm);
662 assert!(result.is_err(), "module missing memory should be rejected");
663 assert!(result.unwrap_err().to_string().contains("memory"));
664 }
665
666 #[test]
668 fn test_missing_process_rejected() {
669 let wasm = wat::parse_str(
670 r#"
671 (module
672 (memory (export "memory") 1)
673 (func (export "alloc") (param i32) (result i32)
674 i32.const 0
675 )
676 )
677 "#,
678 )
679 .expect("valid wat");
680
681 let result = validate_plugin_module(&wasm);
682 assert!(result.is_err(), "module missing process should be rejected");
683 assert!(result.unwrap_err().to_string().contains("process"));
684 }
685
686 #[test]
688 fn test_invalid_wasm_rejected() {
689 let invalid = b"not valid wasm bytes";
690 let result = validate_plugin_module(invalid);
691 assert!(result.is_err(), "invalid WASM should be rejected");
692 }
693
694 #[test]
696 fn test_runtime_config_custom() {
697 let config = RuntimeConfig {
698 max_memory: 512 * 1024 * 1024, max_time_secs: 60,
700 };
701 assert_eq!(config.max_memory, 512 * 1024 * 1024);
702 assert_eq!(config.max_time_secs, 60);
703 }
704
705 #[test]
710 fn test_plugin_manager_new() {
711 let manager = PluginManager::new();
712 assert!(manager.is_empty());
713 assert_eq!(manager.len(), 0);
714 }
715
716 #[test]
717 fn test_plugin_manager_with_config() {
718 let config = RuntimeConfig {
719 max_memory: 128 * 1024 * 1024,
720 max_time_secs: 10,
721 };
722 let manager = PluginManager::with_config(config);
723 assert!(manager.is_empty());
724 }
725
726 #[test]
727 fn test_plugin_manager_default() {
728 let manager = PluginManager::default();
729 assert!(manager.is_empty());
730 assert_eq!(manager.len(), 0);
731 }
732
733 #[test]
734 fn test_watching_plugin_manager_new() {
735 let manager = WatchingPluginManager::new();
736 assert!(manager.is_empty());
737 assert_eq!(manager.len(), 0);
738 assert!(manager.plugin_info().is_empty());
739 }
740
741 #[test]
742 fn test_watching_plugin_manager_with_config() {
743 let config = RuntimeConfig {
744 max_memory: 64 * 1024 * 1024,
745 max_time_secs: 5,
746 };
747 let manager = WatchingPluginManager::with_config(config);
748 assert!(manager.is_empty());
749 }
750
751 #[test]
752 fn test_watching_plugin_manager_default() {
753 let manager = WatchingPluginManager::default();
754 assert!(manager.is_empty());
755 assert_eq!(manager.len(), 0);
756 }
757
758 #[test]
759 fn test_watching_plugin_manager_get_unknown() {
760 let manager = WatchingPluginManager::new();
761 assert!(manager.get("nonexistent").is_none());
762 }
763
764 #[test]
765 fn test_plugin_manager_execute_out_of_bounds() {
766 let manager = PluginManager::new();
767 let input = crate::types::PluginInput {
768 directives: vec![],
769 options: crate::types::PluginOptions::default(),
770 config: None,
771 };
772 let result = manager.execute(0, &input);
773 assert!(result.is_err());
774 assert!(result.unwrap_err().to_string().contains("out of bounds"));
775 }
776
777 #[test]
778 fn test_watching_plugin_manager_execute_out_of_bounds() {
779 let manager = WatchingPluginManager::new();
780 let input = crate::types::PluginInput {
781 directives: vec![],
782 options: crate::types::PluginOptions::default(),
783 config: None,
784 };
785 let result = manager.execute(0, &input);
786 assert!(result.is_err());
787 assert!(result.unwrap_err().to_string().contains("out of bounds"));
788 }
789
790 #[test]
791 fn test_watching_plugin_manager_execute_by_name_unknown() {
792 let manager = WatchingPluginManager::new();
793 let input = crate::types::PluginInput {
794 directives: vec![],
795 options: crate::types::PluginOptions::default(),
796 config: None,
797 };
798 let result = manager.execute_by_name("unknown", &input);
799 assert!(result.is_err());
800 assert!(result.unwrap_err().to_string().contains("not found"));
801 }
802
803 #[test]
804 fn test_plugin_manager_execute_all_empty() {
805 let manager = PluginManager::new();
806 let input = crate::types::PluginInput {
807 directives: vec![],
808 options: crate::types::PluginOptions::default(),
809 config: None,
810 };
811 let result = manager.execute_all(input);
812 assert!(result.is_ok());
813 let output = result.unwrap();
814 assert!(output.directives.is_empty());
815 assert!(output.errors.is_empty());
816 }
817
818 #[test]
819 fn test_watching_plugin_manager_execute_all_empty() {
820 let manager = WatchingPluginManager::new();
821 let input = crate::types::PluginInput {
822 directives: vec![],
823 options: crate::types::PluginOptions::default(),
824 config: None,
825 };
826 let result = manager.execute_all(input);
827 assert!(result.is_ok());
828 let output = result.unwrap();
829 assert!(output.directives.is_empty());
830 assert!(output.errors.is_empty());
831 }
832
833 #[test]
834 fn test_watching_plugin_manager_check_reload_empty() {
835 let mut manager = WatchingPluginManager::new();
836 let result = manager.check_and_reload();
837 assert!(result.is_ok());
838 assert!(!result.unwrap()); }
840
841 #[test]
842 fn test_watching_plugin_manager_reload_all_empty() {
843 let mut manager = WatchingPluginManager::new();
844 let result = manager.reload_all();
845 assert!(result.is_ok()); }
847
848 #[test]
849 fn test_plugin_load_bytes() {
850 let wasm = wat::parse_str(
851 r#"
852 (module
853 (memory (export "memory") 1)
854 (func (export "alloc") (param i32) (result i32)
855 i32.const 0
856 )
857 (func (export "process") (param i32 i32) (result i64)
858 i64.const 0
859 )
860 )
861 "#,
862 )
863 .expect("valid wat");
864
865 let config = RuntimeConfig::default();
866 let result = Plugin::load_bytes("test_plugin", &wasm, &config);
867 assert!(result.is_ok());
868
869 let plugin = result.unwrap();
870 assert_eq!(plugin.name(), "test_plugin");
871 }
872
873 #[test]
874 fn test_plugin_manager_load_bytes() {
875 let wasm = wat::parse_str(
876 r#"
877 (module
878 (memory (export "memory") 1)
879 (func (export "alloc") (param i32) (result i32)
880 i32.const 0
881 )
882 (func (export "process") (param i32 i32) (result i64)
883 i64.const 0
884 )
885 )
886 "#,
887 )
888 .expect("valid wat");
889
890 let mut manager = PluginManager::new();
891 let result = manager.load_bytes("my_plugin", &wasm);
892 assert!(result.is_ok());
893 assert_eq!(result.unwrap(), 0); assert_eq!(manager.len(), 1);
895 assert!(!manager.is_empty());
896 }
897
898 #[test]
899 fn test_plugin_manager_multiple_plugins() {
900 let wasm = wat::parse_str(
901 r#"
902 (module
903 (memory (export "memory") 1)
904 (func (export "alloc") (param i32) (result i32)
905 i32.const 0
906 )
907 (func (export "process") (param i32 i32) (result i64)
908 i64.const 0
909 )
910 )
911 "#,
912 )
913 .expect("valid wat");
914
915 let mut manager = PluginManager::new();
916 manager.load_bytes("plugin1", &wasm).unwrap();
917 manager.load_bytes("plugin2", &wasm).unwrap();
918 manager.load_bytes("plugin3", &wasm).unwrap();
919
920 assert_eq!(manager.len(), 3);
921 }
922
923 #[test]
924 fn test_validate_truncated_wasm() {
925 let truncated = &[0x00, 0x61, 0x73, 0x6d]; let result = validate_plugin_module(truncated);
928 assert!(result.is_err());
929 }
930
931 #[test]
932 fn test_validate_wrong_magic() {
933 let wrong_magic = &[0xFF, 0xFF, 0xFF, 0xFF];
934 let result = validate_plugin_module(wrong_magic);
935 assert!(result.is_err());
936 }
937
938 #[test]
939 fn test_validate_empty_wasm() {
940 let empty: &[u8] = &[];
941 let result = validate_plugin_module(empty);
942 assert!(result.is_err());
943 }
944}