1use crate::{BundleError, Diagnostic, Result};
2use cedar_policy::{EntityTypeName, Schema};
3use regex::{Regex, RegexBuilder, RegexSet, RegexSetBuilder};
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use std::collections::HashSet;
7use std::sync::Arc;
8use treetop_core::{AttrValue, Labeler, RegexLabeler, Resource};
9
10const MAX_LABEL_RULES: usize = 256;
11const MAX_PATTERNS_PER_RULE: usize = 1_024;
12const MAX_TOTAL_PATTERNS: usize = 4_096;
13const MAX_REGEX_BYTES: usize = 16 * 1024;
14const MAX_TOTAL_REGEX_BYTES: usize = 1024 * 1024;
15const REGEX_SET_SIZE_LIMIT: usize = 2 * 1024 * 1024;
16const REGEX_SET_DFA_SIZE_LIMIT: usize = 1024 * 1024;
17const INDIVIDUAL_REGEX_THRESHOLD: usize = 4;
18const INDIVIDUAL_REGEX_SIZE_LIMIT: usize = REGEX_SET_SIZE_LIMIT / INDIVIDUAL_REGEX_THRESHOLD;
19const INDIVIDUAL_REGEX_DFA_SIZE_LIMIT: usize =
20 REGEX_SET_DFA_SIZE_LIMIT / INDIVIDUAL_REGEX_THRESHOLD;
21
22#[derive(Debug, Clone, Serialize, Deserialize)]
23#[serde(deny_unknown_fields)]
24struct RawLabelPattern {
25 name: String,
26 regex: String,
27}
28
29#[derive(Debug, Clone, Serialize)]
31pub struct LabelPattern {
32 name: String,
33 regex: String,
34}
35
36impl PartialEq for LabelPattern {
37 fn eq(&self, other: &Self) -> bool {
38 self.name == other.name && self.regex == other.regex
39 }
40}
41
42impl Eq for LabelPattern {}
43
44impl LabelPattern {
45 pub fn name(&self) -> &str {
46 &self.name
47 }
48
49 pub fn regex(&self) -> &str {
50 &self.regex
51 }
52}
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
55#[serde(deny_unknown_fields)]
56struct RawLabelRule {
57 kind: String,
58 field: String,
59 output: String,
60 patterns: Vec<RawLabelPattern>,
61}
62
63#[derive(Debug, Clone, Serialize)]
65pub struct LabelRule {
66 kind: String,
67 field: String,
68 output: String,
69 patterns: Vec<LabelPattern>,
70 #[serde(skip)]
71 compiled: CompiledPatterns,
72}
73
74#[derive(Debug, Clone)]
75enum CompiledPatterns {
76 Individual(Arc<Vec<Regex>>),
77 Set(Arc<RegexSet>),
78}
79
80impl PartialEq for LabelRule {
81 fn eq(&self, other: &Self) -> bool {
82 self.kind == other.kind
83 && self.field == other.field
84 && self.output == other.output
85 && self.patterns == other.patterns
86 }
87}
88
89impl Eq for LabelRule {}
90
91#[derive(Debug)]
92struct RegexSetLabeler {
93 kind: String,
94 field: String,
95 output: String,
96 names: Vec<String>,
97 compiled: Arc<RegexSet>,
98}
99
100impl Labeler for RegexSetLabeler {
101 fn applies_to(&self, kind: &str) -> bool {
102 self.kind == kind
103 }
104
105 fn apply(&self, resource: &mut Resource) {
106 let Some(AttrValue::String(value)) = resource.attributes().get(&self.field) else {
107 resource.attrs().remove(&self.output);
108 return;
109 };
110 let labels = self
111 .compiled
112 .matches(value)
113 .iter()
114 .map(|index| AttrValue::String(self.names[index].clone()))
115 .collect();
116 resource
117 .attrs()
118 .insert(self.output.clone(), AttrValue::Set(labels));
119 }
120}
121
122impl LabelRule {
123 pub fn kind(&self) -> &str {
124 &self.kind
125 }
126
127 pub fn field(&self) -> &str {
128 &self.field
129 }
130
131 pub fn output(&self) -> &str {
132 &self.output
133 }
134
135 pub fn patterns(&self) -> &[LabelPattern] {
136 &self.patterns
137 }
138}
139
140#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
142#[serde(transparent)]
143pub struct LabelSet(Vec<LabelRule>);
144
145impl LabelSet {
146 pub fn from_json_str(input: &str) -> Result<Self> {
148 let raw: Vec<RawLabelRule> = serde_json::from_str(input).map_err(|error| {
149 let mut diagnostic = Diagnostic::error("labels.invalid_json", error.to_string());
150 diagnostic.line = Some(error.line());
151 diagnostic.column = Some(error.column());
152 BundleError::Validation(vec![diagnostic])
153 })?;
154 Self::from_raw(raw)
155 }
156
157 pub fn validate_schema_json_str(&self, schema_source: &str) -> Result<()> {
159 let schema_json: Value = serde_json::from_str(schema_source).map_err(|error| {
160 BundleError::Validation(vec![Diagnostic::error(
161 "schema.invalid_json",
162 error.to_string(),
163 )])
164 })?;
165 let schema = Schema::from_json_value(schema_json.clone()).map_err(|error| {
166 BundleError::Validation(vec![Diagnostic::error(
167 "schema.aggregate_invalid",
168 error.to_string(),
169 )])
170 })?;
171 let diagnostics = self.validate_schema(&schema, &schema_json);
172 if diagnostics.is_empty() {
173 Ok(())
174 } else {
175 Err(BundleError::Validation(diagnostics))
176 }
177 }
178
179 fn from_raw(raw: Vec<RawLabelRule>) -> Result<Self> {
180 validate_document_limits(&raw)?;
181
182 let mut diagnostics = Vec::new();
183 let mut destinations = HashSet::new();
184 let mut rules = Vec::with_capacity(raw.len());
185
186 for (rule_index, raw_rule) in raw.into_iter().enumerate() {
187 let location = format!("labels[{rule_index}]");
188 if raw_rule.kind.trim().is_empty() {
189 diagnostics.push(Diagnostic::error(
190 "labels.empty_kind",
191 format!("{location}.kind must not be empty"),
192 ));
193 } else if raw_rule.kind.parse::<EntityTypeName>().is_err() {
194 diagnostics.push(Diagnostic::error(
195 "labels.invalid_kind",
196 format!("{location}.kind is not a Cedar entity type"),
197 ));
198 }
199 if raw_rule.field.trim().is_empty() {
200 diagnostics.push(Diagnostic::error(
201 "labels.empty_field",
202 format!("{location}.field must not be empty"),
203 ));
204 }
205 if raw_rule.output.trim().is_empty() {
206 diagnostics.push(Diagnostic::error(
207 "labels.empty_output",
208 format!("{location}.output must not be empty"),
209 ));
210 }
211 if raw_rule.field == raw_rule.output {
212 diagnostics.push(Diagnostic::error(
213 "labels.input_is_output",
214 format!("{location}.field and output must be different"),
215 ));
216 }
217 if !destinations.insert((raw_rule.kind.clone(), raw_rule.output.clone())) {
218 diagnostics.push(Diagnostic::error(
219 "labels.duplicate_destination",
220 format!(
221 "duplicate label destination ({}, {})",
222 raw_rule.kind, raw_rule.output
223 ),
224 ));
225 }
226 if raw_rule.patterns.is_empty() {
227 diagnostics.push(Diagnostic::error(
228 "labels.empty_patterns",
229 format!("{location}.patterns must not be empty"),
230 ));
231 }
232 if raw_rule.patterns.len() > MAX_PATTERNS_PER_RULE {
233 diagnostics.push(Diagnostic::error(
234 "labels.too_many_patterns",
235 format!(
236 "{location}.patterns contains more than {MAX_PATTERNS_PER_RULE} entries"
237 ),
238 ));
239 }
240
241 let mut names = HashSet::new();
242 let mut patterns = Vec::with_capacity(raw_rule.patterns.len());
243 let mut patterns_valid =
244 !raw_rule.patterns.is_empty() && raw_rule.patterns.len() <= MAX_PATTERNS_PER_RULE;
245 for (pattern_index, raw_pattern) in raw_rule.patterns.into_iter().enumerate() {
246 let pattern_location = format!("{location}.patterns[{pattern_index}]");
247 if raw_pattern.name.trim().is_empty() {
248 diagnostics.push(Diagnostic::error(
249 "labels.empty_pattern_name",
250 format!("{pattern_location}.name must not be empty"),
251 ));
252 }
253 if !names.insert(raw_pattern.name.clone()) {
254 diagnostics.push(Diagnostic::error(
255 "labels.duplicate_pattern_name",
256 format!(
257 "duplicate pattern name {:?} in {location}",
258 raw_pattern.name
259 ),
260 ));
261 }
262 if raw_pattern.regex.is_empty() {
263 diagnostics.push(Diagnostic::error(
264 "labels.empty_regex",
265 format!("{pattern_location}.regex must not be empty"),
266 ));
267 patterns_valid = false;
268 } else if raw_pattern.regex.len() > MAX_REGEX_BYTES {
269 diagnostics.push(Diagnostic::error(
270 "labels.regex_too_large",
271 format!("{pattern_location}.regex exceeds {MAX_REGEX_BYTES} bytes"),
272 ));
273 patterns_valid = false;
274 }
275 patterns.push(LabelPattern {
276 name: raw_pattern.name,
277 regex: raw_pattern.regex,
278 });
279 }
280
281 if patterns_valid {
282 match compile_patterns(&patterns) {
283 Ok(compiled) => rules.push(LabelRule {
284 kind: raw_rule.kind,
285 field: raw_rule.field,
286 output: raw_rule.output,
287 patterns,
288 compiled,
289 }),
290 Err(error) => diagnostics.push(Diagnostic::error(
291 "labels.invalid_regex",
292 format!("{location}.patterns cannot be compiled safely: {error}"),
293 )),
294 }
295 }
296 }
297
298 if diagnostics.is_empty() {
299 Ok(Self(rules))
300 } else {
301 Err(BundleError::Validation(diagnostics))
302 }
303 }
304
305 pub(crate) fn combine(sets: impl IntoIterator<Item = Self>) -> Result<Self> {
306 let rules = sets.into_iter().flat_map(|set| set.0).collect::<Vec<_>>();
307 validate_combined_limits(&rules)?;
308 let mut destinations = HashSet::with_capacity(rules.len());
309 let mut diagnostics = Vec::new();
310 for rule in &rules {
311 if !destinations.insert((&rule.kind, &rule.output)) {
312 diagnostics.push(Diagnostic::error(
313 "labels.duplicate_destination",
314 format!(
315 "duplicate label destination ({}, {})",
316 rule.kind, rule.output
317 ),
318 ));
319 }
320 }
321 if diagnostics.is_empty() {
322 Ok(Self(rules))
323 } else {
324 Err(BundleError::Validation(diagnostics))
325 }
326 }
327
328 pub fn rules(&self) -> &[LabelRule] {
329 &self.0
330 }
331
332 pub fn is_empty(&self) -> bool {
333 self.0.is_empty()
334 }
335
336 pub fn to_labelers(&self) -> Vec<Arc<dyn Labeler>> {
338 self.0
339 .iter()
340 .map(|rule| match &rule.compiled {
341 CompiledPatterns::Individual(compiled) => {
342 let patterns = rule
343 .patterns
344 .iter()
345 .zip(compiled.iter())
346 .map(|(pattern, regex)| (pattern.name.clone(), regex.clone()))
347 .collect();
348 Arc::new(RegexLabeler::new(
349 rule.kind.clone(),
350 rule.field.clone(),
351 rule.output.clone(),
352 patterns,
353 )) as Arc<dyn Labeler>
354 }
355 CompiledPatterns::Set(compiled) => {
356 let names = rule
357 .patterns
358 .iter()
359 .map(|pattern| pattern.name.clone())
360 .collect();
361 Arc::new(RegexSetLabeler {
362 kind: rule.kind.clone(),
363 field: rule.field.clone(),
364 output: rule.output.clone(),
365 names,
366 compiled: Arc::clone(compiled),
367 }) as Arc<dyn Labeler>
368 }
369 })
370 .collect()
371 }
372
373 pub(crate) fn validate_schema(&self, schema: &Schema, schema_json: &Value) -> Vec<Diagnostic> {
374 let known_types = schema
375 .entity_types()
376 .map(ToString::to_string)
377 .collect::<HashSet<_>>();
378 let mut diagnostics = Vec::new();
379 for rule in &self.0 {
380 if !known_types.contains(&rule.kind) {
381 diagnostics.push(Diagnostic::error(
382 "labels.unknown_kind",
383 format!("label kind {} is not declared in the schema", rule.kind),
384 ));
385 continue;
386 }
387 let Some(attributes) = entity_attributes(schema_json, &rule.kind) else {
388 diagnostics.push(Diagnostic::error(
389 "labels.missing_shape",
390 format!("label kind {} has no record shape", rule.kind),
391 ));
392 continue;
393 };
394 match attributes.get(&rule.field) {
395 Some(value) if is_string_type(value) => {}
396 Some(value) => diagnostics.push(Diagnostic::error(
397 "labels.field_not_string",
398 format!(
399 "{}.{} must have schema type String, found {value}",
400 rule.kind, rule.field
401 ),
402 )),
403 None => diagnostics.push(Diagnostic::error(
404 "labels.field_missing",
405 format!("{}.{} is not declared in the schema", rule.kind, rule.field),
406 )),
407 }
408 match attributes.get(&rule.output) {
409 Some(value) if is_string_set_type(value) => {}
410 Some(value) => diagnostics.push(Diagnostic::error(
411 "labels.output_not_string_set",
412 format!(
413 "{}.{} must have schema type Set<String>, found {value}",
414 rule.kind, rule.output,
415 ),
416 )),
417 None => diagnostics.push(Diagnostic::error(
418 "labels.output_missing",
419 format!(
420 "{}.{} is not declared in the schema",
421 rule.kind, rule.output
422 ),
423 )),
424 }
425 }
426 diagnostics
427 }
428}
429
430fn compile_patterns(
431 patterns: &[LabelPattern],
432) -> std::result::Result<CompiledPatterns, regex::Error> {
433 if patterns.len() <= INDIVIDUAL_REGEX_THRESHOLD {
434 let compiled = patterns
435 .iter()
436 .map(|pattern| {
437 let mut builder = RegexBuilder::new(&pattern.regex);
438 builder
439 .size_limit(INDIVIDUAL_REGEX_SIZE_LIMIT)
440 .dfa_size_limit(INDIVIDUAL_REGEX_DFA_SIZE_LIMIT);
441 builder.build()
442 })
443 .collect::<std::result::Result<Vec<_>, _>>()?;
444 Ok(CompiledPatterns::Individual(Arc::new(compiled)))
445 } else {
446 let mut builder =
447 RegexSetBuilder::new(patterns.iter().map(|pattern| pattern.regex.as_str()));
448 builder
449 .size_limit(REGEX_SET_SIZE_LIMIT)
450 .dfa_size_limit(REGEX_SET_DFA_SIZE_LIMIT);
451 builder
452 .build()
453 .map(|compiled| CompiledPatterns::Set(Arc::new(compiled)))
454 }
455}
456
457fn validate_document_limits(raw: &[RawLabelRule]) -> Result<()> {
458 if raw.len() > MAX_LABEL_RULES {
459 return Err(BundleError::Validation(vec![Diagnostic::error(
460 "labels.too_many_rules",
461 format!("label document contains more than {MAX_LABEL_RULES} rules"),
462 )]));
463 }
464 let total_patterns = raw
465 .iter()
466 .try_fold(0usize, |total, rule| total.checked_add(rule.patterns.len()))
467 .unwrap_or(usize::MAX);
468 if total_patterns > MAX_TOTAL_PATTERNS {
469 return Err(BundleError::Validation(vec![Diagnostic::error(
470 "labels.too_many_patterns",
471 format!("label document contains more than {MAX_TOTAL_PATTERNS} patterns"),
472 )]));
473 }
474 let total_regex_bytes = raw
475 .iter()
476 .flat_map(|rule| &rule.patterns)
477 .try_fold(0usize, |total, pattern| {
478 total.checked_add(pattern.regex.len())
479 })
480 .unwrap_or(usize::MAX);
481 if total_regex_bytes > MAX_TOTAL_REGEX_BYTES {
482 return Err(BundleError::Validation(vec![Diagnostic::error(
483 "labels.regex_budget_exceeded",
484 format!("label document regex sources exceed {MAX_TOTAL_REGEX_BYTES} total bytes"),
485 )]));
486 }
487 Ok(())
488}
489
490fn validate_combined_limits(rules: &[LabelRule]) -> Result<()> {
491 if rules.len() > MAX_LABEL_RULES {
492 return Err(BundleError::Validation(vec![Diagnostic::error(
493 "labels.too_many_rules",
494 format!("combined label document contains more than {MAX_LABEL_RULES} rules"),
495 )]));
496 }
497 let total_patterns = rules
498 .iter()
499 .try_fold(0usize, |total, rule| total.checked_add(rule.patterns.len()))
500 .unwrap_or(usize::MAX);
501 if total_patterns > MAX_TOTAL_PATTERNS {
502 return Err(BundleError::Validation(vec![Diagnostic::error(
503 "labels.too_many_patterns",
504 format!("combined label document contains more than {MAX_TOTAL_PATTERNS} patterns"),
505 )]));
506 }
507 let total_regex_bytes = rules
508 .iter()
509 .flat_map(|rule| &rule.patterns)
510 .try_fold(0usize, |total, pattern| {
511 total.checked_add(pattern.regex.len())
512 })
513 .unwrap_or(usize::MAX);
514 if total_regex_bytes > MAX_TOTAL_REGEX_BYTES {
515 return Err(BundleError::Validation(vec![Diagnostic::error(
516 "labels.regex_budget_exceeded",
517 format!(
518 "combined label document regex sources exceed {MAX_TOTAL_REGEX_BYTES} total bytes"
519 ),
520 )]));
521 }
522 Ok(())
523}
524
525fn entity_attributes<'a>(
526 schema_json: &'a Value,
527 kind: &str,
528) -> Option<&'a serde_json::Map<String, Value>> {
529 let parsed = kind.parse::<EntityTypeName>().ok()?;
530 let namespace = parsed.namespace().to_string();
531 let namespace_definition = schema_json.as_object()?.get(&namespace)?;
532 let definition = namespace_definition
533 .get("entityTypes")?
534 .get(parsed.basename())?;
535 record_attributes(
536 schema_json,
537 &namespace,
538 definition.get("shape")?,
539 &mut HashSet::new(),
540 )
541}
542
543fn record_attributes<'a>(
544 schema_json: &'a Value,
545 namespace: &str,
546 shape: &'a Value,
547 visited: &mut HashSet<String>,
548) -> Option<&'a serde_json::Map<String, Value>> {
549 if shape.get("type").and_then(Value::as_str) == Some("Record") {
550 return shape.get("attributes")?.as_object();
551 }
552 if shape.get("type").and_then(Value::as_str) != Some("EntityOrCommon") {
553 return None;
554 }
555 let name = shape.get("name")?.as_str()?;
556 let (common_namespace, basename) = name
557 .rsplit_once("::")
558 .map_or((namespace, name), |(namespace, basename)| {
559 (namespace, basename)
560 });
561 let qualified_name = format!("{common_namespace}::{basename}");
562 if !visited.insert(qualified_name) {
563 return None;
564 }
565 let common = schema_json
566 .as_object()?
567 .get(common_namespace)?
568 .get("commonTypes")?
569 .get(basename)?;
570 record_attributes(schema_json, common_namespace, common, visited)
571}
572
573fn is_string_type(value: &Value) -> bool {
574 value.get("type").and_then(Value::as_str) == Some("String")
575 || (value.get("type").and_then(Value::as_str) == Some("EntityOrCommon")
576 && value.get("name").and_then(Value::as_str) == Some("String"))
577}
578
579fn is_string_set_type(value: &Value) -> bool {
580 value.get("type").and_then(Value::as_str) == Some("Set")
581 && value.get("element").is_some_and(is_string_type)
582}
583
584#[cfg(test)]
585mod tests {
586 use super::*;
587
588 #[test]
589 fn strict_label_validation_rejects_unknown_fields() {
590 let error = LabelSet::from_json_str(
591 r#"[{"kind":"App::Host","field":"name","output":"labels","patterns":[{"name":"prod","regex":"prod","extra":true}]}]"#,
592 )
593 .unwrap_err();
594 assert!(error.diagnostics()[0].message.contains("unknown field"));
595 }
596
597 #[test]
598 fn label_set_converts_to_runtime_labelers() {
599 let labels = LabelSet::from_json_str(
600 r#"[{"kind":"App::Host","field":"name","output":"labels","patterns":[{"name":"prod","regex":"^prod"}]}]"#,
601 )
602 .unwrap();
603 assert_eq!(labels.to_labelers().len(), 1);
604 }
605
606 #[test]
607 fn regex_set_labeler_returns_all_matches_and_replaces_untrusted_output() {
608 let labels = LabelSet::from_json_str(
609 r#"[{"kind":"App::Host","field":"name","output":"labels","patterns":[{"name":"prod","regex":"^prod"},{"name":"database","regex":"db$"},{"name":"staging","regex":"^staging"},{"name":"cache","regex":"cache$"},{"name":"worker","regex":"worker"}]}]"#,
610 )
611 .unwrap();
612 let labeler = labels.to_labelers().pop().unwrap();
613 let mut resource = Resource::new("App::Host", "one")
614 .with_attr("name", AttrValue::String("prod-db".to_string()))
615 .with_attr(
616 "labels",
617 AttrValue::Set(vec![AttrValue::String("forged".to_string())]),
618 );
619
620 labeler.apply(&mut resource);
621
622 assert_eq!(
623 resource.attributes().get("labels"),
624 Some(&AttrValue::Set(vec![
625 AttrValue::String("prod".to_string()),
626 AttrValue::String("database".to_string()),
627 ]))
628 );
629 }
630
631 #[test]
632 fn label_document_pattern_budget_is_enforced_before_compilation() {
633 let patterns = (0..=MAX_TOTAL_PATTERNS)
634 .map(|index| {
635 serde_json::json!({
636 "name": format!("pattern-{index}"),
637 "regex": "a",
638 })
639 })
640 .collect::<Vec<_>>();
641 let source = serde_json::json!([{
642 "kind": "App::Host",
643 "field": "name",
644 "output": "labels",
645 "patterns": patterns,
646 }])
647 .to_string();
648
649 let error = LabelSet::from_json_str(&source).unwrap_err();
650
651 assert_eq!(error.diagnostics()[0].code, "labels.too_many_patterns");
652 }
653
654 #[test]
655 fn combined_label_sets_reapply_document_limits() {
656 let labels = LabelSet::from_json_str(
657 r#"[{"kind":"App::Host","field":"name","output":"labels","patterns":[{"name":"prod","regex":"^prod"}]}]"#,
658 )
659 .unwrap();
660
661 let error =
662 LabelSet::combine(std::iter::repeat_n(labels, MAX_LABEL_RULES + 1)).unwrap_err();
663
664 assert_eq!(error.diagnostics()[0].code, "labels.too_many_rules");
665 }
666
667 #[test]
668 fn schema_validation_resolves_common_record_shapes() {
669 let labels = LabelSet::from_json_str(
670 r#"[{"kind":"App::Host","field":"name","output":"labels","patterns":[{"name":"prod","regex":"^prod"}]}]"#,
671 )
672 .unwrap();
673 let schema = r#"{
674 "App": {
675 "commonTypes": {
676 "HostShape": {
677 "type": "Record",
678 "attributes": {
679 "name": {"type": "String", "required": true},
680 "labels": {
681 "type": "Set",
682 "element": {"type": "String"},
683 "required": false
684 }
685 },
686 "additionalAttributes": false
687 }
688 },
689 "entityTypes": {
690 "Host": {
691 "shape": {"type": "EntityOrCommon", "name": "HostShape"}
692 }
693 },
694 "actions": {}
695 }
696 }"#;
697
698 labels.validate_schema_json_str(schema).unwrap();
699 }
700}