From 48f9b742970e2199431066807d3f7a5cd2f31728 Mon Sep 17 00:00:00 2001
From: Anthonios Partheniou <partheniou@google.com>
Date: Mon, 15 Apr 2024 16:40:48 -0400
Subject: [PATCH] feat: Add ability to echo request headers in trailing
 metadata (#1501)

* Echo headers and trailers when the stream ends

* feat: Add ability to echo request headers in response

* lint

* drop usage of slices package

* update comment

* address review feedback

* address review feedback

* address review feedback

* address review feedback

* revert server/genproto
---
 schema/googleapis                    |  2 +-
 server/services/echo_service.go      | 18 ++++++++++--------
 server/services/echo_service_test.go |  9 +++++++--
 3 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/schema/googleapis b/schema/googleapis
index d81d0b9e6..5db2bbe91 160000
--- a/schema/googleapis
+++ b/schema/googleapis
@@ -1 +1 @@
-Subproject commit d81d0b9e6993d6ab425dff4d7c3d05fb2e59fa57
+Subproject commit 5db2bbe911db5a01c6a7564f78e84f8116de91d3
diff --git a/server/services/echo_service.go b/server/services/echo_service.go
index c5cec526d..0e80b469f 100644
--- a/server/services/echo_service.go
+++ b/server/services/echo_service.go
@@ -313,10 +313,11 @@ func echoTrailers(ctx context.Context) {
 		return
 	}
 
-	values := md.Get("showcase-trailer")
-	for _, value := range values {
-		trailer := metadata.Pairs("showcase-trailer", value)
-		grpc.SetTrailer(ctx, trailer)
+	for k, v := range md {
+		for _, value := range v {
+			trailer := metadata.Pairs(k, value)
+			grpc.SetTrailer(ctx, trailer)
+		}
 	}
 }
 
@@ -326,9 +327,10 @@ func echoStreamingTrailers(stream grpc.ServerStream) {
 		return
 	}
 
-	values := md.Get("showcase-trailer")
-	for _, value := range values {
-		trailer := metadata.Pairs("showcase-trailer", value)
-		stream.SetTrailer(trailer)
+	for k, v := range md {
+		for _, value := range v {
+			trailer := metadata.Pairs(k, value)
+			stream.SetTrailer(trailer)
+		}
 	}
 }
diff --git a/server/services/echo_service_test.go b/server/services/echo_service_test.go
index 9aa7fcd10..b38bb3c11 100644
--- a/server/services/echo_service_test.go
+++ b/server/services/echo_service_test.go
@@ -19,6 +19,7 @@ import (
 	"errors"
 	"io"
 	"reflect"
+	"sort"
 	"strings"
 	"testing"
 	"time"
@@ -116,13 +117,16 @@ func (m *mockUnaryStream) Send(resp *pb.EchoResponse) error { return nil }
 func (m *mockUnaryStream) Context() context.Context         { return nil }
 func (m *mockUnaryStream) SetTrailer(md metadata.MD) {
 	m.trail = append(m.trail, md.Get("showcase-trailer")...)
+	m.trail = append(m.trail, md.Get("x-goog-api-version")...)
+	// Sort the trailer values as having a guaranteed order will help with array comparison
+	sort.Strings(m.trail)
 }
 func (m *mockUnaryStream) SetHeader(md metadata.MD) error {
 	m.head = append(m.head, md.Get("x-goog-request-params")...)
 	return nil
 }
 func (m *mockUnaryStream) verify(expectHeadersAndTrailers bool) {
-	if expectHeadersAndTrailers && (!reflect.DeepEqual([]string{"show", "case"}, m.trail) || !reflect.DeepEqual([]string{"showcaseHeader", "anotherHeader"}, m.head)) {
+	if expectHeadersAndTrailers && (!reflect.DeepEqual([]string{"apiVersion", "case", "show"}, m.trail) || !reflect.DeepEqual([]string{"showcaseHeader", "anotherHeader"}, m.head)) {
 		m.t.Errorf("Unary stream did not get all expected headers and trailers.\nGot these headers: %+v\nGot these trailers: %+v", m.head, m.trail)
 	}
 }
@@ -153,6 +157,7 @@ func (m *mockExpandStream) SetTrailer(md metadata.MD) {
 
 func (m *mockExpandStream) SetHeader(md metadata.MD) error {
 	m.head = append(m.head, md.Get("x-goog-request-params")...)
+	m.head = append(m.head, md.Get("x-goog-api-version")...)
 	return nil
 }
 
@@ -898,6 +903,6 @@ func TestBlockError(t *testing.T) {
 
 func appendTestOutgoingMetadata(ctx context.Context, stream grpc.ServerTransportStream) context.Context {
 	ctx = grpc.NewContextWithServerTransportStream(ctx, stream)
-	ctx = metadata.NewIncomingContext(ctx, metadata.Pairs("showcase-trailer", "show", "showcase-trailer", "case", "trailer", "trail", "x-goog-request-params", "showcaseHeader", "x-goog-request-params", "anotherHeader", "header", "head"))
+	ctx = metadata.NewIncomingContext(ctx, metadata.Pairs("showcase-trailer", "show", "showcase-trailer", "case", "trailer", "trail", "x-goog-request-params", "showcaseHeader", "x-goog-request-params", "anotherHeader", "header", "head", "x-goog-api-version", "apiVersion"))
 	return ctx
 }