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
googleapis · Apr 15, 2024 · 48f9b74 · 48f9b74
1 parent a34fc37
commit 48f9b74
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 11 deletions.
diff --git a/schema/googleapis b/schema/googleapis
diff --git 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
@@ -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
 }