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

nimbus/
targeting.rs
