Add support for SGROUP and EGROUP (#13)

* Added support for SGROUP and EGROUP wire types * Added group number tracking for fields sandwiched between SGROUP and EGROUP
Kardbord · Sep 22, 2024 · 4a9c6a4 · 4a9c6a4
1 parent 4e629dc
commit 4a9c6a4
Show file tree

Hide file tree

Showing 5 changed files with 139 additions and 8 deletions.
diff --git a/analyzer/protobuf.spicy b/analyzer/protobuf.spicy
@@ -18,14 +18,31 @@ type WireType = enum {
     VARINT = 0x00,
     I64    = 0x01,
     LEN    = 0x02,
-    SGROUP = 0x03, # deprecated
-    EGROUP = 0x04, # deprecated
+    SGROUP = 0x03, # deprecated, but remains in the wire format.
+    EGROUP = 0x04, # deprecated, but remains in the wire format.
     I32    = 0x05
 };
 
+const UNSET_GROUP: int32 = -1;
+global g_current_group: int32 = UNSET_GROUP;
+
 type TagAndValue = unit {
     tag:   Tag;
     value: Value(self.tag);
+
+    var group: int32 &optional;
+
+    on %done {
+        if (self.value.sgroup) {
+            g_current_group = int32(self.tag.field_num);
+        }
+        if (g_current_group != UNSET_GROUP) {
+            self.group = g_current_group;
+        }
+        if (self.value.egroup) {
+            g_current_group = UNSET_GROUP;
+        }
+    }
 };
 
 # tag := (field << 3) | wire_type
@@ -63,10 +80,13 @@ type Value = unit(tag: Tag) {
         WireType::I32    -> i32    : I32;
         WireType::I64    -> i64    : I64;
         WireType::LEN    -> unimpl : bytes &eod;
-        WireType::SGROUP -> unimpl : bytes &eod;
-        WireType::EGROUP -> unimpl : bytes &eod;
+        WireType::SGROUP ->        : void { self.sgroup = True; } # Group fields are just a tag with no payload.
+        WireType::EGROUP ->        : void { self.egroup = True; } # Group fields are just a tag with no payload.
         # TODO - Finish switch for other types
     };
+
+    var sgroup: bool = False;
+    var egroup: bool = False;
 };
 
 function parse_varint(buf: vector<uint8>) : uint64 {

diff --git a/testing/gtest/SpicyProtobufTest.cpp b/testing/gtest/SpicyProtobufTest.cpp
@@ -17,6 +17,7 @@
 #include <test-data/varint/protobuf_neg_150.h>
 #include <test-data/varint/protobuf_neg_2.h>
 #include <test-data/varint/protobuf_zigzag_150.h>
+#include <test-data/groups/protobuf_groups.h>
 
 namespace { // anonymous namespace
 
@@ -29,7 +30,15 @@ class SpicyProtobufTest : public ::testing::Test {
   static auto SetUpTestSuite() -> void;
   static auto TearDownTestSuite() -> void;
 
-  static auto parseMessage(const char* data, uint64_t size) -> ValueReference<Message>;
+  // Flags to pass to the parseMessage function.
+  // All values should be powers of 2.
+  enum ParseFlag {
+    NONE         = 0,
+    CHECK_GROUPS = 1, // Check that sgroup and egroup are false
+  };
+
+  static auto parseFlagIsSet(uint32_t flags, ParseFlag flag_to_check) -> bool;
+  static auto parseMessage(const char* data, uint64_t size, uint32_t flags = ParseFlag::CHECK_GROUPS) -> ValueReference<Message>;
 };
 
 auto SpicyProtobufTest::SetUpTestSuite() -> void {
@@ -44,14 +53,24 @@ auto SpicyProtobufTest::TearDownTestSuite() -> void {
   ::hilti::rt::done();
 }
 
-auto SpicyProtobufTest::parseMessage(const char* p_data, uint64_t const size) -> ValueReference<Message> {
+auto SpicyProtobufTest::parseFlagIsSet(uint32_t const flags, ParseFlag const flag_to_check) -> bool {
+  return (flags & flag_to_check) > 0;
+}
+
+auto SpicyProtobufTest::parseMessage(const char* p_data, uint64_t const size, uint32_t const flags) -> ValueReference<Message> {
   using namespace ::hilti::rt::reference;
   using ::hilti::rt::Stream;
 
   auto p_input_stream = make_value<Stream>(p_data, size);
   p_input_stream->freeze();
   auto p_msg = make_value<Message>();
   ::hlt::protobuf::Message::parse2(p_msg, p_input_stream, {}, {});
+  if (parseFlagIsSet(flags, ParseFlag::CHECK_GROUPS)) {
+    for (auto const &tag_and_val: *p_msg->message) {
+      EXPECT_FALSE(tag_and_val.value.value()->sgroup);
+      EXPECT_FALSE(tag_and_val.value.value()->egroup);
+    }
+  }
   return p_msg;
 }
 
@@ -840,4 +859,88 @@ TEST_F(SpicyProtobufTest, TestI64) {
 
 }
 
+TEST_F(SpicyProtobufTest, TestGroups) {
+  using namespace __hlt::protobuf;
+
+  // protobuf_groups
+  auto p_data = parseMessage(reinterpret_cast<const char *>(groups_protobuf_groups_binpb), groups_protobuf_groups_binpb_len, ParseFlag::NONE);
+  auto p_msg = p_data->message;
+  ASSERT_EQ(p_msg->size(), 6);
+
+  { // Field 1
+    auto tag_and_val = p_data->message->at(0);
+    EXPECT_EQ(tag_and_val.tag.value()->field_num, 1);
+    EXPECT_EQ(tag_and_val.tag.value()->wire_type.value(), WireType::VARINT);
+    EXPECT_FALSE(tag_and_val.value.value()->sgroup);
+    EXPECT_FALSE(tag_and_val.value.value()->egroup);
+    ASSERT_FALSE(tag_and_val.value.value()->varint->isNull());
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_unsigned, 1);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_zigzag, -1);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_twos_compliment64, 1);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_twos_compliment32, 1);
+    EXPECT_FALSE(tag_and_val.group.has_value());
+  }
+
+  int32_t group_num = -1;
+  { // Field 2, SGROUP
+    auto tag_and_val = p_data->message->at(1);
+    EXPECT_EQ(tag_and_val.tag.value()->field_num, 2);
+    group_num = tag_and_val.tag.value()->field_num;
+    EXPECT_EQ(tag_and_val.tag.value()->wire_type.value(), WireType::SGROUP);
+    EXPECT_TRUE(tag_and_val.value.value()->sgroup);
+    EXPECT_FALSE(tag_and_val.value.value()->egroup);
+    EXPECT_EQ(tag_and_val.group, group_num);
+  }
+
+  { // Field 3
+    auto tag_and_val = p_data->message->at(2);
+    EXPECT_EQ(tag_and_val.tag.value()->field_num, 3);
+    EXPECT_EQ(tag_and_val.tag.value()->wire_type.value(), WireType::VARINT);
+    EXPECT_FALSE(tag_and_val.value.value()->sgroup);
+    EXPECT_FALSE(tag_and_val.value.value()->egroup);
+    ASSERT_FALSE(tag_and_val.value.value()->varint->isNull());
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_unsigned, 2);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_zigzag, 1);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_twos_compliment64, 2);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_twos_compliment32, 2);
+    EXPECT_EQ(tag_and_val.group, group_num);
+  }
+
+  { // Field 4
+    auto tag_and_val = p_data->message->at(3);
+    EXPECT_EQ(tag_and_val.tag.value()->field_num, 4);
+    EXPECT_EQ(tag_and_val.tag.value()->wire_type.value(), WireType::I32);
+    EXPECT_FALSE(tag_and_val.value.value()->sgroup);
+    EXPECT_FALSE(tag_and_val.value.value()->egroup);
+    ASSERT_FALSE(tag_and_val.value.value()->i32->isNull());
+    EXPECT_EQ(tag_and_val.value.value()->i32.value()->as_unsigned, 1078530011);
+    EXPECT_EQ(tag_and_val.value.value()->i32.value()->as_twos_compliment, 1078530011);
+    EXPECT_FLOAT_EQ(tag_and_val.value.value()->i32.value()->as_float, 3.1415927);
+    EXPECT_EQ(tag_and_val.group, group_num);
+  }
+
+  { // Field 2, EGROUP
+    auto tag_and_val = p_data->message->at(4);
+    EXPECT_EQ(tag_and_val.tag.value()->field_num, group_num);
+    EXPECT_EQ(tag_and_val.tag.value()->wire_type.value(), WireType::EGROUP);
+    EXPECT_FALSE(tag_and_val.value.value()->sgroup);
+    EXPECT_TRUE(tag_and_val.value.value()->egroup);
+    EXPECT_EQ(tag_and_val.group, group_num);
+  }
+
+  { // Field 5
+    auto tag_and_val = p_data->message->at(5);
+    EXPECT_EQ(tag_and_val.tag.value()->field_num, 5);
+    EXPECT_EQ(tag_and_val.tag.value()->wire_type.value(), WireType::VARINT);
+    EXPECT_FALSE(tag_and_val.value.value()->sgroup);
+    EXPECT_FALSE(tag_and_val.value.value()->egroup);
+    ASSERT_FALSE(tag_and_val.value.value()->varint->isNull());
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_unsigned, 3);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_zigzag, -2);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_twos_compliment64, 3);
+    EXPECT_EQ(tag_and_val.value.value()->varint.value()->as_twos_compliment32, 3);
+    EXPECT_FALSE(tag_and_val.group.has_value());
+  }
+}
+
 } // anonymous namespace
diff --git a/testing/gtest/test-data/README.md b/testing/gtest/test-data/README.md
@@ -9,12 +9,13 @@ See the `Protoscope` installation instructions located at the reference below.
 Command:
 
 ```bash
+# Generate a .binpb from a .txtpb
 protoscope -s <test-file.txtpb> > <test-file.binpb>
 ```
 
-Example text based format:
+Example text based format (.txtpb):
 
-```bash
+```txt
 field_num: value
 1: 150
 ```

diff --git a/testing/gtest/test-data/groups/protobuf_groups.binpb b/testing/gtest/test-data/groups/protobuf_groups.binpb
@@ -0,0 +1 @@
+%�I@(
diff --git a/testing/gtest/test-data/groups/protobuf_groups.txtpb b/testing/gtest/test-data/groups/protobuf_groups.txtpb
@@ -0,0 +1,6 @@
+1: 1
+2:SGROUP
+  3: 2
+  4: 0x40490fdbi32
+2:EGROUP
+5: 3