Verify that a trailers-only response is well-formed (#685)

Previously, connect-go would accept a trailers-only gRPC response (where
the trailers are in the HTTP headers, signaled by the presence of a
"grpc-status" key in the headers) and assume it was valid w/out further
verification.

However, it should reject trailers-only responses that _also_ include
response body data and/or HTTP trailers, after the HTTP headers as those
are malformed responses.

Author: jhump

connect_ext_test.go

@@ -2307,6 +2307,97 @@ func TestStreamUnexpectedEOF(t *testing.T) {
 	}
 }
 
+func TestTrailersOnlyErrors(t *testing.T) {
+	t.Parallel()
+
+	head := [3]byte{}
+	testcases := []struct {
+		name       string
+		handler    http.HandlerFunc
+		options    []connect.ClientOption
+		expectCode connect.Code
+		expectMsg  string
+	}{{
+		name:    "grpc_body_after_trailers-only",
+		options: []connect.ClientOption{connect.WithGRPC()},
+		handler: func(responseWriter http.ResponseWriter, _ *http.Request) {
+			header := responseWriter.Header()
+			header.Set("Content-Type", "application/grpc")
+			header.Set("Grpc-Status", "3")
+			_, err := responseWriter.Write(head[:])
+			assert.Nil(t, err)
+		},
+		expectCode: connect.CodeInternal,
+		expectMsg:  fmt.Sprintf("internal: corrupt response: %d extra bytes after trailers-only response", len(head)),
+	}, {
+		name:    "grpc-web_body_after_trailers-only",
+		options: []connect.ClientOption{connect.WithGRPCWeb()},
+		handler: func(responseWriter http.ResponseWriter, _ *http.Request) {
+			header := responseWriter.Header()
+			header.Set("Content-Type", "application/grpc-web")
+			header.Set("Grpc-Status", "3")
+			_, err := responseWriter.Write(head[:])
+			assert.Nil(t, err)
+		},
+		expectCode: connect.CodeInternal,
+		expectMsg:  fmt.Sprintf("internal: corrupt response: %d extra bytes after trailers-only response", len(head)),
+	}, {
+		name:    "grpc_trailers_after_trailers-only",
+		options: []connect.ClientOption{connect.WithGRPC()},
+		handler: func(responseWriter http.ResponseWriter, _ *http.Request) {
+			header := responseWriter.Header()
+			header.Set("Content-Type", "application/grpc")
+			header.Set("Grpc-Status", "3")
+			responseWriter.WriteHeader(http.StatusOK)
+			responseWriter.(http.Flusher).Flush() //nolint:forcetypeassert
+			header.Set(http.TrailerPrefix+"Foo", "abc")
+		},
+		expectCode: connect.CodeInternal,
+		expectMsg:  "internal: corrupt response from server: gRPC trailers-only response may not contain HTTP trailers",
+	}, {
+		name:    "grpc-web_trailers_after_trailers-only",
+		options: []connect.ClientOption{connect.WithGRPCWeb()},
+		handler: func(responseWriter http.ResponseWriter, _ *http.Request) {
+			header := responseWriter.Header()
+			header.Set("Content-Type", "application/grpc-web")
+			header.Set("Grpc-Status", "3")
+			responseWriter.WriteHeader(http.StatusOK)
+			responseWriter.(http.Flusher).Flush() //nolint:forcetypeassert
+			header.Set(http.TrailerPrefix+"Foo", "abc")
+		},
+		expectCode: connect.CodeInternal,
+		expectMsg:  "internal: corrupt response from server: gRPC trailers-only response may not contain HTTP trailers",
+	}}
+	for _, testcase := range testcases {
+		testcase := testcase
+		t.Run(testcase.name, func(t *testing.T) {
+			t.Parallel()
+			mux := http.NewServeMux()
+			mux.HandleFunc("/", func(responseWriter http.ResponseWriter, request *http.Request) {
+				_, _ = io.Copy(io.Discard, request.Body)
+				testcase.handler(responseWriter, request)
+			})
+			server := memhttptest.NewServer(t, mux)
+			client := pingv1connect.NewPingServiceClient(
+				server.Client(),
+				server.URL(),
+				testcase.options...,
+			)
+			const upTo = 2
+			request := connect.NewRequest(&pingv1.CountUpRequest{Number: upTo})
+			request.Header().Set("Test-Case", t.Name())
+			stream, err := client.CountUp(context.Background(), request)
+			assert.Nil(t, err)
+			for i := 0; stream.Receive() && i < upTo; i++ {
+				assert.Equal(t, stream.Msg().GetNumber(), 42)
+			}
+			assert.NotNil(t, stream.Err())
+			assert.Equal(t, connect.CodeOf(stream.Err()), testcase.expectCode)
+			assert.Equal(t, stream.Err().Error(), testcase.expectMsg)
+		})
+	}
+}
+
 // TestBlankImportCodeGeneration tests that services.connect.go is generated with
 // blank import statements to services.pb.go so that the service's Descriptor is
 // available in the global proto registry.

envelope.go

@@ -305,9 +305,6 @@ func (r *envelopeReader) Read(env *envelope) *Error {
 			return NewError(CodeUnknown, err)
 		}
 		err = wrapIfContextError(err)
-		if connectErr, ok := asError(err); ok {
-			return connectErr
-		}
 		// Something else has gone wrong - the stream didn't end cleanly.
 		if connectErr, ok := asError(err); ok {
 			return connectErr

protocol_grpc.go

@@ -674,11 +674,16 @@ func grpcValidateResponse(
 		)
 	}
 	// When there's no body, gRPC and gRPC-Web servers may send error information
-	// in the HTTP headers.
+	// in the HTTP headers. When this happens, it's called a "trailers-only" response.
 	if err := grpcErrorFromTrailer(
 		protobuf,
 		response.Header,
 	); err != nil && !errors.Is(err, errTrailersWithoutGRPCStatus) {
+		// Trailers-only responses may not have data in the body or HTTP trailers.
+		if bodyErr := grpcVerifyTrailersOnly(response); bodyErr != nil {
+			return bodyErr
+		}
+
 		// Per the specification, only the HTTP status code and Content-Type should
 		// be treated as headers. The rest should be treated as trailing metadata.
 		if contentType := getHeaderCanonical(response.Header, headerContentType); contentType != "" {
@@ -696,6 +701,39 @@ func grpcValidateResponse(
 	return nil
 }
 
+func grpcVerifyTrailersOnly(response *http.Response) *Error {
+	// Make sure there's nothing in the body.
+	if numBytes, err := discard(response.Body); err != nil {
+		err = wrapIfContextError(err)
+		if connErr, ok := asError(err); ok {
+			return connErr
+		}
+		return errorf(CodeInternal, "corrupt response: I/O error after trailers-only response: %w", err)
+	} else if numBytes > 0 {
+		return errorf(CodeInternal, "corrupt response: %d extra bytes after trailers-only response", numBytes)
+	}
+
+	// Now we know we've reached EOF, so we can look at HTTP trailers.
+	// If headers included "Trailer" key, net/http pre-populates response.Trailer with nil
+	// values. So we need to exclude those to see if there were actually any trailers.
+	var trailerCount int
+	for _, v := range response.Trailer {
+		if len(v) > 0 {
+			trailerCount++
+		}
+	}
+	if trailerCount > 0 {
+		// Invalid response: cannot have both trailers in the header (trailers-only
+		// response) AND trailers after the body.
+		return errorf(
+			CodeInternal,
+			"corrupt response from server: gRPC trailers-only response may not contain HTTP trailers",
+		)
+	}
+
+	return nil
+}
+
 func grpcHTTPToCode(httpCode int) Code {
 	// https://github.com/grpc/grpc/blob/master/doc/http-grpc-status-mapping.md
 	// Note that this is not just the inverse of the gRPC-to-HTTP mapping.