Add gRPC frame error handling (#19)

lizan · web-flow · commit 6c54b75dbd29 · 2017-08-10T11:41:32.000-07:00
diff --git a/src/include/grpc_transcoding/message_reader.h b/src/include/grpc_transcoding/message_reader.h
@@ -72,18 +72,22 @@ class MessageReader {
   //       NextMessage() again.
   std::unique_ptr<::google::protobuf::io::ZeroCopyInputStream> NextMessage();
 
+  ::google::protobuf::util::Status Status() const { return status_; }
+
   // Returns true if the stream has ended (this is permanent); otherwise returns
   // false.
-  bool Finished() const { return finished_; }
+  bool Finished() const { return finished_ || !status_.ok(); }
 
  private:
   TranscoderInputStream* in_;
   // The size of the current message.
-  unsigned int current_message_size_;
+  uint32_t current_message_size_;
   // Whether we have read the current message size or not
   bool have_current_message_size_;
   // Are we all done?
   bool finished_;
+  // Status
+  ::google::protobuf::util::Status status_;
 
   MessageReader(const MessageReader&) = delete;
   MessageReader& operator=(const MessageReader&) = delete;
diff --git a/src/include/grpc_transcoding/transcoder_input_stream.h b/src/include/grpc_transcoding/transcoder_input_stream.h
@@ -26,6 +26,9 @@ class TranscoderInputStream
  public:
   // returns the number of bytes available to read at the moment.
   virtual int64_t BytesAvailable() const = 0;
+
+  // returns if the stream is finished. i.e. No more data will be added.
+  virtual bool Finished() const = 0;
 };
 
 }  // namespace transcoding
diff --git a/src/message_reader.cc b/src/message_reader.cc
@@ -60,21 +60,9 @@ bool ReadStream(pbio::ZeroCopyInputStream* stream, unsigned char* buffer,
   return true;
 }
 
-// Determines whether the stream is finished or not.
-bool IsStreamFinished(pbio::ZeroCopyInputStream* stream) {
-  int size = 0;
-  const void* data = nullptr;
-  if (!stream->Next(&data, &size)) {
-    return true;
-  } else {
-    stream->BackUp(size);
-    return false;
-  }
-}
-
 // A helper function to extract the size from a gRPC wire format message
 // delimiter - see http://www.grpc.io/docs/guides/wire.html.
-unsigned DelimiterToSize(const unsigned char* delimiter) {
+uint32_t DelimiterToSize(const unsigned char* delimiter) {
   unsigned size = 0;
   // Bytes 1-4 are big-endian 32-bit message size
   size = size | static_cast<unsigned>(delimiter[1]);
@@ -90,9 +78,9 @@ unsigned DelimiterToSize(const unsigned char* delimiter) {
 }  // namespace
 
 std::unique_ptr<pbio::ZeroCopyInputStream> MessageReader::NextMessage() {
-  if (finished_) {
+  if (Finished()) {
     // The stream has ended
-    return std::unique_ptr<pbio::ZeroCopyInputStream>();
+    return nullptr;
   }
 
   // Check if we have the current message size. If not try to read it.
@@ -101,35 +89,52 @@ std::unique_ptr<pbio::ZeroCopyInputStream> MessageReader::NextMessage() {
     if (in_->BytesAvailable() < static_cast<pb::int64>(kDelimiterSize)) {
       // We don't have 5 bytes available to read the length of the message.
       // Find out whether the stream is finished and return false.
-      finished_ = IsStreamFinished(in_);
-      return std::unique_ptr<pbio::ZeroCopyInputStream>();
+      finished_ = in_->Finished();
+      if (finished_ && in_->BytesAvailable() != 0) {
+        status_ = google::protobuf::util::Status(
+            google::protobuf::util::error::INTERNAL,
+            "Incomplete gRPC frame header received");
+      }
+      return nullptr;
     }
 
     // Try to read the delimiter
     unsigned char delimiter[kDelimiterSize] = {0};
     if (!ReadStream(in_, delimiter, sizeof(delimiter))) {
       finished_ = true;
-      return std::unique_ptr<pbio::ZeroCopyInputStream>();
+      return nullptr;
+    }
+
+    if (delimiter[0] != 0) {
+      status_ = google::protobuf::util::Status(
+          google::protobuf::util::error::INTERNAL,
+          "Unsupported gRPC frame flag: " + std::to_string(delimiter[0]));
+      return nullptr;
     }
 
     current_message_size_ = DelimiterToSize(delimiter);
     have_current_message_size_ = true;
   }
 
   if (in_->BytesAvailable() < static_cast<pb::int64>(current_message_size_)) {
+    if (in_->Finished()) {
+      status_ = google::protobuf::util::Status(
+          google::protobuf::util::error::INTERNAL,
+          "Incomplete gRPC frame expected size: " +
+              std::to_string(current_message_size_) + " actual size: " +
+              std::to_string(in_->BytesAvailable()));
+    }
     // We don't have a full message
-    return std::unique_ptr<pbio::ZeroCopyInputStream>();
+    return nullptr;
   }
 
-  // We have a message! Use LimitingInputStream to wrap the input stream and
-  // limit it to current_message_size_ bytes to cover only the current message.
-  auto result = std::unique_ptr<pbio::ZeroCopyInputStream>(
-      new pbio::LimitingInputStream(in_, current_message_size_));
-
   // Reset the have_current_message_size_ for the next message
   have_current_message_size_ = false;
 
-  return result;
+  // We have a message! Use LimitingInputStream to wrap the input stream and
+  // limit it to current_message_size_ bytes to cover only the current message.
+  return std::unique_ptr<pbio::ZeroCopyInputStream>(
+      new pbio::LimitingInputStream(in_, current_message_size_));
 }
 
 }  // namespace transcoding
diff --git a/src/message_stream.cc b/src/message_stream.cc
@@ -82,6 +82,8 @@ class InputStreamOverMessageStream : public TranscoderInputStream {
     return static_cast<int64_t>(message_.size() - position_);
   }
 
+  bool Finished() const { return src_->Finished(); }
+
  private:
   // Updates the current message and creates an ArrayInputStream over it.
   void ReadNextMessage() {
diff --git a/src/response_to_json_translator.cc b/src/response_to_json_translator.cc
@@ -44,8 +44,14 @@ bool ResponseToJsonTranslator::NextMessage(std::string* message) {
     // All done
     return false;
   }
+
   // Try to read a message
   auto proto_in = reader_.NextMessage();
+  status_ = reader_.Status();
+  if (!status_.ok()) {
+    return false;
+  }
+
   if (proto_in) {
     std::string json_out;
     if (TranslateMessage(proto_in.get(), &json_out)) {
diff --git a/test/message_reader_test.cc b/test/message_reader_test.cc
@@ -95,6 +95,7 @@ class MessageReaderTestRun {
       }
       // Read the message
       auto stream = reader_->NextMessage();
+      EXPECT_TRUE(reader_->Status().ok());
       if (!stream) {
         ADD_FAILURE() << "No message available" << std::endl;
         return false;
@@ -296,13 +297,15 @@ TEST_F(MessageReaderTest, DirectTest) {
   // Nothing yet
   EXPECT_EQ(nullptr, reader.NextMessage().get());
   EXPECT_FALSE(reader.Finished());
+  EXPECT_TRUE(reader.Status().ok());
 
   // Add the delimiter for the first message
   input_stream.AddChunk(SizeToDelimiter(message1.size()));
 
   // Still nothing
   EXPECT_EQ(nullptr, reader.NextMessage().get());
   EXPECT_FALSE(reader.Finished());
+  EXPECT_TRUE(reader.Status().ok());
 
   // Add the message itself
   input_stream.AddChunk(message1);
@@ -313,6 +316,7 @@ TEST_F(MessageReaderTest, DirectTest) {
   EXPECT_EQ(message1, ReadAllFromStream(message1_stream.get()));
   EXPECT_EQ(nullptr, reader.NextMessage().get());
   EXPECT_FALSE(reader.Finished());
+  EXPECT_TRUE(reader.Status().ok());
 
   // Add part of the message2
   input_stream.AddChunk(SizeToDelimiter(message2.size()));
@@ -321,6 +325,7 @@ TEST_F(MessageReaderTest, DirectTest) {
   // No message should be available
   EXPECT_EQ(nullptr, reader.NextMessage().get());
   EXPECT_FALSE(reader.Finished());
+  EXPECT_TRUE(reader.Status().ok());
 
   // Add the rest of the second message
   input_stream.AddChunk(message2.substr(10));
@@ -331,6 +336,7 @@ TEST_F(MessageReaderTest, DirectTest) {
   EXPECT_EQ(message2, ReadAllFromStream(message2_stream.get()));
   EXPECT_EQ(nullptr, reader.NextMessage().get());
   EXPECT_FALSE(reader.Finished());
+  EXPECT_TRUE(reader.Status().ok());
 
   // Adding both message3 & message4 in one shot and Finish the stream
   input_stream.AddChunk(SizeToDelimiter(message3.size()));
@@ -341,6 +347,7 @@ TEST_F(MessageReaderTest, DirectTest) {
 
   // Not finished yet as we still have messages to read
   EXPECT_FALSE(reader.Finished());
+  EXPECT_TRUE(reader.Status().ok());
 
   // Now both message3 & message4 must be available
   auto message3_stream = reader.NextMessage();
@@ -355,6 +362,45 @@ TEST_F(MessageReaderTest, DirectTest) {
   // All done, the reader must be finished
   EXPECT_EQ(nullptr, reader.NextMessage().get());
   EXPECT_TRUE(reader.Finished());
+  EXPECT_TRUE(reader.Status().ok());
+}
+
+TEST_F(MessageReaderTest, IncompleteFrameHeader) {
+  TestZeroCopyInputStream input_stream;
+  MessageReader reader(&input_stream);
+
+  input_stream.AddChunk("\x0a");
+  input_stream.Finish();
+
+  EXPECT_EQ(nullptr, reader.NextMessage().get());
+  EXPECT_FALSE(reader.Status().ok());
+  EXPECT_EQ(reader.Status().error_message(),
+            "Incomplete gRPC frame header received");
+}
+
+TEST_F(MessageReaderTest, InvalidFrameFlag) {
+  TestZeroCopyInputStream input_stream;
+  MessageReader reader(&input_stream);
+
+  input_stream.AddChunk(std::string("\x0A\x00\x00\x00\x00", 5));
+  input_stream.Finish();
+
+  EXPECT_EQ(nullptr, reader.NextMessage().get());
+  EXPECT_FALSE(reader.Status().ok());
+  EXPECT_EQ(reader.Status().error_message(), "Unsupported gRPC frame flag: 10");
+}
+
+TEST_F(MessageReaderTest, IncompleteFrame) {
+  TestZeroCopyInputStream input_stream;
+  MessageReader reader(&input_stream);
+
+  input_stream.AddChunk(std::string("\x00\x00\x00\x00\x05\x00", 6));
+  input_stream.Finish();
+
+  EXPECT_EQ(nullptr, reader.NextMessage().get());
+  EXPECT_FALSE(reader.Status().ok());
+  EXPECT_EQ(reader.Status().error_message(),
+            "Incomplete gRPC frame expected size: 5 actual size: 1");
 }
 
 }  // namespace
diff --git a/test/response_to_json_translator_test.cc b/test/response_to_json_translator_test.cc
@@ -825,6 +825,75 @@ TEST_F(ResponseToJsonTranslatorTest, Streaming5KMessages) {
   EXPECT_TRUE(ExpectJsonArrayEq(expected_json_array, actual_json_array));
 }
 
+TEST_F(ResponseToJsonTranslatorTest, IncompleteFrameHeader) {
+  // Load the service config
+  ::google::api::Service service;
+  ASSERT_TRUE(
+      transcoding::testing::LoadService("bookstore_service.pb.txt", &service));
+
+  // Create a TypeHelper using the service config
+  TypeHelper type_helper(service.types(), service.enums());
+
+  TestZeroCopyInputStream input_stream;
+  ResponseToJsonTranslator translator(
+      type_helper.Resolver(), "type.googleapis.com/Shelf", true, &input_stream);
+
+  input_stream.AddChunk(std::string("\x0A", 1));
+  input_stream.Finish();
+
+  std::string actual;
+  EXPECT_FALSE(translator.NextMessage(&actual));
+  EXPECT_FALSE(translator.Status().ok());
+  EXPECT_EQ(translator.Status().error_message(),
+            "Incomplete gRPC frame header received");
+}
+
+TEST_F(ResponseToJsonTranslatorTest, InvalidFrameFlag) {
+  // Load the service config
+  ::google::api::Service service;
+  ASSERT_TRUE(
+      transcoding::testing::LoadService("bookstore_service.pb.txt", &service));
+
+  // Create a TypeHelper using the service config
+  TypeHelper type_helper(service.types(), service.enums());
+
+  TestZeroCopyInputStream input_stream;
+  ResponseToJsonTranslator translator(
+      type_helper.Resolver(), "type.googleapis.com/Shelf", true, &input_stream);
+
+  input_stream.AddChunk(std::string("\x0A\x00\x00\x00\x00", 5));
+  input_stream.Finish();
+
+  std::string actual;
+  EXPECT_FALSE(translator.NextMessage(&actual));
+  EXPECT_FALSE(translator.Status().ok());
+  EXPECT_EQ(translator.Status().error_message(),
+            "Unsupported gRPC frame flag: 10");
+}
+
+TEST_F(ResponseToJsonTranslatorTest, IncompleteFrame) {
+  // Load the service config
+  ::google::api::Service service;
+  ASSERT_TRUE(
+      transcoding::testing::LoadService("bookstore_service.pb.txt", &service));
+
+  // Create a TypeHelper using the service config
+  TypeHelper type_helper(service.types(), service.enums());
+
+  TestZeroCopyInputStream input_stream;
+  ResponseToJsonTranslator translator(
+      type_helper.Resolver(), "type.googleapis.com/Shelf", true, &input_stream);
+
+  input_stream.AddChunk(std::string("\x00\x00\x00\x00\x05\x00", 6));
+  input_stream.Finish();
+
+  std::string actual;
+  EXPECT_FALSE(translator.NextMessage(&actual));
+  EXPECT_FALSE(translator.Status().ok());
+  EXPECT_EQ(translator.Status().error_message(),
+            "Incomplete gRPC frame expected size: 5 actual size: 1");
+}
+
 }  // namespace
 }  // namespace testing
 }  // namespace transcoding

Original file line number	Diff line number	Diff line change
`@@ -82,6 +82,8 @@ class InputStreamOverMessageStream : public TranscoderInputStream {`
`82`	`82`	`return static_cast<int64_t>(message_.size() - position_);`
`83`	`83`	`}`
`84`	`84`
	`85`	`+ bool Finished() const { return src_->Finished(); }`
	`86`	`+`
`85`	`87`	`private:`
`86`	`88`	`// Updates the current message and creates an ArrayInputStream over it.`
`87`	`89`	`void ReadNextMessage() {`