1use std::collections::{BTreeMap, BTreeSet};
41use std::fmt;
42
43use zenkey::origin::HostId;
44
45use crate::model::facts::{KeyFacts, KeyShape, OriginKind};
46use crate::model::snapshot_diff::structural_of;
47use crate::report::{Asked, MapEvidence, OriginPair, Side, Snapshot, SnapshotRow, Unmapped};
48
49#[derive(Debug, Clone, PartialEq, Eq)]
52pub struct Label {
53 pub source: String,
55 pub verified: bool,
59}
60
61#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct OriginProfile {
65 pub origin: String,
67 pub producers: BTreeSet<String>,
70 pub label: Asked<Option<Label>>,
76}
77
78impl OriginProfile {
79 fn verified_label(&self) -> Option<&str> {
81 match &self.label {
82 Asked::Asked(Some(l)) if l.verified => Some(&l.source),
83 _ => None,
84 }
85 }
86}
87
88fn is_bridge_document(facts: &KeyFacts) -> bool {
92 let KeyShape::V1(f) = &facts.shape else {
93 return false;
94 };
95 f.origin_kind == OriginKind::Host
96 && f.class == "state"
97 && f.producer.is_some()
98 && matches!(f.subject.as_slice(), [s] if s == "health" || s == "sensor")
99}
100
101fn bridge_claim(row: &SnapshotRow) -> Option<(Option<String>, Option<String>)> {
104 let doc = structural_of(row)?;
105 let field = |name: &str| doc.get(name).and_then(|v| v.as_str()).map(str::to_string);
106 Some((field("source"), field("host_id")))
107}
108
109pub fn origin_profiles(snapshot: &Snapshot) -> Vec<OriginProfile> {
111 struct Acc {
112 producers: BTreeSet<String>,
113 claims: Vec<(Option<String>, Option<String>)>,
115 bridge_rows: usize,
116 }
117 let base = snapshot.header.base.as_str();
118 let mut acc: BTreeMap<String, Acc> = BTreeMap::new();
119 for row in &snapshot.rows {
120 let facts = KeyFacts::project(base, &row.key);
121 let KeyShape::V1(f) = &facts.shape else {
122 continue;
123 };
124 if f.origin_kind != OriginKind::Host {
125 continue;
126 }
127 let entry = acc.entry(f.origin.clone()).or_insert_with(|| Acc {
128 producers: BTreeSet::new(),
129 claims: Vec::new(),
130 bridge_rows: 0,
131 });
132 if let Some(p) = &f.producer {
133 entry.producers.insert(p.clone());
134 }
135 if is_bridge_document(&facts) && !row.delete {
137 entry.bridge_rows += 1;
138 if let Some(claim) = bridge_claim(row) {
139 entry.claims.push(claim);
140 }
141 }
142 }
143 acc.into_iter()
144 .map(|(origin, a)| {
145 let label = if a.bridge_rows == 0 {
146 Asked::NotAsked
147 } else {
148 let sources: BTreeSet<&str> =
149 a.claims.iter().filter_map(|(s, _)| s.as_deref()).collect();
150 match sources.into_iter().collect::<Vec<_>>().as_slice() {
151 [source] => {
154 let verified = a
155 .claims
156 .iter()
157 .filter(|(s, _)| s.as_deref() == Some(source))
158 .all(|(_, h)| h.as_deref() == Some(origin.as_str()));
159 Asked::Asked(Some(Label {
160 source: (*source).to_string(),
161 verified,
162 }))
163 }
164 _ => Asked::Asked(None),
167 }
168 };
169 OriginProfile {
170 origin,
171 producers: a.producers,
172 label,
173 }
174 })
175 .collect()
176}
177
178#[derive(Debug, Clone, PartialEq, Eq, Default)]
181pub struct MapPlan {
182 pub pairs: Vec<OriginPair>,
183 pub unmapped: Vec<Unmapped>,
184}
185
186impl MapPlan {
187 pub fn is_complete(&self) -> bool {
189 self.unmapped.is_empty()
190 }
191
192 pub fn b_to_a(&self) -> BTreeMap<&str, &str> {
194 self.pairs
195 .iter()
196 .map(|p| (p.b.as_str(), p.a.as_str()))
197 .collect()
198 }
199}
200
201#[derive(Debug, Clone, PartialEq, Eq)]
204pub enum MapError {
205 UnknownOrigin { origin: String, side: Side },
207 PairedTwice { origin: String, side: Side },
209}
210
211impl fmt::Display for MapError {
212 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
213 match self {
214 MapError::UnknownOrigin { origin, side } => write!(
215 f,
216 "--map names {origin}, which is not a host origin in {}",
217 side_name(*side)
218 ),
219 MapError::PairedTwice { origin, side } => write!(
220 f,
221 "--map names {origin} (in {}) twice; an origin pairs once",
222 side_name(*side)
223 ),
224 }
225 }
226}
227
228impl std::error::Error for MapError {}
229
230fn side_name(side: Side) -> &'static str {
231 match side {
232 Side::A => "a",
233 Side::B => "b",
234 }
235}
236
237fn producer_set(p: &OriginProfile) -> String {
239 format!(
240 "{{{}}}",
241 p.producers.iter().cloned().collect::<Vec<_>>().join(", ")
242 )
243}
244
245pub fn plan_map(
249 a: &[OriginProfile],
250 b: &[OriginProfile],
251 explicit: &[(HostId, HostId)],
252) -> Result<MapPlan, MapError> {
253 let (ia, ib) = (index(a), index(b));
254 let mut used_a: BTreeSet<&str> = BTreeSet::new();
255 let mut used_b: BTreeSet<&str> = BTreeSet::new();
256 let mut pairs = Vec::new();
257
258 for (x, y) in explicit {
260 let (x, y) = (x.as_str(), y.as_str());
261 if !ia.contains_key(x) {
262 return Err(MapError::UnknownOrigin {
263 origin: x.to_string(),
264 side: Side::A,
265 });
266 }
267 if !ib.contains_key(y) {
268 return Err(MapError::UnknownOrigin {
269 origin: y.to_string(),
270 side: Side::B,
271 });
272 }
273 if !used_a.insert(x) {
274 return Err(MapError::PairedTwice {
275 origin: x.to_string(),
276 side: Side::A,
277 });
278 }
279 if !used_b.insert(y) {
280 return Err(MapError::PairedTwice {
281 origin: y.to_string(),
282 side: Side::B,
283 });
284 }
285 pairs.push(OriginPair {
286 a: x.to_string(),
287 b: y.to_string(),
288 evidence: MapEvidence::Explicit,
289 });
290 }
291
292 {
296 let (fa, fb) = (free(a, &used_a), free(b, &used_b));
297 let (la, lb) = (by_label(&fa), by_label(&fb));
298 for p in &fa {
299 let Some(label) = p.verified_label() else {
300 continue;
301 };
302 if la[label].len() != 1 {
303 continue;
304 }
305 let Some([q]) = lb.get(label).map(Vec::as_slice) else {
306 continue;
307 };
308 used_a.insert(&p.origin);
309 used_b.insert(&q.origin);
310 pairs.push(OriginPair {
311 a: p.origin.clone(),
312 b: q.origin.clone(),
313 evidence: MapEvidence::Label {
314 source: label.to_string(),
315 },
316 });
317 }
318 }
319
320 {
322 let (fa, fb) = (free(a, &used_a), free(b, &used_b));
323 let (sa, sb) = (by_set(&fa), by_set(&fb));
324 for p in &fa {
325 if sa[&p.producers].len() != 1 {
326 continue;
327 }
328 let Some([q]) = sb.get(&p.producers).map(Vec::as_slice) else {
329 continue;
330 };
331 used_a.insert(&p.origin);
332 used_b.insert(&q.origin);
333 pairs.push(OriginPair {
334 a: p.origin.clone(),
335 b: q.origin.clone(),
336 evidence: MapEvidence::ProducerSet,
337 });
338 }
339 }
340
341 let (fa, fb) = (free(a, &used_a), free(b, &used_b));
345 let mut unmapped = Vec::new();
346 for p in &fa {
347 unmapped.push(Unmapped {
348 origin: p.origin.clone(),
349 side: Side::A,
350 reason: why_unpaired(p, &fa, &fb, Side::A),
351 });
352 }
353 for p in &fb {
354 unmapped.push(Unmapped {
355 origin: p.origin.clone(),
356 side: Side::B,
357 reason: why_unpaired(p, &fb, &fa, Side::B),
358 });
359 }
360 Ok(MapPlan { pairs, unmapped })
361}
362
363fn index(side: &[OriginProfile]) -> BTreeMap<&str, &OriginProfile> {
364 side.iter().map(|p| (p.origin.as_str(), p)).collect()
365}
366
367fn free<'p>(side: &'p [OriginProfile], used: &BTreeSet<&str>) -> Vec<&'p OriginProfile> {
369 side.iter()
370 .filter(|p| !used.contains(p.origin.as_str()))
371 .collect()
372}
373
374fn by_label<'p>(side: &[&'p OriginProfile]) -> BTreeMap<&'p str, Vec<&'p OriginProfile>> {
376 let mut m: BTreeMap<&str, Vec<&OriginProfile>> = BTreeMap::new();
377 for p in side {
378 if let Some(l) = p.verified_label() {
379 m.entry(l).or_default().push(p);
380 }
381 }
382 m
383}
384
385fn by_set<'p>(
387 side: &[&'p OriginProfile],
388) -> BTreeMap<&'p BTreeSet<String>, Vec<&'p OriginProfile>> {
389 let mut m: BTreeMap<&BTreeSet<String>, Vec<&OriginProfile>> = BTreeMap::new();
390 for p in side {
391 m.entry(&p.producers).or_default().push(p);
392 }
393 m
394}
395
396fn why_unpaired(
400 p: &OriginProfile,
401 this: &[&OriginProfile],
402 other: &[&OriginProfile],
403 side: Side,
404) -> String {
405 let (here, there) = match side {
406 Side::A => ("a", "b"),
407 Side::B => ("b", "a"),
408 };
409 let label = match &p.label {
410 Asked::NotAsked => "no health/sensor row: label not asked".to_string(),
411 Asked::Asked(None) => "health/sensor row carries no usable `source`".to_string(),
412 Asked::Asked(Some(l)) if !l.verified => format!(
413 "label `{}` not verified: `host_id` absent or not this origin",
414 l.source
415 ),
416 Asked::Asked(Some(l)) => {
417 let claims = |side: &[&OriginProfile]| {
418 side.iter()
419 .filter(|q| q.verified_label() == Some(l.source.as_str()))
420 .count()
421 };
422 let (n_here, n_there) = (claims(this), claims(other));
423 if n_here > 1 {
424 format!("label `{}` claimed by {n_here} origins in {here}", l.source)
425 } else if n_there == 0 {
426 let unverified = other.iter().any(|q| {
427 matches!(&q.label, Asked::Asked(Some(m)) if m.source == l.source && !m.verified)
428 });
429 if unverified {
430 format!(
431 "label `{}` claimed in {there} only by an unverified document",
432 l.source
433 )
434 } else {
435 format!("label `{}` claimed by no origin in {there}", l.source)
436 }
437 } else {
438 format!(
439 "label `{}` claimed by {n_there} origins in {there}",
440 l.source
441 )
442 }
443 }
444 };
445 let matches =
446 |side: &[&OriginProfile]| side.iter().filter(|q| q.producers == p.producers).count();
447 let (n_here, n_there) = (matches(this), matches(other));
448 let set = producer_set(p);
449 let producers = if n_there == 0 {
450 format!("producer set {set} matches no origin in {there}")
451 } else if n_here > 1 {
452 format!("producer set {set} shared by {n_here} origins in {here}")
453 } else {
454 format!("producer set {set} matches {n_there} origins in {there}")
455 };
456 format!("{label}; {producers}")
457}
458
459#[cfg(test)]
460pub(crate) mod tests {
461 use super::*;
462 use crate::report::{AnsweredBy, Holder, RegistrationWire, VerdictWire, ZsnapHeader};
463
464 fn header(base: &str) -> ZsnapHeader {
465 ZsnapHeader {
466 zsnap: 1,
467 selectors: vec![zenkey::grammar::with_base(base, "v1/**")],
468 base: base.into(),
469 collected_at: "2026-09-06T00:00:00Z".into(),
470 collection_span_s: 0.5,
471 asked: 1,
472 answered: 0,
473 elided: 0,
474 errors: 0,
475 superseded: 0,
476 roster: Asked::NotAsked,
477 }
478 }
479
480 pub(crate) fn row(key: &str, body: &str) -> SnapshotRow {
481 use base64::Engine as _;
482 let origin = key
483 .split('/')
484 .skip_while(|c| *c != "v1")
485 .nth(1)
486 .unwrap_or("")
487 .to_string();
488 SnapshotRow {
489 key: key.into(),
490 delete: false,
491 bytes: Some(base64::engine::general_purpose::STANDARD.encode(body)),
492 encoding: Some("application/json".into()),
493 timestamp: None,
494 stamper: None,
495 source: None,
496 source_zid: None,
497 registration: RegistrationWire::RegistryNotLoaded,
498 verdict: VerdictWire::NotValidated {
499 reason: "no_registry".into(),
500 },
501 holder: Holder::Live {
502 origin,
503 answered_by: AnsweredBy::Stamper,
504 },
505 }
506 }
507
508 pub(crate) fn host(base: &str, origin: &str, label: &str, extra: &[&str]) -> Vec<SnapshotRow> {
510 let k = |rel: &str| zenkey::grammar::with_base(base, format!("v1/{origin}/{rel}"));
511 let mut rows = vec![
512 row(
513 &k("state/sysinfo/health"),
514 &format!(r#"{{"host_id":"{origin}","source":"{label}","status":"ok"}}"#),
515 ),
516 row(&k("telemetry/sysinfo/disk/root/used"), r#"{"value":41.0}"#),
517 ];
518 for p in extra {
519 rows.push(row(&k(&format!("state/{p}/rotated")), r#"{"n":3}"#));
520 }
521 rows
522 }
523
524 pub(crate) fn snap(base: &str, rows: Vec<SnapshotRow>) -> Snapshot {
525 let mut rows = rows;
526 rows.sort_by(|x, y| x.key.cmp(&y.key));
527 Snapshot {
528 header: header(base),
529 rows,
530 }
531 }
532
533 const A1: &str = "h-aaaaaaaaaaa1";
534 const A2: &str = "h-aaaaaaaaaaa2";
535 const B1: &str = "h-bbbbbbbbbbb1";
536 const B2: &str = "h-bbbbbbbbbbb2";
537
538 fn hid(s: &str) -> HostId {
539 HostId::parse(s).unwrap()
540 }
541
542 #[test]
545 fn a_renamed_fleet_pairs_every_origin_on_its_label() {
546 let a = snap(
547 "acme",
548 [
549 host("acme", A1, "web", &[]),
550 host("acme", A2, "db", &["logs"]),
551 ]
552 .concat(),
553 );
554 let b = snap(
555 "acme",
556 [
557 host("acme", B1, "web", &[]),
558 host("acme", B2, "db", &["logs"]),
559 ]
560 .concat(),
561 );
562 let (pa, pb) = (origin_profiles(&a), origin_profiles(&b));
563 assert_eq!(pa.len(), 2);
564 assert_eq!(
565 pa[1].producers,
566 ["logs", "sysinfo"].into_iter().map(String::from).collect()
567 );
568 assert_eq!(
569 pa[0].label,
570 Asked::Asked(Some(Label {
571 source: "web".into(),
572 verified: true
573 }))
574 );
575 let plan = plan_map(&pa, &pb, &[]).unwrap();
576 assert!(plan.is_complete());
577 assert_eq!(
578 plan.pairs,
579 vec![
580 OriginPair {
581 a: A1.into(),
582 b: B1.into(),
583 evidence: MapEvidence::Label {
584 source: "web".into()
585 }
586 },
587 OriginPair {
588 a: A2.into(),
589 b: B2.into(),
590 evidence: MapEvidence::Label {
591 source: "db".into()
592 }
593 },
594 ]
595 );
596 }
597
598 #[test]
602 fn an_ambiguous_label_leaves_both_claimants_unpaired_and_says_why() {
603 let a = snap("acme", host("acme", A1, "node", &[]));
604 let b = snap(
605 "acme",
606 [host("acme", B1, "node", &[]), host("acme", B2, "node", &[])].concat(),
607 );
608 let plan = plan_map(&origin_profiles(&a), &origin_profiles(&b), &[]).unwrap();
609 assert!(plan.pairs.is_empty());
610 let reasons: Vec<(&str, Side, &str)> = plan
611 .unmapped
612 .iter()
613 .map(|u| (u.origin.as_str(), u.side, u.reason.as_str()))
614 .collect();
615 assert_eq!(
616 reasons,
617 vec![
618 (
619 A1,
620 Side::A,
621 "label `node` claimed by 2 origins in b; producer set {sysinfo} matches 2 origins in b"
622 ),
623 (
624 B1,
625 Side::B,
626 "label `node` claimed by 2 origins in b; producer set {sysinfo} shared by 2 origins in b"
627 ),
628 (
629 B2,
630 Side::B,
631 "label `node` claimed by 2 origins in b; producer set {sysinfo} shared by 2 origins in b"
632 ),
633 ]
634 );
635 }
636
637 #[test]
641 fn an_explicit_pair_beats_a_conflicting_label() {
642 let a = snap(
643 "acme",
644 [
645 host("acme", A1, "web", &[]),
646 host("acme", A2, "db", &["logs"]),
647 ]
648 .concat(),
649 );
650 let b = snap(
651 "acme",
652 [
653 host("acme", B1, "web", &[]),
654 host("acme", B2, "db", &["logs"]),
655 ]
656 .concat(),
657 );
658 let plan = plan_map(
659 &origin_profiles(&a),
660 &origin_profiles(&b),
661 &[(hid(A1), hid(B2))],
662 )
663 .unwrap();
664 assert_eq!(plan.pairs[0].evidence, MapEvidence::Explicit);
665 assert_eq!(
666 (plan.pairs[0].a.as_str(), plan.pairs[0].b.as_str()),
667 (A1, B2)
668 );
669 assert_eq!(plan.pairs.len(), 1);
673 assert_eq!(plan.unmapped.len(), 2);
674 assert_eq!(
675 plan.unmapped[0].reason,
676 "label `db` claimed by no origin in b; producer set {logs, sysinfo} matches no origin in b"
677 );
678 }
679
680 #[test]
684 fn an_origin_only_in_b_is_unmapped_on_side_b() {
685 let a = snap("acme", host("acme", A1, "web", &[]));
686 let b = snap(
687 "acme",
688 [
689 host("acme", B1, "web", &[]),
690 host("acme", B2, "db", &["logs"]),
691 ]
692 .concat(),
693 );
694 let plan = plan_map(&origin_profiles(&a), &origin_profiles(&b), &[]).unwrap();
695 assert_eq!(plan.pairs.len(), 1);
696 assert_eq!(
697 plan.unmapped,
698 vec![Unmapped {
699 origin: B2.into(),
700 side: Side::B,
701 reason: "label `db` claimed by no origin in a; producer set {logs, sysinfo} matches no origin in a".into(),
702 }]
703 );
704 }
705
706 #[test]
709 fn an_unlabelled_fleet_pairs_on_producer_sets_and_says_the_label_was_not_asked() {
710 let strip = |rows: Vec<SnapshotRow>| -> Vec<SnapshotRow> {
711 rows.into_iter()
712 .filter(|r| !r.key.ends_with("/health"))
713 .collect()
714 };
715 let a = snap(
716 "acme",
717 strip(
718 [
719 host("acme", A1, "web", &[]),
720 host("acme", A2, "db", &["logs"]),
721 ]
722 .concat(),
723 ),
724 );
725 let b = snap(
726 "acme",
727 strip(
728 [
729 host("acme", B1, "web", &[]),
730 host("acme", B2, "db", &["logs"]),
731 ]
732 .concat(),
733 ),
734 );
735 let pa = origin_profiles(&a);
736 assert_eq!(pa[0].label, Asked::NotAsked);
737 let plan = plan_map(&pa, &origin_profiles(&b), &[]).unwrap();
738 assert!(plan.is_complete());
739 assert!(
740 plan.pairs
741 .iter()
742 .all(|p| p.evidence == MapEvidence::ProducerSet)
743 );
744 assert_eq!(plan.b_to_a()[B2], A2);
745
746 let a = snap(
749 "acme",
750 strip([host("acme", A1, "x", &[]), host("acme", A2, "y", &[])].concat()),
751 );
752 let plan = plan_map(&origin_profiles(&a), &origin_profiles(&b), &[]).unwrap();
753 assert_eq!(
754 plan.unmapped[0].reason,
755 "no health/sensor row: label not asked; producer set {sysinfo} shared by 2 origins in a"
756 );
757 }
758
759 #[test]
762 fn an_unverified_label_does_not_pair() {
763 let mut a_rows = host("acme", A1, "web", &[]);
764 a_rows[0] = row(
765 &format!("acme/v1/{A1}/state/sysinfo/health"),
766 r#"{"host_id":"h-000000000000","source":"web"}"#,
767 );
768 let a = snap("acme", a_rows);
769 let b = snap(
770 "acme",
771 [host("acme", B1, "web", &[]), host("acme", B2, "web", &[])].concat(),
772 );
773 let pa = origin_profiles(&a);
774 assert_eq!(
775 pa[0].label,
776 Asked::Asked(Some(Label {
777 source: "web".into(),
778 verified: false
779 }))
780 );
781 let plan = plan_map(&pa, &origin_profiles(&b), &[]).unwrap();
782 assert!(plan.pairs.is_empty());
783 assert!(
784 plan.unmapped[0]
785 .reason
786 .starts_with("label `web` not verified")
787 );
788 assert!(
789 plan.unmapped[1]
790 .reason
791 .starts_with("label `web` claimed by 2 origins in b"),
792 "{}",
793 plan.unmapped[1].reason
794 );
795 }
796
797 #[test]
800 fn an_explicit_pair_must_name_origins_both_snapshots_hold() {
801 let a = snap("acme", host("acme", A1, "web", &[]));
802 let b = snap("acme", host("acme", B1, "web", &[]));
803 let (pa, pb) = (origin_profiles(&a), origin_profiles(&b));
804 assert_eq!(
805 plan_map(&pa, &pb, &[(hid(A2), hid(B1))]),
806 Err(MapError::UnknownOrigin {
807 origin: A2.into(),
808 side: Side::A
809 })
810 );
811 assert_eq!(
812 plan_map(&pa, &pb, &[(hid(A1), hid(B2))])
813 .unwrap_err()
814 .to_string(),
815 "--map names h-bbbbbbbbbbb2, which is not a host origin in b"
816 );
817 let twice = snap(
818 "acme",
819 [host("acme", A1, "web", &[]), host("acme", A2, "db", &[])].concat(),
820 );
821 assert_eq!(
822 plan_map(
823 &origin_profiles(&twice),
824 &pb,
825 &[(hid(A1), hid(B1)), (hid(A2), hid(B1))]
826 ),
827 Err(MapError::PairedTwice {
828 origin: B1.into(),
829 side: Side::B
830 })
831 );
832 }
833
834 #[test]
837 fn a_service_origin_is_not_profiled() {
838 let mut rows = host("acme", A1, "web", &[]);
839 rows.push(row("acme/v1/@catalog/state/entity/x", "{}"));
840 let p = origin_profiles(&snap("acme", rows));
841 assert_eq!(p.len(), 1);
842 assert_eq!(p[0].origin, A1);
843 }
844}