From fadddbc57c8e58e0e3f022743a22b68a3ccde94c Mon Sep 17 00:00:00 2001
From: Scott James Remnant <keybuk@google.com>
Date: Thu, 28 Apr 2022 17:41:22 -0700
Subject: [PATCH] pw_protobuf: Check fields don't exceed message bounds

Change-Id: Ic809af6d5d8903a8896d5cb3073505cddcbba8c0
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/92762
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Pigweed-Auto-Submit: Scott James Remnant <keybuk@google.com>
Reviewed-by: Keir Mierle <keir@google.com>
---
 pw_protobuf/codegen_message_test.cc | 81 +++++++++++++++++++++++++++++
 pw_protobuf/encoder.cc              |  4 +-
 pw_protobuf/stream_decoder.cc       |  3 +-
 3 files changed, 86 insertions(+), 2 deletions(-)

diff --git a/pw_protobuf/codegen_message_test.cc b/pw_protobuf/codegen_message_test.cc
index a5c0d72daf..be026e978d 100644
--- a/pw_protobuf/codegen_message_test.cc
+++ b/pw_protobuf/codegen_message_test.cc
@@ -17,6 +17,7 @@
 
 #include "gtest/gtest.h"
 #include "pw_preprocessor/compiler.h"
+#include "pw_protobuf/internal/codegen.h"
 #include "pw_status/status.h"
 #include "pw_status/status_with_size.h"
 #include "pw_stream/memory_stream.h"
@@ -943,6 +944,49 @@ TEST(CodegenMessage, ReadNestedForcedCallback) {
   ASSERT_EQ(status, OkStatus());
 }
 
+class BreakableDecoder : public KeyValuePair::StreamDecoder {
+ public:
+  constexpr BreakableDecoder(stream::Reader& reader) : StreamDecoder(reader) {}
+
+  Status Read(KeyValuePair::Message& message,
+              std::span<const MessageField> table) {
+    return ::pw::protobuf::StreamDecoder::Read(
+        std::as_writable_bytes(std::span(&message, 1)), table);
+  }
+};
+
+TEST(CodegenMessage, DISABLED_ReadDoesNotOverrun) {
+  // Deliberately construct a message table that attempts to violate the bounds
+  // of the structure. We're not testing that a developer can't do this, rather
+  // that the protobuf decoder can't be exploited in this way.
+  constexpr MessageField kMessageFields[] = {
+      {1,
+       WireType::kDelimited,
+       sizeof(std::byte),
+       static_cast<VarintType>(0),
+       false,
+       false,
+       false,
+       0,
+       sizeof(KeyValuePair::Message) * 2,
+       {}},
+  };
+
+  // clang-format off
+  constexpr uint8_t proto_data[] = {
+    // id=1, len=9,
+    0x0a, 0x08, 'd', 'o', 'n', 't', 'e', 'a', 't', 'm', 'e',
+  };
+  // clang-format on
+
+  stream::MemoryReader reader(std::as_bytes(std::span(proto_data)));
+  BreakableDecoder decoder(reader);
+
+  KeyValuePair::Message message{};
+  // ASSERT_CRASH
+  std::ignore = decoder.Read(message, kMessageFields);
+}
+
 TEST(CodegenMessage, Write) {
   constexpr uint8_t pigweed_data[] = {
       0x10, 0x20, 0x30, 0x40, 0x50, 0x60, 0x70, 0x80};
@@ -1561,6 +1605,43 @@ TEST(CodegenMessage, WriteNestedForcedCallback) {
             0);
 }
 
+class BreakableEncoder : public KeyValuePair::MemoryEncoder {
+ public:
+  constexpr BreakableEncoder(ByteSpan buffer)
+      : KeyValuePair::MemoryEncoder(buffer) {}
+
+  Status Write(const KeyValuePair::Message& message,
+               std::span<const MessageField> table) {
+    return ::pw::protobuf::StreamEncoder::Write(
+        std::as_bytes(std::span(&message, 1)), table);
+  }
+};
+
+TEST(CodegenMessage, DISABLED_WriteDoesNotOverrun) {
+  // Deliberately construct a message table that attempts to violate the bounds
+  // of the structure. We're not testing that a developer can't do this, rather
+  // that the protobuf encoder can't be exploited in this way.
+  constexpr MessageField kMessageFields[] = {
+      {1,
+       WireType::kDelimited,
+       sizeof(std::byte),
+       static_cast<VarintType>(0),
+       false,
+       false,
+       false,
+       0,
+       sizeof(KeyValuePair::Message) * 2,
+       {}},
+  };
+
+  std::byte encode_buffer[64];
+
+  BreakableEncoder encoder(encode_buffer);
+  KeyValuePair::Message message{};
+  // ASSERT_CRASH
+  std::ignore = encoder.Write(message, kMessageFields);
+}
+
 // The following tests cover using the codegen struct Message and callbacks in
 // different ways.
 
diff --git a/pw_protobuf/encoder.cc b/pw_protobuf/encoder.cc
index 7c99b77b02..c027f34993 100644
--- a/pw_protobuf/encoder.cc
+++ b/pw_protobuf/encoder.cc
@@ -241,7 +241,9 @@ Status StreamEncoder::Write(std::span<const std::byte> message,
     // Calculate the span of bytes corresponding to the structure field to
     // read from.
     const auto values =
-        std::span(message.data() + field.field_offset(), field.field_size());
+        message.subspan(field.field_offset(), field.field_size());
+    PW_CHECK(values.begin() >= message.begin() &&
+             values.end() <= message.end());
 
     // If the field is using callbacks, interpret the input field accordingly
     // and allow the caller to provide custom handling.
diff --git a/pw_protobuf/stream_decoder.cc b/pw_protobuf/stream_decoder.cc
index 9432cb65f5..29f42e19c8 100644
--- a/pw_protobuf/stream_decoder.cc
+++ b/pw_protobuf/stream_decoder.cc
@@ -529,7 +529,8 @@ Status StreamDecoder::Read(std::span<std::byte> message,
     // Calculate the span of bytes corresponding to the structure field to
     // output into.
     const auto out =
-        std::span(message.data() + field->field_offset(), field->field_size());
+        message.subspan(field->field_offset(), field->field_size());
+    PW_CHECK(out.begin() >= message.begin() && out.end() <= message.end());
 
     // If the field is using callbacks, interpret the output field accordingly
     // and allow the caller to provide custom handling.