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    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
145// This is the common entry point to JEXL evaluation.
146// The targeting attributes and additional context should have been merged and calculated before
147// getting here.
148// Any additional transforms should be added here.
149pub 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}