use kaynine_core::error::PolicyError;
use kaynine_core::policy::{Decision, Policy, PolicyRequest};
use kaynine_core::store::SessionStore;
use std::sync::Arc;
pub struct PolicyChain {
pub layers: Vec<Arc<dyn Policy>>,
}
impl PolicyChain {
pub fn new(layers: Vec<Arc<dyn Policy>>) -> Self {
Self { layers }
}
}
#[async_trait::async_trait]
impl Policy for PolicyChain {
async fn evaluate(&self, request: PolicyRequest) -> Result<Decision, PolicyError> {
let mut first_ask: Option<Decision> = None;
for layer in &self.layers {
match layer.evaluate(request.clone()).await? {
Decision::Deny { reason } => return Ok(Decision::Deny { reason }),
Decision::Ask => {
first_ask.get_or_insert(Decision::Ask);
}
Decision::Allow => {}
}
}
Ok(first_ask.unwrap_or(Decision::Allow))
}
}
pub struct GrantPolicy {
pub store: Arc<dyn SessionStore>,
pub namespace: String,
}
#[async_trait::async_trait]
impl Policy for GrantPolicy {
async fn evaluate(&self, request: PolicyRequest) -> Result<Decision, PolicyError> {
let grants = self
.store
.list_grants(&self.namespace)
.await
.map_err(|e| PolicyError::Evaluate(e.to_string()))?;
let call_caps: Vec<serde_json::Value> = request
.call
.capabilities
.iter()
.map(serde_json::to_value)
.collect::<Result<_, _>>()
.map_err(|e| PolicyError::Evaluate(e.to_string()))?;
for grant in &grants {
let Some(granted) = grant.get("capability") else {
continue;
};
if call_caps.iter().any(|cap| cap == granted) {
return Ok(Decision::Allow);
}
}
Ok(Decision::Ask)
}
}
#[cfg(test)]
mod tests {
use super::*;
use kaynine_core::ids::{RunId, SessionId, ToolCallId};
use kaynine_core::policy::DenyAllPolicy;
use kaynine_core::store::SessionStore;
use kaynine_core::testing::ScriptedPolicy;
use kaynine_core::tool::{Capability, PreparedToolCall};
use kaynine_store::sqlite::SqliteStore;
use std::collections::VecDeque;
use std::sync::Mutex;
struct FixedPolicy(Decision);
#[async_trait::async_trait]
impl Policy for FixedPolicy {
async fn evaluate(&self, _request: PolicyRequest) -> Result<Decision, PolicyError> {
Ok(self.0.clone())
}
}
fn request_with_caps(caps: Vec<Capability>) -> PolicyRequest {
PolicyRequest {
call: PreparedToolCall {
call_id: ToolCallId::from("call-1"),
name: "bash".into(),
arguments: serde_json::json!({}),
capabilities: caps,
},
session_id: SessionId::from("s"),
run_id: RunId::from("r"),
turn: 1,
}
}
#[tokio::test]
async fn policy_chain_priority() {
let chain = PolicyChain::new(vec![
Arc::new(FixedPolicy(Decision::Ask)),
Arc::new(DenyAllPolicy {
reason: "机密".into(),
}),
]);
assert_eq!(
chain.evaluate(request_with_caps(vec![])).await.unwrap(),
Decision::Deny {
reason: "机密".into()
}
);
let chain = PolicyChain::new(vec![
Arc::new(FixedPolicy(Decision::Allow)),
Arc::new(FixedPolicy(Decision::Ask)),
Arc::new(FixedPolicy(Decision::Allow)),
]);
assert_eq!(
chain.evaluate(request_with_caps(vec![])).await.unwrap(),
Decision::Ask
);
let chain = PolicyChain::new(vec![
Arc::new(FixedPolicy(Decision::Allow)),
Arc::new(FixedPolicy(Decision::Allow)),
]);
assert_eq!(
chain.evaluate(request_with_caps(vec![])).await.unwrap(),
Decision::Allow
);
}
#[tokio::test]
async fn policy_chain_defers_to_scripted_order() {
let script = ScriptedPolicy {
decisions: Mutex::new(VecDeque::from(vec![Decision::Ask])),
};
let chain = PolicyChain::new(vec![
Arc::new(script),
Arc::new(FixedPolicy(Decision::Allow)),
]);
assert_eq!(
chain.evaluate(request_with_caps(vec![])).await.unwrap(),
Decision::Ask
);
}
async fn store_with_grant(capability: &Capability) -> Arc<SqliteStore> {
let path =
std::env::temp_dir().join(format!("kaynine-grant-test-{}.db", uuid::Uuid::new_v4()));
let store = SqliteStore::open(&path).expect("open store");
store.migrate().await.expect("migrate");
store
.insert_grant(
"grant-1",
"prod",
serde_json::json!({
"grant_id": "grant-1",
"namespace": "prod",
"capability": capability,
}),
)
.await
.expect("insert grant");
Arc::new(store)
}
#[tokio::test]
async fn grant_policy_allows_matched_capability() {
let cap = Capability::FileRead {
path: "/tmp/data".into(),
};
let store = store_with_grant(&cap).await;
let policy = GrantPolicy {
store: store as Arc<dyn SessionStore>,
namespace: "prod".into(),
};
assert_eq!(
policy
.evaluate(request_with_caps(vec![
Capability::NetworkAccess {
origin: "example.com".into()
},
cap.clone()
]))
.await
.unwrap(),
Decision::Allow
);
assert_eq!(
policy
.evaluate(request_with_caps(vec![Capability::FileWrite {
path: "/tmp/data".into()
}]))
.await
.unwrap(),
Decision::Ask
);
}
}