chio_guards/
memory_governance.rs

1//! MemoryGovernanceGuard -- enforce memory store allowlist, retention
2//! TTL ceilings, and per-session memory-entry counts on
3//! [`ToolAction::MemoryWrite`] and [`ToolAction::MemoryRead`] actions.
4//!
5//! See `docs/protocols/STRUCTURAL-SECURITY-FIXES.md` section 3. The guard
6//! sources its policy from two places:
7//!
8//! 1. **Capability constraints** on the matched grant
9//!    ([`Constraint::MemoryStoreAllowlist`]): when present, writes and
10//!    reads targeting a store outside the allowlist are denied.
11//! 2. **Guard configuration** ([`MemoryGovernanceConfig`]): provides
12//!    deployment-wide defaults for `max_memory_entries`,
13//!    `max_retention_ttl_secs`, and per-store overrides.  Operators can
14//!    use these even when the capability grammar does not surface the
15//!    equivalent constraints.
16//!
17//! The guard keeps an in-memory per-session counter of memory writes so
18//! it can enforce [`MemoryGovernanceConfig::max_memory_entries`]
19//! deterministically without touching shared kernel state.
20//!
21//! # Fail-closed semantics
22//!
23//! - memory writes without a parseable store key are denied when the
24//!   matched grant carries a non-empty `MemoryStoreAllowlist`;
25//! - malformed deny-pattern regex input causes
26//!   [`MemoryGovernanceGuard::with_config`] to return
27//!   [`MemoryGovernanceError::InvalidPattern`];
28//! - writes with an explicit retention TTL above `max_retention_ttl_secs`
29//!   are denied;
30//! - writes whose total matches / exceeds `max_memory_entries` are
31//!   denied (fail-closed on counter mutex poisoning).
32
33use std::collections::HashMap;
34use std::sync::Mutex;
35
36use regex::Regex;
37use serde::{Deserialize, Serialize};
38use serde_json::Value;
39
40use chio_core::capability::scope::Constraint;
41use chio_kernel::{Guard, GuardContext, GuardDecision, KernelError};
42
43use crate::action::{extract_action_checked, ToolAction};
44
45/// Errors produced when building a [`MemoryGovernanceGuard`].
46#[derive(Debug, thiserror::Error)]
47pub enum MemoryGovernanceError {
48    /// A `deny_patterns` entry was not a valid regex.
49    #[error("invalid deny pattern `{pattern}`: {source}")]
50    InvalidPattern {
51        pattern: String,
52        #[source]
53        source: regex::Error,
54    },
55}
56
57/// Configuration for [`MemoryGovernanceGuard`].
58#[derive(Clone, Debug, Deserialize, Serialize)]
59#[serde(deny_unknown_fields)]
60pub struct MemoryGovernanceConfig {
61    /// Enable/disable the guard entirely.
62    #[serde(default = "default_true")]
63    pub enabled: bool,
64    /// Hard-coded store allowlist applied on top of the capability-level
65    /// [`Constraint::MemoryStoreAllowlist`].  Empty means "no additional
66    /// allowlist" (capability-level list still applies).
67    #[serde(default)]
68    pub store_allowlist: Vec<String>,
69    /// Maximum memory-entry count per agent + session combination.  When
70    /// `Some(n)`, the `n`-th write is denied.
71    #[serde(default, skip_serializing_if = "Option::is_none")]
72    pub max_memory_entries: Option<u64>,
73    /// Maximum retention TTL (seconds) allowed on a single write.  When
74    /// `Some(ttl)`, writes requesting a larger TTL -- or indefinite
75    /// retention (missing TTL) -- are denied.
76    #[serde(default, skip_serializing_if = "Option::is_none")]
77    pub max_retention_ttl_secs: Option<u64>,
78    /// Maximum content size (bytes) for a single memory write.  `None`
79    /// disables the check.
80    #[serde(default, skip_serializing_if = "Option::is_none")]
81    pub max_content_size_bytes: Option<u64>,
82    /// Extra regex patterns that deny a write when the content matches.
83    #[serde(default)]
84    pub deny_patterns: Vec<String>,
85}
86
87fn default_true() -> bool {
88    true
89}
90
91impl Default for MemoryGovernanceConfig {
92    fn default() -> Self {
93        Self {
94            enabled: true,
95            store_allowlist: Vec::new(),
96            max_memory_entries: None,
97            max_retention_ttl_secs: None,
98            max_content_size_bytes: None,
99            deny_patterns: Vec::new(),
100        }
101    }
102}
103
104/// Session key used for per-session memory-entry counting.
105type SessionKey = (String, String); // (agent_id, capability_id)
106
107/// Guard implementing memory governance.
108pub struct MemoryGovernanceGuard {
109    enabled: bool,
110    store_allowlist: Vec<String>,
111    max_memory_entries: Option<u64>,
112    max_retention_ttl_secs: Option<u64>,
113    max_content_size_bytes: Option<u64>,
114    deny_patterns: Vec<Regex>,
115    counters: Mutex<HashMap<SessionKey, u64>>,
116}
117
118impl MemoryGovernanceGuard {
119    /// Build a guard with default configuration (no limits).  Non-guard
120    /// code paths remain fully permissive until a capability constraint
121    /// or config field is supplied.
122    pub fn new() -> Self {
123        Self::with_config(MemoryGovernanceConfig::default()).unwrap_or_else(|_| Self {
124            enabled: true,
125            store_allowlist: Vec::new(),
126            max_memory_entries: None,
127            max_retention_ttl_secs: None,
128            max_content_size_bytes: None,
129            deny_patterns: Vec::new(),
130            counters: Mutex::new(HashMap::new()),
131        })
132    }
133
134    /// Build a guard with explicit configuration.
135    pub fn with_config(config: MemoryGovernanceConfig) -> Result<Self, MemoryGovernanceError> {
136        let mut deny_patterns = Vec::with_capacity(config.deny_patterns.len());
137        for pat in &config.deny_patterns {
138            let re = Regex::new(pat).map_err(|e| MemoryGovernanceError::InvalidPattern {
139                pattern: pat.clone(),
140                source: e,
141            })?;
142            deny_patterns.push(re);
143        }
144        Ok(Self {
145            enabled: config.enabled,
146            store_allowlist: config.store_allowlist,
147            max_memory_entries: config.max_memory_entries,
148            max_retention_ttl_secs: config.max_retention_ttl_secs,
149            max_content_size_bytes: config.max_content_size_bytes,
150            deny_patterns,
151            counters: Mutex::new(HashMap::new()),
152        })
153    }
154
155    /// Current counter value for a session (test / observability helper).
156    pub fn session_count(&self, agent_id: &str, capability_id: &str) -> u64 {
157        self.counters
158            .lock()
159            .ok()
160            .and_then(|g| {
161                g.get(&(agent_id.to_string(), capability_id.to_string()))
162                    .copied()
163            })
164            .unwrap_or(0)
165    }
166
167    /// Gather the effective store allowlist from the matched grant plus
168    /// the guard-level config.  Returns `None` if neither source supplies
169    /// a non-empty allowlist.
170    fn effective_store_allowlist<'a>(&'a self, ctx: &'a GuardContext<'a>) -> Option<Vec<String>> {
171        let mut combined: Vec<String> = self.store_allowlist.clone();
172        if let Some(grant) = ctx
173            .matched_grant_index
174            .and_then(|i| ctx.scope.grants.get(i))
175        {
176            for c in &grant.constraints {
177                if let Constraint::MemoryStoreAllowlist(list) = c {
178                    combined.extend(list.iter().cloned());
179                }
180            }
181        }
182        if combined.is_empty() {
183            None
184        } else {
185            Some(combined)
186        }
187    }
188
189    /// Increment the per-session write counter and return the new value.
190    /// Fails closed (treats poisoning as "over limit") on mutex poisoning.
191    fn bump_counter(&self, key: SessionKey) -> Result<u64, KernelError> {
192        let mut guard = self.counters.lock().map_err(|_| {
193            KernelError::Internal("memory-governance guard counter mutex poisoned".to_string())
194        })?;
195        let entry = guard.entry(key).or_insert(0);
196        *entry = entry.saturating_add(1);
197        Ok(*entry)
198    }
199}
200
201impl Default for MemoryGovernanceGuard {
202    fn default() -> Self {
203        Self::new()
204    }
205}
206
207impl Guard for MemoryGovernanceGuard {
208    fn name(&self) -> &str {
209        "memory-governance"
210    }
211
212    fn evaluate(&self, ctx: &GuardContext) -> Result<GuardDecision, KernelError> {
213        if !self.enabled {
214            return Ok(GuardDecision::allow());
215        }
216
217        let action = match extract_action_checked(&ctx.request.tool_name, &ctx.request.arguments) {
218            Ok(action) => action,
219            Err(_) => return Ok(GuardDecision::deny(Vec::new())),
220        };
221
222        match action {
223            ToolAction::MemoryWrite { store, .. } => self.evaluate_write(ctx, &store),
224            ToolAction::MemoryRead { store, .. } => self.evaluate_read(ctx, &store),
225            _ => Ok(GuardDecision::allow()),
226        }
227    }
228}
229
230impl MemoryGovernanceGuard {
231    fn evaluate_write(
232        &self,
233        ctx: &GuardContext,
234        store: &str,
235    ) -> Result<GuardDecision, KernelError> {
236        // 1. Store allowlist (capability + guard config).
237        if let Some(allow) = self.effective_store_allowlist(ctx) {
238            if !allow.iter().any(|s| store_matches(s, store)) {
239                return Ok(GuardDecision::deny(Vec::new()));
240            }
241        }
242
243        // 2. Retention TTL ceiling.
244        if let Some(max_ttl) = self.max_retention_ttl_secs {
245            let requested = extract_retention_ttl(&ctx.request.arguments);
246            match requested {
247                None => {
248                    // Missing TTL with a configured ceiling is treated
249                    // as a request for indefinite retention and denied.
250                    return Ok(GuardDecision::deny(Vec::new()));
251                }
252                Some(ttl) if ttl > max_ttl => {
253                    return Ok(GuardDecision::deny(Vec::new()));
254                }
255                Some(_) => {}
256            }
257        }
258
259        // 3. Content size.
260        if let Some(max_bytes) = self.max_content_size_bytes {
261            match extract_content_size_bytes(&ctx.request.arguments) {
262                Some(size) => {
263                    if size > max_bytes {
264                        return Ok(GuardDecision::deny(Vec::new()));
265                    }
266                }
267                None => {
268                    return Ok(GuardDecision::deny(Vec::new()));
269                }
270            }
271        }
272
273        // 4. Deny patterns on content.
274        if !self.deny_patterns.is_empty() {
275            if let Some(content) = extract_content_text(&ctx.request.arguments) {
276                for re in &self.deny_patterns {
277                    if re.is_match(&content) {
278                        return Ok(GuardDecision::deny(Vec::new()));
279                    }
280                }
281            }
282        }
283
284        // 5. Per-session entry limit.  We bump the counter only after
285        //    the previous gates pass; denials do not consume quota.
286        if let Some(max_entries) = self.max_memory_entries {
287            let key = (ctx.agent_id.to_string(), ctx.request.capability.id.clone());
288            let count = self.bump_counter(key)?;
289            if count > max_entries {
290                return Ok(GuardDecision::deny(Vec::new()));
291            }
292        }
293
294        Ok(GuardDecision::allow())
295    }
296
297    fn evaluate_read(&self, ctx: &GuardContext, store: &str) -> Result<GuardDecision, KernelError> {
298        // Reads respect the store allowlist so an agent cannot read from
299        // a forbidden store even when the write path is blocked.
300        if let Some(allow) = self.effective_store_allowlist(ctx) {
301            if !allow.iter().any(|s| store_matches(s, store)) {
302                return Ok(GuardDecision::deny(Vec::new()));
303            }
304        }
305        Ok(GuardDecision::allow())
306    }
307}
308
309/// Store allowlist match: supports exact match and `*` wildcard.
310fn store_matches(pattern: &str, store: &str) -> bool {
311    if pattern == "*" {
312        return true;
313    }
314    if let Some(prefix) = pattern.strip_suffix('*') {
315        return store.starts_with(prefix);
316    }
317    pattern == store
318}
319
320/// Read an explicit retention TTL (seconds) from the arguments.
321fn extract_retention_ttl(arguments: &Value) -> Option<u64> {
322    for key in [
323        "retention_ttl",
324        "retentionTtl",
325        "retention_ttl_secs",
326        "retentionTtlSecs",
327        "ttl",
328        "ttl_secs",
329        "expires_in",
330        "expiresIn",
331    ] {
332        if let Some(v) = arguments.get(key).and_then(|v| v.as_u64()) {
333            return Some(v);
334        }
335    }
336    None
337}
338
339/// Read an explicit content byte size from the arguments, falling back
340/// to the length of the `content` / `text` string when present.
341fn extract_content_size_bytes(arguments: &Value) -> Option<u64> {
342    for key in ["content_size", "contentSize", "content_bytes", "size"] {
343        if let Some(v) = arguments.get(key).and_then(|v| v.as_u64()) {
344            return Some(v);
345        }
346    }
347    extract_content_text(arguments).map(|s| s.len() as u64)
348}
349
350/// Extract the text body of a memory write for regex / size checks.
351fn extract_content_text(arguments: &Value) -> Option<String> {
352    for key in ["content", "text", "value", "vector_text", "payload"] {
353        if let Some(v) = arguments.get(key).and_then(|v| v.as_str()) {
354            if !v.is_empty() {
355                return Some(v.to_string());
356            }
357        }
358    }
359    None
360}
361
362#[cfg(test)]
363mod tests {
364    use super::*;
365
366    #[test]
367    fn store_matches_wildcards() {
368        assert!(store_matches("*", "anything"));
369        assert!(store_matches("agent-*", "agent-notes"));
370        assert!(!store_matches("agent-*", "other"));
371        assert!(store_matches("agent-notes", "agent-notes"));
372    }
373
374    #[test]
375    fn extract_retention_ttl_reads_common_keys() {
376        let args = serde_json::json!({"ttl": 600});
377        assert_eq!(extract_retention_ttl(&args), Some(600));
378        let camel = serde_json::json!({"retentionTtl": 120});
379        assert_eq!(extract_retention_ttl(&camel), Some(120));
380        let none = serde_json::json!({});
381        assert_eq!(extract_retention_ttl(&none), None);
382    }
383
384    #[test]
385    fn content_size_falls_back_to_text_length() {
386        let args = serde_json::json!({"content": "hello"});
387        assert_eq!(extract_content_size_bytes(&args), Some(5));
388    }
389}
chio_guards/memory_governance.rs

chio_guards/
memory_governance.rs
