From ebd098392a8ba5a611762fbfd06de8fa06bca8ae Mon Sep 17 00:00:00 2001
From: Menghan Li <menghanl@google.com>
Date: Fri, 8 Apr 2022 15:41:56 -0700
Subject: [PATCH] xds/eds: reject EDS resources with multiple instances of the
 same locality in the same priority (#5303)

---
 .../xdsclient/xdsresource/unmarshal_eds.go    | 13 ++++-
 .../xdsresource/unmarshal_eds_test.go         | 53 +++++++++++++++++++
 2 files changed, 64 insertions(+), 2 deletions(-)

diff --git a/xds/internal/xdsclient/xdsresource/unmarshal_eds.go b/xds/internal/xdsclient/xdsresource/unmarshal_eds.go
index 147870cdf6b..9eb7117d9a2 100644
--- a/xds/internal/xdsclient/xdsresource/unmarshal_eds.go
+++ b/xds/internal/xdsclient/xdsresource/unmarshal_eds.go
@@ -107,7 +107,7 @@ func parseEDSRespProto(m *v3endpointpb.ClusterLoadAssignment) (EndpointsUpdate,
 	for _, dropPolicy := range m.GetPolicy().GetDropOverloads() {
 		ret.Drops = append(ret.Drops, parseDropPolicy(dropPolicy))
 	}
-	priorities := make(map[uint32]struct{})
+	priorities := make(map[uint32]map[string]bool)
 	for _, locality := range m.Endpoints {
 		l := locality.GetLocality()
 		if l == nil {
@@ -119,7 +119,16 @@ func parseEDSRespProto(m *v3endpointpb.ClusterLoadAssignment) (EndpointsUpdate,
 			SubZone: l.SubZone,
 		}
 		priority := locality.GetPriority()
-		priorities[priority] = struct{}{}
+		localitiesWithPriority := priorities[priority]
+		if localitiesWithPriority == nil {
+			localitiesWithPriority = make(map[string]bool)
+			priorities[priority] = localitiesWithPriority
+		}
+		lidStr, _ := lid.ToString()
+		if localitiesWithPriority[lidStr] {
+			return EndpointsUpdate{}, fmt.Errorf("duplicate locality %s with the same priority %v", lidStr, priority)
+		}
+		localitiesWithPriority[lidStr] = true
 		ret.Localities = append(ret.Localities, Locality{
 			ID:        lid,
 			Endpoints: parseEndpoints(locality.GetLbEndpoints()),
diff --git a/xds/internal/xdsclient/xdsresource/unmarshal_eds_test.go b/xds/internal/xdsclient/xdsresource/unmarshal_eds_test.go
index 5c6118d4e72..db9d4f52896 100644
--- a/xds/internal/xdsclient/xdsresource/unmarshal_eds_test.go
+++ b/xds/internal/xdsclient/xdsresource/unmarshal_eds_test.go
@@ -64,6 +64,17 @@ func (s) TestEDSParseRespProto(t *testing.T) {
 			want:    EndpointsUpdate{},
 			wantErr: true,
 		},
+		{
+			name: "duplicate-locality-in-the-same-priority",
+			m: func() *v3endpointpb.ClusterLoadAssignment {
+				clab0 := newClaBuilder("test", nil)
+				clab0.addLocality("locality-0", 1, 0, []string{"addr1:314"}, nil)
+				clab0.addLocality("locality-0", 1, 0, []string{"addr1:314"}, nil) // Duplicate locality with the same priority.
+				return clab0.Build()
+			}(),
+			want:    EndpointsUpdate{},
+			wantErr: true,
+		},
 		{
 			name: "good",
 			m: func() *v3endpointpb.ClusterLoadAssignment {
@@ -105,6 +116,48 @@ func (s) TestEDSParseRespProto(t *testing.T) {
 			},
 			wantErr: false,
 		},
+		{
+			name: "good duplicate locality with different priority",
+			m: func() *v3endpointpb.ClusterLoadAssignment {
+				clab0 := newClaBuilder("test", nil)
+				clab0.addLocality("locality-1", 1, 1, []string{"addr1:314"}, &addLocalityOptions{
+					Health: []v3corepb.HealthStatus{v3corepb.HealthStatus_UNHEALTHY},
+					Weight: []uint32{271},
+				})
+				// Same locality name, but with different priority.
+				clab0.addLocality("locality-1", 1, 0, []string{"addr2:159"}, &addLocalityOptions{
+					Health: []v3corepb.HealthStatus{v3corepb.HealthStatus_DRAINING},
+					Weight: []uint32{828},
+				})
+				return clab0.Build()
+			}(),
+			want: EndpointsUpdate{
+				Drops: nil,
+				Localities: []Locality{
+					{
+						Endpoints: []Endpoint{{
+							Address:      "addr1:314",
+							HealthStatus: EndpointHealthStatusUnhealthy,
+							Weight:       271,
+						}},
+						ID:       internal.LocalityID{SubZone: "locality-1"},
+						Priority: 1,
+						Weight:   1,
+					},
+					{
+						Endpoints: []Endpoint{{
+							Address:      "addr2:159",
+							HealthStatus: EndpointHealthStatusDraining,
+							Weight:       828,
+						}},
+						ID:       internal.LocalityID{SubZone: "locality-1"},
+						Priority: 0,
+						Weight:   1,
+					},
+				},
+			},
+			wantErr: false,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {