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

vtcode_bash_runner/
policy.rs
