Skip to main content

telos_agent/integrations/plugin/
policy_loader.rs

1//! Command-backed plugin policy with a controlled working directory and environment.
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5
6use async_trait::async_trait;
7use tokio::io::AsyncWriteExt;
8use tokio::process::Command;
9
10use crate::agent::policies::{Policy, PolicyContext, PolicyOutcome};
11use crate::error::AgentError;
12
13pub struct CommandPolicy {
14    name: String,
15    command: String,
16    args: Vec<String>,
17    timeout_ms: u64,
18    plugin_root: PathBuf,
19    env: HashMap<String, String>,
20}
21
22impl CommandPolicy {
23    pub fn new(
24        name: String,
25        command: String,
26        args: Vec<String>,
27        timeout_ms: u64,
28        plugin_root: PathBuf,
29        env: HashMap<String, String>,
30    ) -> Self {
31        Self { name, command, args, timeout_ms, plugin_root, env }
32    }
33
34    fn resolve(&self, value: &str) -> String {
35        value.replace("${PLUGIN_ROOT}", &self.plugin_root.to_string_lossy())
36    }
37
38    fn command_path(&self) -> String {
39        let command = self.resolve(&self.command);
40        let path = Path::new(&command);
41        if path.is_relative()
42            && (command.starts_with('.') || command.contains('/') || command.contains('\\'))
43        {
44            self.plugin_root.join(path).to_string_lossy().into_owned()
45        } else {
46            command
47        }
48    }
49}
50
51#[async_trait]
52impl Policy for CommandPolicy {
53    fn name(&self) -> &str {
54        &self.name
55    }
56
57    async fn evaluate(&self, context: &PolicyContext) -> Result<PolicyOutcome, AgentError> {
58        let input = serde_json::to_vec(context).map_err(|error| AgentError::ToolExecution {
59            tool: self.name.clone(),
60            message: format!("policy serialization error: {error}"),
61        })?;
62        let mut command = Command::new(self.command_path());
63        command
64            .args(self.args.iter().map(|arg| self.resolve(arg)))
65            .current_dir(&self.plugin_root)
66            .env_clear()
67            .envs(&self.env)
68            .env("PLUGIN_ROOT", &self.plugin_root)
69            .stdin(std::process::Stdio::piped())
70            .stdout(std::process::Stdio::piped())
71            .stderr(std::process::Stdio::piped())
72            .kill_on_drop(true);
73        let mut child = command.spawn().map_err(|error| AgentError::ToolExecution {
74            tool: self.name.clone(),
75            message: format!("policy failed to spawn: {error}"),
76        })?;
77        if let Some(mut stdin) = child.stdin.take() {
78            stdin.write_all(&input).await.map_err(|error| AgentError::ToolExecution {
79                tool: self.name.clone(),
80                message: format!("policy stdin failed: {error}"),
81            })?;
82        }
83        let output = tokio::time::timeout(
84            std::time::Duration::from_millis(self.timeout_ms),
85            child.wait_with_output(),
86        )
87        .await
88        .map_err(|_| AgentError::ToolExecution {
89            tool: self.name.clone(),
90            message: "policy timed out".into(),
91        })?
92        .map_err(|error| AgentError::ToolExecution {
93            tool: self.name.clone(),
94            message: format!("policy failed: {error}"),
95        })?;
96        if !output.status.success() {
97            return Err(AgentError::ToolExecution {
98                tool: self.name.clone(),
99                message: String::from_utf8_lossy(&output.stderr).trim().to_string(),
100            });
101        }
102        serde_json::from_slice(&output.stdout).map_err(|error| AgentError::ToolExecution {
103            tool: self.name.clone(),
104            message: format!("invalid policy outcome: {error}"),
105        })
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn resolves_plugin_root_in_commands_and_arguments() {
115        let root = PathBuf::from("example-plugin");
116        let policy = CommandPolicy::new(
117            "p".into(),
118            "./bin/check".into(),
119            vec!["${PLUGIN_ROOT}/config.json".into()],
120            1000,
121            root.clone(),
122            HashMap::new(),
123        );
124        assert_eq!(PathBuf::from(policy.command_path()), root.join("./bin/check"));
125        assert_eq!(PathBuf::from(policy.resolve(&policy.args[0])), root.join("config.json"));
126    }
127}