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