From c672451950653990bd607c8ba08733d6f36d85fc Mon Sep 17 00:00:00 2001
From: Ernest Nguyen <58267404+erni27@users.noreply.github.com>
Date: Mon, 10 Oct 2022 21:48:01 +0200
Subject: [PATCH] xds/xdsclient: add sum of EDS locality weights check (#5703)

---
 .../xdsclient/xdsresource/unmarshal_eds.go       | 16 +++++++++++-----
 .../xdsclient/xdsresource/unmarshal_eds_test.go  | 12 ++++++++++++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/xds/internal/xdsclient/xdsresource/unmarshal_eds.go b/xds/internal/xdsclient/xdsresource/unmarshal_eds.go
index 7d4b89dc9dc..15b0d88f9f1 100644
--- a/xds/internal/xdsclient/xdsresource/unmarshal_eds.go
+++ b/xds/internal/xdsclient/xdsresource/unmarshal_eds.go
@@ -19,6 +19,7 @@ package xdsresource
 
 import (
 	"fmt"
+	"math"
 	"net"
 	"strconv"
 
@@ -118,6 +119,7 @@ func parseEDSRespProto(m *v3endpointpb.ClusterLoadAssignment, logger *grpclog.Pr
 		ret.Drops = append(ret.Drops, parseDropPolicy(dropPolicy))
 	}
 	priorities := make(map[uint32]map[string]bool)
+	sumOfWeights := make(map[uint32]uint64)
 	for _, locality := range m.Endpoints {
 		l := locality.GetLocality()
 		if l == nil {
@@ -128,17 +130,21 @@ func parseEDSRespProto(m *v3endpointpb.ClusterLoadAssignment, logger *grpclog.Pr
 			logger.Warningf("Ignoring locality %s with weight 0", pretty.ToJSON(l))
 			continue
 		}
-		lid := internal.LocalityID{
-			Region:  l.Region,
-			Zone:    l.Zone,
-			SubZone: l.SubZone,
-		}
 		priority := locality.GetPriority()
+		sumOfWeights[priority] += uint64(weight)
+		if sumOfWeights[priority] > math.MaxUint32 {
+			return EndpointsUpdate{}, fmt.Errorf("sum of weights of localities at the same priority %d exceeded maximal value", priority)
+		}
 		localitiesWithPriority := priorities[priority]
 		if localitiesWithPriority == nil {
 			localitiesWithPriority = make(map[string]bool)
 			priorities[priority] = localitiesWithPriority
 		}
+		lid := internal.LocalityID{
+			Region:  l.Region,
+			Zone:    l.Zone,
+			SubZone: l.SubZone,
+		}
 		lidStr, _ := lid.ToString()
 		if localitiesWithPriority[lidStr] {
 			return EndpointsUpdate{}, fmt.Errorf("duplicate locality %s with the same priority %v", lidStr, priority)
diff --git a/xds/internal/xdsclient/xdsresource/unmarshal_eds_test.go b/xds/internal/xdsclient/xdsresource/unmarshal_eds_test.go
index 28ceb11c6e1..02718e09dda 100644
--- a/xds/internal/xdsclient/xdsresource/unmarshal_eds_test.go
+++ b/xds/internal/xdsclient/xdsresource/unmarshal_eds_test.go
@@ -99,6 +99,18 @@ func (s) TestEDSParseRespProto(t *testing.T) {
 			}(),
 			want: EndpointsUpdate{},
 		},
+		{
+			name: "max sum of weights at the same priority exceeded",
+			m: func() *v3endpointpb.ClusterLoadAssignment {
+				clab0 := newClaBuilder("test", nil)
+				clab0.addLocality("locality-1", 1, 0, []string{"addr1:314"}, nil)
+				clab0.addLocality("locality-2", 4294967295, 1, []string{"addr2:159"}, nil)
+				clab0.addLocality("locality-3", 1, 1, []string{"addr2:88"}, nil)
+				return clab0.Build()
+			}(),
+			want:    EndpointsUpdate{},
+			wantErr: true,
+		},
 		{
 			name: "good",
 			m: func() *v3endpointpb.ClusterLoadAssignment {