Skip to main content

albert_runtime/
permissions.rs

1use std::collections::BTreeMap;
2use serde_json::Value;
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
5pub enum PermissionMode {
6    ReadOnly,
7    WorkspaceWrite,
8    DangerFullAccess,
9    Prompt,
10    Allow,
11}
12
13impl PermissionMode {
14    #[must_use]
15    pub fn as_str(self) -> &'static str {
16        match self {
17            Self::ReadOnly => "read-only",
18            Self::WorkspaceWrite => "workspace-write",
19            Self::DangerFullAccess => "danger-full-access",
20            Self::Prompt => "prompt",
21            Self::Allow => "allow",
22        }
23    }
24}
25
26#[derive(Debug, Clone, PartialEq, Eq)]
27pub struct PermissionRequest {
28    pub tool_name: String,
29    pub input: String,
30    pub current_mode: PermissionMode,
31    pub required_mode: PermissionMode,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub enum PermissionPromptDecision {
36    Allow,
37    AllowWithEdits { new_input: String },
38    Deny { reason: String },
39}
40
41pub trait PermissionPrompter {
42    fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision;
43}
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum PermissionOutcome {
47    Allow,
48    AllowWithEdits { new_input: String },
49    Deny { reason: String },
50}
51
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct PermissionPolicy {
54    active_mode: PermissionMode,
55    tool_requirements: BTreeMap<String, PermissionMode>,
56}
57
58impl PermissionPolicy {
59    #[must_use]
60    pub fn new(active_mode: PermissionMode) -> Self {
61        Self {
62            active_mode,
63            tool_requirements: BTreeMap::new(),
64        }
65    }
66
67    #[must_use]
68    pub fn with_tool_requirement(
69        mut self,
70        tool_name: impl Into<String>,
71        required_mode: PermissionMode,
72    ) -> Self {
73        self.tool_requirements
74            .insert(tool_name.into(), required_mode);
75        self
76    }
77
78    #[must_use]
79    pub fn active_mode(&self) -> PermissionMode {
80        self.active_mode
81    }
82
83    #[must_use]
84    pub fn required_mode_for(&self, tool_name: &str) -> PermissionMode {
85        self.tool_requirements
86            .get(tool_name)
87            .copied()
88            .unwrap_or(PermissionMode::DangerFullAccess)
89    }
90
91    /// Authorize a tool call and sanitize input to prevent dynamic escalation
92    #[must_use]
93    pub fn authorize(
94        &self,
95        tool_name: &str,
96        input: &str,
97        mut prompter: Option<&mut dyn PermissionPrompter>,
98    ) -> PermissionOutcome {
99        let current_mode = self.active_mode();
100        let required_mode = self.required_mode_for(tool_name);
101        
102        // Security sanitization: strip LLM-controlled sandbox disabling
103        let sanitized_input = sanitize_tool_input(input);
104
105        if current_mode == PermissionMode::Allow || current_mode >= required_mode {
106            return PermissionOutcome::Allow;
107        }
108
109        let request = PermissionRequest {
110            tool_name: tool_name.to_string(),
111            input: sanitized_input,
112            current_mode,
113            required_mode,
114        };
115
116        if current_mode == PermissionMode::Prompt
117            || (current_mode == PermissionMode::WorkspaceWrite
118                && required_mode == PermissionMode::DangerFullAccess)
119        {
120            return match prompter.as_mut() {
121                Some(prompter) => match prompter.decide(&request) {
122                    PermissionPromptDecision::Allow => PermissionOutcome::Allow,
123                    PermissionPromptDecision::AllowWithEdits { new_input } => {
124                        PermissionOutcome::AllowWithEdits { new_input }
125                    }
126                    PermissionPromptDecision::Deny { reason } => PermissionOutcome::Deny { reason },
127                },
128                None => PermissionOutcome::Deny {
129                    reason: format!(
130                        "tool '{tool_name}' requires approval to escalate from {} to {}",
131                        current_mode.as_str(),
132                        required_mode.as_str()
133                    ),
134                },
135            };
136        }
137
138        PermissionOutcome::Deny {
139            reason: format!(
140                "tool '{tool_name}' requires {} permission; current mode is {}",
141                required_mode.as_str(),
142                current_mode.as_str()
143            ),
144        }
145    }
146}
147
148/// Strip restricted flags from tool input JSON
149fn sanitize_tool_input(input: &str) -> String {
150    if let Ok(mut val) = serde_json::from_str::<Value>(input) {
151        if let Some(obj) = val.as_object_mut() {
152            // Revoke dynamic access to sandbox disabling
153            obj.remove("dangerouslyDisableSandbox");
154            return serde_json::to_string(obj).unwrap_or_else(|_| input.to_string());
155        }
156    }
157    input.to_string()
158}
159
160#[cfg(test)]
161mod tests {
162    use super::{
163        PermissionMode, PermissionOutcome, PermissionPolicy, PermissionPromptDecision,
164        PermissionPrompter, PermissionRequest, sanitize_tool_input,
165    };
166
167    struct RecordingPrompter {
168        seen: Vec<PermissionRequest>,
169        allow: bool,
170    }
171
172    impl PermissionPrompter for RecordingPrompter {
173        fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision {
174            self.seen.push(request.clone());
175            if self.allow {
176                PermissionPromptDecision::Allow
177            } else {
178                PermissionPromptDecision::Deny {
179                    reason: "not now".to_string(),
180                }
181            }
182        }
183    }
184
185    #[test]
186    fn allows_tools_when_active_mode_meets_requirement() {
187        let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
188            .with_tool_requirement("read_file", PermissionMode::ReadOnly)
189            .with_tool_requirement("write_file", PermissionMode::WorkspaceWrite);
190
191        assert_eq!(
192            policy.authorize("read_file", "{}", None),
193            PermissionOutcome::Allow
194        );
195    }
196
197    #[test]
198    fn sanitizes_dangerously_disable_sandbox_flag() {
199        let input = r#"{"command":"ls","dangerouslyDisableSandbox":true}"#;
200        let sanitized = sanitize_tool_input(input);
201        assert!(!sanitized.contains("dangerouslyDisableSandbox"));
202        assert!(sanitized.contains("ls"));
203    }
204
205    #[test]
206    fn authorize_uses_sanitized_input_for_prompts() {
207        let policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite)
208            .with_tool_requirement("bash", PermissionMode::DangerFullAccess);
209        let mut prompter = RecordingPrompter {
210            seen: Vec::new(),
211            allow: true,
212        };
213
214        let input = r#"{"command":"rm -rf /","dangerouslyDisableSandbox":true}"#;
215        policy.authorize("bash", input, Some(&mut prompter));
216
217        assert_eq!(prompter.seen.len(), 1);
218        assert!(!prompter.seen[0].input.contains("dangerouslyDisableSandbox"));
219    }
220}