From f385139abcfb2ea4ed1ea19df18720f078cabb3a Mon Sep 17 00:00:00 2001
From: Antoine Bardoux <antoine.b@padoa-group.com>
Date: Fri, 12 May 2023 12:16:08 +0200
Subject: [PATCH] Allow multiple hostnames in internal annotation for pods

---
 source/pod.go      | 51 +++++++++++++++++++++++++++++-----------------
 source/pod_test.go | 40 ++++++++++++++++++++++++++++++++++++
 source/source.go   |  8 ++++++--
 3 files changed, 78 insertions(+), 21 deletions(-)

diff --git a/source/pod.go b/source/pod.go
index 123468539..912879b95 100644
--- a/source/pod.go
+++ b/source/pod.go
@@ -89,27 +89,17 @@ func (ps *podSource) Endpoints(ctx context.Context) ([]*endpoint.Endpoint, error
 			continue
 		}
 
-		if domain, ok := pod.Annotations[internalHostnameAnnotationKey]; ok {
-			addToEndpointMap(endpointMap, domain, suitableType(pod.Status.PodIP), pod.Status.PodIP)
-		}
-
-		if domain, ok := pod.Annotations[hostnameAnnotationKey]; ok {
-			node, _ := ps.nodeInformer.Lister().Get(pod.Spec.NodeName)
-			for _, address := range node.Status.Addresses {
-				recordType := suitableType(address.Address)
-				// IPv6 addresses are labeled as NodeInternalIP despite being usable externally as well.
-				if address.Type == corev1.NodeExternalIP || (address.Type == corev1.NodeInternalIP && recordType == endpoint.RecordTypeAAAA) {
-					addToEndpointMap(endpointMap, domain, recordType, address.Address)
-				}
-			}
-		}
-
-		if ps.compatibility == "kops-dns-controller" {
-			if domain, ok := pod.Annotations[kopsDNSControllerInternalHostnameAnnotationKey]; ok {
+		if domainAnnotation, ok := pod.Annotations[internalHostnameAnnotationKey]; ok {
+			domainList := splitHostnameAnnotation(domainAnnotation)
+			for _, domain := range domainList {
 				addToEndpointMap(endpointMap, domain, suitableType(pod.Status.PodIP), pod.Status.PodIP)
-			}
 
-			if domain, ok := pod.Annotations[kopsDNSControllerHostnameAnnotationKey]; ok {
+			}
+		}
+
+		if domainAnnotation, ok := pod.Annotations[hostnameAnnotationKey]; ok {
+			domainList := splitHostnameAnnotation(domainAnnotation)
+			for _, domain := range domainList {
 				node, _ := ps.nodeInformer.Lister().Get(pod.Spec.NodeName)
 				for _, address := range node.Status.Addresses {
 					recordType := suitableType(address.Address)
@@ -120,6 +110,29 @@ func (ps *podSource) Endpoints(ctx context.Context) ([]*endpoint.Endpoint, error
 				}
 			}
 		}
+
+		if ps.compatibility == "kops-dns-controller" {
+			if domainAnnotation, ok := pod.Annotations[kopsDNSControllerInternalHostnameAnnotationKey]; ok {
+				domainList := splitHostnameAnnotation(domainAnnotation)
+				for _, domain := range domainList {
+					addToEndpointMap(endpointMap, domain, suitableType(pod.Status.PodIP), pod.Status.PodIP)
+				}
+			}
+
+			if domainAnnotation, ok := pod.Annotations[kopsDNSControllerHostnameAnnotationKey]; ok {
+				domainList := splitHostnameAnnotation(domainAnnotation)
+				for _, domain := range domainList {
+					node, _ := ps.nodeInformer.Lister().Get(pod.Spec.NodeName)
+					for _, address := range node.Status.Addresses {
+						recordType := suitableType(address.Address)
+						// IPv6 addresses are labeled as NodeInternalIP despite being usable externally as well.
+						if address.Type == corev1.NodeExternalIP || (address.Type == corev1.NodeInternalIP && recordType == endpoint.RecordTypeAAAA) {
+							addToEndpointMap(endpointMap, domain, recordType, address.Address)
+						}
+					}
+				}
+			}
+		}
 	}
 	endpoints := []*endpoint.Endpoint{}
 	for key, targets := range endpointMap {
diff --git a/source/pod_test.go b/source/pod_test.go
index 5a57aedc2..549a9ebf9 100644
--- a/source/pod_test.go
+++ b/source/pod_test.go
@@ -526,6 +526,46 @@ func TestPodSource(t *testing.T) {
 				},
 			},
 		},
+		{
+			"split record for internal hostname annotation",
+			"",
+			"",
+			[]*endpoint.Endpoint{
+				{DNSName: "internal.a.foo.example.org", Targets: endpoint.Targets{"10.0.1.1"}, RecordType: endpoint.RecordTypeA},
+				{DNSName: "internal.b.foo.example.org", Targets: endpoint.Targets{"10.0.1.1"}, RecordType: endpoint.RecordTypeA},
+			},
+			false,
+			[]*corev1.Node{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "my-node1",
+					},
+					Status: corev1.NodeStatus{
+						Addresses: []corev1.NodeAddress{
+							{Type: corev1.NodeInternalIP, Address: "10.0.1.1"},
+						},
+					},
+				},
+			},
+			[]*corev1.Pod{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name:      "my-pod1",
+						Namespace: "kube-system",
+						Annotations: map[string]string{
+							internalHostnameAnnotationKey: "internal.a.foo.example.org,internal.b.foo.example.org",
+						},
+					},
+					Spec: corev1.PodSpec{
+						HostNetwork: true,
+						NodeName:    "my-node1",
+					},
+					Status: corev1.PodStatus{
+						PodIP: "10.0.1.1",
+					},
+				},
+			},
+		},
 	} {
 		tc := tc
 		t.Run(tc.title, func(t *testing.T) {
diff --git a/source/source.go b/source/source.go
index 91b83bb4f..e63c511ba 100644
--- a/source/source.go
+++ b/source/source.go
@@ -157,7 +157,7 @@ func getHostnamesFromAnnotations(annotations map[string]string) []string {
 	if !exists {
 		return nil
 	}
-	return strings.Split(strings.Replace(hostnameAnnotation, " ", "", -1), ",")
+	return splitHostnameAnnotation(hostnameAnnotation)
 }
 
 func getAccessFromAnnotations(annotations map[string]string) string {
@@ -173,7 +173,11 @@ func getInternalHostnamesFromAnnotations(annotations map[string]string) []string
 	if !exists {
 		return nil
 	}
-	return strings.Split(strings.Replace(internalHostnameAnnotation, " ", "", -1), ",")
+	return splitHostnameAnnotation(internalHostnameAnnotation)
+}
+
+func splitHostnameAnnotation(annotation string) []string {
+	return strings.Split(strings.Replace(annotation, " ", "", -1), ",")
 }
 
 func getAliasFromAnnotations(annotations map[string]string) bool {