chio_guards/
data_flow.rs

1//! Data flow guard -- enforces cumulative bytes-read/written limits via session journal.
2//!
3//! This guard reads cumulative data flow statistics from the session journal
4//! and denies requests that would cause the session to exceed configured
5//! byte limits for reads, writes, or combined I/O.
6//!
7//! The guard fails closed: if the session journal is unavailable or returns
8//! an error, the request is denied.
9
10use std::sync::Arc;
11
12use chio_http_session::SessionJournal;
13#[cfg(test)]
14use chio_kernel::Verdict;
15use chio_kernel::{Guard, GuardContext, GuardDecision, KernelError};
16
17// ---------------------------------------------------------------------------
18// DataFlowConfig
19// ---------------------------------------------------------------------------
20
21/// Configuration for cumulative data flow limits.
22#[derive(Clone, Debug, Default)]
23pub struct DataFlowConfig {
24    /// Maximum cumulative bytes read per session. None means unlimited.
25    pub max_bytes_read: Option<u64>,
26    /// Maximum cumulative bytes written per session. None means unlimited.
27    pub max_bytes_written: Option<u64>,
28    /// Maximum cumulative bytes (read + written) per session. None means unlimited.
29    pub max_bytes_total: Option<u64>,
30}
31
32// ---------------------------------------------------------------------------
33// DataFlowGuard
34// ---------------------------------------------------------------------------
35
36/// Guard that enforces cumulative data flow limits using the session journal.
37///
38/// Reads the journal's cumulative data flow statistics and denies requests
39/// if any configured limit has been reached.
40pub struct DataFlowGuard {
41    journal: Arc<SessionJournal>,
42    config: DataFlowConfig,
43}
44
45impl DataFlowGuard {
46    /// Create a new guard with the given journal and configuration.
47    pub fn new(journal: Arc<SessionJournal>, config: DataFlowConfig) -> Self {
48        Self { journal, config }
49    }
50}
51
52impl Guard for DataFlowGuard {
53    fn name(&self) -> &str {
54        "data-flow"
55    }
56
57    fn evaluate(&self, _ctx: &GuardContext) -> Result<GuardDecision, KernelError> {
58        let snapshot = self.journal.snapshot().map_err(|e| {
59            KernelError::Internal(format!("data-flow guard journal error (fail-closed): {e}"))
60        })?;
61        let flow = snapshot.data_flow;
62
63        // Check bytes read limit.
64        if let Some(max_read) = self.config.max_bytes_read {
65            if flow.total_bytes_read >= max_read {
66                return Ok(GuardDecision::deny(Vec::new()));
67            }
68        }
69
70        // Check bytes written limit.
71        if let Some(max_written) = self.config.max_bytes_written {
72            if flow.total_bytes_written >= max_written {
73                return Ok(GuardDecision::deny(Vec::new()));
74            }
75        }
76
77        // Check total I/O limit.
78        if let Some(max_total) = self.config.max_bytes_total {
79            let total = flow
80                .total_bytes_read
81                .saturating_add(flow.total_bytes_written);
82            if total >= max_total {
83                return Ok(GuardDecision::deny(Vec::new()));
84            }
85        }
86
87        Ok(GuardDecision::allow())
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use chio_http_session::RecordParams;
95
96    fn make_journal(session_id: &str) -> Arc<SessionJournal> {
97        Arc::new(SessionJournal::new(session_id.to_string()))
98    }
99
100    fn make_ctx() -> (
101        chio_kernel::ToolCallRequest,
102        chio_core::capability::scope::ChioScope,
103        String,
104        String,
105    ) {
106        let kp = chio_core::crypto::Keypair::generate();
107        let scope = chio_core::capability::scope::ChioScope::default();
108        let agent_id = kp.public_key().to_hex();
109        let server_id = "srv-test".to_string();
110
111        let cap_body = chio_core::capability::token::CapabilityTokenBody {
112            id: "cap-test".to_string(),
113            issuer: kp.public_key(),
114            subject: kp.public_key(),
115            scope: scope.clone(),
116            issued_at: 0,
117            expires_at: u64::MAX,
118            delegation_chain: vec![],
119            aggregate_invocation_budget: None,
120        };
121        let cap =
122            chio_core::capability::token::CapabilityToken::sign(cap_body, &kp).expect("sign cap");
123
124        let request = chio_kernel::ToolCallRequest {
125            request_id: "req-test".to_string(),
126            capability: cap,
127            tool_name: "read_file".to_string(),
128            server_id: server_id.clone(),
129            agent_id: agent_id.clone(),
130            arguments: serde_json::json!({"path": "/app/src/main.rs"}),
131            dpop_proof: None,
132            execution_nonce: None,
133            governed_intent: None,
134            approval_token: None,
135            approval_tokens: Vec::new(),
136            threshold_approval_proposal: None,
137            supplemental_authorization: None,
138            model_metadata: None,
139            federated_origin_kernel_id: None,
140        };
141
142        (request, scope, agent_id, server_id)
143    }
144
145    fn guard_ctx<'a>(
146        request: &'a chio_kernel::ToolCallRequest,
147        scope: &'a chio_core::capability::scope::ChioScope,
148        agent_id: &'a String,
149        server_id: &'a String,
150    ) -> chio_kernel::GuardContext<'a> {
151        chio_kernel::GuardContext {
152            request,
153            scope,
154            agent_id,
155            server_id,
156            session_filesystem_roots: None,
157            matched_grant_index: None,
158        }
159    }
160
161    #[test]
162    fn guard_name() {
163        let journal = make_journal("sess-1");
164        let guard = DataFlowGuard::new(journal, DataFlowConfig::default());
165        assert_eq!(guard.name(), "data-flow");
166    }
167
168    #[test]
169    fn unlimited_allows_all() {
170        let journal = make_journal("sess-1");
171        // Add some data flow.
172        journal
173            .record(RecordParams {
174                tool_name: "read_file".to_string(),
175                server_id: "srv".to_string(),
176                agent_id: "agent".to_string(),
177                bytes_read: 1_000_000,
178                bytes_written: 500_000,
179                delegation_depth: 0,
180                allowed: true,
181            })
182            .expect("record");
183
184        let guard = DataFlowGuard::new(journal, DataFlowConfig::default());
185        let (request, scope, agent_id, server_id) = make_ctx();
186        let ctx = guard_ctx(&request, &scope, &agent_id, &server_id);
187        assert_eq!(guard.evaluate(&ctx).expect("ok"), Verdict::Allow);
188    }
189
190    #[test]
191    fn denies_when_bytes_read_exceeded() {
192        let journal = make_journal("sess-read");
193        journal
194            .record(RecordParams {
195                tool_name: "read_file".to_string(),
196                server_id: "srv".to_string(),
197                agent_id: "agent".to_string(),
198                bytes_read: 500,
199                bytes_written: 0,
200                delegation_depth: 0,
201                allowed: true,
202            })
203            .expect("record");
204
205        let guard = DataFlowGuard::new(
206            journal,
207            DataFlowConfig {
208                max_bytes_read: Some(500),
209                ..DataFlowConfig::default()
210            },
211        );
212
213        let (request, scope, agent_id, server_id) = make_ctx();
214        let ctx = guard_ctx(&request, &scope, &agent_id, &server_id);
215        assert_eq!(guard.evaluate(&ctx).expect("ok"), Verdict::Deny);
216    }
217
218    #[test]
219    fn denies_when_bytes_written_exceeded() {
220        let journal = make_journal("sess-write");
221        journal
222            .record(RecordParams {
223                tool_name: "write_file".to_string(),
224                server_id: "srv".to_string(),
225                agent_id: "agent".to_string(),
226                bytes_read: 0,
227                bytes_written: 1000,
228                delegation_depth: 0,
229                allowed: true,
230            })
231            .expect("record");
232
233        let guard = DataFlowGuard::new(
234            journal,
235            DataFlowConfig {
236                max_bytes_written: Some(999),
237                ..DataFlowConfig::default()
238            },
239        );
240
241        let (request, scope, agent_id, server_id) = make_ctx();
242        let ctx = guard_ctx(&request, &scope, &agent_id, &server_id);
243        assert_eq!(guard.evaluate(&ctx).expect("ok"), Verdict::Deny);
244    }
245
246    #[test]
247    fn denies_when_total_exceeded() {
248        let journal = make_journal("sess-total");
249        journal
250            .record(RecordParams {
251                tool_name: "read_file".to_string(),
252                server_id: "srv".to_string(),
253                agent_id: "agent".to_string(),
254                bytes_read: 300,
255                bytes_written: 200,
256                delegation_depth: 0,
257                allowed: true,
258            })
259            .expect("record");
260
261        let guard = DataFlowGuard::new(
262            journal,
263            DataFlowConfig {
264                max_bytes_total: Some(500),
265                ..DataFlowConfig::default()
266            },
267        );
268
269        let (request, scope, agent_id, server_id) = make_ctx();
270        let ctx = guard_ctx(&request, &scope, &agent_id, &server_id);
271        assert_eq!(guard.evaluate(&ctx).expect("ok"), Verdict::Deny);
272    }
273
274    #[test]
275    fn allows_when_under_limit() {
276        let journal = make_journal("sess-under");
277        journal
278            .record(RecordParams {
279                tool_name: "read_file".to_string(),
280                server_id: "srv".to_string(),
281                agent_id: "agent".to_string(),
282                bytes_read: 100,
283                bytes_written: 50,
284                delegation_depth: 0,
285                allowed: true,
286            })
287            .expect("record");
288
289        let guard = DataFlowGuard::new(
290            journal,
291            DataFlowConfig {
292                max_bytes_read: Some(1000),
293                max_bytes_written: Some(1000),
294                max_bytes_total: Some(2000),
295            },
296        );
297
298        let (request, scope, agent_id, server_id) = make_ctx();
299        let ctx = guard_ctx(&request, &scope, &agent_id, &server_id);
300        assert_eq!(guard.evaluate(&ctx).expect("ok"), Verdict::Allow);
301    }
302
303    #[test]
304    fn cumulative_across_multiple_entries() {
305        let journal = make_journal("sess-cumulative");
306        for _ in 0..5 {
307            journal
308                .record(RecordParams {
309                    tool_name: "read_file".to_string(),
310                    server_id: "srv".to_string(),
311                    agent_id: "agent".to_string(),
312                    bytes_read: 200,
313                    bytes_written: 0,
314                    delegation_depth: 0,
315                    allowed: true,
316                })
317                .expect("record");
318        }
319        // Total bytes_read = 1000.
320
321        let guard = DataFlowGuard::new(
322            journal,
323            DataFlowConfig {
324                max_bytes_read: Some(999),
325                ..DataFlowConfig::default()
326            },
327        );
328
329        let (request, scope, agent_id, server_id) = make_ctx();
330        let ctx = guard_ctx(&request, &scope, &agent_id, &server_id);
331        assert_eq!(guard.evaluate(&ctx).expect("ok"), Verdict::Deny);
332    }
333}
chio_guards/data_flow.rs

chio_guards/
data_flow.rs
