From cc872ef7547877e14e9432ce2a1949b9e69b0538 Mon Sep 17 00:00:00 2001
From: enxebre <alberto.garcial@hotmail.com>
Date: Tue, 20 Sep 2022 10:46:07 +0200
Subject: [PATCH] Add support to reconcile labels from Machines to Nodes

---
 api/v1beta1/machine_types.go                  |  7 +++
 .../machine/machine_controller_noderef.go     | 44 +++++++++++++++++++
 .../machine_controller_noderef_test.go        | 26 +++++++++++
 3 files changed, 77 insertions(+)

diff --git a/api/v1beta1/machine_types.go b/api/v1beta1/machine_types.go
index 73bb663a044c..c10bd2e0a4b7 100644
--- a/api/v1beta1/machine_types.go
+++ b/api/v1beta1/machine_types.go
@@ -62,6 +62,13 @@ const (
 	// This annotation can be set on BootstrapConfig or Machine objects. The value set on the Machine object takes precedence.
 	// This annotation can only be used on Control Plane Machines.
 	MachineCertificatesExpiryDateAnnotation = "machine.cluster.x-k8s.io/certificates-expiry"
+
+	// NodeRoleLabelPrefix is one of the CAPI managed Node label domains.
+	NodeRoleLabelPrefix = "node-role.kubernetes.io"
+	// NodeRestrictionLabelDomain is one of the CAPI managed Node label domains.
+	NodeRestrictionLabelDomain = "node-restriction.kubernetes.io"
+	// ManagedNodeLabelDomain is one of the CAPI managed domain Node label domains.
+	ManagedNodeLabelDomain = "node.cluster.x-k8s.io"
 )
 
 // ANCHOR: MachineSpec
diff --git a/internal/controllers/machine/machine_controller_noderef.go b/internal/controllers/machine/machine_controller_noderef.go
index 29b5a397679f..e18f2dfb13e1 100644
--- a/internal/controllers/machine/machine_controller_noderef.go
+++ b/internal/controllers/machine/machine_controller_noderef.go
@@ -19,10 +19,12 @@ package machine
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"github.com/pkg/errors"
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/klog/v2"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -116,6 +118,17 @@ func (r *Reconciler) reconcileNode(ctx context.Context, cluster *clusterv1.Clust
 		}
 	}
 
+	options := []client.PatchOption{
+		client.FieldOwner("capi-machine"),
+		client.ForceOwnership,
+	}
+	nodeCopy := node.DeepCopy()
+	nodeCopy.Labels = getManagedLabels(machine.Labels)
+	nodePatch := unstructuredNode(nodeCopy)
+	if err := remoteClient.Patch(ctx, nodePatch, client.Apply, options...); err != nil {
+		return ctrl.Result{}, errors.Wrap(err, "failed to apply patch label to the node")
+	}
+
 	// Do the remaining node health checks, then set the node health to true if all checks pass.
 	status, message := summarizeNodeConditions(node)
 	if status == corev1.ConditionFalse {
@@ -131,6 +144,37 @@ func (r *Reconciler) reconcileNode(ctx context.Context, cluster *clusterv1.Clust
 	return ctrl.Result{}, nil
 }
 
+// unstructuredNode returns a raw unstructured from a Node.
+func unstructuredNode(node *corev1.Node) *unstructured.Unstructured {
+	obj := &unstructured.Unstructured{}
+	obj.SetAPIVersion("v1")
+	obj.SetKind("Node")
+	obj.SetName(node.Name)
+	obj.SetUID(node.UID)
+	obj.SetLabels(node.Labels)
+	return obj
+}
+
+// getManagedLabels gets a map[string]string and returns another map[string]string
+// filtering out labels not managed by CAPI.
+func getManagedLabels(labels map[string]string) map[string]string {
+	managedLabels := make(map[string]string)
+	for key, value := range labels {
+		dnsSubdomainOrName := strings.Split(key, "/")[0]
+		if dnsSubdomainOrName == clusterv1.NodeRoleLabelPrefix {
+			managedLabels[key] = value
+		}
+		if dnsSubdomainOrName == clusterv1.NodeRestrictionLabelDomain || strings.HasSuffix(dnsSubdomainOrName, "."+clusterv1.NodeRestrictionLabelDomain) {
+			managedLabels[key] = value
+		}
+		if dnsSubdomainOrName == clusterv1.ManagedNodeLabelDomain || strings.HasSuffix(dnsSubdomainOrName, "."+clusterv1.ManagedNodeLabelDomain) {
+			managedLabels[key] = value
+		}
+	}
+
+	return managedLabels
+}
+
 // summarizeNodeConditions summarizes a Node's conditions and returns the summary of condition statuses and concatenate failed condition messages:
 // if there is at least 1 semantically-negative condition, summarized status = False;
 // if there is at least 1 semantically-positive condition when there is 0 semantically negative condition, summarized status = True;
diff --git a/internal/controllers/machine/machine_controller_noderef_test.go b/internal/controllers/machine/machine_controller_noderef_test.go
index 6eedb805bf64..610475b5f586 100644
--- a/internal/controllers/machine/machine_controller_noderef_test.go
+++ b/internal/controllers/machine/machine_controller_noderef_test.go
@@ -227,3 +227,29 @@ func TestSummarizeNodeConditions(t *testing.T) {
 		})
 	}
 }
+
+func TestGetManagedLabels(t *testing.T) {
+	// Create managedLabels map from known managed prefixes.
+	managedLabels := map[string]string{}
+	managedLabels[clusterv1.ManagedNodeLabelDomain] = ""
+	managedLabels["custom-prefix."+clusterv1.NodeRestrictionLabelDomain] = ""
+	managedLabels["custom-prefix."+clusterv1.NodeRestrictionLabelDomain+"/anything"] = ""
+	managedLabels[clusterv1.NodeRoleLabelPrefix+"/anything"] = ""
+
+	// Append arbitrary labels.
+	allLabels := map[string]string{
+		"foo":                               "",
+		"bar":                               "",
+		"company.xyz/node.cluster.x-k8s.io": "not-managed",
+		"gpu-node.cluster.x-k8s.io":         "not-managed",
+		"company.xyz/node-restriction.kubernetes.io": "not-managed",
+		"gpu-node-restriction.kubernetes.io":         "not-managed",
+	}
+	for k, v := range managedLabels {
+		allLabels[k] = v
+	}
+
+	g := NewWithT(t)
+	got := getManagedLabels(allLabels)
+	g.Expect(got).To(BeEquivalentTo(managedLabels))
+}