loong_kernel/
policy.rs

1use std::{
2    collections::BTreeSet,
3    sync::{
4        Mutex,
5        atomic::{AtomicU64, Ordering},
6    },
7};
8
9// Re-export data types from contracts
10pub use loong_contracts::{PolicyContext, PolicyDecision, PolicyRequest};
11
12use crate::{contracts::CapabilityToken, errors::PolicyError, pack::VerticalPackManifest};
13
14pub trait PolicyEngine: Send + Sync {
15    fn issue_token(
16        &self,
17        pack: &VerticalPackManifest,
18        agent_id: &str,
19        now_epoch_s: u64,
20        ttl_s: u64,
21    ) -> Result<CapabilityToken, PolicyError>;
22
23    fn authorize(
24        &self,
25        token: &CapabilityToken,
26        runtime_pack_id: &str,
27        now_epoch_s: u64,
28        required: &std::collections::BTreeSet<crate::contracts::Capability>,
29    ) -> Result<(), PolicyError>;
30
31    fn revoke_token(&self, token_id: &str) -> Result<(), PolicyError>;
32
33    fn revoke_generation(&self, _below: u64) {
34        // Default no-op
35    }
36
37    fn check_tool_call(&self, _request: &PolicyRequest) -> PolicyDecision {
38        PolicyDecision::Allow
39    }
40}
41
42#[derive(Debug, Default)]
43pub struct StaticPolicyEngine {
44    token_seq: AtomicU64,
45    revoked_tokens: Mutex<BTreeSet<String>>,
46    generation: AtomicU64,
47    revoked_below_generation: AtomicU64,
48}
49
50impl StaticPolicyEngine {
51    fn next_token_id(&self) -> String {
52        let seq = self.token_seq.fetch_add(1, Ordering::Relaxed) + 1;
53        format!("tok-{seq:016x}")
54    }
55
56    /// Revoke all tokens with generation <= `below`.
57    ///
58    /// Note: tokens issued concurrently during this call may land in the
59    /// revoked range. This is acceptable for StaticPolicyEngine (test/dev).
60    /// A production engine should use a lock or AcqRel ordering.
61    pub fn revoke_generation(&self, below: u64) {
62        self.revoked_below_generation
63            .fetch_max(below, Ordering::Relaxed);
64        // Fast-forward generation so newly issued tokens won't be immediately revoked.
65        self.generation.fetch_max(below, Ordering::Relaxed);
66    }
67
68    pub fn current_generation(&self) -> u64 {
69        self.generation.load(Ordering::Relaxed)
70    }
71}
72
73impl PolicyEngine for StaticPolicyEngine {
74    fn issue_token(
75        &self,
76        pack: &VerticalPackManifest,
77        agent_id: &str,
78        now_epoch_s: u64,
79        ttl_s: u64,
80    ) -> Result<CapabilityToken, PolicyError> {
81        let generation = self.generation.fetch_add(1, Ordering::Relaxed) + 1;
82        Ok(CapabilityToken {
83            token_id: self.next_token_id(),
84            pack_id: pack.pack_id.clone(),
85            agent_id: agent_id.to_owned(),
86            allowed_capabilities: pack.granted_capabilities.clone(),
87            issued_at_epoch_s: now_epoch_s,
88            expires_at_epoch_s: now_epoch_s.saturating_add(ttl_s),
89            generation,
90        })
91    }
92
93    fn authorize(
94        &self,
95        token: &CapabilityToken,
96        runtime_pack_id: &str,
97        now_epoch_s: u64,
98        required: &std::collections::BTreeSet<crate::contracts::Capability>,
99    ) -> Result<(), PolicyError> {
100        if self
101            .revoked_tokens
102            .lock()
103            .map_err(|_err| PolicyError::RevokedToken {
104                token_id: token.token_id.clone(),
105            })?
106            .contains(&token.token_id)
107        {
108            return Err(PolicyError::RevokedToken {
109                token_id: token.token_id.clone(),
110            });
111        }
112
113        let threshold = self.revoked_below_generation.load(Ordering::Relaxed);
114        if token.generation > 0 && token.generation <= threshold {
115            return Err(PolicyError::RevokedToken {
116                token_id: token.token_id.clone(),
117            });
118        }
119
120        if token.pack_id != runtime_pack_id {
121            return Err(PolicyError::PackMismatch {
122                token_pack_id: token.pack_id.clone(),
123                runtime_pack_id: runtime_pack_id.to_owned(),
124            });
125        }
126
127        if now_epoch_s > token.expires_at_epoch_s {
128            return Err(PolicyError::ExpiredToken {
129                token_id: token.token_id.clone(),
130                expires_at_epoch_s: token.expires_at_epoch_s,
131            });
132        }
133
134        for capability in required {
135            if !token.allowed_capabilities.contains(capability) {
136                return Err(PolicyError::MissingCapability {
137                    token_id: token.token_id.clone(),
138                    capability: *capability,
139                });
140            }
141        }
142
143        Ok(())
144    }
145
146    fn revoke_token(&self, token_id: &str) -> Result<(), PolicyError> {
147        let mut revoked = self
148            .revoked_tokens
149            .lock()
150            .map_err(|_err| PolicyError::RevokedToken {
151                token_id: token_id.to_owned(),
152            })?;
153        revoked.insert(token_id.to_owned());
154        Ok(())
155    }
156
157    fn revoke_generation(&self, below: u64) {
158        self.revoke_generation(below);
159    }
160
161    // Deprecated: Tool policy is now enforced via PolicyExtensionChain.
162    fn check_tool_call(&self, _request: &PolicyRequest) -> PolicyDecision {
163        PolicyDecision::Allow
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use std::collections::BTreeSet;
170
171    use serde_json::json;
172
173    use super::*;
174
175    fn policy_request(tool_name: &str, parameters: serde_json::Value) -> PolicyRequest {
176        PolicyRequest {
177            tool_name: tool_name.to_owned(),
178            parameters,
179            pack_id: "test-pack".to_owned(),
180            agent_id: "test-agent".to_owned(),
181            capabilities_used: BTreeSet::new(),
182            context: PolicyContext::default(),
183        }
184    }
185
186    #[test]
187    fn deprecated_check_tool_call_always_allows() {
188        let engine = StaticPolicyEngine::default();
189        let request = policy_request("shell.exec", json!({"command": "rm", "args": ["-rf", "/"]}));
190        assert_eq!(engine.check_tool_call(&request), PolicyDecision::Allow);
191    }
192}
loong_kernel/policy.rs

loong_kernel/
policy.rs
