src/protozero/message_unittest.cc

/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "perfetto/protozero/message.h"

#include <limits>
#include <memory>
#include <random>
#include <string>
#include <utility>
#include <vector>

#include "perfetto/base/logging.h"
#include "perfetto/ext/base/string_utils.h"
#include "perfetto/protozero/message_handle.h"
#include "perfetto/protozero/proto_utils.h"
#include "perfetto/protozero/root_message.h"
#include "src/base/test/utils.h"
#include "src/protozero/test/example_proto/test_messages.pbzero.h"
#include "src/protozero/test/fake_scattered_buffer.h"
#include "test/gtest_and_gmock.h"

namespace protozero {

namespace {

constexpr size_t kChunkSize = 16;
constexpr uint8_t kTestBytes[] = {0, 0, 0, 0, 0x42, 1, 0x42, 0xff, 0x42, 0};
constexpr const char kStartWatermark[] = {'a', 'b', 'c', 'd',
                                          '1', '2', '3', '\0'};
constexpr const char kEndWatermark[] = {'9', '8', '7', '6',
                                        'z', 'w', 'y', '\0'};

class FakeRootMessage : public RootMessage<Message> {};
class FakeChildMessage : public Message {};

uint32_t SimpleHash(const std::string& str) {
  uint32_t hash = 5381;
  for (char c : str)
    hash = 33 * hash + static_cast<uint32_t>(c);
  return hash;
}

class MessageTest : public ::testing::Test {
 public:
  void SetUp() override { SetChunkSize(kChunkSize); }

  void TearDown() override {
    // Check that none of the messages created by the text fixtures below did
    // under/overflow their heap boundaries.
    for (std::unique_ptr<uint8_t[]>& mem : messages_) {
      EXPECT_STREQ(kStartWatermark, reinterpret_cast<char*>(mem.get()));
      EXPECT_STREQ(kEndWatermark,
                   reinterpret_cast<char*>(mem.get() + sizeof(kStartWatermark) +
                                           sizeof(FakeRootMessage)));
      FakeRootMessage* msg = reinterpret_cast<FakeRootMessage*>(
          mem.get() + sizeof(kStartWatermark));
      msg->~FakeRootMessage();
      mem.reset();
    }
    messages_.clear();
    stream_writer_.reset();
    buffer_.reset();
  }

  void ResetMessage(FakeRootMessage* msg) { msg->Reset(stream_writer_.get()); }

  void SetChunkSize(size_t size) {
    buffer_.reset(new FakeScatteredBuffer(size));
    stream_writer_.reset(new ScatteredStreamWriter(buffer_.get()));
    chunk_size_ = size;
    readback_pos_ = 0;
  }

  FakeRootMessage* NewMessage() {
    std::unique_ptr<uint8_t[]> mem(
        new uint8_t[sizeof(kStartWatermark) + sizeof(FakeRootMessage) +
                    sizeof(kEndWatermark)]);
    uint8_t* msg_start = mem.get() + sizeof(kStartWatermark);
    memcpy(mem.get(), kStartWatermark, sizeof(kStartWatermark));
    memset(msg_start, 0, sizeof(FakeRootMessage));
    memcpy(msg_start + sizeof(FakeRootMessage), kEndWatermark,
           sizeof(kEndWatermark));
    messages_.push_back(std::move(mem));
    FakeRootMessage* msg = new (msg_start) FakeRootMessage();
    msg->Reset(stream_writer_.get());
    return msg;
  }

  FakeRootMessage* NewMessageWithSizeField() {
    FakeRootMessage* msg = NewMessage();
    uint8_t* size_field =
        stream_writer_->ReserveBytes(proto_utils::kMessageLengthFieldSize);
    memset(size_field, 0u, proto_utils::kMessageLengthFieldSize);
    msg->set_size_field(size_field);
    return msg;
  }

  size_t GetNumSerializedBytes() {
    if (buffer_->chunks().empty())
      return 0;
    return buffer_->chunks().size() * chunk_size_ -
           stream_writer_->bytes_available();
  }

  std::string GetNextSerializedBytes(size_t num_bytes) {
    size_t old_readback_pos = readback_pos_;
    readback_pos_ += num_bytes;
    return buffer_->GetBytesAsString(old_readback_pos, num_bytes);
  }

  static void BuildNestedMessages(Message* msg,
                                  uint32_t max_depth,
                                  bool empty = false,
                                  uint32_t depth = 0) {
    if (!empty) {
      for (uint32_t i = 1; i <= 128; ++i)
        msg->AppendBytes(i, kTestBytes, sizeof(kTestBytes));
    }

    if (depth < max_depth) {
      auto* nested_msg =
          msg->BeginNestedMessage<FakeChildMessage>(1 + depth * 10);
      BuildNestedMessages(nested_msg, max_depth, empty, depth + 1);
    }

    if (!empty) {
      for (uint32_t i = 129; i <= 256; ++i)
        msg->AppendVarInt(i, 42);
    }

    if ((depth & 2) == 0)
      msg->Finalize();
  }

 protected:
  std::unique_ptr<FakeScatteredBuffer> buffer_;
  std::unique_ptr<ScatteredStreamWriter> stream_writer_;
  std::vector<std::unique_ptr<uint8_t[]>> messages_;
  size_t chunk_size_{};
  size_t readback_pos_{};
};

TEST_F(MessageTest, ZeroLengthArraysAndStrings) {
  Message* msg = NewMessage();
  msg->AppendBytes(1 /* field_id */, nullptr, 0);
  msg->AppendString(2 /* field_id */, "");

  EXPECT_EQ(4u, msg->Finalize());
  EXPECT_EQ(4u, GetNumSerializedBytes());

  // These lines match the serialization of the Append* calls above.
  ASSERT_EQ("0A00", GetNextSerializedBytes(2));
  ASSERT_EQ("1200", GetNextSerializedBytes(2));
}

TEST_F(MessageTest, BasicTypesNoNesting) {
  Message* msg = NewMessage();
  msg->AppendVarInt(1 /* field_id */, 0);
  msg->AppendVarInt(2 /* field_id */, std::numeric_limits<uint32_t>::max());
  msg->AppendVarInt(3 /* field_id */, 42);
  msg->AppendVarInt(4 /* field_id */, std::numeric_limits<uint64_t>::max());
  msg->AppendFixed(5 /* field_id */, 3.1415f /* float */);
  msg->AppendFixed(6 /* field_id */, 3.14159265358979323846 /* double */);
  msg->AppendBytes(7 /* field_id */, kTestBytes, sizeof(kTestBytes));

  // Field ids > 16 are expected to be varint encoded (preamble > 1 byte)
  msg->AppendString(257 /* field_id */, "0123456789abcdefABCDEF");
  msg->AppendSignedVarInt(3 /* field_id */, -21);

  EXPECT_EQ(74u, msg->Finalize());
  EXPECT_EQ(74u, GetNumSerializedBytes());

  // These lines match the serialization of the Append* calls above.
  ASSERT_EQ("0800", GetNextSerializedBytes(2));
  ASSERT_EQ("10FFFFFFFF0F", GetNextSerializedBytes(6));
  ASSERT_EQ("182A", GetNextSerializedBytes(2));
  ASSERT_EQ("20FFFFFFFFFFFFFFFFFF01", GetNextSerializedBytes(11));
  ASSERT_EQ("2D560E4940", GetNextSerializedBytes(5));
  ASSERT_EQ("31182D4454FB210940", GetNextSerializedBytes(9));
  ASSERT_EQ("3A0A00000000420142FF4200", GetNextSerializedBytes(12));
  ASSERT_EQ("8A101630313233343536373839616263646566414243444546",
            GetNextSerializedBytes(25));
  ASSERT_EQ("1829", GetNextSerializedBytes(2));
}

TEST_F(MessageTest, NestedMessagesSimple) {
  Message* root_msg = NewMessage();
  root_msg->AppendVarInt(1 /* field_id */, 1);

  FakeChildMessage* nested_msg =
      root_msg->BeginNestedMessage<FakeChildMessage>(128 /* field_id */);
  ASSERT_EQ(0u, reinterpret_cast<uintptr_t>(nested_msg) % sizeof(void*));
  nested_msg->AppendVarInt(2 /* field_id */, 2);

  nested_msg =
      root_msg->BeginNestedMessage<FakeChildMessage>(129 /* field_id */);
  nested_msg->AppendVarInt(4 /* field_id */, 2);

  root_msg->AppendVarInt(5 /* field_id */, 3);

  // The expected size of the root message is supposed to be 14 bytes:
  //   2 bytes for the varint field (id: 1) (1 for preamble and one for payload)
  //   3 bytes for the preamble of the 1st nested message (2 for id, 1 for size)
  //   2 bytes for the varint field (id: 2) of the 1st nested message
  //   3 bytes for the premable of the 2nd nested message
  //   2 bytes for the varint field (id: 4) of the 2nd nested message.
  //   2 bytes for the last varint (id : 5) field of the root message.
  // Test also that finalization is idempontent and Finalize() can be safely
  // called more than once without side effects.
  for (int i = 0; i < 3; ++i) {
    EXPECT_EQ(14u, root_msg->Finalize());
    EXPECT_EQ(14u, GetNumSerializedBytes());
  }

  ASSERT_EQ("0801", GetNextSerializedBytes(2));

  ASSERT_EQ("820802", GetNextSerializedBytes(3));
  ASSERT_EQ("1002", GetNextSerializedBytes(2));

  ASSERT_EQ("8A0802", GetNextSerializedBytes(3));
  ASSERT_EQ("2002", GetNextSerializedBytes(2));

  ASSERT_EQ("2803", GetNextSerializedBytes(2));
}

// Tests using a AppendScatteredBytes to append raw bytes to
// a message using multiple individual buffers.
TEST_F(MessageTest, AppendScatteredBytes) {
  Message* root_msg = NewMessage();

  uint8_t buffer[42];
  memset(buffer, 0x42, sizeof(buffer));

  ContiguousMemoryRange ranges[] = {{buffer, buffer + sizeof(buffer)},
                                    {buffer, buffer + sizeof(buffer)}};
  root_msg->AppendScatteredBytes(1 /* field_id */, ranges, 2);
  EXPECT_EQ(86u, root_msg->Finalize());
  EXPECT_EQ(86u, GetNumSerializedBytes());

  // field_id
  EXPECT_EQ("0A", GetNextSerializedBytes(1));
  // field length
  EXPECT_EQ("54", GetNextSerializedBytes(1));
  // start of contents
  EXPECT_EQ("42424242", GetNextSerializedBytes(4));
}

TEST_F(MessageTest, AppendRawProtoBytesFinalizesNestedMessage) {
  Message* root_msg = NewMessage();

  uint8_t buffer[42];
  memset(buffer, 0x42, sizeof(buffer));

  FakeChildMessage* nested_msg =
      root_msg->BeginNestedMessage<FakeChildMessage>(9001 /* field_id */);
  nested_msg->AppendVarInt(4 /* field_id */, 2);
  uint8_t* nested_msg_size_field = nested_msg->size_field();

  EXPECT_FALSE(nested_msg->is_finalized());
  EXPECT_EQ(0u, *nested_msg_size_field);

  root_msg->AppendRawProtoBytes(buffer, sizeof(buffer));

  // Nested message should have been finalized as a side effect of appending
  // raw bytes.
  EXPECT_EQ(0x2u, *nested_msg_size_field);
}

TEST_F(MessageTest, AppendScatteredBytesFinalizesNestedMessage) {
  Message* root_msg = NewMessage();

  uint8_t buffer[42];
  memset(buffer, 0x42, sizeof(buffer));

  FakeChildMessage* nested_msg =
      root_msg->BeginNestedMessage<FakeChildMessage>(9001 /* field_id */);
  nested_msg->AppendVarInt(4 /* field_id */, 2);
  uint8_t* nested_msg_size_field = nested_msg->size_field();

  EXPECT_FALSE(nested_msg->is_finalized());
  EXPECT_EQ(0u, *nested_msg_size_field);

  ContiguousMemoryRange ranges[] = {{buffer, buffer + sizeof(buffer)}};
  root_msg->AppendScatteredBytes(1 /* field_id */, ranges, 1);

  // Nested message should have been finalized as a side effect of appending
  // scattered bytes.
  EXPECT_EQ(0x2u, *nested_msg_size_field);
}

TEST_F(MessageTest, StressTest) {
  std::vector<Message*> nested_msgs;

  Message* root_msg = NewMessage();
  BuildNestedMessages(root_msg, /*max_depth=*/10);
  root_msg->Finalize();

  // The main point of this test is to stress the code paths and test for
  // unexpected crashes of the production code. The actual serialization is
  // already covered in the other text fixtures. Keeping just a final smoke test
  // here on the full buffer hash.
  std::string full_buf = GetNextSerializedBytes(GetNumSerializedBytes());
  size_t buf_hash = SimpleHash(full_buf);
  EXPECT_EQ(0xf9e32b65, buf_hash);
}

TEST_F(MessageTest, DeeplyNested) {
  Message* root_msg = NewMessage();
  BuildNestedMessages(root_msg, /*max_depth=*/1000);
  root_msg->Finalize();

  std::string full_buf = GetNextSerializedBytes(GetNumSerializedBytes());
  size_t buf_hash = SimpleHash(full_buf);
  EXPECT_EQ(0xc0fde419, buf_hash);
}

TEST_F(MessageTest, DeeplyNestedEmptyMessages) {
  // Stress test writing deeply nested empty messages, many of which will be
  // packed into the protobuf length field.

  // Use a larger chunk size for this test so there is more opportunity to pack
  // messages.
  SetChunkSize(4096u);

  Message* root_msg = NewMessage();
  BuildNestedMessages(root_msg, /*max_depth=*/1000, /*empty=*/true);
  root_msg->Finalize();

  std::string full_buf = GetNextSerializedBytes(GetNumSerializedBytes());
  size_t buf_hash = SimpleHash(full_buf);
  EXPECT_EQ(0x9371fe8eu, buf_hash);
}

TEST_F(MessageTest, DestructInvalidMessageHandle) {
  FakeRootMessage* msg = NewMessage();
  EXPECT_DCHECK_DEATH({
    MessageHandle<FakeRootMessage> handle(msg);
    ResetMessage(msg);
  });
}

TEST_F(MessageTest, MessageHandle) {
  FakeRootMessage* msg3 = NewMessageWithSizeField();
  FakeRootMessage* msg2 = NewMessageWithSizeField();
  FakeRootMessage* msg1 = NewMessageWithSizeField();
  FakeRootMessage* ignored_msg = NewMessage();
  uint8_t* msg1_size = msg1->size_field();
  uint8_t* msg2_size = msg2->size_field();
  uint8_t* msg3_size = msg3->size_field();

  // Test that the handle going out of scope causes the finalization of the
  // target message and triggers the optional callback.
  {
    MessageHandle<FakeRootMessage> handle1(msg1);
    handle1->AppendBytes(1 /* field_id */, kTestBytes, 1 /* size */);
    ASSERT_EQ(0u, msg1_size[0]);
  }
  ASSERT_EQ(0x3u, msg1_size[0]);

  // Test that the handle can be late initialized.
  MessageHandle<FakeRootMessage> handle2(ignored_msg);
  handle2 = MessageHandle<FakeRootMessage>(msg2);
  handle2->AppendBytes(1 /* field_id */, kTestBytes, 2 /* size */);
  ASSERT_EQ(0u, msg2_size[0]);  // |msg2| should not be finalized yet.

  // Test that std::move works and does NOT cause finalization of the moved
  // message.
  MessageHandle<FakeRootMessage> handle_swp(ignored_msg);
  handle_swp = std::move(handle2);
  ASSERT_EQ(0u, msg2_size[0]);  // msg2 should be NOT finalized yet.
  handle_swp->AppendBytes(2 /* field_id */, kTestBytes, 3 /* size */);

  MessageHandle<FakeRootMessage> handle3(msg3);
  handle3->AppendBytes(1 /* field_id */, kTestBytes, 4 /* size */);
  ASSERT_EQ(0u, msg3_size[0]);  // msg2 should be NOT finalized yet.

  // Both |handle3| and |handle_swp| point to a valid message (respectively,
  // |msg3| and |msg2|). Now move |handle3| into |handle_swp|.
  handle_swp = std::move(handle3);
  ASSERT_EQ(0x89u, msg2_size[0]);  // |msg2| should be finalized at this point.

  // At this point writing into handle_swp should actually write into |msg3|.
  ASSERT_EQ(msg3, &*handle_swp);
  handle_swp->AppendBytes(2 /* field_id */, kTestBytes, 8 /* size */);
  MessageHandle<FakeRootMessage> another_handle(ignored_msg);
  handle_swp = std::move(another_handle);
  ASSERT_EQ(0x90u, msg3_size[0]);  // |msg3| should be finalized at this point.

#if PERFETTO_DCHECK_IS_ON()
  // In developer builds w/ PERFETTO_DCHECK on a finalized message should
  // invalidate the handle, in order to early catch bugs in the client code.
  FakeRootMessage* msg4 = NewMessage();
  MessageHandle<FakeRootMessage> handle4(msg4);
  ASSERT_EQ(msg4, &*handle4);
  msg4->Finalize();
  ASSERT_THAT(handle4.get(), testing::IsNull());
#endif

  // Test also the behavior of handle with non-root (nested) messages.

  uint8_t* size_msg_2;
  {
    auto* nested_msg_1 = NewMessage()->BeginNestedMessage<FakeChildMessage>(3);
    MessageHandle<FakeChildMessage> child_handle_1(nested_msg_1);
    uint8_t* size_msg_1 = nested_msg_1->size_field();
    memset(size_msg_1, 0, proto_utils::kMessageLengthFieldSize);
    child_handle_1->AppendVarInt(1, 0x11);

    auto* nested_msg_2 = NewMessage()->BeginNestedMessage<FakeChildMessage>(2);
    size_msg_2 = nested_msg_2->size_field();
    memset(size_msg_2, 0, proto_utils::kMessageLengthFieldSize);
    MessageHandle<FakeChildMessage> child_handle_2(nested_msg_2);
    child_handle_2->AppendVarInt(2, 0xFF);

    // |nested_msg_1| should not be finalized yet.
    ASSERT_EQ(0u, size_msg_1[0]);

    // This move should cause |nested_msg_1| to be finalized, but not
    // |nested_msg_2|, which will be finalized only after the current scope.
    child_handle_1 = std::move(child_handle_2);
    ASSERT_EQ(0x82u, size_msg_1[0]);
    ASSERT_EQ(0u, size_msg_2[0]);
  }
  ASSERT_EQ(0x3u, size_msg_2[0]);
}

TEST_F(MessageTest, MoveMessageHandle) {
  FakeRootMessage* msg = NewMessageWithSizeField();
  uint8_t* msg_size = msg->size_field();

  // Test that the handle going out of scope causes the finalization of the
  // target message.
  {
    MessageHandle<FakeRootMessage> handle1(msg);
    MessageHandle<FakeRootMessage> handle2{};
    handle1->AppendBytes(1 /* field_id */, kTestBytes, 1 /* size */);
    handle2 = std::move(handle1);
    ASSERT_EQ(0u, msg_size[0]);
  }
  ASSERT_EQ(0x3u, msg_size[0]);
}

TEST_F(MessageTest, FinalizeWithCompaction) {
  FakeRootMessage* msg = NewMessageWithSizeField();

  msg->AppendBytes(1 /* field_id */, kTestBytes, 5 /* size */);
  uint32_t size = msg->Finalize();
  EXPECT_EQ(7u, size);
  EXPECT_EQ(8u, GetNumSerializedBytes());
}

TEST_F(MessageTest, FinalizeWithoutCompaction) {
  FakeRootMessage* msg = NewMessageWithSizeField();

  // This message doesn't fit into a single chunk, so it won't be compacted.
  msg->AppendBytes(1 /* field_id */, kTestBytes, sizeof(kTestBytes) /* size */);
  msg->AppendBytes(1 /* field_id */, kTestBytes, sizeof(kTestBytes) /* size */);
  uint32_t size = msg->Finalize();
  EXPECT_EQ(24u, size);
  EXPECT_EQ(28u, GetNumSerializedBytes());
}

TEST_F(MessageTest, ProtoGroupGeneratedMessages) {
  RootMessage<protozero::test::protos::pbzero::EveryField> root;
  root.Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  auto* outer = root.add_field_nested();
  auto* inner = outer->add_field_nested();
  auto* leaf = inner->add_field_nested();
  leaf->set_field_int32(150);
  root.set_field_int32(42);
  EXPECT_EQ(11u, root.Finalize());
  EXPECT_EQ(11u, root.Finalize());
  EXPECT_EQ("737373089601040404082A", GetNextSerializedBytes(11));
}

TEST_F(MessageTest, ProtoGroupEmptyAndReset) {
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  EXPECT_EQ(0u, root->Finalize());
  EXPECT_EQ(0u, root->Finalize());
  EXPECT_EQ(0u, GetNumSerializedBytes());
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  root->BeginNestedMessage<Message>(1);
  EXPECT_EQ(2u, root->Finalize());
  root->Reset(stream_writer_.get());
  root->BeginNestedMessage<Message>(1);
  EXPECT_EQ(2u, root->Finalize());
  EXPECT_EQ("0B040A00", GetNextSerializedBytes(4));
}

TEST_F(MessageTest, LengthDelimitedRootWithErasedType) {
  Message* root = NewMessage();
  root->BeginNestedMessage<Message>(1)->AppendVarInt(1, 1);
  EXPECT_EQ(root->Finalize(), 4u);
  EXPECT_EQ(root->Finalize(), 4u);
  EXPECT_EQ("0A020801", GetNextSerializedBytes(4));
}

TEST_F(MessageTest, ProtoGroupHighFieldIdAndSiblings) {
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  auto* child = root->BeginNestedMessage<Message>(100);
  child->AppendVarInt(1, 0);

  // The child's size excludes its closing byte. The parent writes that byte.
  EXPECT_EQ(2u, child->Finalize());
  EXPECT_EQ(2u, child->Finalize());

  root->BeginNestedMessage<Message>(1)->AppendVarInt(2, 150);
  root->AppendVarInt(2, 42);
  EXPECT_EQ(12u, root->Finalize());
  EXPECT_EQ("A3060800040B10960104102A", GetNextSerializedBytes(12));
}

TEST_F(MessageTest, ProtoGroupCloseAtChunkBoundary) {
  SetChunkSize(4);
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  auto* child = root->BeginNestedMessage<Message>(1);
  child->AppendVarInt(2, 150);
  ASSERT_EQ(4u, GetNumSerializedBytes());
  EXPECT_EQ(nullptr, child->size_field());

  // The closing byte goes into the next chunk. The first chunk is unchanged.
  const auto published = buffer_->GetChunkAsString(0);
  EXPECT_EQ(5u, root->Finalize());
  EXPECT_EQ(published, buffer_->GetChunkAsString(0));
  EXPECT_EQ("0B10960104", GetNextSerializedBytes(5));
}

TEST_F(MessageTest, ProtoGroupBytesAcrossChunks) {
  SetChunkSize(4);
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  auto* child = root->BeginNestedMessage<Message>(1);
  child->AppendBytes(2, kTestBytes, sizeof(kTestBytes));
  EXPECT_EQ(14u, root->Finalize());
  EXPECT_EQ("0B120A00000000420142FF420004", GetNextSerializedBytes(14));
}

// Roots and nested messages share one Finalize(). A root gets no closing byte,
// even through a plain Message*, because it has no parent.
TEST_F(MessageTest, ProtoGroupRootThroughMessagePointerHasNoClosingByte) {
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  root->BeginNestedMessage<Message>(1)->AppendVarInt(2, 7);
  Message* as_plain_message = root;
  EXPECT_EQ(4u, as_plain_message->Finalize());
  EXPECT_EQ(4u, GetNumSerializedBytes());
  EXPECT_EQ("0B100704", GetNextSerializedBytes(4));
}

// A child finalized early, even twice, gets exactly one closing byte.
// The parent writes it right after the child, before its own next field.
TEST_F(MessageTest, ProtoGroupEarlyChildFinalizeClosesOnce) {
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  auto* child = root->BeginNestedMessage<Message>(1);
  child->AppendVarInt(2, 7);
  EXPECT_EQ(2u, child->Finalize());
  EXPECT_EQ(2u, child->Finalize());

  // The parent has not written the closing byte yet.
  EXPECT_EQ(3u, GetNumSerializedBytes());

  root->AppendVarInt(3, 1);
  EXPECT_EQ(6u, root->Finalize());
  EXPECT_EQ(6u, root->Finalize());
  EXPECT_EQ("0B1007041801", GetNextSerializedBytes(6));
}

// A handle on a nested message finalizes the child.
// The closing byte still comes from the parent, once.
TEST_F(MessageTest, ProtoGroupNestedHandleClosesInParent) {
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  {
    MessageHandle<Message> handle(root->BeginNestedMessage<Message>(1));
    handle->AppendVarInt(2, 7);
  }
  EXPECT_EQ(3u, GetNumSerializedBytes());
  EXPECT_EQ(4u, root->Finalize());
  EXPECT_EQ("0B100704", GetNextSerializedBytes(4));
}

// Deep nesting closes innermost first, one byte per level.
TEST_F(MessageTest, ProtoGroupDeepNesting) {
  constexpr size_t kDepth = 100;
  SetChunkSize(4096u);
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  Message* msg = root;
  for (size_t i = 0; i < kDepth; ++i)
    msg = msg->BeginNestedMessage<Message>(1);
  EXPECT_EQ(2 * kDepth, root->Finalize());

  std::string expected;
  for (size_t i = 0; i < kDepth; ++i)
    expected += "0B";
  for (size_t i = 0; i < kDepth; ++i)
    expected += "04";
  EXPECT_EQ(expected, GetNextSerializedBytes(2 * kDepth));
}

// The encoding does not leak into a later length-delimited Reset(), and a
// length-delimited child still gets its length field.
TEST_F(MessageTest, LengthDelimitedAfterProtoGroupReset) {
  auto* root = NewMessage();
  root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
  root->BeginNestedMessage<Message>(1)->AppendVarInt(2, 7);
  EXPECT_EQ(4u, root->Finalize());

  root->Reset(stream_writer_.get());
  EXPECT_EQ(Message::Encoding::kLengthDelimited, root->encoding());
  root->BeginNestedMessage<Message>(1)->AppendVarInt(2, 7);
  EXPECT_EQ(4u, root->Finalize());

  EXPECT_EQ("0B100704", GetNextSerializedBytes(4));
  EXPECT_EQ("0A021007", GetNextSerializedBytes(4));
}

// A random message tree: varint fields and nested messages.
struct TreeNode {
  struct Item {
    uint32_t field_id = 0;
    uint64_t value = 0;                // Varint fields only.
    std::unique_ptr<TreeNode> nested;  // Set for a nested message.
  };
  std::vector<Item> items;
};

std::unique_ptr<TreeNode> RandomTree(std::minstd_rand* rnd, int depth) {
  auto node = std::make_unique<TreeNode>();
  const uint32_t num_items = (*rnd)() % 4;
  for (uint32_t i = 0; i < num_items; ++i) {
    TreeNode::Item item;
    item.field_id = 1 + (*rnd)() % 300;  // One- and two-byte tags.
    if (depth > 0 && (*rnd)() % 2)
      item.nested = RandomTree(rnd, depth - 1);
    else
      item.value = (*rnd)();
    node->items.push_back(std::move(item));
  }
  return node;
}

// The reference proto group encoding: a nested message is its start tag, its
// content and a closing byte. The root has no framing.
void EncodeReference(const TreeNode& node, std::vector<uint8_t>* out) {
  uint8_t varint[proto_utils::kMaxSimpleFieldEncodedSize];
  for (const TreeNode::Item& item : node.items) {
    if (item.nested) {
      uint8_t* end = proto_utils::WriteVarInt(
          proto_utils::MakeTagStartGroup(item.field_id), varint);
      out->insert(out->end(), varint, end);
      EncodeReference(*item.nested, out);
      out->push_back(proto_utils::kProtoGroupEndByte);
    } else {
      uint8_t* end = proto_utils::WriteVarInt(
          proto_utils::MakeTagVarInt(item.field_id), varint);
      end = proto_utils::WriteVarInt(item.value, end);
      out->insert(out->end(), varint, end);
    }
  }
}

// Writes |node| through protozero. Each nested message is finalized early at
// random, the way a nested MessageHandle finalizes it.
void WriteTree(const TreeNode& node, Message* msg, std::minstd_rand* rnd) {
  for (const TreeNode::Item& item : node.items) {
    if (item.nested) {
      Message* child = msg->BeginNestedMessage<Message>(item.field_id);
      WriteTree(*item.nested, child, rnd);
      if ((*rnd)() % 2)
        child->Finalize();
    } else {
      msg->AppendVarInt(item.field_id, item.value);
    }
  }
}

// The closing byte lands in the right place for any tree, any early
// finalization and any chunk size, including closes at a chunk boundary.
TEST_F(MessageTest, ProtoGroupRandomTreesMatchReferenceEncoding) {
  std::minstd_rand rnd(0x2545f491);
  for (int i = 0; i < 500; ++i) {
    SCOPED_TRACE(i);
    SetChunkSize(4 + static_cast<size_t>(i % 29));
    const std::unique_ptr<TreeNode> tree = RandomTree(&rnd, /*depth=*/5);
    std::vector<uint8_t> expected;
    EncodeReference(*tree, &expected);

    auto* root = NewMessage();
    root->Reset(stream_writer_.get(), Message::Encoding::kProtoGroup);
    WriteTree(*tree, root, &rnd);
    ASSERT_EQ(expected.size(), root->Finalize());
    ASSERT_EQ(expected.size(), GetNumSerializedBytes());

    const char* expected_data = reinterpret_cast<const char*>(expected.data());
    const std::string expected_hex = perfetto::base::ToUpper(
        perfetto::base::ToHex(expected_data, expected.size()));
    ASSERT_EQ(expected_hex, GetNextSerializedBytes(expected.size()));
  }
}

}  // namespace
}  // namespace protozero