1use std::collections::HashMap;
2
3use rhai::{Dynamic, Engine, EvalAltResult, Scope};
4use thiserror::Error;
5use tracing::instrument;
6
7#[derive(Debug, Error)]
8pub enum ScriptError {
9 #[error("script error: {0}")]
10 Eval(Box<EvalAltResult>),
11 #[error("type error: {0}")]
12 Type(String),
13}
14
15pub type Result<T> = std::result::Result<T, ScriptError>;
16
17pub struct ScriptEngine {
20 engine: Engine,
21}
22
23impl Default for ScriptEngine {
24 fn default() -> Self {
25 Self::new()
26 }
27}
28
29impl ScriptEngine {
30 pub fn new() -> Self {
31 let mut engine = Engine::new();
32 engine.set_max_expr_depths(64, 64);
33 engine.set_max_operations(1_000_000);
34 register_date_fns(&mut engine);
35 Self { engine }
36 }
37
38 pub fn eval_bool(&self, expr: &str, vars: &HashMap<String, serde_json::Value>) -> Result<bool> {
40 let (result, _) = self.run(expr, vars)?;
41 result
42 .as_bool()
43 .map_err(|_| ScriptError::Type(format!("condition '{expr}' did not return bool")))
44 }
45
46 #[instrument(name = "script_run", skip(self, source, vars))]
49 pub fn run(
50 &self,
51 source: &str,
52 vars: &HashMap<String, serde_json::Value>,
53 ) -> Result<(Dynamic, HashMap<String, Dynamic>)> {
54 let mut scope = Scope::new();
55 for (k, v) in vars {
56 scope.push_dynamic(k.clone(), json_to_dynamic(v));
57 }
58
59 let result = self
60 .engine
61 .eval_with_scope::<Dynamic>(&mut scope, source)
62 .map_err(ScriptError::Eval)?;
63
64 let exports: HashMap<String, Dynamic> = scope
65 .iter_raw()
66 .map(|(name, _, val)| (name.to_owned(), val.clone()))
67 .collect();
68
69 Ok((result, exports))
70 }
71}
72
73fn register_date_fns(engine: &mut Engine) {
76 use chrono::{Datelike, Local, NaiveDate};
77
78 engine.register_fn("now_year", || Local::now().year() as i64);
79 engine.register_fn("now_month", || Local::now().month() as i64);
80 engine.register_fn("now_day", || Local::now().day() as i64);
81 engine.register_fn("today", || Local::now().format("%Y-%m-%d").to_string());
82
83 fn next_month(year: i64, month: i64) -> Option<(i32, u32)> {
86 if !(1..=12).contains(&month) {
87 return None;
88 }
89 if year < i32::MIN as i64 || year > i32::MAX as i64 {
90 return None;
91 }
92 if month == 12 {
93 let y2 = year.checked_add(1)?;
94 if y2 > i32::MAX as i64 {
95 return None;
96 }
97 Some((y2 as i32, 1))
98 } else {
99 Some((year as i32, (month + 1) as u32))
100 }
101 }
102
103 fn valid_ym(year: i64, month: i64) -> Option<(i32, u32)> {
104 if !(1..=12).contains(&month) {
105 return None;
106 }
107 if year < i32::MIN as i64 || year > i32::MAX as i64 {
108 return None;
109 }
110 Some((year as i32, month as u32))
111 }
112
113 engine.register_fn("end_of_month", |year: i64, month: i64| -> i64 {
114 next_month(year, month)
115 .and_then(|(y2, m2)| NaiveDate::from_ymd_opt(y2, m2, 1))
116 .and_then(|d| d.pred_opt())
117 .map(|d| d.day() as i64)
118 .unwrap_or(0)
119 });
120
121 engine.register_fn(
122 "end_of_month_str",
123 |year: i64, month: i64, fmt: &str| -> String {
124 next_month(year, month)
125 .and_then(|(y2, m2)| NaiveDate::from_ymd_opt(y2, m2, 1))
126 .and_then(|d| d.pred_opt())
127 .map(|d| d.format(fmt).to_string())
128 .unwrap_or_default()
129 },
130 );
131
132 engine.register_fn("start_of_month", |year: i64, month: i64| -> i64 {
133 valid_ym(year, month)
134 .and_then(|(y, m)| NaiveDate::from_ymd_opt(y, m, 1))
135 .map(|d| d.day() as i64)
136 .unwrap_or(0)
137 });
138
139 engine.register_fn(
140 "start_of_month_str",
141 |year: i64, month: i64, fmt: &str| -> String {
142 valid_ym(year, month)
143 .and_then(|(y, m)| NaiveDate::from_ymd_opt(y, m, 1))
144 .map(|d| d.format(fmt).to_string())
145 .unwrap_or_default()
146 },
147 );
148}
149
150fn json_to_dynamic(v: &serde_json::Value) -> Dynamic {
151 match v {
152 serde_json::Value::Null => Dynamic::UNIT,
153 serde_json::Value::Bool(b) => Dynamic::from(*b),
154 serde_json::Value::Number(n) => {
155 if let Some(i) = n.as_i64() {
156 Dynamic::from(i)
157 } else {
158 Dynamic::from(n.as_f64().unwrap_or(0.0))
159 }
160 }
161 serde_json::Value::String(s) => Dynamic::from(s.clone()),
162 serde_json::Value::Array(arr) => {
163 Dynamic::from(arr.iter().map(json_to_dynamic).collect::<Vec<_>>())
164 }
165 serde_json::Value::Object(map) => {
166 let m: rhai::Map = map
167 .iter()
168 .map(|(k, v)| (k.clone().into(), json_to_dynamic(v)))
169 .collect();
170 Dynamic::from(m)
171 }
172 }
173}