agentguard_core/policy/
store.rs

1//! Policy store: loads/saves/validates Cedar policies and schemas.
2
3use crate::error::Result;
4use crate::policy::types::{PolicySource, Severity, ValidationIssue, ValidationReport};
5use crate::schema::SchemaParsed;
6use cedar_policy::{Policy, PolicyId, PolicySet, ValidationMode, Validator};
7use std::io::Read;
8use std::path::{Path, PathBuf};
9use std::str::FromStr;
10use walkdir::WalkDir;
11
12const MAX_POLICY_FILES: usize = 1024;
13const MAX_POLICY_TREE_ENTRIES: usize = 4096;
14const MAX_POLICY_FILE_BYTES: u64 = 1_048_576;
15const MAX_TOTAL_POLICY_BYTES: u64 = 16_777_216;
16const MAX_SCHEMA_BYTES: u64 = 1_048_576;
17
18#[derive(Debug, Clone)]
19pub struct PolicyStore {
20    pub root: PathBuf,
21}
22
23impl PolicyStore {
24    pub fn open(root: impl AsRef<Path>) -> Result<Self> {
25        let root = root.as_ref().to_path_buf();
26        if !root.exists() {
27            std::fs::create_dir_all(root.join("policies"))?;
28        }
29        Ok(Self { root })
30    }
31
32    pub fn default_root() -> PathBuf {
33        PathBuf::from(".agentguard")
34    }
35
36    pub fn schema_path(&self) -> PathBuf {
37        self.root.join("schema.cedarschema")
38    }
39
40    pub fn policies_dir(&self) -> PathBuf {
41        self.root.join("policies")
42    }
43
44    pub fn load_policies(&self) -> Result<(PolicySet, Vec<PolicySource>)> {
45        let mut set = PolicySet::new();
46        let mut sources = Vec::new();
47        let dir = self.policies_dir();
48
49        if !dir.exists() {
50            return Ok((set, sources));
51        }
52
53        let mut policy_count = 0usize;
54        let mut total_bytes = 0u64;
55        let mut entries_seen = 0usize;
56        for entry in WalkDir::new(&dir) {
57            let entry = entry.map_err(|e| crate::error::Error::Walk(e.to_string()))?;
58            entries_seen += 1;
59            if entries_seen > MAX_POLICY_TREE_ENTRIES {
60                return Err(crate::error::Error::PolicyParse {
61                    message: format!(
62                        "policy store tree exceeds the limit of {MAX_POLICY_TREE_ENTRIES} entries"
63                    ),
64                    file: entry.path().display().to_string(),
65                });
66            }
67            if entry.file_type().is_file()
68                && entry.path().extension().and_then(|s| s.to_str()) == Some("cedar")
69            {
70                policy_count += 1;
71                if policy_count > MAX_POLICY_FILES {
72                    return Err(crate::error::Error::PolicyParse {
73                        message: format!("policy count exceeds the limit of {MAX_POLICY_FILES}"),
74                        file: entry.path().display().to_string(),
75                    });
76                }
77                let src = read_bounded_text(entry.path(), MAX_POLICY_FILE_BYTES, "policy")?;
78                total_bytes = total_bytes.checked_add(src.len() as u64).ok_or_else(|| {
79                    crate::error::Error::PolicyParse {
80                        message: "total policy source size overflow".into(),
81                        file: entry.path().display().to_string(),
82                    }
83                })?;
84                if total_bytes > MAX_TOTAL_POLICY_BYTES {
85                    return Err(crate::error::Error::PolicyParse {
86                        message: format!(
87                            "total policy source size exceeds the limit of {MAX_TOTAL_POLICY_BYTES} bytes"
88                        ),
89                        file: entry.path().display().to_string(),
90                    });
91                }
92                sources.push(PolicySource {
93                    path: entry.path().to_path_buf(),
94                    text: src,
95                });
96            }
97        }
98
99        // Parse only after the complete input set is within its bounds. This
100        // prevents a directory with many large, valid prefix files from
101        // allocating policy ASTs before the total-size cap is enforced.
102        // Sort after the bounded walk to preserve deterministic source order
103        // without buffering an unbounded directory for WalkDir's sorter.
104        sources.sort_by(|left, right| left.path.cmp(&right.path));
105        for src in &sources {
106            let file_set =
107                PolicySet::from_str(&src.text).map_err(|e| crate::error::Error::PolicyParse {
108                    message: e.to_string(),
109                    file: src.path.display().to_string(),
110                })?;
111            set.merge(&file_set, true)
112                .map_err(|e| crate::error::Error::PolicyParse {
113                    message: e.to_string(),
114                    file: src.path.display().to_string(),
115                })?;
116        }
117
118        Ok((set, sources))
119    }
120
121    pub fn load_schema(&self) -> Result<Option<SchemaParsed>> {
122        let p = self.schema_path();
123        if !p.exists() {
124            return Ok(None);
125        }
126        let text = read_bounded_text(&p, MAX_SCHEMA_BYTES, "schema")?;
127        let (schema, _warnings) = cedar_policy::Schema::from_cedarschema_str(&text)
128            .map_err(|e| crate::error::Error::Schema(e.to_string()))?;
129        Ok(Some(SchemaParsed {
130            schema,
131            source: text,
132        }))
133    }
134
135    pub fn validate(&self) -> Result<ValidationReport> {
136        let (policies, sources) = self.load_policies()?;
137        let schema = self.load_schema()?;
138
139        let mut errors = Vec::new();
140        let mut warnings = Vec::new();
141
142        for src in &sources {
143            let pid = PolicyId::new(src.path.to_string_lossy());
144            if let Err(e) = Policy::parse(Some(pid), src.text.clone()) {
145                errors.push(ValidationIssue {
146                    policy: src.path.to_string_lossy().to_string(),
147                    severity: Severity::Error,
148                    message: e.to_string(),
149                });
150            }
151        }
152
153        if let Some(s) = &schema {
154            let validator = Validator::new(s.schema.clone());
155            let result = validator.validate(&policies, ValidationMode::Strict);
156            for err in result.validation_errors() {
157                errors.push(ValidationIssue {
158                    policy: err.policy_id().to_string(),
159                    severity: Severity::Error,
160                    message: err.to_string(),
161                });
162            }
163            for warn in result.validation_warnings() {
164                warnings.push(ValidationIssue {
165                    policy: warn.policy_id().to_string(),
166                    severity: Severity::Warning,
167                    message: warn.to_string(),
168                });
169            }
170        } else {
171            warnings.push(ValidationIssue {
172                policy: "<store>".into(),
173                severity: Severity::Warning,
174                message: "no schema present; skipping type validation".into(),
175            });
176        }
177
178        Ok(ValidationReport {
179            policy_count: policies.policies().count(),
180            errors,
181            warnings,
182        })
183    }
184
185    /// Write a policy file under `policies/`. `name` is sanitized to a
186    /// single filename component: path separators and `..` segments are
187    /// stripped, and the result is rejected if it would be empty or a
188    /// parent-directory reference.
189    pub fn write_policy(&self, name: &str, text: &str) -> Result<PathBuf> {
190        // Strip path separators, NULs, and the like. We keep the
191        // sanitization simple: anything that's not [A-Za-z0-9._-] is
192        // replaced with '_'. This blocks '..' and '/'.
193        let safe: String = name
194            .chars()
195            .map(|c| {
196                if c.is_ascii_alphanumeric() || c == '.' || c == '_' || c == '-' {
197                    c
198                } else {
199                    '_'
200                }
201            })
202            .collect();
203        // Reject any name that, after sanitization, is empty or starts
204        // with '.' (a hidden file or '.' / '..').
205        if safe.is_empty() || safe.starts_with('.') {
206            return Err(crate::error::Error::PolicyParse {
207                message: format!("invalid policy name: {:?}", name),
208                file: name.to_string(),
209            });
210        }
211        let path = self.policies_dir().join(format!("{}.cedar", safe));
212        // Defense in depth: confirm the resolved path stays inside the
213        // policies directory (guards against symlink races).
214        let policies_dir = self.policies_dir();
215        let canonical_policies = std::fs::canonicalize(&policies_dir).unwrap_or(policies_dir);
216        let resolved = path.clone();
217        if let Ok(canonical) = std::fs::canonicalize(&resolved) {
218            if !canonical.starts_with(&canonical_policies) {
219                return Err(crate::error::Error::PolicyParse {
220                    message: "policy path escapes policies dir".into(),
221                    file: path.display().to_string(),
222                });
223            }
224        }
225        std::fs::create_dir_all(self.policies_dir())?;
226        std::fs::write(&path, text)?;
227        Ok(path)
228    }
229
230    pub fn write_schema(&self, text: &str) -> Result<()> {
231        std::fs::write(self.schema_path(), text)?;
232        Ok(())
233    }
234}
235
236fn read_bounded_text(path: &Path, max_bytes: u64, label: &str) -> Result<String> {
237    let mut bytes = Vec::new();
238    std::fs::File::open(path)?
239        .take(max_bytes + 1)
240        .read_to_end(&mut bytes)?;
241    if bytes.len() as u64 > max_bytes {
242        return Err(crate::error::Error::Other(format!(
243            "{} exceeds the {label} size limit of {max_bytes} bytes",
244            path.display()
245        )));
246    }
247    String::from_utf8(bytes).map_err(|error| {
248        crate::error::Error::Other(format!("{} is not valid UTF-8: {error}", path.display()))
249    })
250}
251
252#[cfg(test)]
253mod tests {
254    use super::*;
255    use std::fs;
256    use tempfile::tempdir;
257
258    #[test]
259    fn policy_file_count_over_limit_fails_instead_of_silently_truncating() {
260        let dir = tempdir().unwrap();
261        let store = PolicyStore::open(dir.path()).unwrap();
262        fs::create_dir_all(store.policies_dir()).unwrap();
263        for index in 0..=MAX_POLICY_FILES {
264            fs::write(store.policies_dir().join(format!("{index:04}.cedar")), "").unwrap();
265        }
266
267        let error = store.load_policies().unwrap_err();
268        assert!(error.to_string().contains("policy count exceeds"));
269    }
270
271    #[test]
272    fn policy_tree_entry_count_is_bounded() {
273        let dir = tempdir().unwrap();
274        let store = PolicyStore::open(dir.path()).unwrap();
275        fs::create_dir_all(store.policies_dir()).unwrap();
276        for index in 0..MAX_POLICY_TREE_ENTRIES {
277            fs::write(
278                store.policies_dir().join(format!("{index:04}.txt")),
279                "ignored non-policy file",
280            )
281            .unwrap();
282        }
283
284        let error = store.load_policies().unwrap_err();
285        assert!(error.to_string().contains("policy store tree exceeds"));
286    }
287
288    #[test]
289    fn total_policy_source_size_is_bounded_before_policy_parsing() {
290        let dir = tempdir().unwrap();
291        let store = PolicyStore::open(dir.path()).unwrap();
292        fs::create_dir_all(store.policies_dir()).unwrap();
293        let comment = format!("// {}\n", "x".repeat(MAX_POLICY_FILE_BYTES as usize - 4));
294        for index in 0..=MAX_TOTAL_POLICY_BYTES / MAX_POLICY_FILE_BYTES {
295            fs::write(
296                store.policies_dir().join(format!("{index:04}.cedar")),
297                &comment,
298            )
299            .unwrap();
300        }
301
302        let error = store.load_policies().unwrap_err();
303        assert!(error
304            .to_string()
305            .contains("total policy source size exceeds"));
306    }
307
308    #[test]
309    fn individual_policy_file_size_is_bounded() {
310        let dir = tempdir().unwrap();
311        let store = PolicyStore::open(dir.path()).unwrap();
312        fs::create_dir_all(store.policies_dir()).unwrap();
313        fs::write(
314            store.policies_dir().join("oversized.cedar"),
315            vec![b'x'; MAX_POLICY_FILE_BYTES as usize + 1],
316        )
317        .unwrap();
318
319        let error = store.load_policies().unwrap_err();
320        assert!(error.to_string().contains("policy size limit"));
321    }
322
323    #[test]
324    fn schema_size_is_bounded_before_schema_parsing() {
325        let dir = tempdir().unwrap();
326        let store = PolicyStore::open(dir.path()).unwrap();
327        fs::write(
328            store.schema_path(),
329            vec![b'x'; MAX_SCHEMA_BYTES as usize + 1],
330        )
331        .unwrap();
332
333        let error = store.load_schema().unwrap_err();
334        assert!(error.to_string().contains("schema size limit"));
335    }
336
337    #[test]
338    fn policy_directories_do_not_consume_the_policy_file_limit() {
339        let dir = tempdir().unwrap();
340        let store = PolicyStore::open(dir.path()).unwrap();
341        let nested = store.policies_dir().join("nested");
342        fs::create_dir_all(&nested).unwrap();
343        fs::write(
344            nested.join("allow.cedar"),
345            "permit(principal, action, resource);",
346        )
347        .unwrap();
348
349        let (policies, sources) = store.load_policies().unwrap();
350        assert_eq!(policies.policies().count(), 1);
351        assert_eq!(sources.len(), 1);
352    }
353}