From e7603197f028d034c70abba5380b15de35de7bb2 Mon Sep 17 00:00:00 2001
From: Brenna N Epp <brennae@google.com>
Date: Mon, 10 Jun 2024 13:50:58 -0700
Subject: [PATCH] chore(gensupport): merge x-goog-api-client vals into a single
 header (#2612)

---
 internal/gensupport/send.go      | 24 +++++++++++++--
 internal/gensupport/send_test.go | 53 ++++++++++++++++++++++++++++----
 2 files changed, 68 insertions(+), 9 deletions(-)

diff --git a/internal/gensupport/send.go b/internal/gensupport/send.go
index f39dd00d99f..f6716134ebf 100644
--- a/internal/gensupport/send.go
+++ b/internal/gensupport/send.go
@@ -48,8 +48,24 @@ func SendRequest(ctx context.Context, client *http.Client, req *http.Request) (*
 	if ctx != nil {
 		headers := callctx.HeadersFromContext(ctx)
 		for k, vals := range headers {
-			for _, v := range vals {
-				req.Header.Add(k, v)
+			if k == "x-goog-api-client" {
+				// Merge all values into a single "x-goog-api-client" header.
+				var mergedVal strings.Builder
+				baseXGoogHeader := req.Header.Get("X-Goog-Api-Client")
+				if baseXGoogHeader != "" {
+					mergedVal.WriteString(baseXGoogHeader)
+					mergedVal.WriteRune(' ')
+				}
+				for _, v := range vals {
+					mergedVal.WriteString(v)
+					mergedVal.WriteRune(' ')
+				}
+				// Remove the last space and replace the header on the request.
+				req.Header.Set(k, mergedVal.String()[:mergedVal.Len()-1])
+			} else {
+				for _, v := range vals {
+					req.Header.Add(k, v)
+				}
 			}
 		}
 	}
@@ -118,7 +134,9 @@ func sendAndRetry(ctx context.Context, client *http.Client, req *http.Request, r
 	var err error
 	attempts := 1
 	invocationID := uuid.New().String()
-	baseXGoogHeader := req.Header.Get("X-Goog-Api-Client")
+
+	xGoogHeaderVals := req.Header.Values("X-Goog-Api-Client")
+	baseXGoogHeader := strings.Join(xGoogHeaderVals, " ")
 
 	// Loop to retry the request, up to the context deadline.
 	var pause time.Duration
diff --git a/internal/gensupport/send_test.go b/internal/gensupport/send_test.go
index 59534408393..502ab352ebd 100644
--- a/internal/gensupport/send_test.go
+++ b/internal/gensupport/send_test.go
@@ -9,6 +9,7 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"regexp"
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
@@ -36,13 +37,24 @@ func TestSendRequestWithRetry(t *testing.T) {
 }
 
 type headerRoundTripper struct {
-	wantHeader http.Header
+	wantHeader        http.Header
+	wantXgoogAPIRegex string // test x-goog-api-client separately
 }
 
 func (rt *headerRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
+	// Test x-goog-api-client with a regex, since invocation ids are randomly generated
+	match, err := regexp.MatchString(rt.wantXgoogAPIRegex, r.Header.Get("X-Goog-Api-Client"))
+	if err != nil {
+		return nil, fmt.Errorf("compiling regexp: %v", err)
+	}
+	if !match {
+		return nil, fmt.Errorf("X-Goog-Api-Client header has wrong format\ngot %v\nwant regex matching %v", r.Header.Get("X-Goog-Api-Client"), rt.wantXgoogAPIRegex)
+	}
+
 	// Ignore x-goog headers sent by SendRequestWithRetry
-	r.Header.Del("X-Goog-Api-Client")
 	r.Header.Del("X-Goog-Gcs-Idempotency-Token")
+	r.Header.Del("X-Goog-Api-Client") // this was tested above already
+
 	if diff := cmp.Diff(r.Header, rt.wantHeader); diff != "" {
 		return nil, fmt.Errorf("headers don't match: %v", diff)
 	}
@@ -53,19 +65,48 @@ func (rt *headerRoundTripper) RoundTrip(r *http.Request) (*http.Response, error)
 func TestSendRequestHeader(t *testing.T) {
 	ctx := context.Background()
 	ctx = callctx.SetHeaders(ctx, "foo", "100", "bar", "200")
-	client := http.Client{
-		Transport: &headerRoundTripper{
-			wantHeader: map[string][]string{"Foo": {"100"}, "Bar": {"200"}},
-		},
+	transport := &headerRoundTripper{
+		wantHeader: map[string][]string{"Foo": {"100"}, "Bar": {"200"}},
 	}
+	client := http.Client{Transport: transport}
+
 	req, _ := http.NewRequest("GET", "url", nil)
 	if _, err := SendRequest(ctx, &client, req); err != nil {
 		t.Errorf("SendRequest: %v", err)
 	}
+
+	// SendRequestWithRetry adds retry and idempotency headers
+	transport.wantXgoogAPIRegex = "^gccl-invocation-id/.{36} gccl-attempt-count/[0-9]+ $"
 	req2, _ := http.NewRequest("GET", "url", nil)
 	if _, err := SendRequestWithRetry(ctx, &client, req2, nil); err != nil {
+		t.Errorf("SendRequestWithRetry: %v", err)
+	}
+}
+
+// Ensure that x-goog-api-client headers set via the context are merged properly
+// and passed through to the request as expected.
+func TestSendRequestXgoogHeaderxxx(t *testing.T) {
+	ctx := context.Background()
+	ctx = callctx.SetHeaders(ctx, "x-goog-api-client", "val/1", "bar", "200", "x-goog-api-client", "val/2")
+	ctx = callctx.SetHeaders(ctx, "x-goog-api-client", "val/11 val/22")
+
+	transport := &headerRoundTripper{
+		wantHeader:        map[string][]string{"Bar": {"200"}},
+		wantXgoogAPIRegex: "^val/1 val/2 val/11 val/22$",
+	}
+	client := http.Client{Transport: transport}
+
+	req, _ := http.NewRequest("GET", "url", nil)
+	if _, err := SendRequest(ctx, &client, req); err != nil {
 		t.Errorf("SendRequest: %v", err)
 	}
+
+	// SendRequestWithRetry adds retry and idempotency headers
+	transport.wantXgoogAPIRegex = fmt.Sprintf("^gccl-invocation-id/.{36} gccl-attempt-count/[0-9]+ %s$", transport.wantXgoogAPIRegex[1:len(transport.wantXgoogAPIRegex)-1])
+	req2, _ := http.NewRequest("GET", "url", nil)
+	if _, err := SendRequestWithRetry(ctx, &client, req2, nil); err != nil {
+		t.Errorf("SendRequestWithRetry: %v", err)
+	}
 }
 
 type brokenRoundTripper struct{}