Skip to main content

llm_kernel/dlp/
policy.rs

1//! Policy resolution — map a [`DataPolicy`] + [`Sensitivity`] to a [`PolicyAction`].
2//!
3//! ```
4//! use llm_kernel::dlp::{DataPolicy, PolicyAction, PolicyThreshold, Sensitivity, lookup};
5//!
6//! let mut policy = DataPolicy::default();
7//! policy.thresholds.push(PolicyThreshold {
8//!     min_sensitivity: Sensitivity::Restricted,
9//!     action: PolicyAction::Block,
10//! });
11//! assert_eq!(lookup(&policy, Sensitivity::Restricted), PolicyAction::Block);
12//! assert_eq!(lookup(&policy, Sensitivity::Internal), PolicyAction::Allow);
13//! ```
14
15pub use crate::provider::policy::{
16    DataPolicy, ImagePolicy, PolicyAction, PolicyThreshold, Sensitivity,
17};
18
19/// Resolve the action for `sensitivity` under `policy`.
20///
21/// Picks the action of the **highest** threshold whose `min_sensitivity` is
22/// satisfied; [`PolicyAction::Allow`] when no threshold matches (the
23/// permissive default). Works regardless of `thresholds` ordering.
24pub fn lookup(policy: &DataPolicy, sensitivity: Sensitivity) -> PolicyAction {
25    policy
26        .thresholds
27        .iter()
28        .filter(|t| sensitivity >= t.min_sensitivity)
29        .max_by_key(|t| t.min_sensitivity)
30        .map(|t| t.action.clone())
31        .unwrap_or(PolicyAction::Allow)
32}
33
34#[cfg(test)]
35mod tests {
36    use super::*;
37
38    fn cloud_policy() -> DataPolicy {
39        DataPolicy {
40            image: ImagePolicy::Allow,
41            thresholds: vec![
42                PolicyThreshold {
43                    min_sensitivity: Sensitivity::Internal,
44                    action: PolicyAction::Warn,
45                },
46                PolicyThreshold {
47                    min_sensitivity: Sensitivity::Restricted,
48                    action: PolicyAction::Block,
49                },
50                PolicyThreshold {
51                    min_sensitivity: Sensitivity::Confidential,
52                    action: PolicyAction::ReRoute {
53                        provider: "ollama".to_string(),
54                    },
55                },
56            ],
57        }
58    }
59
60    #[test]
61    fn highest_matching_threshold_wins_regardless_of_order() {
62        let policy = cloud_policy(); // deliberately unordered
63        assert_eq!(lookup(&policy, Sensitivity::Public), PolicyAction::Allow);
64        assert_eq!(lookup(&policy, Sensitivity::Internal), PolicyAction::Warn);
65        assert_eq!(
66            lookup(&policy, Sensitivity::Confidential),
67            PolicyAction::ReRoute {
68                provider: "ollama".to_string()
69            }
70        );
71        assert_eq!(
72            lookup(&policy, Sensitivity::Restricted),
73            PolicyAction::Block
74        );
75    }
76
77    #[test]
78    fn below_all_thresholds_allows() {
79        assert_eq!(
80            lookup(&cloud_policy(), Sensitivity::Public),
81            PolicyAction::Allow
82        );
83    }
84
85    #[test]
86    fn empty_policy_always_allows() {
87        let policy = DataPolicy::default();
88        assert_eq!(
89            lookup(&policy, Sensitivity::Restricted),
90            PolicyAction::Allow
91        );
92    }
93}