Fix gRPC Web spec violation related to header/trailer names (improbable-eng#271)

Ensures in-body trailer keys are lower case as per the spec.

Author: ktr0731

go/grpcweb/grpc_web_response.go

@@ -114,9 +114,9 @@ func (w *grpcWebResponse) copyTrailersToPayload() {
 	w.wrapped.(http.Flusher).Flush()
 }
 
-func (w *grpcWebResponse) extractTrailerHeaders() http.Header {
+func (w *grpcWebResponse) extractTrailerHeaders() trailer {
 	flushedHeaders := w.wrapped.Header()
-	trailerHeaders := make(http.Header)
+	trailerHeaders := trailer{make(http.Header)}
 	for k, vv := range w.headers {
 		// Skip the pre-annoucement of Trailer headers. Don't add them to the response headers.
 		if strings.ToLower(k) == "trailer" {

go/grpcweb/trailer.go

@@ -0,0 +1,32 @@
+//Copyright 2018 Improbable. All Rights Reserved.
+// See LICENSE for licensing terms.
+
+package grpcweb
+
+import (
+	"net/http"
+	"strings"
+)
+
+// gRPC-Web spec says that must use lower-case header/trailer names.
+// See "HTTP wire protocols" section in
+// https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-WEB.md#protocol-differences-vs-grpc-over-http2
+type trailer struct {
+	http.Header
+}
+
+func (t trailer) Add(key, value string) {
+	key = strings.ToLower(key)
+	t.Header[key] = append(t.Header[key], value)
+}
+
+func (t trailer) Get(key string) string {
+	if t.Header == nil {
+		return ""
+	}
+	v := t.Header[key]
+	if len(v) == 0 {
+		return ""
+	}
+	return v[0]
+}

go/grpcweb/trailer_test.go

@@ -0,0 +1,11 @@
+package grpcweb
+
+import "net/http"
+
+type Trailer struct {
+	trailer
+}
+
+func HTTPTrailerToGrpcWebTrailer(httpTrailer http.Header) Trailer {
+	return Trailer{trailer{httpTrailer}}
+}

go/grpcweb/wrapper.go

@@ -4,10 +4,9 @@
 package grpcweb
 
 import (
+	"context"
 	"net/http"
 	"net/url"
-
-	"context"
 	"strings"
 	"time"
 

go/grpcweb/wrapper_test.go

@@ -22,6 +22,10 @@ import (
 	"time"
 
 	"github.com/golang/protobuf/proto"
+	google_protobuf "github.com/golang/protobuf/ptypes/empty"
+	"github.com/improbable-eng/grpc-web/go/grpcweb"
+	testproto "github.com/improbable-eng/grpc-web/test/go/_proto/improbable/grpcweb/test"
+	"github.com/mwitkow/go-conntrack/connhelpers"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
@@ -32,11 +36,6 @@ import (
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/grpclog"
 	"google.golang.org/grpc/metadata"
-
-	google_protobuf "github.com/golang/protobuf/ptypes/empty"
-	"github.com/improbable-eng/grpc-web/go/grpcweb"
-	testproto "github.com/improbable-eng/grpc-web/test/go/_proto/improbable/grpcweb/test"
-	"github.com/mwitkow/go-conntrack/connhelpers"
 )
 
 var (
@@ -112,7 +111,7 @@ func (s *GrpcWebWrapperTestSuite) makeRequest(verb string, method string, header
 	return resp, err
 }
 
-func (s *GrpcWebWrapperTestSuite) makeGrpcRequest(method string, reqHeaders http.Header, requestMessages [][]byte) (headers http.Header, trailers http.Header, responseMessages [][]byte, err error) {
+func (s *GrpcWebWrapperTestSuite) makeGrpcRequest(method string, reqHeaders http.Header, requestMessages [][]byte) (headers http.Header, trailers grpcweb.Trailer, responseMessages [][]byte, err error) {
 	writer := new(bytes.Buffer)
 	for _, msgBytes := range requestMessages {
 		grpcPreamble := []byte{0, 0, 0, 0, 0}
@@ -122,33 +121,32 @@ func (s *GrpcWebWrapperTestSuite) makeGrpcRequest(method string, reqHeaders http
 	}
 	resp, err := s.makeRequest("POST", method, reqHeaders, writer)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, grpcweb.Trailer{}, nil, err
 	}
 	defer resp.Body.Close()
 	contents, err := ioutil.ReadAll(resp.Body)
 	if err != nil {
-		return nil, nil, nil, err
+		return nil, grpcweb.Trailer{}, nil, err
 	}
 	reader := bytes.NewReader(contents)
-	trailers = make(http.Header)
 	for {
 		grpcPreamble := []byte{0, 0, 0, 0, 0}
 		readCount, err := reader.Read(grpcPreamble)
 		if err == io.EOF {
 			break
 		}
 		if readCount != 5 || err != nil {
-			return nil, nil, nil, fmt.Errorf("Unexpected end of body in preamble: %v", err)
+			return nil, grpcweb.Trailer{}, nil, fmt.Errorf("Unexpected end of body in preamble: %v", err)
 		}
 		payloadLength := binary.BigEndian.Uint32(grpcPreamble[1:])
 		payloadBytes := make([]byte, payloadLength)
 
 		readCount, err = reader.Read(payloadBytes)
 		if uint32(readCount) != payloadLength || err != nil {
-			return nil, nil, nil, fmt.Errorf("Unexpected end of msg: %v", err)
+			return nil, grpcweb.Trailer{}, nil, fmt.Errorf("Unexpected end of msg: %v", err)
 		}
 		if grpcPreamble[0]&(1<<7) == (1 << 7) { // MSB signifies the trailer parser
-			trailers = readHeadersFromBytes(payloadBytes)
+			trailers = readTrailersFromBytes(s.T(), payloadBytes)
 		} else {
 			responseMessages = append(responseMessages, payloadBytes)
 		}
@@ -166,7 +164,7 @@ func (s *GrpcWebWrapperTestSuite) TestPingEmpty() {
 	assert.Equal(s.T(), 1, len(responses), "PingEmpty is an unary response")
 	s.assertTrailerGrpcCode(trailers, codes.OK, "")
 	s.assertHeadersContainMetadata(headers, expectedHeaders)
-	s.assertHeadersContainMetadata(trailers, expectedTrailers)
+	s.assertTrailersContainMetadata(trailers, expectedTrailers)
 }
 
 func (s *GrpcWebWrapperTestSuite) TestPing() {
@@ -179,7 +177,7 @@ func (s *GrpcWebWrapperTestSuite) TestPing() {
 	assert.Equal(s.T(), 1, len(responses), "PingEmpty is an unary response")
 	s.assertTrailerGrpcCode(trailers, codes.OK, "")
 	s.assertHeadersContainMetadata(headers, expectedHeaders)
-	s.assertHeadersContainMetadata(trailers, expectedTrailers)
+	s.assertTrailersContainMetadata(trailers, expectedTrailers)
 	s.assertHeadersContainCorsExpectedHeaders(headers, expectedHeaders)
 }
 
@@ -195,7 +193,7 @@ func (s *GrpcWebWrapperTestSuite) TestPingError_WithTrailersInData() {
 	assert.Equal(s.T(), 0, len(responses), "PingError is an unary response that has no payload")
 	s.assertTrailerGrpcCode(trailers, codes.Unimplemented, "Not implemented PingError")
 	s.assertHeadersContainMetadata(headers, expectedHeaders)
-	s.assertHeadersContainMetadata(trailers, expectedTrailers)
+	s.assertTrailersContainMetadata(trailers, expectedTrailers)
 	s.assertHeadersContainCorsExpectedHeaders(headers, expectedHeaders)
 }
 
@@ -209,7 +207,7 @@ func (s *GrpcWebWrapperTestSuite) TestPingError_WithTrailersInHeaders() {
 	require.NoError(s.T(), err, "No error on making request")
 
 	assert.Equal(s.T(), 0, len(responses), "PingError is an unary response that has no payload")
-	s.assertTrailerGrpcCode(headers, codes.Unimplemented, "Not implemented PingError")
+	s.assertHeadersGrpcCode(headers, codes.Unimplemented, "Not implemented PingError")
 	// s.assertHeadersContainMetadata(headers, expectedHeaders) // TODO(mwitkow): There is a bug in gRPC where headers don't get added if no payload exists.
 	s.assertHeadersContainMetadata(headers, expectedTrailers)
 	s.assertHeadersContainCorsExpectedHeaders(headers, expectedTrailers)
@@ -224,7 +222,7 @@ func (s *GrpcWebWrapperTestSuite) TestPingList() {
 	assert.Equal(s.T(), expectedListResponses, len(responses), "the number of expected proto fields shouold match")
 	s.assertTrailerGrpcCode(trailers, codes.OK, "")
 	s.assertHeadersContainMetadata(headers, expectedHeaders)
-	s.assertHeadersContainMetadata(trailers, expectedTrailers)
+	s.assertTrailersContainMetadata(trailers, expectedTrailers)
 	s.assertHeadersContainCorsExpectedHeaders(headers, expectedHeaders)
 }
 
@@ -312,6 +310,14 @@ func (s *GrpcWebWrapperTestSuite) assertHeadersContainMetadata(headers http.Head
 	}
 }
 
+func (s *GrpcWebWrapperTestSuite) assertTrailersContainMetadata(trailers grpcweb.Trailer, meta metadata.MD) {
+	for k, v := range meta {
+		for _, vv := range v {
+			assert.Equal(s.T(), trailers.Get(k), vv, "Expected there to be %v=%v", k, vv)
+		}
+	}
+}
+
 func (s *GrpcWebWrapperTestSuite) assertHeadersContainCorsExpectedHeaders(headers http.Header, meta metadata.MD) {
 	value := headers.Get("Access-Control-Expose-Headers")
 	assert.NotEmpty(s.T(), value, "cors: access control expose headers should not be empty")
@@ -323,7 +329,15 @@ func (s *GrpcWebWrapperTestSuite) assertHeadersContainCorsExpectedHeaders(header
 	}
 }
 
-func (s *GrpcWebWrapperTestSuite) assertTrailerGrpcCode(trailers http.Header, code codes.Code, desc string) {
+func (s *GrpcWebWrapperTestSuite) assertHeadersGrpcCode(headers http.Header, code codes.Code, desc string) {
+	require.NotEmpty(s.T(), headers.Get("grpc-status"), "grpc-status must not be empty in trailers")
+	statusCode, err := strconv.Atoi(headers.Get("grpc-status"))
+	require.NoError(s.T(), err, "no error parsing grpc-status")
+	assert.EqualValues(s.T(), code, statusCode, "grpc-status must match expected code")
+	assert.EqualValues(s.T(), desc, headers.Get("grpc-message"), "grpc-message is expected to match")
+}
+
+func (s *GrpcWebWrapperTestSuite) assertTrailerGrpcCode(trailers grpcweb.Trailer, code codes.Code, desc string) {
 	require.NotEmpty(s.T(), trailers.Get("grpc-status"), "grpc-status must not be empty in trailers")
 	statusCode, err := strconv.Atoi(trailers.Get("grpc-status"))
 	require.NoError(s.T(), err, "no error parsing grpc-status")
@@ -340,14 +354,40 @@ func serializeProtoMessages(messages []proto.Message) [][]byte {
 	return out
 }
 
-func readHeadersFromBytes(dataBytes []byte) http.Header {
+func readTrailersFromBytes(t *testing.T, dataBytes []byte) grpcweb.Trailer {
 	bufferReader := bytes.NewBuffer(dataBytes)
 	tp := textproto.NewReader(bufio.NewReader(bufferReader))
+
+	// First, read bytes as MIME headers.
+	// However, it normalizes header names by textproto.CanonicalMIMEHeaderKey.
+	// In the next step, replace header names by raw one.
 	mimeHeader, err := tp.ReadMIMEHeader()
 	if err == nil {
-		return make(http.Header)
+		return grpcweb.Trailer{}
+	}
+
+	trailers := make(http.Header)
+	bufferReader = bytes.NewBuffer(dataBytes)
+	tp = textproto.NewReader(bufio.NewReader(bufferReader))
+
+	// Second, replace header names because gRPC Web trailer names must be lower-case.
+	for {
+		line, err := tp.ReadLine()
+		if err == io.EOF {
+			break
+		}
+		require.NoError(t, err, "failed to read header line")
+
+		i := strings.IndexByte(line, ':')
+		if i == -1 {
+			require.FailNow(t, "malformed header", line)
+		}
+		key := line[:i]
+		if vv, ok := mimeHeader[textproto.CanonicalMIMEHeaderKey(key)]; ok {
+			trailers[key] = vv
+		}
 	}
-	return http.Header(mimeHeader)
+	return grpcweb.HTTPTrailerToGrpcWebTrailer(trailers)
 }
 
 func headerWithFlag(flags ...string) http.Header {