vtcode_bash_runner/
policy.rs

1use crate::executor::{CommandCategory, CommandInvocation};
2use anyhow::{Result, bail};
3use std::collections::HashSet;
4use std::path::Path;
5use std::sync::Arc;
6use vtcode_commons::WorkspacePaths;
7
8pub trait CommandPolicy: Send + Sync {
9    fn check(&self, invocation: &CommandInvocation) -> Result<()>;
10}
11
12pub struct AllowAllPolicy;
13
14impl CommandPolicy for AllowAllPolicy {
15    fn check(&self, _invocation: &CommandInvocation) -> Result<()> {
16        Ok(())
17    }
18}
19
20#[derive(Clone)]
21pub struct WorkspaceGuardPolicy {
22    workspace: Arc<dyn WorkspacePaths>,
23    allowed_commands: Option<HashSet<CommandCategory>>,
24}
25
26impl WorkspaceGuardPolicy {
27    pub fn new(workspace: Arc<dyn WorkspacePaths>) -> Self {
28        Self {
29            workspace,
30            allowed_commands: None,
31        }
32    }
33
34    pub fn with_allowed_commands(
35        mut self,
36        commands: impl IntoIterator<Item = CommandCategory>,
37    ) -> Self {
38        self.allowed_commands = Some(commands.into_iter().collect());
39        self
40    }
41
42    fn ensure_within_workspace(&self, path: &Path) -> Result<()> {
43        let root = self.workspace.workspace_root();
44        if !path.starts_with(root) {
45            bail!(
46                "path `{}` escapes the workspace root `{}`",
47                path.display(),
48                root.display()
49            );
50        }
51        Ok(())
52    }
53}
54
55impl CommandPolicy for WorkspaceGuardPolicy {
56    fn check(&self, invocation: &CommandInvocation) -> Result<()> {
57        if let Some(allowed) = &self.allowed_commands
58            && !allowed.contains(&invocation.category)
59        {
60            bail!(
61                "command category {:?} is not permitted",
62                invocation.category
63            );
64        }
65
66        self.ensure_within_workspace(&invocation.working_dir)?;
67
68        for path in &invocation.touched_paths {
69            self.ensure_within_workspace(path)?;
70        }
71
72        Ok(())
73    }
74}