permissions.rs 6.7 KB

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