treetop_core/
labels.rs

1use arc_swap::ArcSwap;
2use regex::Regex;
3use std::sync::Arc;
4
5use crate::types::{AttrValue, Resource};
6
7/// Trait for objects that can label resources based on their attributes.
8///
9/// Implementations should be fast and side-effect free beyond mutating the
10/// provided `Resource`'s attributes. Repeated application should be safe and
11/// ideally idempotent.
12pub trait Labeler: Send + Sync {
13    /// Returns true if this labeler applies to resources of the given kind.
14    ///
15    /// e.g. "Host", "Database::Table"; you can also support wildcard/globs if you want.
16    fn applies_to(&self, kind: &str) -> bool;
17
18    /// Mutates the resource by injecting derived attributes (e.g., sets of labels).
19    fn apply(&self, res: &mut Resource);
20}
21
22/// A labeler that uses regular expressions for matching on resource attributes.
23#[allow(dead_code)]
24#[derive(Debug, Clone)]
25pub struct RegexLabeler {
26    /// The kind of resource this labeler applies to, e.g. "Host"
27    kind: String,
28    /// attribute to read from, e.g. "name"
29    field: String,
30    /// attribute to write to, e.g. "nameLabels"
31    output: String,
32    /// Rulesets for matching resource attributes
33    table: Vec<(String, Regex)>,
34}
35
36impl RegexLabeler {
37    /// Create a regex-based labeler.
38    ///
39    /// - `kind`: resource kind this applies to (e.g., "Host")
40    /// - `field`: attribute to read from (e.g., "name")
41    /// - `output`: attribute to write labels to (e.g., "nameLabels")
42    /// - `table`: vector of `(label, regex)` pairs
43    #[allow(dead_code)]
44    pub fn new(
45        kind: impl Into<String>,
46        field: impl Into<String>,
47        output: impl Into<String>,
48        table: Vec<(String, Regex)>,
49    ) -> Self {
50        Self {
51            kind: kind.into(),
52            field: field.into(),
53            output: output.into(),
54            table,
55        }
56    }
57}
58
59impl Labeler for RegexLabeler {
60    fn applies_to(&self, kind: &str) -> bool {
61        self.kind == kind
62    }
63
64    fn apply(&self, res: &mut Resource) {
65        let Some(AttrValue::String(value)) = res.attrs().get(&self.field).cloned() else {
66            return;
67        };
68        let mut out: Vec<AttrValue> = Vec::new();
69        for (label, re) in &self.table {
70            if re.is_match(&value) {
71                out.push(AttrValue::String(label.clone()));
72            }
73        }
74        if !out.is_empty() {
75            match res.attrs().get_mut(&self.output) {
76                Some(AttrValue::Set(existing)) => existing.extend(out),
77                _ => {
78                    res.attrs().insert(self.output.clone(), AttrValue::Set(out));
79                }
80            }
81        }
82    }
83}
84
85/// Implementation of the LabelRegistry.
86///
87/// Consumption of this registry goes through the static `LABEL_REGISTRY`.
88pub struct LabelRegistry {
89    inner: ArcSwap<Vec<Arc<dyn Labeler>>>,
90}
91impl LabelRegistry {
92    /// Applies all labelers in the registry to the given resource.
93    ///
94    /// Labelers run in insertion order. If multiple labelers target the same
95    /// output set attribute, values are appended.
96    pub fn apply(&self, res: &mut Resource) {
97        let snapshot = self.inner.load();
98        let kind_owned = res.kind().to_owned();
99        for l in snapshot.iter() {
100            if l.applies_to(&kind_owned) {
101                l.apply(res);
102            }
103        }
104    }
105
106    /// Loads a set of labelers into the registry, atomically.
107    ///
108    /// Replaces all prior labelers in a single swap. New `evaluate()` calls use
109    /// the new set immediately; in-flight evaluations continue with the old set.
110    pub fn reload(&self, labelers: Vec<Arc<dyn Labeler>>) {
111        self.inner.store(Arc::new(labelers));
112    }
113}
114
115/// Builder for creating a LabelRegistry with labelers.
116///
117/// This uses a builder pattern to ensure labelers are properly initialized
118/// before the registry is used.
119///
120/// # Example
121///
122/// ```rust
123/// use std::sync::Arc;
124/// use treetop_core::{LabelRegistryBuilder, RegexLabeler};
125/// use regex::Regex;
126///
127/// let registry = LabelRegistryBuilder::new()
128///     .add_labeler(Arc::new(RegexLabeler::new(
129///         "Host",
130///         "name",
131///         "nameLabels",
132///         vec![("prod".to_string(), Regex::new(r"\.prod\.").unwrap())],
133///     )))
134///     .build();
135/// ```
136pub struct LabelRegistryBuilder {
137    labelers: Vec<Arc<dyn Labeler>>,
138}
139
140impl LabelRegistryBuilder {
141    /// Create a new, empty label registry builder.
142    pub fn new() -> Self {
143        Self {
144            labelers: Vec::new(),
145        }
146    }
147
148    /// Add a labeler to the registry.
149    ///
150    /// This can be called repeatedly to build up a registry before `build()`.
151    pub fn add_labeler(mut self, labeler: Arc<dyn Labeler>) -> Self {
152        self.labelers.push(labeler);
153        self
154    }
155
156    /// Build the label registry.
157    ///
158    /// Consumes the builder and returns an initialized registry ready to use
159    /// with `PolicyEngine::with_label_registry()`.
160    pub fn build(self) -> LabelRegistry {
161        LabelRegistry {
162            inner: ArcSwap::from_pointee(self.labelers),
163        }
164    }
165}
166
167impl Default for LabelRegistryBuilder {
168    fn default() -> Self {
169        Self::new()
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use std::collections::BTreeSet;
177    use yare::parameterized;
178
179    fn compile(rules: Vec<(&str, &str)>) -> Vec<(String, Regex)> {
180        rules
181            .into_iter()
182            .map(|(l, p)| (l.to_string(), Regex::new(p).unwrap()))
183            .collect()
184    }
185
186    fn get_label_strings(res: &mut Resource, key: &str) -> BTreeSet<String> {
187        match res.attrs().get(key) {
188            Some(AttrValue::Set(v)) => v
189                .iter()
190                .filter_map(|a| {
191                    if let AttrValue::String(s) = a {
192                        Some(s.clone())
193                    } else {
194                        None
195                    }
196                })
197                .collect(),
198            _ => BTreeSet::new(),
199        }
200    }
201
202    #[parameterized(
203        simple_match = {
204            "Host", "name", "nameLabels",
205            vec![("prod", r"(^|\.)prod\.example\.com$")],
206            "db12.prod.example.com",
207            &["prod"]
208        },
209        no_match = {
210            "Host", "name", "nameLabels",
211            vec![("corp", r"(^|\.)corp\.example\.com$")],
212            "web.dev.example.com",
213            &[]
214        },
215        multi_match = {
216            "Host", "name", "nameLabels",
217            vec![("prod", r"(^|\.)prod\."), ("db", r"(^|\.)db\d+\.")],
218            "db42.prod.example.com",
219            &["db","prod"]
220        }
221    )]
222    fn regex_labeler_apply_basic(
223        kind: &str,
224        field: &str,
225        output: &str,
226        rules: Vec<(&str, &str)>,
227        input: &str,
228        expected: &[&str],
229    ) {
230        let labeler = RegexLabeler::new(kind, field, output, compile(rules));
231
232        let mut res = Resource::new(kind, input);
233        res.attrs()
234            .insert(field.to_string(), AttrValue::String(input.to_string()));
235
236        labeler.apply(&mut res);
237
238        let got = get_label_strings(&mut res, output);
239        let want: BTreeSet<String> = expected.iter().map(|s| s.to_string()).collect();
240        assert_eq!(got, want);
241    }
242
243    #[test]
244    fn regex_labeler_missing_input_field_is_noop() {
245        let labeler = RegexLabeler::new(
246            "Host",
247            "name",
248            "nameLabels",
249            compile(vec![("prod", r"(^|\.)prod\.")]),
250        );
251
252        let mut res = Resource::new("Host", "db99.prod.example.com");
253        // no "name" inserted
254
255        labeler.apply(&mut res);
256        assert!(res.attrs().get("nameLabels").is_none());
257    }
258
259    #[test]
260    fn regex_labeler_appends_to_existing_set() {
261        let labeler = RegexLabeler::new(
262            "Host",
263            "name",
264            "nameLabels",
265            compile(vec![("prod", r"(^|\.)prod\."), ("db", r"(^|\.)db\d+\.")]),
266        );
267
268        let mut res = Resource::new("Host", "db99.prod.example.com");
269        res.attrs().insert(
270            "name".into(),
271            AttrValue::String("db99.prod.example.com".into()),
272        );
273        res.attrs().insert(
274            "nameLabels".into(),
275            AttrValue::Set(vec![AttrValue::String("pre".into())]),
276        );
277
278        labeler.apply(&mut res);
279
280        let labels = get_label_strings(&mut res, "nameLabels");
281        assert!(labels.contains("pre"));
282        assert!(labels.contains("prod"));
283        assert!(labels.contains("db"));
284    }
285}
treetop_core/labels.rs

treetop_core/
labels.rs
