From b0a946cf0cf1adb094d7da9dbf2d054905702c73 Mon Sep 17 00:00:00 2001
From: Doug Fawley <dfawley@google.com>
Date: Thu, 7 Sep 2023 16:54:08 -0700
Subject: [PATCH] xds: fix hash policy header to skip bin headers and use extra
 metadata (#6609)

---
 xds/internal/resolver/serviceconfig.go      | 22 +++++++++++----
 xds/internal/resolver/serviceconfig_test.go | 30 +++++++++++++++++++++
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/xds/internal/resolver/serviceconfig.go b/xds/internal/resolver/serviceconfig.go
index 02470ddca5e4..06f6a47519c4 100644
--- a/xds/internal/resolver/serviceconfig.go
+++ b/xds/internal/resolver/serviceconfig.go
@@ -31,6 +31,7 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/internal/envconfig"
 	"google.golang.org/grpc/internal/grpcrand"
+	"google.golang.org/grpc/internal/grpcutil"
 	iresolver "google.golang.org/grpc/internal/resolver"
 	"google.golang.org/grpc/internal/serviceconfig"
 	"google.golang.org/grpc/internal/wrr"
@@ -229,19 +230,30 @@ func retryConfigToPolicy(config *xdsresource.RetryConfig) *serviceconfig.RetryPo
 func (cs *configSelector) generateHash(rpcInfo iresolver.RPCInfo, hashPolicies []*xdsresource.HashPolicy) uint64 {
 	var hash uint64
 	var generatedHash bool
+	var md, emd metadata.MD
+	var mdRead bool
 	for _, policy := range hashPolicies {
 		var policyHash uint64
 		var generatedPolicyHash bool
 		switch policy.HashPolicyType {
 		case xdsresource.HashPolicyTypeHeader:
-			md, ok := metadata.FromOutgoingContext(rpcInfo.Context)
-			if !ok {
+			if strings.HasSuffix(policy.HeaderName, "-bin") {
 				continue
 			}
-			values := md.Get(policy.HeaderName)
-			// If the header isn't present, no-op.
+			if !mdRead {
+				md, _ = metadata.FromOutgoingContext(rpcInfo.Context)
+				emd, _ = grpcutil.ExtraMetadata(rpcInfo.Context)
+				mdRead = true
+			}
+			values := emd.Get(policy.HeaderName)
 			if len(values) == 0 {
-				continue
+				// Extra metadata (e.g. the "content-type" header) takes
+				// precedence over the user's metadata.
+				values = md.Get(policy.HeaderName)
+				if len(values) == 0 {
+					// If the header isn't present at all, this policy is a no-op.
+					continue
+				}
 			}
 			joinedValues := strings.Join(values, ",")
 			if policy.Regex != nil {
diff --git a/xds/internal/resolver/serviceconfig_test.go b/xds/internal/resolver/serviceconfig_test.go
index 786b003c0154..a75bdb73ebe0 100644
--- a/xds/internal/resolver/serviceconfig_test.go
+++ b/xds/internal/resolver/serviceconfig_test.go
@@ -25,6 +25,7 @@ import (
 
 	xxhash "github.com/cespare/xxhash/v2"
 	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/internal/grpcutil"
 	iresolver "google.golang.org/grpc/internal/resolver"
 	"google.golang.org/grpc/metadata"
 	_ "google.golang.org/grpc/xds/internal/balancer/cdsbalancer" // To parse LB config
@@ -106,6 +107,35 @@ func (s) TestGenerateRequestHash(t *testing.T) {
 				Method:  "/some-method",
 			},
 		},
+		// Tests that bin headers are skipped.
+		{
+			name: "skip-bin",
+			hashPolicies: []*xdsresource.HashPolicy{{
+				HashPolicyType: xdsresource.HashPolicyTypeHeader,
+				HeaderName:     "something-bin",
+			}, {
+				HashPolicyType: xdsresource.HashPolicyTypeChannelID,
+			}},
+			requestHashWant: channelID,
+			rpcInfo: iresolver.RPCInfo{
+				Context: metadata.NewOutgoingContext(context.Background(), metadata.Pairs("something-bin", "xyz")),
+			},
+		},
+		// Tests that extra metadata takes precedence over the user's metadata.
+		{
+			name: "extra-metadata",
+			hashPolicies: []*xdsresource.HashPolicy{{
+				HashPolicyType: xdsresource.HashPolicyTypeHeader,
+				HeaderName:     "content-type",
+			}},
+			requestHashWant: xxhash.Sum64String("grpc value"),
+			rpcInfo: iresolver.RPCInfo{
+				Context: grpcutil.WithExtraMetadata(
+					metadata.NewOutgoingContext(context.Background(), metadata.Pairs("content-type", "user value")),
+					metadata.Pairs("content-type", "grpc value"),
+				),
+			},
+		},
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {