vtcode_bash_runner/
policy.rs

1use crate::executor::{CommandCategory, CommandInvocation};
2use anyhow::{Result, bail};
3use hashbrown::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 { workspace, allowed_commands: None }
29    }
30
31    pub fn with_allowed_commands(mut self, commands: impl IntoIterator<Item = CommandCategory>) -> Self {
32        self.allowed_commands = Some(commands.into_iter().collect());
33        self
34    }
35
36    fn ensure_within_workspace(&self, path: &Path) -> Result<()> {
37        let root = self.workspace.workspace_root();
38        vtcode_commons::paths::ensure_path_within_workspace(path, root).map_err(|error| {
39            error.context(format!("path `{}` escapes the workspace root `{}`", path.display(), root.display()))
40        })?;
41        Ok(())
42    }
43}
44
45impl CommandPolicy for WorkspaceGuardPolicy {
46    fn check(&self, invocation: &CommandInvocation) -> Result<()> {
47        if let Some(allowed) = &self.allowed_commands
48            && !allowed.contains(&invocation.category)
49        {
50            bail!("command category {:?} is not permitted", invocation.category);
51        }
52
53        self.ensure_within_workspace(&invocation.working_dir)?;
54
55        for path in &invocation.touched_paths {
56            self.ensure_within_workspace(path)?;
57        }
58
59        Ok(())
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66    use crate::executor::ShellKind;
67    use std::path::PathBuf;
68
69    struct StaticWorkspace {
70        root: PathBuf,
71    }
72
73    impl WorkspacePaths for StaticWorkspace {
74        fn workspace_root(&self) -> &Path {
75            &self.root
76        }
77
78        fn config_dir(&self) -> PathBuf {
79            self.root.join("config")
80        }
81    }
82
83    fn policy() -> WorkspaceGuardPolicy {
84        WorkspaceGuardPolicy::new(Arc::new(StaticWorkspace { root: PathBuf::from("/tmp/workspace") }))
85    }
86
87    fn invocation(working_dir: &str, touched: &[&str]) -> CommandInvocation {
88        CommandInvocation::new(
89            ShellKind::Unix,
90            "true".to_string(),
91            CommandCategory::ListDirectory,
92            PathBuf::from(working_dir),
93        )
94        .with_paths(touched.iter().map(PathBuf::from).collect())
95    }
96
97    #[test]
98    fn accepts_paths_inside_workspace() {
99        let invocation = invocation("/tmp/workspace/src", &["/tmp/workspace/file.txt"]);
100        assert!(policy().check(&invocation).is_ok());
101    }
102
103    #[test]
104    fn rejects_working_dir_traversal_escape() {
105        let invocation = invocation("/tmp/workspace/../etc", &[]);
106        assert!(policy().check(&invocation).is_err());
107    }
108
109    #[test]
110    fn rejects_touched_path_traversal_escape() {
111        let invocation = invocation("/tmp/workspace", &["/tmp/workspace/../../etc/passwd"]);
112        assert!(policy().check(&invocation).is_err());
113    }
114
115    #[test]
116    fn accepts_traversal_that_stays_inside_workspace() {
117        let invocation = invocation("/tmp/workspace/src/../src", &[]);
118        assert!(policy().check(&invocation).is_ok());
119    }
120
121    #[test]
122    fn rejects_path_outside_workspace() {
123        let invocation = invocation("/tmp/other", &[]);
124        assert!(policy().check(&invocation).is_err());
125    }
126}
vtcode_bash_runner/policy.rs

vtcode_bash_runner/
policy.rs
