Fix propagation of headers with grpc server errors

.travis.yml

@@ -1,8 +1,7 @@
 sudo: false
 language: go
 go:
-  - 1.7
-  - 1.8
+  - "1.10"
 
 install:
   - go get google.golang.org/grpc

proxy/handler.go

@@ -116,13 +116,9 @@ func (s *handler) forwardClientToServer(src grpc.ClientStream, dst grpc.ServerSt
 	go func() {
 		f := &frame{}
 		for i := 0; ; i++ {
-			if err := src.RecvMsg(f); err != nil {
-				ret <- err // this can be io.EOF which is happy case
-				break
-			}
 			if i == 0 {
-				// This is a bit of a hack, but client to server headers are only readable after first client msg is
-				// received but must be written to server stream before the first msg is flushed.
+				// This is a bit of a hack, but client to server headers must be forwarded even in case of errors.
+				// and must be written to server stream before the first msg is flushed.
 				// This is the only place to do it nicely.
 				md, err := src.Header()
 				if err != nil {
@@ -134,6 +130,12 @@ func (s *handler) forwardClientToServer(src grpc.ClientStream, dst grpc.ServerSt
 					break
 				}
 			}
+
+			if err := src.RecvMsg(f); err != nil {
+				ret <- err // this can be io.EOF which is happy case
+				break
+			}
+
 			if err := dst.SendMsg(f); err != nil {
 				ret <- err
 				break

proxy/handler_test.go

@@ -61,6 +61,8 @@ func (s *assertingService) Ping(ctx context.Context, ping *pb.PingRequest) (*pb.
 }
 
 func (s *assertingService) PingError(ctx context.Context, ping *pb.PingRequest) (*pb.Empty, error) {
+	// Send user headers.
+	grpc.SendHeader(ctx, metadata.Pairs(serverHeaderMdKey, "I like turtles."))
 	return nil, grpc.Errorf(codes.FailedPrecondition, "Userspace error.")
 }
 
@@ -95,6 +97,12 @@ func (s *assertingService) PingStream(stream pb.TestService_PingStreamServer) er
 	return nil
 }
 
+func (s *assertingService) PingStreamError(stream pb.TestService_PingStreamErrorServer) error {
+	stream.SendHeader(metadata.Pairs(serverHeaderMdKey, "I like turtles."))
+
+	return grpc.Errorf(codes.FailedPrecondition, "Userspace error.")
+}
+
 // ProxyHappySuite tests the "happy" path of handling: that everything works in absence of connection issues.
 type ProxyHappySuite struct {
 	suite.Suite
@@ -139,6 +147,15 @@ func (s *ProxyHappySuite) TestPingCarriesServerHeadersAndTrailers() {
 	assert.Len(s.T(), trailerMd, 1, "server response trailers must contain server data")
 }
 
+func (s *ProxyHappySuite) TestPingCarriesServerHeadersWhenServerError() {
+	headerMd := metadata.Pairs("bar", "foo")
+	trailerMd := make(metadata.MD)
+	_, err := s.testClient.PingError(s.ctx(), &pb.PingRequest{Value: "foo"}, grpc.Header(&headerMd), grpc.Trailer(&trailerMd))
+	require.Error(s.T(), err, "PingError should never succeed")
+	assert.Equal(s.T(), codes.FailedPrecondition, grpc.Code(err))
+	assert.Contains(s.T(), headerMd, serverHeaderMdKey, "server response headers must contain server data")
+}
+
 func (s *ProxyHappySuite) TestPingErrorPropagatesAppError() {
 	_, err := s.testClient.PingError(s.ctx(), &pb.PingRequest{Value: "foo"})
 	require.Error(s.T(), err, "PingError should never succeed")
@@ -182,6 +199,23 @@ func (s *ProxyHappySuite) TestPingStream_FullDuplexWorks() {
 	assert.Len(s.T(), trailerMd, 1, "PingList trailer headers user contain metadata")
 }
 
+func (s *ProxyHappySuite) TestPingStream_FullDuplexWithErrorsAndMetadata() {
+	stream, err := s.testClient.PingStreamError(s.ctx())
+	require.NoError(s.T(), err, "PingStream request should be successful.")
+
+	ping := &pb.PingRequest{Value: "foo:1"}
+	require.NoError(s.T(), stream.Send(ping), "sending to PingStreamError must not fail")
+	// Check that the header arrives.
+	headerMd, err := stream.Header()
+	require.NoError(s.T(), err, "PingStreamError headers should not error.")
+	assert.Contains(s.T(), headerMd, serverHeaderMdKey, "PingStreamError response headers user should contain metadata")
+
+	_, err = stream.Recv()
+	assert.Equal(s.T(), codes.FailedPrecondition, grpc.Code(err))
+
+	require.NoError(s.T(), stream.CloseSend(), "no error on close send")
+}
+
 func (s *ProxyHappySuite) TestPingStream_StressTest() {
 	for i := 0; i < 50; i++ {
 		s.TestPingStream_FullDuplexWorks()