From a3ac555a5e37fcd2519955ca461b680376e05a9e Mon Sep 17 00:00:00 2001
From: John Maguire <john.maguire@hashicorp.com>
Date: Thu, 12 Sep 2024 11:37:36 -0400
Subject: [PATCH] [NET-10952] fix cluster dns lookup family to gracefully
 handle ipv6 (#21703)

* update jwks cluster creation to gracefully handle ipv6

* update unit tests for dns lookup family

* Add changelog
---
 .changelog/21703.txt       |  3 +++
 agent/xds/clusters.go      | 22 ++++++++++++++++-
 agent/xds/clusters_test.go | 50 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 .changelog/21703.txt

diff --git a/.changelog/21703.txt b/.changelog/21703.txt
new file mode 100644
index 0000000000..41d226e489
--- /dev/null
+++ b/.changelog/21703.txt
@@ -0,0 +1,3 @@
+```release-note:bug
+jwt-provider: change dns lookup family from the default of AUTO which would prefer ipv6 to ALL if LOGICAL_DNS is used or PREFER_IPV4 if STRICT_DNS is used to gracefully handle transitions to ipv6.
+```
diff --git a/agent/xds/clusters.go b/agent/xds/clusters.go
index ba04dada22..1d69f804c7 100644
--- a/agent/xds/clusters.go
+++ b/agent/xds/clusters.go
@@ -214,9 +214,12 @@ func makeJWTProviderCluster(p *structs.JWTProviderConfigEntry) (*envoy_cluster_v
 		return nil, err
 	}
 
+	discoveryType := makeJWKSDiscoveryClusterType(p.JSONWebKeySet.Remote)
+	lookupFamily := makeJWKSClusterDNSLookupFamilyType(discoveryType)
 	cluster := &envoy_cluster_v3.Cluster{
 		Name:                 makeJWKSClusterName(p.Name),
-		ClusterDiscoveryType: makeJWKSDiscoveryClusterType(p.JSONWebKeySet.Remote),
+		ClusterDiscoveryType: discoveryType,
+		DnsLookupFamily:      lookupFamily,
 		LoadAssignment: &envoy_endpoint_v3.ClusterLoadAssignment{
 			ClusterName: makeJWKSClusterName(p.Name),
 			Endpoints: []*envoy_endpoint_v3.LocalityLbEndpoints{
@@ -278,6 +281,23 @@ func makeJWKSDiscoveryClusterType(r *structs.RemoteJWKS) *envoy_cluster_v3.Clust
 	return ct
 }
 
+func makeJWKSClusterDNSLookupFamilyType(r *envoy_cluster_v3.Cluster_Type) envoy_cluster_v3.Cluster_DnsLookupFamily {
+	// When using LOGICAL_DNS we want to use the Cluster_ALL lookup family which will fetch all the ip addresses for a given hostname and then
+	// try to connect to each one and will create the cluster based on the first one that passes.
+	// When using STRICT_DNS we want to use the CLUSTER_V4_PREFERRED lookup family which will prefer
+	// creating clusters using ipv4 addresses if those are available.
+	// Otherwise we fallback to Cluser_AUTO which will use the default behavior, and will be ignored as per the documentation.
+	// https://www.envoyproxy.io/docs/envoy/latest/api-v3/config/cluster/v3/cluster.proto#envoy-v3-api-enum-config-cluster-v3-cluster-dnslookupfamily
+	switch r.Type {
+	case envoy_cluster_v3.Cluster_LOGICAL_DNS:
+		return envoy_cluster_v3.Cluster_ALL
+	case envoy_cluster_v3.Cluster_STRICT_DNS:
+		return envoy_cluster_v3.Cluster_V4_PREFERRED
+	default:
+		return envoy_cluster_v3.Cluster_AUTO
+	}
+}
+
 func makeJWTCertValidationContext(p *structs.JWKSCluster) *envoy_tls_v3.CertificateValidationContext {
 	vc := &envoy_tls_v3.CertificateValidationContext{}
 	if p == nil || p.TLSCertificates == nil {
diff --git a/agent/xds/clusters_test.go b/agent/xds/clusters_test.go
index b5105bb698..75ef066931 100644
--- a/agent/xds/clusters_test.go
+++ b/agent/xds/clusters_test.go
@@ -380,6 +380,56 @@ func TestMakeJWKSDiscoveryClusterType(t *testing.T) {
 	}
 }
 
+func TestMakeJWKSClusterDNSLookupFamilyType(t *testing.T) {
+	tests := map[string]struct {
+		clusterType             *envoy_cluster_v3.Cluster_Type
+		expectedDNSLookupFamily envoy_cluster_v3.Cluster_DnsLookupFamily
+	}{
+		// strict dns and logical dns are the only ones that are different
+		"jwks with strict dns": {
+			clusterType: &envoy_cluster_v3.Cluster_Type{
+				Type: envoy_cluster_v3.Cluster_STRICT_DNS,
+			},
+			expectedDNSLookupFamily: envoy_cluster_v3.Cluster_V4_PREFERRED,
+		},
+		"jwks with logical dns": {
+			clusterType: &envoy_cluster_v3.Cluster_Type{
+				Type: envoy_cluster_v3.Cluster_LOGICAL_DNS,
+			},
+			expectedDNSLookupFamily: envoy_cluster_v3.Cluster_ALL,
+		},
+		// all should be auto from here down
+		"jwks with cluster EDS": {
+			clusterType: &envoy_cluster_v3.Cluster_Type{
+				Type: envoy_cluster_v3.Cluster_EDS,
+			},
+			expectedDNSLookupFamily: envoy_cluster_v3.Cluster_AUTO,
+		},
+		"jwks with static dns": {
+			clusterType: &envoy_cluster_v3.Cluster_Type{
+				Type: envoy_cluster_v3.Cluster_STATIC,
+			},
+			expectedDNSLookupFamily: envoy_cluster_v3.Cluster_AUTO,
+		},
+
+		"jwks with original dst": {
+			clusterType: &envoy_cluster_v3.Cluster_Type{
+				Type: envoy_cluster_v3.Cluster_ORIGINAL_DST,
+			},
+			expectedDNSLookupFamily: envoy_cluster_v3.Cluster_AUTO,
+		},
+	}
+
+	for name, tt := range tests {
+		tt := tt
+		t.Run(name, func(t *testing.T) {
+			actualDNSLookupFamily := makeJWKSClusterDNSLookupFamilyType(tt.clusterType)
+
+			require.Equal(t, tt.expectedDNSLookupFamily, actualDNSLookupFamily)
+		})
+	}
+}
+
 func TestParseJWTRemoteURL(t *testing.T) {
 	tests := map[string]struct {
 		uri            string