permissions.rs 6.9 KB

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