diff --git a/include/protozero/basic_pbf_writer.hpp b/include/protozero/basic_pbf_writer.hpp index 9e0af4ea..266e82ac 100644 --- a/include/protozero/basic_pbf_writer.hpp +++ b/include/protozero/basic_pbf_writer.hpp @@ -75,6 +75,8 @@ class basic_pbf_writer { // This is usually 0. If there is an open submessage, this is set in the // parent to the position where the data of the submessage is written to. + // If the size of the submessage was known when it was opened, this is + // set to the position where the data of the submessage must end instead. std::size_t m_pos = 0; void add_varint(uint64_t value) { @@ -179,12 +181,13 @@ class basic_pbf_writer { m_rollback_pos = buffer_customization::size(m_data); add_field(tag, pbf_wire_type::length_delimited); buffer_customization::append_zeros(m_data, std::size_t(reserve_bytes)); + m_pos = buffer_customization::size(m_data); } else { m_rollback_pos = size_is_known; add_length_varint(tag, static_cast(size)); reserve(size); + m_pos = buffer_customization::size(m_data) + size; } - m_pos = buffer_customization::size(m_data); } void rollback_submessage() { @@ -210,7 +213,14 @@ class basic_pbf_writer { void close_submessage() { protozero_assert(m_data); - if (m_pos == 0 || m_rollback_pos == size_is_known) { + if (m_pos == 0) { + return; + } + if (m_rollback_pos == size_is_known) { + const auto end_pos = m_pos; + m_pos = 0; + protozero_assert(buffer_customization::size(m_data) == end_pos && "submessage with known size must have exactly that many bytes written"); + (void)end_pos; return; } if (buffer_customization::size(m_data) - m_pos == 0) { diff --git a/test/t/repeated_packed_fixed32/reader_test_cases.cpp b/test/t/repeated_packed_fixed32/reader_test_cases.cpp index fc2a179d..6bfdda0d 100644 --- a/test/t/repeated_packed_fixed32/reader_test_cases.cpp +++ b/test/t/repeated_packed_fixed32/reader_test_cases.cpp @@ -7,3 +7,56 @@ using cpp_type = uint32_t; #include + +TEST_CASE("write repeated packed fixed32 field using packed field with known size") { + std::string buffer; + protozero::pbf_writer pw{buffer}; + + SECTION("correct number of elements") { + { + protozero::packed_field_fixed32 field{pw, 1, 2}; + field.add_element(17); + field.add_element(1); + } + pw.add_uint32(2, 5); + + REQUIRE(buffer == std::string{"\x0a\x08\x11\x00\x00\x00\x01\x00\x00\x00\x10\x05", 12}); + } + + SECTION("correct number of elements with commit") { + protozero::packed_field_fixed32 field{pw, 1, 2}; + field.add_element(17); + field.add_element(1); + field.commit(); + pw.add_uint32(2, 5); + + REQUIRE(buffer == std::string{"\x0a\x08\x11\x00\x00\x00\x01\x00\x00\x00\x10\x05", 12}); + } + + SECTION("too few elements") { + protozero::packed_field_fixed32 field{pw, 1, 2}; + field.add_element(17); + REQUIRE_THROWS_AS(field.commit(), assert_error); + } + + SECTION("too many elements") { + protozero::packed_field_fixed32 field{pw, 1, 2}; + field.add_element(17); + field.add_element(1); + field.add_element(2); + REQUIRE_THROWS_AS(field.commit(), assert_error); + } +} + +TEST_CASE("write submessage with known size") { + std::string buffer; + protozero::pbf_writer pw{buffer}; + + { + protozero::pbf_writer sw{pw, 1, 2}; + sw.add_uint32(1, 5); + } + pw.add_uint32(2, 5); + + REQUIRE(buffer == std::string{"\x0a\x02\x08\x05\x10\x05", 6}); +}