Expose last parsed delimiter in MessageReader (#58)

Currently, `MessageReader::NextMessage` will advance the buffer past the gRPC delimiter.

However, we need the delimiter if we wish to preserve the original message. Instead of having our extraction libraries reform the delimiter, just rely on `MessageReader` to expose the last parsed message.

Author: nareddyt

script/ci.sh

@@ -17,4 +17,4 @@
 ################################################################################
 
 bazel build //...
-bazel test //...
+bazel test //... --test_output=errors

src/BUILD

@@ -214,7 +214,7 @@ cc_library(
 
 cc_library(
     name = "transcoder_input_stream",
-    srcs = [
+    hdrs = [
         "include/grpc_transcoding/transcoder_input_stream.h",
     ],
     includes = [

src/include/grpc_transcoding/message_reader.h

@@ -25,6 +25,17 @@ namespace grpc {
 
 namespace transcoding {
 
+// The number of bytes in the delimiter for gRPC wire format's
+// `Length-Prefixed-Message`.
+constexpr size_t kGrpcDelimiterByteSize = 5;
+
+// Return type that contains both the proto message and the preceding gRPC data
+// frame.
+struct MessageAndGrpcFrame {
+  std::unique_ptr<::google::protobuf::io::ZeroCopyInputStream> message;
+  unsigned char grpc_frame[kGrpcDelimiterByteSize];
+};
+
 // MessageReader helps extract full messages from a ZeroCopyInputStream of
 // messages in gRPC wire format (http://www.grpc.io/docs/guides/wire.html). Each
 // message is returned in a ZeroCopyInputStream. MessageReader doesn't advance
@@ -70,8 +81,15 @@ class MessageReader {
   //       of the original ZeroCopyInputStream and the MessageReader relies on
   //       the caller to advance the stream to the next message before calling
   //       NextMessage() again.
+  // NOTE: the caller should check `Status()` is OK after calling this method.
   std::unique_ptr<::google::protobuf::io::ZeroCopyInputStream> NextMessage();
 
+  // An overload that also outputs the gRPC message delimiter for the parsed
+  // message. The caller is free to take ownership of contents in `grpc_frame`.
+  // NOTE: the caller must check the `message` is NOT nullptr and the `Status()`
+  //       is OK before consuming the `grpc_frame`.
+  MessageAndGrpcFrame NextMessageAndGrpcFrame();
+
   ::google::protobuf::util::Status Status() const { return status_; }
 
   // Returns true if the stream has ended (this is permanent); otherwise returns
@@ -88,6 +106,8 @@ class MessageReader {
   bool finished_;
   // Status
   ::google::protobuf::util::Status status_;
+  // Buffer to store the current delimiter value.
+  unsigned char delimiter_[kGrpcDelimiterByteSize];
 
   MessageReader(const MessageReader&) = delete;
   MessageReader& operator=(const MessageReader&) = delete;

src/message_reader.cc

@@ -85,8 +85,8 @@ std::unique_ptr<pbio::ZeroCopyInputStream> MessageReader::NextMessage() {
 
   // Check if we have the current message size. If not try to read it.
   if (!have_current_message_size_) {
-    const size_t kDelimiterSize = 5;
-    if (in_->BytesAvailable() < static_cast<pb::int64>(kDelimiterSize)) {
+    if (in_->BytesAvailable()
+        < static_cast<pb::int64>(kGrpcDelimiterByteSize)) {
       // 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_ = in_->Finished();
@@ -98,21 +98,21 @@ std::unique_ptr<pbio::ZeroCopyInputStream> MessageReader::NextMessage() {
       return nullptr;
     }
 
-    // Try to read the delimiter
-    unsigned char delimiter[kDelimiterSize] = {0};
-    if (!ReadStream(in_, delimiter, sizeof(delimiter))) {
+    // Try to read the delimiter.
+    memset(delimiter_, 0, kGrpcDelimiterByteSize);
+    if (!ReadStream(in_, delimiter_, kGrpcDelimiterByteSize)) {
       finished_ = true;
       return nullptr;
     }
 
-    if (delimiter[0] != 0) {
+    if (delimiter_[0] != 0) {
       status_ = google::protobuf::util::Status(
           google::protobuf::util::error::INTERNAL,
-          "Unsupported gRPC frame flag: " + std::to_string(delimiter[0]));
+          "Unsupported gRPC frame flag: " + std::to_string(delimiter_[0]));
       return nullptr;
     }
 
-    current_message_size_ = DelimiterToSize(delimiter);
+    current_message_size_ = DelimiterToSize(delimiter_);
     have_current_message_size_ = true;
   }
 
@@ -137,6 +137,13 @@ std::unique_ptr<pbio::ZeroCopyInputStream> MessageReader::NextMessage() {
       new pbio::LimitingInputStream(in_, current_message_size_));
 }
 
+MessageAndGrpcFrame MessageReader::NextMessageAndGrpcFrame() {
+  MessageAndGrpcFrame out;
+  out.message = NextMessage();
+  memcpy(out.grpc_frame, delimiter_, kGrpcDelimiterByteSize);
+  return out;
+}
+
 }  // namespace transcoding
 
 }  // namespace grpc

test/message_reader_test.cc

@@ -47,6 +47,7 @@ struct ExpectedAt {
   // The position in the input, after which this message is expected
   size_t at;
   std::string message;
+  std::string delimiter;
 };
 
 // MessageReaderTestRun tests a single MessageReader processing the input as
@@ -87,25 +88,33 @@ class MessageReaderTestRun {
     // While we still have expected messages before or at the current position
     // try to match.
     while (next_expected_ != std::end(expected_) &&
-           next_expected_->at <= position_) {
+        next_expected_->at <= position_) {
       // Must not be finished as we expect a message
       if (reader_->Finished()) {
         ADD_FAILURE() << "Finished unexpectedly" << std::endl;
         return false;
       }
       // Read the message
-      auto stream = reader_->NextMessage();
+      MessageAndGrpcFrame result = reader_->NextMessageAndGrpcFrame();
       EXPECT_TRUE(reader_->Status().ok());
-      if (!stream) {
+      if (!result.message) {
         ADD_FAILURE() << "No message available" << std::endl;
         return false;
       }
       // Match the message with the expected message
-      auto message = ReadAllFromStream(stream.get());
+      auto message = ReadAllFromStream(result.message.get());
       if (next_expected_->message != message) {
         EXPECT_EQ(next_expected_->message, message);
         return false;
       }
+      // Match the delimiter.
+      std::string delimiter_string =
+          std::string(reinterpret_cast<const char*>(result.grpc_frame),
+                      kGrpcDelimiterByteSize);
+      if (delimiter_string != next_expected_->delimiter) {
+        EXPECT_EQ(delimiter_string, next_expected_->delimiter);
+        return false;
+      }
       // Move to the next expected message
       ++next_expected_;
     }
@@ -150,7 +159,8 @@ class MessageReaderTestCase {
       input_ += message;
       // Remember that we should expect this message after input_.size() bytes
       // are processed.
-      expected_.emplace_back(ExpectedAt{input_.size(), message});
+      expected_.emplace_back(ExpectedAt{input_.size(), message,
+                                        SizeToDelimiter(message.size())});
     }
   }