permissions.rs 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117
  1. use std::collections::BTreeMap;
  2. #[derive(Debug, Clone, Copy, PartialEq, Eq)]
  3. pub enum PermissionMode {
  4. Allow,
  5. Deny,
  6. Prompt,
  7. }
  8. #[derive(Debug, Clone, PartialEq, Eq)]
  9. pub struct PermissionRequest {
  10. pub tool_name: String,
  11. pub input: String,
  12. }
  13. #[derive(Debug, Clone, PartialEq, Eq)]
  14. pub enum PermissionPromptDecision {
  15. Allow,
  16. Deny { reason: String },
  17. }
  18. pub trait PermissionPrompter {
  19. fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision;
  20. }
  21. #[derive(Debug, Clone, PartialEq, Eq)]
  22. pub enum PermissionOutcome {
  23. Allow,
  24. Deny { reason: String },
  25. }
  26. #[derive(Debug, Clone, PartialEq, Eq)]
  27. pub struct PermissionPolicy {
  28. default_mode: PermissionMode,
  29. tool_modes: BTreeMap<String, PermissionMode>,
  30. }
  31. impl PermissionPolicy {
  32. #[must_use]
  33. pub fn new(default_mode: PermissionMode) -> Self {
  34. Self {
  35. default_mode,
  36. tool_modes: BTreeMap::new(),
  37. }
  38. }
  39. #[must_use]
  40. pub fn with_tool_mode(mut self, tool_name: impl Into<String>, mode: PermissionMode) -> Self {
  41. self.tool_modes.insert(tool_name.into(), mode);
  42. self
  43. }
  44. #[must_use]
  45. pub fn mode_for(&self, tool_name: &str) -> PermissionMode {
  46. self.tool_modes
  47. .get(tool_name)
  48. .copied()
  49. .unwrap_or(self.default_mode)
  50. }
  51. #[must_use]
  52. pub fn authorize(
  53. &self,
  54. tool_name: &str,
  55. input: &str,
  56. mut prompter: Option<&mut dyn PermissionPrompter>,
  57. ) -> PermissionOutcome {
  58. match self.mode_for(tool_name) {
  59. PermissionMode::Allow => PermissionOutcome::Allow,
  60. PermissionMode::Deny => PermissionOutcome::Deny {
  61. reason: format!("tool '{tool_name}' denied by permission policy"),
  62. },
  63. PermissionMode::Prompt => match prompter.as_mut() {
  64. Some(prompter) => match prompter.decide(&PermissionRequest {
  65. tool_name: tool_name.to_string(),
  66. input: input.to_string(),
  67. }) {
  68. PermissionPromptDecision::Allow => PermissionOutcome::Allow,
  69. PermissionPromptDecision::Deny { reason } => PermissionOutcome::Deny { reason },
  70. },
  71. None => PermissionOutcome::Deny {
  72. reason: format!("tool '{tool_name}' requires interactive approval"),
  73. },
  74. },
  75. }
  76. }
  77. }
  78. #[cfg(test)]
  79. mod tests {
  80. use super::{
  81. PermissionMode, PermissionOutcome, PermissionPolicy, PermissionPromptDecision,
  82. PermissionPrompter, PermissionRequest,
  83. };
  84. struct AllowPrompter;
  85. impl PermissionPrompter for AllowPrompter {
  86. fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision {
  87. assert_eq!(request.tool_name, "bash");
  88. PermissionPromptDecision::Allow
  89. }
  90. }
  91. #[test]
  92. fn uses_tool_specific_overrides() {
  93. let policy = PermissionPolicy::new(PermissionMode::Deny)
  94. .with_tool_mode("bash", PermissionMode::Prompt);
  95. let outcome = policy.authorize("bash", "echo hi", Some(&mut AllowPrompter));
  96. assert_eq!(outcome, PermissionOutcome::Allow);
  97. assert!(matches!(
  98. policy.authorize("edit", "x", None),
  99. PermissionOutcome::Deny { .. }
  100. ));
  101. }
  102. }