update tree with propagate logic and tests

Signed-off-by: Phil Brookes <[email protected]>

internal/common/tree.go

@@ -13,6 +13,7 @@ type DNSTreeNode struct {
 	Name     string
 	Children []*DNSTreeNode
 	DataSets []DNSTreeNodeData
+	Parent   *DNSTreeNode
 }
 
 // DNSTreeNodeData holds a data for the enpoint(s) that correspond to this node
@@ -25,6 +26,171 @@ type DNSTreeNodeData struct {
 	Targets          []string
 }
 
+// PropagateStoppableLabel takes a propLabel (and value) to propagate throughout a tree, and a stopLabel
+// whenever the label is propagated to a dataset in a node which also has the stopLabel, this node and
+// all of the children of this node and any parents will have the propLabel removed.
+//
+// N.B. Any node with no parents is assumed to have the stopLabel, even when not present, to prevent the
+// propLabel propagating to the entire tree (if this is required, use `AddLabelToBranch` on the root node).
+//
+// The overview of the logic of this function is as follows:
+// - Spread propLabels as greedily as possible
+//   - Any labelled node labels all of it's children
+//   - Any node with all children labelled get's the propLabel too
+//
+// - Resolve the stopLabels
+//   - Any node with the stopLabel and the propLabel:
+//   - Has the propLabel removed from itself and all it's children
+//   - Has the propLabel removed from any parent (or parent's parent) that has the label
+
+func PropagateStoppableLabel(node *DNSTreeNode, propLabel, value, stopLabel string) {
+	//propagate labels regardless of stop labels
+	propagateLabel(node, propLabel, value)
+
+	//propagate stop labels
+	resolveStops(node, propLabel, value, stopLabel)
+}
+
+func resolveStops(node *DNSTreeNode, label, value, stopLabel string) {
+	if isRoot(node) && allChildrenHaveLabel(node, label, value) {
+		RemoveLabelFromTree(node, label)
+		//entire tree cleaned so no need for any further checks
+		return
+	}
+
+	//remove label from stop labelled children
+	for _, c := range node.Children {
+		d := findDataSetForChild(node, c.Name)
+		//has label and stop label = remove label from this dataset and all children
+		if d != nil && d.Labels[stopLabel] != "" && d.Labels[label] != "" {
+			delete(d.Labels, label)
+			RemoveLabelFromTree(c, label)
+			RemoveLabelFromParents(node, label)
+		}
+
+	}
+
+}
+
+func propagateLabel(node *DNSTreeNode, label, value string) {
+	for _, c := range node.Children {
+		d := findDataSetForChild(node, c.Name)
+		if d != nil && d.Labels[label] != "" {
+			//this child is labelled, indiscriminately label entire tree under this child
+			AddLabelToTree(c, label, value)
+		} else {
+			// this child is not labelled, continue descending to propagate label
+			propagateLabel(c, label, value)
+		}
+	}
+
+	// if all children are labelled, label this branch in parent node
+	if len(node.Children) > 0 && allChildrenHaveLabel(node, label, value) && node.Parent != nil {
+		AddLabelToBranch(node.Parent, node.Name, label, value)
+	}
+}
+
+func isRoot(node *DNSTreeNode) bool {
+	return node.Parent == nil
+}
+
+func allChildrenHaveLabel(node *DNSTreeNode, label, value string) bool {
+	for _, c := range node.Children {
+		if !HasLabelForBranch(node, c.Name, label, value) {
+			return false
+		}
+	}
+	return true
+}
+
+func findDataSetForChild(node *DNSTreeNode, name string) *DNSTreeNodeData {
+	for _, d := range node.DataSets {
+		if slices.Contains(d.Targets, name) {
+			return &d
+		}
+	}
+	return nil
+}
+
+func AddLabelToBranch(node *DNSTreeNode, branch, label, value string) {
+	d := findDataSetForChild(node, branch)
+	if d == nil {
+		node.DataSets = append(node.DataSets, DNSTreeNodeData{
+			Labels: endpoint.Labels{
+				label: value,
+			},
+			Targets: []string{
+				branch,
+			},
+		})
+	} else {
+		if len(d.Targets) == 1 {
+			d.Labels[label] = value
+		} else {
+			//remove target from shared dataset and recreate uniquely
+			for i, t := range d.Targets {
+				if t == branch {
+					d.Targets = append(d.Targets[:i], d.Targets[i+1:]...)
+					newDS := DNSTreeNodeData{
+						Labels:  d.Labels.DeepCopy(),
+						Targets: []string{branch},
+					}
+					newDS.Labels[label] = value
+					node.DataSets = append(node.DataSets, newDS)
+				}
+			}
+		}
+	}
+}
+
+func AddChild(parent *DNSTreeNode, child *DNSTreeNode) {
+	parent.Children = append(parent.Children, child)
+	child.Parent = parent
+}
+
+func RemoveLabelFromParents(node *DNSTreeNode, label string) {
+	if isRoot(node) {
+		return
+	}
+
+	d := findDataSetForChild(node.Parent, node.Name)
+	if d == nil {
+		return
+	}
+
+	delete(d.Labels, label)
+
+	RemoveLabelFromParents(node.Parent, label)
+}
+
+func RemoveLabelFromTree(node *DNSTreeNode, label string) {
+	for _, d := range node.DataSets {
+		delete(d.Labels, label)
+	}
+
+	for _, c := range node.Children {
+		RemoveLabelFromTree(c, label)
+	}
+}
+
+func AddLabelToTree(node *DNSTreeNode, label, value string) {
+	for _, c := range node.Children {
+		AddLabelToBranch(node, c.Name, label, value)
+		AddLabelToTree(c, label, value)
+	}
+}
+
+func HasLabelForBranch(node *DNSTreeNode, branch, label, value string) bool {
+	d := findDataSetForChild(node, branch)
+	if d == nil {
+		return false
+	}
+	if v, ok := d.Labels[label]; ok {
+		return value == v
+	}
+	return false
+}
+
 // RemoveNode removes a node from a tree.
 // If the node was the only child of the parent node,
 // the parent will be removed as well unless the parent is a root node
@@ -115,21 +281,11 @@ func ToEndpoints(node *DNSTreeNode, endpoints *[]*endpoint.Endpoint) *[]*endpoin
 	if isALeafNode(node) {
 		return endpoints
 	}
-	targets := []string{}
+
 	for _, child := range node.Children {
-		targets = append(targets, child.Name)
 		ToEndpoints(child, endpoints)
 	}
 
-	// this should not happen. the node is either leaf or has datasets (unless the cree was made manually)
-	if node.DataSets == nil {
-		*endpoints = append(*endpoints, &endpoint.Endpoint{
-			DNSName: node.Name,
-			Targets: targets,
-		})
-		return endpoints
-	}
-
 	for _, data := range node.DataSets {
 		*endpoints = append(*endpoints, &endpoint.Endpoint{
 			DNSName:          node.Name,
@@ -162,6 +318,7 @@ func populateNode(node *DNSTreeNode, record *v1alpha1.DNSRecord) {
 	}
 
 	for _, c := range children {
+		c.Parent = node
 		populateNode(c, record)
 	}
 	node.Children = children
@@ -204,5 +361,5 @@ func findDataSets(name string, record *v1alpha1.DNSRecord) []DNSTreeNodeData {
 // isALeafNode check if this is the last node in a tree
 func isALeafNode(node *DNSTreeNode) bool {
 	// no children means this is pointing to an IP or a host outside of the DNS Record
-	return node.Children == nil || len(node.Children) == 0
+	return len(node.Children) == 0
 }