grpc-json: no buffering at end_stream or non-gRPC response

source/common/grpc/common.cc

@@ -48,6 +48,17 @@ bool Common::hasGrpcContentType(const Http::HeaderMap& headers) {
   return false;
 }
 
+bool Common::isGrpcResponseHeader(const Http::HeaderMap& headers, bool end_stream) {
+  if (end_stream) {
+    // Trailers-only response, only grpc-status is required.
+    return headers.GrpcStatus() != nullptr;
+  }
+  if (Http::Utility::getResponseStatus(headers) != enumToInt(Http::Code::OK)) {
+    return false;
+  }
+  return hasGrpcContentType(headers);
+}
+
 void Common::chargeStat(const Upstream::ClusterInfo& cluster, const std::string& protocol,
                         const std::string& grpc_service, const std::string& grpc_method,
                         const Http::HeaderEntry* grpc_status) {

source/common/grpc/common.h

@@ -32,6 +32,13 @@ class Common {
    */
   static bool hasGrpcContentType(const Http::HeaderMap& headers);
 
+  /**
+   * @param headers the headers to parse.
+   * @param bool indicating wether the header is at end_stream.
+   * @return bool indicating whether the header is a gRPC reseponse header
+   */
+  static bool isGrpcResponseHeader(const Http::HeaderMap& headers, bool end_stream);
+
   /**
    * Returns the GrpcStatus code from a given set of trailers, if present.
    * @param trailers the trailers to parse.

source/common/grpc/json_transcoder_filter.cc

@@ -295,6 +295,10 @@ void JsonTranscoderFilter::setDecoderFilterCallbacks(
 
 Http::FilterHeadersStatus JsonTranscoderFilter::encodeHeaders(Http::HeaderMap& headers,
                                                               bool end_stream) {
+  if (!Common::isGrpcResponseHeader(headers, end_stream)) {
+    error_ = true;
+  }
+
   if (error_ || !transcoder_) {
     return Http::FilterHeadersStatus::Continue;
   }
@@ -329,7 +333,7 @@ Http::FilterDataStatus JsonTranscoderFilter::encodeData(Buffer::Instance& data,
 
   readToBuffer(*transcoder_->ResponseOutput(), data);
 
-  if (!method_->server_streaming()) {
+  if (!method_->server_streaming() && !end_stream) {
     // Buffer until the response is complete.
     return Http::FilterDataStatus::StopIterationAndBuffer;
   }

test/common/grpc/common_test.cc

@@ -143,5 +143,21 @@ TEST(GrpcCommonTest, HasGrpcContentType) {
   EXPECT_FALSE(isGrpcContentType("application/grpc-web+foo"));
 }
 
+TEST(GrpcCommonTest, IsGrpcResponseHeader) {
+  Http::TestHeaderMapImpl grpc_status_only{{":status", "500"}, {"grpc-status", "14"}};
+  EXPECT_TRUE(Common::isGrpcResponseHeader(grpc_status_only, true));
+  EXPECT_FALSE(Common::isGrpcResponseHeader(grpc_status_only, false));
+
+  Http::TestHeaderMapImpl grpc_response_header{{":status", "200"},
+                                               {"content-type", "application/grpc"}};
+  EXPECT_FALSE(Common::isGrpcResponseHeader(grpc_response_header, true));
+  EXPECT_TRUE(Common::isGrpcResponseHeader(grpc_response_header, false));
+
+  Http::TestHeaderMapImpl json_response_header{{":status", "200"},
+                                               {"content-type", "application/json"}};
+  EXPECT_FALSE(Common::isGrpcResponseHeader(json_response_header, true));
+  EXPECT_FALSE(Common::isGrpcResponseHeader(json_response_header, false));
+}
+
 } // namespace Grpc
 } // namespace Envoy

test/common/grpc/json_transcoder_filter_test.cc

@@ -243,8 +243,16 @@ TEST_F(GrpcJsonTranscoderFilterTest, NoTranscoding) {
   Http::TestHeaderMapImpl request_trailers;
   EXPECT_EQ(Http::FilterTrailersStatus::Continue, filter_.decodeTrailers(request_trailers));
 
-  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_.encodeHeaders(request_headers, false));
-  EXPECT_EQ(expected_request_headers, request_headers);
+  Http::TestHeaderMapImpl response_headers{{"content-type", "application/grpc"},
+                                           {":status", "200"},
+                                           {":path", "/grpc.service/UnknownGrpcMethod"}};
+
+  Http::TestHeaderMapImpl expected_response_headers{{"content-type", "application/grpc"},
+                                                    {":status", "200"},
+                                                    {":path", "/grpc.service/UnknownGrpcMethod"}};
+
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_.encodeHeaders(response_headers, false));
+  EXPECT_EQ(expected_response_headers, response_headers);
 
   Buffer::OwnedImpl response_data{"{}"};
   EXPECT_EQ(Http::FilterDataStatus::Continue, filter_.encodeData(response_data, false));
@@ -334,6 +342,44 @@ TEST_F(GrpcJsonTranscoderFilterTest, TranscodingUnaryError) {
   EXPECT_EQ(0, request_data.length());
 }
 
+TEST_F(GrpcJsonTranscoderFilterTest, TranscodingUnaryTimeout) {
+  Http::TestHeaderMapImpl request_headers{
+      {"content-type", "application/json"}, {":method", "POST"}, {":path", "/shelf"}};
+
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_.decodeHeaders(request_headers, false));
+  EXPECT_EQ("application/grpc", request_headers.get_("content-type"));
+  EXPECT_EQ("/bookstore.Bookstore/CreateShelf", request_headers.get_(":path"));
+  EXPECT_EQ("trailers", request_headers.get_("te"));
+
+  Buffer::OwnedImpl request_data{"{\"theme\": \"Children\"}"};
+
+  EXPECT_EQ(Http::FilterDataStatus::Continue, filter_.decodeData(request_data, true));
+
+  Http::TestHeaderMapImpl response_headers{
+      {":status", "504"}, {"content-length", "24"}, {"content-type", "text/plain"}};
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_.encodeHeaders(response_headers, false));
+  EXPECT_EQ(Http::FilterDataStatus::Continue, filter_.encodeData(request_data, true));
+}
+
+TEST_F(GrpcJsonTranscoderFilterTest, TranscodingUnaryNotGrpcResponse) {
+  Http::TestHeaderMapImpl request_headers{
+      {"content-type", "application/json"}, {":method", "POST"}, {":path", "/shelf"}};
+
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_.decodeHeaders(request_headers, false));
+  EXPECT_EQ("application/grpc", request_headers.get_("content-type"));
+  EXPECT_EQ("/bookstore.Bookstore/CreateShelf", request_headers.get_(":path"));
+  EXPECT_EQ("trailers", request_headers.get_("te"));
+
+  Buffer::OwnedImpl request_data{"{\"theme\": \"Children\"}"};
+
+  EXPECT_EQ(Http::FilterDataStatus::Continue, filter_.decodeData(request_data, true));
+
+  Http::TestHeaderMapImpl response_headers{
+      {":status", "200"}, {"content-length", "24"}, {"content-type", "text/plain"}};
+  EXPECT_EQ(Http::FilterHeadersStatus::Continue, filter_.encodeHeaders(response_headers, false));
+  EXPECT_EQ(Http::FilterDataStatus::Continue, filter_.encodeData(request_data, true));
+}
+
 struct GrpcJsonTranscoderFilterPrintTestParam {
   std::string config_json_;
   std::string expected_response_;