allow setting expected size

fandreuz · fandreuz · commit 84d0ce1c178d · 2025-06-10T20:04:32.000Z
diff --git a/src/protobuf.cpp b/src/protobuf.cpp
@@ -69,20 +69,35 @@ void ProtoBuffer::field(protobuf_index_t index, const unsigned char* s, size_t l
     _offset += len;
 }
 
-protobuf_mark_t ProtoBuffer::startMessage(protobuf_index_t index) {
+protobuf_mark_t ProtoBuffer::startMessage(protobuf_index_t index, size_t len_byte_count) {
     tag(index, LEN);
 
-    ensureCapacity(NESTED_FIELD_BYTE_COUNT);
-    _offset += NESTED_FIELD_BYTE_COUNT;
-    return _offset;
+    ensureCapacity(len_byte_count);
+
+    protobuf_mark_t mark{_offset, len_byte_count};
+    _offset += len_byte_count;
+    return mark;
 }
 
-void ProtoBuffer::commitMessage(protobuf_mark_t mark) {
-    size_t message_length = _offset - mark;
-    for (size_t i = 0; i < NESTED_FIELD_BYTE_COUNT - 1; ++i) {
-        size_t idx = mark - NESTED_FIELD_BYTE_COUNT + i;
-        _data[idx] = (unsigned char) (0x80 | (message_length & 0x7f));
-        message_length >>= 7;
+void ProtoBuffer::commitMessage(const protobuf_mark_t& mark) {
+    size_t actual_len = _offset - (mark.message_start + mark.expected_len_byte_count);
+    size_t actual_len_byte_count = varIntSize(actual_len);
+    if (actual_len_byte_count > mark.expected_len_byte_count) {
+        _offset = mark.message_start;
+        // The tag is already accounted for before mark.message_start
+        ensureCapacity(actual_len_byte_count + actual_len);
+        memmove(_data + mark.message_start + actual_len_byte_count,
+                _data + mark.message_start + mark.expected_len_byte_count,
+                actual_len);
+        _offset = mark.message_start + actual_len_byte_count + actual_len;
+    } else {
+        actual_len_byte_count = mark.expected_len_byte_count;
+    }
+
+    for (size_t i = 0; i < actual_len_byte_count - 1; ++i) {
+        size_t idx = mark.message_start + i;
+        _data[idx] = (unsigned char) (0x80 | (actual_len & 0x7f));
+        actual_len >>= 7;
     }
-    _data[mark - 1] = (unsigned char) message_length;
+    _data[mark.message_start + actual_len_byte_count - 1] = (unsigned char) actual_len;
 }
diff --git a/src/protobuf.h b/src/protobuf.h
@@ -14,7 +14,10 @@ protobuf_t VARINT = 0;
 protobuf_t LEN = 2;
 
 typedef u32 protobuf_index_t;
-typedef u32 protobuf_mark_t;
+struct protobuf_mark_t {
+    size_t message_start;
+    size_t expected_len_byte_count;
+};
 
 // We assume the length of a nested field can be represented with 3 varint bytes.
 const size_t NESTED_FIELD_BYTE_COUNT = 3;
@@ -46,8 +49,8 @@ class ProtoBuffer {
     void field(protobuf_index_t index, const char* s, size_t len);
     void field(protobuf_index_t index, const unsigned char* s, size_t len);
 
-    protobuf_mark_t startMessage(protobuf_index_t index);
-    void commitMessage(protobuf_mark_t mark);
+    protobuf_mark_t startMessage(protobuf_index_t index, size_t len_byte_count = NESTED_FIELD_BYTE_COUNT);
+    void commitMessage(const protobuf_mark_t& mark);
 
     void putVarInt(u64 n);
     static size_t varIntSize(u64 value);
diff --git a/test/native/protobufBufferTest.cpp b/test/native/protobufBufferTest.cpp
@@ -7,42 +7,44 @@
 #include "testRunner.hpp"
 #include <string.h>
 
-TEST_CASE(Buffer_test_var32_0) {
+static u64 readVarInt(const unsigned char* data) {
+    size_t idx = 0;
+    u64 output = 0;
+    while (data[idx] & 0x80) {
+        output += ((u64) data[idx] & 0x7F) << (7 * idx);
+        idx++;
+    }
+    return output + ((u64) data[idx] << (7 * idx));
+}
+
+TEST_CASE(Buffer_test_varint_0) {
     ProtoBuffer buf(100);
 
     buf.field(3, (u64)0);
 
     CHECK_EQ(buf.offset(), 2);
     CHECK_EQ(buf.data()[0], (3 << 3) | VARINT);
-    CHECK_EQ(buf.data()[1], 0);
+    CHECK_EQ(readVarInt(buf.data() + 1), 0);
 }
 
-TEST_CASE(Buffer_test_var32_150) {
+TEST_CASE(Buffer_test_varint_150) {
     ProtoBuffer buf(100);
 
     buf.field(3, (u64)150);
 
     CHECK_EQ(buf.offset(), 3);
     CHECK_EQ(buf.data()[0], (3 << 3) | VARINT);
-    CHECK_EQ(buf.data()[1], 150);
-    CHECK_EQ(buf.data()[2], 1);
+    CHECK_EQ(readVarInt(buf.data() + 1), 150);
 }
 
-TEST_CASE(Buffer_test_var64_LargeNumber) {
+TEST_CASE(Buffer_test_varint_LargeNumber) {
     ProtoBuffer buf(100);
 
     buf.field(1, (u64)17574838338834838);
 
     CHECK_EQ(buf.offset(), 9);
     CHECK_EQ(buf.data()[0], (1 << 3) | VARINT);
-    CHECK_EQ(buf.data()[1], 150);
-    CHECK_EQ(buf.data()[2], 163);
-    CHECK_EQ(buf.data()[3], 175);
-    CHECK_EQ(buf.data()[4], 225);
-    CHECK_EQ(buf.data()[5], 142);
-    CHECK_EQ(buf.data()[6], 135);
-    CHECK_EQ(buf.data()[7], 156);
-    CHECK_EQ(buf.data()[8], 31);
+    CHECK_EQ(readVarInt(buf.data() + 1), 17574838338834838);
 }
 
 TEST_CASE(Buffer_test_var32_bool) {
@@ -127,11 +129,49 @@ TEST_CASE(Buffer_test_nestedMessageWithString) {
     CHECK_EQ(strncmp((const char*) buf.data() + 10, "ciao", 4), 0);
 }
 
+TEST_CASE(Buffer_test_wrongExpectedLenByteCount) {
+    ProtoBuffer buf(100);
+
+    protobuf_mark_t mark = buf.startMessage(3, 1);
+    unsigned char s[128];
+    for (unsigned char i = 0; i < 128; ++i) s[i] = i;
+    buf.field(5, s, 128);
+    buf.commitMessage(mark);
+
+    CHECK_EQ(buf.offset(), 1 + 2 + 1 + 2 + 128);
+    CHECK_EQ(buf.data()[0], (3 << 3) | LEN);
+    CHECK_EQ(readVarInt(buf.data() + 1), 1 + 2 + 128);
+    CHECK_EQ(buf.data()[3], (5 << 3) | LEN);
+    CHECK_EQ(readVarInt(buf.data() + 4), 128);
+    for (unsigned char i = 0; i < 128; ++i) CHECK_EQ(buf.data()[i + 6], i);
+}
+
+TEST_CASE(Buffer_test_wrongExpectedLenByteCountWithNestedMessage) {
+    ProtoBuffer buf(100);
+
+    protobuf_mark_t mark1 = buf.startMessage(3, 1);
+    protobuf_mark_t mark2 = buf.startMessage(4, 1);
+    unsigned char s[128];
+    for (unsigned char i = 0; i < 128; ++i) s[i] = i;
+    buf.field(5, s, 128);
+    buf.commitMessage(mark2);
+    buf.commitMessage(mark1);
+
+    CHECK_EQ(buf.offset(), 1 + 2 + 1 + 2 + 1 + 2 + 128);
+    CHECK_EQ(buf.data()[0], (3 << 3) | LEN);
+    CHECK_EQ(readVarInt(buf.data() + 1), 1 + 2 + 1 + 2 + 128);
+    CHECK_EQ(buf.data()[3], (4 << 3) | LEN);
+    CHECK_EQ(readVarInt(buf.data() + 4), 1 + 2 + 128);
+    CHECK_EQ(buf.data()[6], (5 << 3) | LEN);
+    CHECK_EQ(readVarInt(buf.data() + 7), 128);
+    for (unsigned char i = 0; i < 128; ++i) CHECK_EQ(buf.data()[i + 9], i);
+}
+
 TEST_CASE(Buffer_test_maxTag) {
     ProtoBuffer buf(100);
 
     // https://protobuf.dev/programming-guides/proto3/#assigning-field-numbers
-    const protobuf_mark_t max_tag = 536870911;
+    const protobuf_index_t max_tag = 536870911;
     buf.field(max_tag, (u64) 3);
 
     CHECK_EQ(buf.offset(), 6);
