use crate::{BundleError, Diagnostic, Result};
use cedar_policy::{EntityTypeName, Schema};
use regex::Regex;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashSet;
use std::sync::Arc;
use treetop_core::{Labeler, RegexLabeler};
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct RawLabelPattern {
name: String,
regex: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct LabelPattern {
name: String,
regex: String,
}
impl LabelPattern {
pub fn name(&self) -> &str {
&self.name
}
pub fn regex(&self) -> &str {
&self.regex
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct RawLabelRule {
kind: String,
field: String,
output: String,
patterns: Vec<RawLabelPattern>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct LabelRule {
kind: String,
field: String,
output: String,
patterns: Vec<LabelPattern>,
}
impl LabelRule {
pub fn kind(&self) -> &str {
&self.kind
}
pub fn field(&self) -> &str {
&self.field
}
pub fn output(&self) -> &str {
&self.output
}
pub fn patterns(&self) -> &[LabelPattern] {
&self.patterns
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize)]
#[serde(transparent)]
pub struct LabelSet(Vec<LabelRule>);
impl LabelSet {
pub fn from_json_str(input: &str) -> Result<Self> {
let raw: Vec<RawLabelRule> = serde_json::from_str(input).map_err(|error| {
let mut diagnostic = Diagnostic::error("labels.invalid_json", error.to_string());
diagnostic.line = Some(error.line());
diagnostic.column = Some(error.column());
BundleError::Validation(vec![diagnostic])
})?;
Self::from_raw(raw)
}
pub fn validate_schema_json_str(&self, schema_source: &str) -> Result<()> {
let schema_json: Value = serde_json::from_str(schema_source).map_err(|error| {
BundleError::Validation(vec![Diagnostic::error(
"schema.invalid_json",
error.to_string(),
)])
})?;
let schema = Schema::from_json_value(schema_json.clone()).map_err(|error| {
BundleError::Validation(vec![Diagnostic::error(
"schema.aggregate_invalid",
error.to_string(),
)])
})?;
let diagnostics = self.validate_schema(&schema, &schema_json);
if diagnostics.is_empty() {
Ok(())
} else {
Err(BundleError::Validation(diagnostics))
}
}
fn from_raw(raw: Vec<RawLabelRule>) -> Result<Self> {
let mut diagnostics = Vec::new();
let mut destinations = HashSet::new();
let mut rules = Vec::with_capacity(raw.len());
for (rule_index, raw_rule) in raw.into_iter().enumerate() {
let location = format!("labels[{rule_index}]");
if raw_rule.kind.trim().is_empty() {
diagnostics.push(Diagnostic::error(
"labels.empty_kind",
format!("{location}.kind must not be empty"),
));
} else if raw_rule.kind.parse::<EntityTypeName>().is_err() {
diagnostics.push(Diagnostic::error(
"labels.invalid_kind",
format!("{location}.kind is not a Cedar entity type"),
));
}
if raw_rule.field.trim().is_empty() {
diagnostics.push(Diagnostic::error(
"labels.empty_field",
format!("{location}.field must not be empty"),
));
}
if raw_rule.output.trim().is_empty() {
diagnostics.push(Diagnostic::error(
"labels.empty_output",
format!("{location}.output must not be empty"),
));
}
if raw_rule.field == raw_rule.output {
diagnostics.push(Diagnostic::error(
"labels.input_is_output",
format!("{location}.field and output must be different"),
));
}
if !destinations.insert((raw_rule.kind.clone(), raw_rule.output.clone())) {
diagnostics.push(Diagnostic::error(
"labels.duplicate_destination",
format!(
"duplicate label destination ({}, {})",
raw_rule.kind, raw_rule.output
),
));
}
if raw_rule.patterns.is_empty() {
diagnostics.push(Diagnostic::error(
"labels.empty_patterns",
format!("{location}.patterns must not be empty"),
));
}
let mut names = HashSet::new();
let mut patterns = Vec::with_capacity(raw_rule.patterns.len());
for (pattern_index, raw_pattern) in raw_rule.patterns.into_iter().enumerate() {
let pattern_location = format!("{location}.patterns[{pattern_index}]");
if raw_pattern.name.trim().is_empty() {
diagnostics.push(Diagnostic::error(
"labels.empty_pattern_name",
format!("{pattern_location}.name must not be empty"),
));
}
if !names.insert(raw_pattern.name.clone()) {
diagnostics.push(Diagnostic::error(
"labels.duplicate_pattern_name",
format!(
"duplicate pattern name {:?} in {location}",
raw_pattern.name
),
));
}
if raw_pattern.regex.is_empty() {
diagnostics.push(Diagnostic::error(
"labels.empty_regex",
format!("{pattern_location}.regex must not be empty"),
));
} else if let Err(error) = Regex::new(&raw_pattern.regex) {
diagnostics.push(Diagnostic::error(
"labels.invalid_regex",
format!("{pattern_location}.regex is invalid: {error}"),
));
}
patterns.push(LabelPattern {
name: raw_pattern.name,
regex: raw_pattern.regex,
});
}
rules.push(LabelRule {
kind: raw_rule.kind,
field: raw_rule.field,
output: raw_rule.output,
patterns,
});
}
if diagnostics.is_empty() {
Ok(Self(rules))
} else {
Err(BundleError::Validation(diagnostics))
}
}
pub(crate) fn combine(sets: impl IntoIterator<Item = Self>) -> Result<Self> {
let raw = sets
.into_iter()
.flat_map(|set| set.0)
.map(|rule| RawLabelRule {
kind: rule.kind,
field: rule.field,
output: rule.output,
patterns: rule
.patterns
.into_iter()
.map(|pattern| RawLabelPattern {
name: pattern.name,
regex: pattern.regex,
})
.collect(),
})
.collect();
Self::from_raw(raw)
}
pub fn rules(&self) -> &[LabelRule] {
&self.0
}
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
pub fn to_labelers(&self) -> Vec<Arc<dyn Labeler>> {
self.0
.iter()
.map(|rule| {
let patterns = rule
.patterns
.iter()
.map(|pattern| {
(
pattern.name.clone(),
Regex::new(&pattern.regex)
.expect("LabelSet invariant: regular expressions are compiled"),
)
})
.collect();
Arc::new(RegexLabeler::new(
rule.kind.clone(),
rule.field.clone(),
rule.output.clone(),
patterns,
)) as Arc<dyn Labeler>
})
.collect()
}
pub(crate) fn validate_schema(&self, schema: &Schema, schema_json: &Value) -> Vec<Diagnostic> {
let known_types = schema
.entity_types()
.map(ToString::to_string)
.collect::<HashSet<_>>();
let mut diagnostics = Vec::new();
for rule in &self.0 {
if !known_types.contains(&rule.kind) {
diagnostics.push(Diagnostic::error(
"labels.unknown_kind",
format!("label kind {} is not declared in the schema", rule.kind),
));
continue;
}
let Some(attributes) = entity_attributes(schema_json, &rule.kind) else {
diagnostics.push(Diagnostic::error(
"labels.missing_shape",
format!("label kind {} has no record shape", rule.kind),
));
continue;
};
match attributes.get(&rule.field) {
Some(value) if is_string_type(value) => {}
Some(value) => diagnostics.push(Diagnostic::error(
"labels.field_not_string",
format!(
"{}.{} must have schema type String, found {value}",
rule.kind, rule.field
),
)),
None => diagnostics.push(Diagnostic::error(
"labels.field_missing",
format!("{}.{} is not declared in the schema", rule.kind, rule.field),
)),
}
match attributes.get(&rule.output) {
Some(value) if is_string_set_type(value) => {}
Some(value) => diagnostics.push(Diagnostic::error(
"labels.output_not_string_set",
format!(
"{}.{} must have schema type Set<String>, found {value}",
rule.kind, rule.output,
),
)),
None => diagnostics.push(Diagnostic::error(
"labels.output_missing",
format!(
"{}.{} is not declared in the schema",
rule.kind, rule.output
),
)),
}
}
diagnostics
}
}
fn entity_attributes<'a>(
schema_json: &'a Value,
kind: &str,
) -> Option<&'a serde_json::Map<String, Value>> {
let parsed = kind.parse::<EntityTypeName>().ok()?;
let namespace = parsed.namespace().to_string();
let namespace_definition = schema_json.as_object()?.get(&namespace)?;
let definition = namespace_definition
.get("entityTypes")?
.get(parsed.basename())?;
record_attributes(
schema_json,
&namespace,
definition.get("shape")?,
&mut HashSet::new(),
)
}
fn record_attributes<'a>(
schema_json: &'a Value,
namespace: &str,
shape: &'a Value,
visited: &mut HashSet<String>,
) -> Option<&'a serde_json::Map<String, Value>> {
if shape.get("type").and_then(Value::as_str) == Some("Record") {
return shape.get("attributes")?.as_object();
}
if shape.get("type").and_then(Value::as_str) != Some("EntityOrCommon") {
return None;
}
let name = shape.get("name")?.as_str()?;
let (common_namespace, basename) = name
.rsplit_once("::")
.map_or((namespace, name), |(namespace, basename)| {
(namespace, basename)
});
let qualified_name = format!("{common_namespace}::{basename}");
if !visited.insert(qualified_name) {
return None;
}
let common = schema_json
.as_object()?
.get(common_namespace)?
.get("commonTypes")?
.get(basename)?;
record_attributes(schema_json, common_namespace, common, visited)
}
fn is_string_type(value: &Value) -> bool {
value.get("type").and_then(Value::as_str) == Some("String")
|| (value.get("type").and_then(Value::as_str) == Some("EntityOrCommon")
&& value.get("name").and_then(Value::as_str) == Some("String"))
}
fn is_string_set_type(value: &Value) -> bool {
value.get("type").and_then(Value::as_str) == Some("Set")
&& value.get("element").is_some_and(is_string_type)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn strict_label_validation_rejects_unknown_fields() {
let error = LabelSet::from_json_str(
r#"[{"kind":"App::Host","field":"name","output":"labels","patterns":[{"name":"prod","regex":"prod","extra":true}]}]"#,
)
.unwrap_err();
assert!(error.diagnostics()[0].message.contains("unknown field"));
}
#[test]
fn label_set_converts_to_runtime_labelers() {
let labels = LabelSet::from_json_str(
r#"[{"kind":"App::Host","field":"name","output":"labels","patterns":[{"name":"prod","regex":"^prod"}]}]"#,
)
.unwrap();
assert_eq!(labels.to_labelers().len(), 1);
}
#[test]
fn schema_validation_resolves_common_record_shapes() {
let labels = LabelSet::from_json_str(
r#"[{"kind":"App::Host","field":"name","output":"labels","patterns":[{"name":"prod","regex":"^prod"}]}]"#,
)
.unwrap();
let schema = r#"{
"App": {
"commonTypes": {
"HostShape": {
"type": "Record",
"attributes": {
"name": {"type": "String", "required": true},
"labels": {
"type": "Set",
"element": {"type": "String"},
"required": false
}
},
"additionalAttributes": false
}
},
"entityTypes": {
"Host": {
"shape": {"type": "EntityOrCommon", "name": "HostShape"}
}
},
"actions": {}
}
}"#;
labels.validate_schema_json_str(schema).unwrap();
}
}