Skip to main content

sentri_library/
loader.rs

1//! Library loader for TOML-based invariants.
2
3use sentri_core::model::Invariant;
4use sentri_core::Result;
5use sentri_dsl_parser::parse_invariant;
6use std::path::Path;
7use tracing::info;
8
9/// Loads invariants from TOML files.
10pub struct LibraryLoader;
11
12impl LibraryLoader {
13    /// Load invariants from a TOML file.
14    ///
15    /// Expects TOML structure like:
16    /// ```toml
17    /// [[invariants]]
18    /// name = "balance_conservation"
19    /// expression = "sum_balances == total_supply"
20    /// severity = "critical"
21    /// ```
22    pub fn load_from_toml(path: &Path) -> Result<Vec<Invariant>> {
23        info!("Loading invariants from {:?}", path);
24
25        let content = std::fs::read_to_string(path).map_err(sentri_core::InvarError::IoError)?;
26
27        // Parse TOML
28        let table: toml::Table = toml::from_str(&content)
29            .map_err(|e| sentri_core::InvarError::ConfigError(e.to_string()))?;
30
31        let mut invariants = Vec::new();
32
33        // Extract invariants from table
34        if let Some(inv_array) = table.get("invariants").and_then(|v| v.as_array()) {
35            for (idx, inv_table) in inv_array.iter().enumerate() {
36                match parse_invariant_table(inv_table) {
37                    Ok(inv) => {
38                        info!("Loaded invariant: {}", inv.name);
39                        invariants.push(inv);
40                    }
41                    Err(e) => {
42                        tracing::warn!("Failed to parse invariant at index {}: {}", idx, e);
43                    }
44                }
45            }
46        }
47
48        info!(
49            "Loaded {} invariants from {}",
50            invariants.len(),
51            path.display()
52        );
53        Ok(invariants)
54    }
55
56    /// Load all invariants from a directory.
57    pub fn load_from_dir(dir: &Path) -> Result<Vec<Invariant>> {
58        let mut all_invariants = Vec::new();
59
60        // Read all .toml files in directory
61        let entries = std::fs::read_dir(dir).map_err(sentri_core::InvarError::IoError)?;
62
63        for entry in entries {
64            let entry = entry.map_err(sentri_core::InvarError::IoError)?;
65            let path = entry.path();
66
67            if path.extension().is_some_and(|ext| ext == "toml") {
68                let invariants = Self::load_from_toml(&path)?;
69                all_invariants.extend(invariants);
70            }
71        }
72
73        Ok(all_invariants)
74    }
75}
76
77/// Parse an invariant from a TOML table value.
78///
79/// Creates a complete invariant by parsing the expression string using the DSL parser.
80fn parse_invariant_table(table: &toml::Value) -> Result<Invariant> {
81    let table = table.as_table().ok_or_else(|| {
82        sentri_core::InvarError::ConfigError("Invariant must be a table".to_string())
83    })?;
84
85    let name = table
86        .get("name")
87        .and_then(|v| v.as_str())
88        .ok_or_else(|| {
89            sentri_core::InvarError::ConfigError("Invariant must have a 'name' field".to_string())
90        })?
91        .to_string();
92
93    let expression_str = table
94        .get("expression")
95        .and_then(|v| v.as_str())
96        .ok_or_else(|| {
97            sentri_core::InvarError::ConfigError(
98                "Invariant must have an 'expression' field".to_string(),
99            )
100        })?;
101
102    let severity = table
103        .get("severity")
104        .and_then(|v| v.as_str())
105        .unwrap_or("medium")
106        .to_string();
107
108    let category = table
109        .get("category")
110        .and_then(|v| v.as_str())
111        .unwrap_or("general")
112        .to_string();
113
114    let description = table
115        .get("description")
116        .and_then(|v| v.as_str())
117        .map(|s| s.to_string());
118
119    // Parse the expression string using the DSL parser
120    // Construct full invariant format for the parser
121    let full_invariant_str = format!(r#"invariant {} {{ {} }}"#, name, expression_str);
122
123    let mut parsed_invariant = parse_invariant(&full_invariant_str)?;
124
125    // Override with TOML-provided severity and category
126    parsed_invariant.severity = severity;
127    parsed_invariant.category = category;
128    if description.is_some() {
129        parsed_invariant.description = description;
130    }
131
132    info!(
133        "Parsed invariant '{}' with expression '{}' (severity: {}, category: {})",
134        name, expression_str, parsed_invariant.severity, parsed_invariant.category
135    );
136
137    Ok(parsed_invariant)
138}