fleetforge_policy/packs/
prompt_injection.rs

1use std::path::Path;
2use std::sync::Arc;
3
4#[cfg(feature = "wasm")]
5use anyhow::Context;
6use anyhow::{anyhow, Result};
7#[cfg(feature = "wasm")]
8use once_cell::sync::OnceCell;
9
10use crate::PolicyEngine;
11#[cfg(feature = "wasm")]
12use crate::PolicyModule;
13
14#[cfg(feature = "wasm")]
15use crate::WasmPolicyEngine;
16
17#[cfg(feature = "wasm")]
18const ENTRYPOINT: &str = "fleetforge/prompt_injection/decision";
19
20#[cfg(feature = "wasm")]
21static EMBEDDED_ENGINE: OnceCell<Arc<dyn PolicyEngine>> = OnceCell::new();
22
23/// Loads the embedded prompt injection policy pack compiled to Wasm.
24pub fn embedded() -> Result<Arc<dyn PolicyEngine>> {
25    #[cfg(feature = "wasm")]
26    {
27        EMBEDDED_ENGINE
28            .get_or_try_init(|| {
29                let bytes = include_bytes!("../../packs/prompt_injection/policy.wasm");
30                build_engine_from_bytes(bytes)
31            })
32            .map(Arc::clone)
33    }
34
35    #[cfg(not(feature = "wasm"))]
36    {
37        Err(anyhow!(
38            "prompt injection wasm pack requires the 'wasm' feature to be enabled"
39        ))
40    }
41}
42
43/// Loads the prompt injection pack from a Wasm file on disk.
44pub fn from_path(path: impl AsRef<Path>) -> Result<Arc<dyn PolicyEngine>> {
45    #[cfg(feature = "wasm")]
46    {
47        let bytes = std::fs::read(path.as_ref()).with_context(|| {
48            format!(
49                "failed to read prompt injection policy at {}",
50                path.as_ref().display()
51            )
52        })?;
53        build_engine_from_bytes(&bytes)
54    }
55
56    #[cfg(not(feature = "wasm"))]
57    {
58        let _ = path.as_ref();
59        Err(anyhow!(
60            "prompt injection wasm pack requires the 'wasm' feature to be enabled"
61        ))
62    }
63}
64
65#[cfg(feature = "wasm")]
66fn build_engine_from_bytes(bytes: &[u8]) -> Result<Arc<dyn PolicyEngine>> {
67    let module = PolicyModule::from_bytes(bytes).with_entrypoint(ENTRYPOINT);
68    let engine = WasmPolicyEngine::new(&module)
69        .context("failed to initialise prompt injection Wasm policy")?;
70    Ok(Arc::new(engine))
71}
fleetforge_policy/packs/prompt_injection.rs

fleetforge_policy/packs/
prompt_injection.rs
