pw_protobuf: Support providing known length to StreamDecoder

When a protobuf is encoded on a stream, the wire format may include
advance details of the length of the message.

Allow this to be specified to StreamDecoder so that it does not read
(and consume) bytes beyond the bounds of the message to be decoded.

Change-Id: Ifabed2ab4f39dda50c3945f9cc943852aee898d9
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/74722
Pigweed-Auto-Submit: Scott James Remnant <keybuk@google.com>
Reviewed-by: Yecheng Zhao <zyecheng@google.com>
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>

4 files changed, 73 insertions, 1 deletions

diff --git a/pw_protobuf/docs.rst b/pw_protobuf/docs.rst
index 260eb26e8..d605eade7 100644
--- a/pw_protobuf/docs.rst
+++ b/pw_protobuf/docs.rst
@@ -378,6 +378,18 @@ of the stream into a provided buffer.
     return status.IsOutOfRange() ? OkStatus() : status;
   }
 
+Where the length of the protobuf message is known in advance, the decoder can
+be prevented from reading from the stream beyond the known bounds by specifying
+the known length to the decoder:
+
+.. code-block:: c++
+
+  pw::protobuf::StreamDecoder decoder(reader, message_length);
+
+When a decoder constructed in this way goes out of scope, it will consume any
+remaining bytes in `message_length` allowing the next `Read` to be data after
+the protobuf, even when it was not fully parsed.
+
 The ``StreamDecoder`` can also return a ``StreamDecoder::BytesReader`` for
 reading bytes fields, avoiding the need to copy data out directly.
 
diff --git a/pw_protobuf/public/pw_protobuf/stream_decoder.h b/pw_protobuf/public/pw_protobuf/stream_decoder.h
index f26f18179..fda6f4eb2 100644
--- a/pw_protobuf/public/pw_protobuf/stream_decoder.h
+++ b/pw_protobuf/public/pw_protobuf/stream_decoder.h
@@ -98,8 +98,14 @@ class StreamDecoder {
   };
 
   constexpr StreamDecoder(stream::Reader& reader)
+      : StreamDecoder(reader, std::numeric_limits<size_t>::max()) {}
+
+  // Allow the maximum length of the protobuf to be specified to the decoder
+  // for streaming situations. When constructed in this way, the decoder will
+  // consume any remaining bytes when it goes out of scope.
+  constexpr StreamDecoder(stream::Reader& reader, size_t length)
       : reader_(reader),
-        stream_bounds_({0, std::numeric_limits<size_t>::max()}),
+        stream_bounds_({0, length}),
         position_(0),
         current_field_(kInitialFieldKey),
         delimited_field_size_(0),
diff --git a/pw_protobuf/stream_decoder.cc b/pw_protobuf/stream_decoder.cc
index 22f64cc72..8df6f2860 100644
--- a/pw_protobuf/stream_decoder.cc
+++ b/pw_protobuf/stream_decoder.cc
@@ -14,6 +14,8 @@
 
 #include "pw_protobuf/stream_decoder.h"
 
+#include <limits>
+
 #include "pw_assert/check.h"
 #include "pw_status/status.h"
 #include "pw_status/status_with_size.h"
@@ -84,6 +86,11 @@ StatusWithSize StreamDecoder::BytesReader::DoRead(ByteSpan destination) {
 StreamDecoder::~StreamDecoder() {
   if (parent_ != nullptr) {
     parent_->CloseNestedDecoder(*this);
+  } else if (stream_bounds_.high < std::numeric_limits<size_t>::max()) {
+    if (status_.ok()) {
+      // Advance the stream to the end of the bounds.
+      PW_CHECK(Advance(stream_bounds_.high).ok());
+    }
   }
 }
 
diff --git a/pw_protobuf/stream_decoder_test.cc b/pw_protobuf/stream_decoder_test.cc
index d083b9065..f1385c4ca 100644
--- a/pw_protobuf/stream_decoder_test.cc
+++ b/pw_protobuf/stream_decoder_test.cc
@@ -714,5 +714,52 @@ TEST(StreamDecoder, ReadDelimitedField_DoesntOverConsume) {
   EXPECT_EQ(result.value(), 42);
 }
 
+TEST(StreamDecoder, Decode_WithLength) {
+  // clang-format off
+  constexpr uint8_t encoded_proto[] = {
+    // type=int32, k=1, v=42
+    0x08, 0x2a,
+    // This field is beyond the range of the protobuf:
+    // type=sint32, k=2, v=-13
+    0x10, 0x19,
+  };
+  // clang-format on
+
+  stream::MemoryReader reader(std::as_bytes(std::span(encoded_proto)));
+  StreamDecoder decoder(reader, /*length=*/2u);
+
+  EXPECT_EQ(decoder.Next(), OkStatus());
+  ASSERT_EQ(decoder.FieldNumber().value(), 1u);
+  Result<int32_t> int32 = decoder.ReadInt32();
+  ASSERT_EQ(int32.status(), OkStatus());
+  EXPECT_EQ(int32.value(), 42);
+
+  EXPECT_EQ(decoder.Next(), Status::OutOfRange());
+}
+
+TEST(StreamDecoder, Decode_WithLength_SkipsToEnd) {
+  // clang-format off
+  constexpr uint8_t encoded_proto[] = {
+    // type=string, k=1, v="Hello world"
+    0x08, 0x0b, 'H', 'e', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd',
+    // This field is beyond the range of the protobuf:
+    // type=sint32, k=2, v=-13
+    0x10, 0x19,
+  };
+  // clang-format on
+
+  stream::MemoryReader reader(std::as_bytes(std::span(encoded_proto)));
+  {
+    StreamDecoder decoder(reader, /*length=*/13u);
+
+    EXPECT_EQ(decoder.Next(), OkStatus());
+    ASSERT_EQ(decoder.FieldNumber().value(), 1u);
+    // Don't read the value out, or advance further. Destructing the object
+    // should advance to the end of the length given.
+  }
+
+  EXPECT_EQ(reader.Tell(), 13u);
+}
+
 }  // namespace
 }  // namespace pw::protobuf