nimbus/
targeting.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at https://mozilla.org/MPL/2.0/.
4
5use 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    #[cfg(feature = "stateful")]
59    pub fn eval_jexl_debug(&self, expression: String) -> Result<String> {
60        let eval_result = jexl_eval_raw(
61            &expression,
62            &self.context,
63            self.event_store.clone(),
64            self.gecko_pref_store.clone(),
65        );
66
67        let response = match eval_result {
68            Ok(value) => {
69                serde_json::json!({
70                    "success": true,
71                    "result": value
72                })
73            }
74            Err(e) => {
75                serde_json::json!({
76                    "success": false,
77                    "error": e.to_string()
78                })
79            }
80        };
81
82        serde_json::to_string_pretty(&response).map_err(|e| {
83            NimbusError::JSONError("Failed to serialize JEXL result".to_string(), e.to_string())
84        })
85    }
86
87    pub fn evaluate_jexl_raw_value(&self, expr: &str) -> Result<Value> {
88        cfg_if::cfg_if! {
89            if #[cfg(feature = "stateful")] {
90                jexl_eval_raw(expr, &self.context, self.event_store.clone(), self.gecko_pref_store.clone())
91            } else {
92                jexl_eval_raw(expr, &self.context)
93            }
94        }
95    }
96
97    pub(crate) fn put(&self, key: &str, value: bool) -> Self {
98        let context = if let Value::Object(map) = &self.context {
99            let mut map = map.clone();
100            map.insert(key.to_string(), Value::Bool(value));
101            Value::Object(map)
102        } else {
103            self.context.clone()
104        };
105
106        Self {
107            context,
108            #[cfg(feature = "stateful")]
109            event_store: self.event_store.clone(),
110            #[cfg(feature = "stateful")]
111            gecko_pref_store: self.gecko_pref_store.clone(),
112            #[cfg(feature = "stateful")]
113            targeting_attributes: self.targeting_attributes.clone(),
114        }
115    }
116}
117
118pub fn jexl_eval_raw<Context: serde::Serialize>(
119    expression_statement: &str,
120    context: &Context,
121    #[cfg(feature = "stateful")] event_store: Arc<Mutex<EventStore>>,
122    #[cfg(feature = "stateful")] gecko_pref_store: Option<Arc<GeckoPrefStore>>,
123) -> Result<Value> {
124    let evaluator = Evaluator::new();
125
126    #[cfg(feature = "stateful")]
127    let evaluator = evaluator
128        .with_transform("versionCompare", |args| Ok(version_compare(args)?))
129        .with_transform("eventSum", |args| {
130            Ok(query_event_store(
131                event_store.clone(),
132                EventQueryType::Sum,
133                args,
134            )?)
135        })
136        .with_transform("eventCountNonZero", |args| {
137            Ok(query_event_store(
138                event_store.clone(),
139                EventQueryType::CountNonZero,
140                args,
141            )?)
142        })
143        .with_transform("eventAveragePerInterval", |args| {
144            Ok(query_event_store(
145                event_store.clone(),
146                EventQueryType::AveragePerInterval,
147                args,
148            )?)
149        })
150        .with_transform("eventAveragePerNonZeroInterval", |args| {
151            Ok(query_event_store(
152                event_store.clone(),
153                EventQueryType::AveragePerNonZeroInterval,
154                args,
155            )?)
156        })
157        .with_transform("eventLastSeen", |args| {
158            Ok(query_event_store(
159                event_store.clone(),
160                EventQueryType::LastSeen,
161                args,
162            )?)
163        })
164        .with_transform("preferenceIsUserSet", |args| {
165            Ok(query_gecko_pref_store(gecko_pref_store.clone(), args)?)
166        })
167        .with_transform("bucketSample", bucket_sample);
168
169    evaluator
170        .eval_in_context(expression_statement, context)
171        .map_err(|err| NimbusError::EvaluationError(err.to_string()))
172}
173
174// This is the common entry point to JEXL evaluation.
175// The targeting attributes and additional context should have been merged and calculated before
176// getting here.
177// Any additional transforms should be added here.
178pub fn jexl_eval<Context: serde::Serialize>(
179    expression_statement: &str,
180    context: &Context,
181    #[cfg(feature = "stateful")] event_store: Arc<Mutex<EventStore>>,
182    #[cfg(feature = "stateful")] gecko_pref_store: Option<Arc<GeckoPrefStore>>,
183) -> Result<bool> {
184    let res = jexl_eval_raw(
185        expression_statement,
186        context,
187        #[cfg(feature = "stateful")]
188        event_store,
189        #[cfg(feature = "stateful")]
190        gecko_pref_store,
191    )?;
192    match res.as_bool() {
193        Some(v) => Ok(v),
194        None => Err(NimbusError::InvalidExpression),
195    }
196}
197
198#[cfg(feature = "stateful")]
199fn bucket_sample(args: &[Value]) -> anyhow::Result<Value> {
200    fn get_arg_as_u32(args: &[Value], idx: usize, name: &str) -> anyhow::Result<u32> {
201        match args.get(idx) {
202            None => Err(anyhow!("{} doesn't exist in jexl transform", name)),
203            Some(Value::Number(n)) => {
204                let n: f64 = if let Some(n) = n.as_u64() {
205                    n as f64
206                } else if let Some(n) = n.as_i64() {
207                    n as f64
208                } else if let Some(n) = n.as_f64() {
209                    n
210                } else {
211                    unreachable!();
212                };
213
214                debug_assert!(n >= 0.0, "JEXL parser does not support negative values");
215                if n > u32::MAX as f64 {
216                    Err(anyhow!("{} is out of range", name))
217                } else {
218                    Ok(n as u32)
219                }
220            }
221            Some(_) => Err(anyhow!("{} is not a number", name)),
222        }
223    }
224
225    let input = args
226        .first()
227        .ok_or_else(|| anyhow!("input doesn't exist in jexl transform"))?;
228    let start = get_arg_as_u32(args, 1, "start")?;
229    let count = get_arg_as_u32(args, 2, "count")?;
230    let total = get_arg_as_u32(args, 3, "total")?;
231
232    let result = crate::sampling::bucket_sample(input, start, count, total)?;
233
234    Ok(Value::Bool(result))
235}
nimbus/targeting.rs

nimbus/
targeting.rs
