1use crate::{NimbusError, Result};
6
7use jexl_eval::Evaluator;
8use serde::Serialize;
9use serde_json::Value;
10
11cfg_if::cfg_if! {
12 if #[cfg(feature = "stateful")] {
13 use anyhow::anyhow;
14 use crate::{TargetingAttributes, stateful::{behavior::{EventStore, EventQueryType, query_event_store}, gecko_prefs::{GeckoPrefStore, query_gecko_pref_store}}};
15 use std::sync::{Arc, Mutex};
16 use firefox_versioning::compare::version_compare;
17 }
18}
19
20#[derive(Clone)]
21pub struct NimbusTargetingHelper {
22 pub(crate) context: Value,
23 #[cfg(feature = "stateful")]
24 pub(crate) event_store: Arc<Mutex<EventStore>>,
25 #[cfg(feature = "stateful")]
26 pub(crate) gecko_pref_store: Option<Arc<GeckoPrefStore>>,
27 #[cfg(feature = "stateful")]
28 pub(crate) targeting_attributes: Option<TargetingAttributes>,
29}
30
31impl NimbusTargetingHelper {
32 pub fn new<C: Serialize>(
33 context: C,
34 #[cfg(feature = "stateful")] event_store: Arc<Mutex<EventStore>>,
35 #[cfg(feature = "stateful")] gecko_pref_store: Option<Arc<GeckoPrefStore>>,
36 ) -> Self {
37 Self {
38 context: serde_json::to_value(context).unwrap(),
39 #[cfg(feature = "stateful")]
40 event_store,
41 #[cfg(feature = "stateful")]
42 gecko_pref_store,
43 #[cfg(feature = "stateful")]
44 targeting_attributes: None,
45 }
46 }
47
48 pub fn eval_jexl(&self, expr: String) -> Result<bool> {
49 cfg_if::cfg_if! {
50 if #[cfg(feature = "stateful")] {
51 jexl_eval(&expr, &self.context, self.event_store.clone(), self.gecko_pref_store.clone())
52 } else {
53 jexl_eval(&expr, &self.context)
54 }
55 }
56 }
57
58 pub fn evaluate_jexl_raw_value(&self, expr: &str) -> Result<Value> {
59 cfg_if::cfg_if! {
60 if #[cfg(feature = "stateful")] {
61 jexl_eval_raw(expr, &self.context, self.event_store.clone(), self.gecko_pref_store.clone())
62 } else {
63 jexl_eval_raw(expr, &self.context)
64 }
65 }
66 }
67
68 pub(crate) fn put(&self, key: &str, value: bool) -> Self {
69 let context = if let Value::Object(map) = &self.context {
70 let mut map = map.clone();
71 map.insert(key.to_string(), Value::Bool(value));
72 Value::Object(map)
73 } else {
74 self.context.clone()
75 };
76
77 Self {
78 context,
79 #[cfg(feature = "stateful")]
80 event_store: self.event_store.clone(),
81 #[cfg(feature = "stateful")]
82 gecko_pref_store: self.gecko_pref_store.clone(),
83 #[cfg(feature = "stateful")]
84 targeting_attributes: self.targeting_attributes.clone(),
85 }
86 }
87}
88
89pub fn jexl_eval_raw<Context: serde::Serialize>(
90 expression_statement: &str,
91 context: &Context,
92 #[cfg(feature = "stateful")] event_store: Arc<Mutex<EventStore>>,
93 #[cfg(feature = "stateful")] gecko_pref_store: Option<Arc<GeckoPrefStore>>,
94) -> Result<Value> {
95 let evaluator = Evaluator::new();
96
97 #[cfg(feature = "stateful")]
98 let evaluator = evaluator
99 .with_transform("versionCompare", |args| Ok(version_compare(args)?))
100 .with_transform("eventSum", |args| {
101 Ok(query_event_store(
102 event_store.clone(),
103 EventQueryType::Sum,
104 args,
105 )?)
106 })
107 .with_transform("eventCountNonZero", |args| {
108 Ok(query_event_store(
109 event_store.clone(),
110 EventQueryType::CountNonZero,
111 args,
112 )?)
113 })
114 .with_transform("eventAveragePerInterval", |args| {
115 Ok(query_event_store(
116 event_store.clone(),
117 EventQueryType::AveragePerInterval,
118 args,
119 )?)
120 })
121 .with_transform("eventAveragePerNonZeroInterval", |args| {
122 Ok(query_event_store(
123 event_store.clone(),
124 EventQueryType::AveragePerNonZeroInterval,
125 args,
126 )?)
127 })
128 .with_transform("eventLastSeen", |args| {
129 Ok(query_event_store(
130 event_store.clone(),
131 EventQueryType::LastSeen,
132 args,
133 )?)
134 })
135 .with_transform("preferenceIsUserSet", |args| {
136 Ok(query_gecko_pref_store(gecko_pref_store.clone(), args)?)
137 })
138 .with_transform("bucketSample", bucket_sample);
139
140 evaluator
141 .eval_in_context(expression_statement, context)
142 .map_err(|err| NimbusError::EvaluationError(err.to_string()))
143}
144
145pub fn jexl_eval<Context: serde::Serialize>(
150 expression_statement: &str,
151 context: &Context,
152 #[cfg(feature = "stateful")] event_store: Arc<Mutex<EventStore>>,
153 #[cfg(feature = "stateful")] gecko_pref_store: Option<Arc<GeckoPrefStore>>,
154) -> Result<bool> {
155 let res = jexl_eval_raw(
156 expression_statement,
157 context,
158 #[cfg(feature = "stateful")]
159 event_store,
160 #[cfg(feature = "stateful")]
161 gecko_pref_store,
162 )?;
163 match res.as_bool() {
164 Some(v) => Ok(v),
165 None => Err(NimbusError::InvalidExpression),
166 }
167}
168
169#[cfg(feature = "stateful")]
170fn bucket_sample(args: &[Value]) -> anyhow::Result<Value> {
171 fn get_arg_as_u32(args: &[Value], idx: usize, name: &str) -> anyhow::Result<u32> {
172 match args.get(idx) {
173 None => Err(anyhow!("{} doesn't exist in jexl transform", name)),
174 Some(Value::Number(n)) => {
175 let n: f64 = if let Some(n) = n.as_u64() {
176 n as f64
177 } else if let Some(n) = n.as_i64() {
178 n as f64
179 } else if let Some(n) = n.as_f64() {
180 n
181 } else {
182 unreachable!();
183 };
184
185 debug_assert!(n >= 0.0, "JEXL parser does not support negative values");
186 if n > u32::MAX as f64 {
187 Err(anyhow!("{} is out of range", name))
188 } else {
189 Ok(n as u32)
190 }
191 }
192 Some(_) => Err(anyhow!("{} is not a number", name)),
193 }
194 }
195
196 let input = args
197 .first()
198 .ok_or_else(|| anyhow!("input doesn't exist in jexl transform"))?;
199 let start = get_arg_as_u32(args, 1, "start")?;
200 let count = get_arg_as_u32(args, 2, "count")?;
201 let total = get_arg_as_u32(args, 3, "total")?;
202
203 let result = crate::sampling::bucket_sample(input, start, count, total)?;
204
205 Ok(Value::Bool(result))
206}