Add protocol buffer message encoding utility class

PiperOrigin-RevId: 550006265

Author: jcking

internal/BUILD

@@ -328,6 +328,7 @@ cc_test(
     deps = [
         ":proto_wire",
         ":testing",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:cord",
     ],
 )

internal/proto_wire.cc

@@ -14,9 +14,13 @@
 
 #include "internal/proto_wire.h"
 
+#include <limits>
 #include <string>
+#include <utility>
 
 #include "absl/base/optimization.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 
 namespace cel::internal {
 
@@ -110,4 +114,36 @@ absl::StatusOr<absl::Cord> ProtoWireDecoder::ReadLengthDelimited() {
   return result;
 }
 
+absl::Status ProtoWireEncoder::WriteTag(ProtoWireTag tag) {
+  ABSL_DCHECK(!tag_.has_value());
+  if (ABSL_PREDICT_FALSE(tag.field_number() == 0)) {
+    // Cannot easily add test coverage as we assert during debug builds that
+    // ProtoWireTag is valid upon construction.
+    return absl::InvalidArgumentError(
+        absl::StrCat("invalid field number encountered encoding ", message_));
+  }
+  if (ABSL_PREDICT_FALSE(!ProtoWireTypeIsValid(tag.type()))) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("invalid wire type encountered encoding field ",
+                     tag.field_number(), " of ", message_));
+  }
+  VarintEncode(static_cast<uint32_t>(tag), data_);
+  tag_.emplace(tag);
+  return absl::OkStatus();
+}
+
+absl::Status ProtoWireEncoder::WriteLengthDelimited(absl::Cord data) {
+  ABSL_DCHECK(tag_.has_value() &&
+              tag_->type() == ProtoWireType::kLengthDelimited);
+  if (ABSL_PREDICT_FALSE(data.size() > std::numeric_limits<uint32_t>::max())) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("out of range length encountered encoding field ",
+                     tag_->field_number(), " of ", message_));
+  }
+  VarintEncode(static_cast<uint32_t>(data.size()), data_);
+  data_.Append(std::move(data));
+  tag_.reset();
+  return absl::OkStatus();
+}
+
 }  // namespace cel::internal

internal/proto_wire.h

@@ -88,19 +88,27 @@ enum class ProtoWireType : uint32_t {
   kFixed32 = 5,
 };
 
+inline constexpr uint32_t kProtoWireTypeMask = uint32_t{0x7};
+inline constexpr int kFieldNumberShift = 3;
+
 class ProtoWireTag final {
  public:
+  static constexpr uint32_t kTypeMask = uint32_t{0x7};
+  static constexpr int kFieldNumberShift = 3;
+
   constexpr explicit ProtoWireTag(uint32_t tag) : tag_(tag) {}
 
   constexpr ProtoWireTag(uint32_t field_number, ProtoWireType type)
-      : ProtoWireTag((field_number << 3) | static_cast<uint32_t>(type)) {
-    ABSL_ASSERT(((field_number << 3) >> 3) == field_number);
+      : ProtoWireTag((field_number << kFieldNumberShift) |
+                     static_cast<uint32_t>(type)) {
+    ABSL_ASSERT(((field_number << kFieldNumberShift) >> kFieldNumberShift) ==
+                field_number);
   }
 
-  constexpr uint32_t field_number() const { return tag_ >> 3; }
+  constexpr uint32_t field_number() const { return tag_ >> kFieldNumberShift; }
 
   constexpr ProtoWireType type() const {
-    return static_cast<ProtoWireType>(tag_ & uint32_t{0x7});
+    return static_cast<ProtoWireType>(tag_ & kTypeMask);
   }
 
   // NOLINTNEXTLINE(google-explicit-constructor)
@@ -110,6 +118,14 @@ class ProtoWireTag final {
   uint32_t tag_;
 };
 
+inline constexpr bool ProtoWireTypeIsValid(ProtoWireType type) {
+  // Ensure `type` is only [0-5]. The bitmask for `type` is 0x7 which allows 6
+  // to exist, but that is not used and invalid. We detect that here.
+  return (static_cast<uint32_t>(type) & uint32_t{0x7}) ==
+             static_cast<uint32_t>(type) &&
+         static_cast<uint32_t>(type) != uint32_t{0x6};
+}
+
 // Creates the "tag" of a record, see
 // https://protobuf.dev/programming-guides/encoding/#structure.
 inline constexpr uint32_t MakeProtoWireTag(uint32_t field_number,
@@ -168,6 +184,30 @@ inline void VarintEncode(bool value, absl::Cord& buffer) {
   buffer.Append(absl::string_view(&scratch, 1));
 }
 
+inline void Fixed32EncodeUnsafe(uint64_t value, char* buffer) {
+  buffer[0] = static_cast<char>(static_cast<uint8_t>(value));
+  buffer[1] = static_cast<char>(static_cast<uint8_t>(value >> 8));
+  buffer[2] = static_cast<char>(static_cast<uint8_t>(value >> 16));
+  buffer[3] = static_cast<char>(static_cast<uint8_t>(value >> 24));
+}
+
+// Encodes `value` as a fixed-size number, see
+// https://protobuf.dev/programming-guides/encoding/#non-varint-numbers.
+inline void Fixed32Encode(uint32_t value, absl::Cord& buffer) {
+  // `absl::Cord::GetAppendBuffer` will allocate a block regardless of whether
+  // `buffer` has enough inline storage space left. To take advantage of inline
+  // storage space, we need to just do a plain append.
+  char scratch[4];
+  Fixed32EncodeUnsafe(value, scratch);
+  buffer.Append(absl::string_view(scratch, ABSL_ARRAYSIZE(scratch)));
+}
+
+// Encodes `value` as a fixed-size number, see
+// https://protobuf.dev/programming-guides/encoding/#non-varint-numbers.
+inline void Fixed32Encode(float value, absl::Cord& buffer) {
+  Fixed32Encode(absl::bit_cast<uint32_t>(value), buffer);
+}
+
 inline void Fixed64EncodeUnsafe(uint64_t value, char* buffer) {
   buffer[0] = static_cast<char>(static_cast<uint8_t>(value));
   buffer[1] = static_cast<char>(static_cast<uint8_t>(value >> 8));
@@ -308,11 +348,12 @@ Fixed32Decode(const absl::Cord& data) {
 }
 
 inline absl::optional<ProtoWireTag> DecodeProtoWireTag(uint32_t value) {
-  if (ABSL_PREDICT_FALSE((value >> 3) == 0)) {
+  if (ABSL_PREDICT_FALSE((value >> ProtoWireTag::kFieldNumberShift) == 0)) {
     // Field number is 0.
     return absl::nullopt;
   }
-  if (ABSL_PREDICT_FALSE((value & uint32_t{0x7}) == uint32_t{0x6})) {
+  if (ABSL_PREDICT_FALSE(!ProtoWireTypeIsValid(
+          static_cast<ProtoWireType>(value & ProtoWireTag::kTypeMask)))) {
     // Wire type is 6, only 0-5 are used.
     return absl::nullopt;
   }
@@ -409,6 +450,60 @@ class ProtoWireDecoder {
   absl::optional<ProtoWireTag> tag_;
 };
 
+class ProtoWireEncoder final {
+ public:
+  explicit ProtoWireEncoder(absl::string_view message
+                                ABSL_ATTRIBUTE_LIFETIME_BOUND,
+                            absl::Cord& data ABSL_ATTRIBUTE_LIFETIME_BOUND)
+      : message_(message), data_(data), original_data_size_(data_.size()) {}
+
+  bool empty() const { return size() == 0; }
+
+  size_t size() const { return data_.size() - original_data_size_; }
+
+  absl::Status WriteTag(ProtoWireTag tag);
+
+  template <typename T>
+  std::enable_if_t<std::is_integral_v<T>, absl::Status> WriteVarint(T value) {
+    ABSL_DCHECK(tag_.has_value() && tag_->type() == ProtoWireType::kVarint);
+    VarintEncode(value, data_);
+    tag_.reset();
+    return absl::OkStatus();
+  }
+
+  template <typename T>
+  std::enable_if_t<sizeof(T) == 4 &&
+                       (std::is_integral_v<T> || std::is_floating_point_v<T>),
+                   absl::Status>
+  WriteFixed32(T value) {
+    ABSL_DCHECK(tag_.has_value() && tag_->type() == ProtoWireType::kFixed32);
+    Fixed32Encode(value, data_);
+    tag_.reset();
+    return absl::OkStatus();
+  }
+
+  template <typename T>
+  std::enable_if_t<sizeof(T) == 8 &&
+                       (std::is_integral_v<T> || std::is_floating_point_v<T>),
+                   absl::Status>
+  WriteFixed64(T value) {
+    ABSL_DCHECK(tag_.has_value() && tag_->type() == ProtoWireType::kFixed64);
+    Fixed64Encode(value, data_);
+    tag_.reset();
+    return absl::OkStatus();
+  }
+
+  absl::Status WriteLengthDelimited(absl::Cord data);
+
+  void EnsureFullyEncoded() { ABSL_DCHECK(!tag_.has_value()); }
+
+ private:
+  absl::string_view message_;
+  absl::Cord& data_;
+  const size_t original_data_size_;
+  absl::optional<ProtoWireTag> tag_;
+};
+
 }  // namespace cel::internal
 
 #endif  // THIRD_PARTY_CEL_CPP_INTERNAL_PROTO_WIRE_H_

internal/proto_wire_test.cc

@@ -17,6 +17,7 @@
 #include <limits>
 
 #include "absl/strings/cord.h"
+#include "absl/strings/string_view.h"
 #include "internal/testing.h"
 
 namespace cel::internal {
@@ -112,6 +113,13 @@ absl::Cord Fixed64Encode(T value) {
   return cord;
 }
 
+template <typename T>
+absl::Cord Fixed32Encode(T value) {
+  absl::Cord cord;
+  internal::Fixed32Encode(value, cord);
+  return cord;
+}
+
 }  // namespace
 
 TEST(Fixed64, Encode) {
@@ -122,6 +130,10 @@ TEST(Fixed64, Decode) {
   EXPECT_THAT(Fixed64Decode<double>(Fixed64Encode(0.0)), Optional(Eq(0.0)));
 }
 
+TEST(Fixed32, Encode) {
+  EXPECT_EQ(Fixed32Encode(0.0f), Fixed32Encode(uint32_t{0}));
+}
+
 TEST(Fixed32, Decode) {
   EXPECT_THAT(Fixed32Decode<float>(
                   absl::Cord(absl::string_view("\x00\x00\x00\x00", 4))),
@@ -206,6 +218,73 @@ TEST(SkipLengthValue, Decoder) {
   }
 }
 
+TEST(ProtoWireEncoder, BadTag) {
+  absl::Cord data;
+  ProtoWireEncoder encoder("foo.Bar", data);
+  EXPECT_TRUE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 0);
+  EXPECT_OK(encoder.WriteTag(ProtoWireTag(1, ProtoWireType::kVarint)));
+  EXPECT_OK(encoder.WriteVarint(1));
+  encoder.EnsureFullyEncoded();
+  EXPECT_FALSE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 2);
+  EXPECT_EQ(data, "\x08\x01");
+}
+
+TEST(ProtoWireEncoder, Varint) {
+  absl::Cord data;
+  ProtoWireEncoder encoder("foo.Bar", data);
+  EXPECT_TRUE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 0);
+  EXPECT_OK(encoder.WriteTag(ProtoWireTag(1, ProtoWireType::kVarint)));
+  EXPECT_OK(encoder.WriteVarint(1));
+  encoder.EnsureFullyEncoded();
+  EXPECT_FALSE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 2);
+  EXPECT_EQ(data, "\x08\x01");
+}
+
+TEST(ProtoWireEncoder, Fixed32) {
+  absl::Cord data;
+  ProtoWireEncoder encoder("foo.Bar", data);
+  EXPECT_TRUE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 0);
+  EXPECT_OK(encoder.WriteTag(ProtoWireTag(1, ProtoWireType::kFixed32)));
+  EXPECT_OK(encoder.WriteFixed32(0.0f));
+  encoder.EnsureFullyEncoded();
+  EXPECT_FALSE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 5);
+  EXPECT_EQ(data, absl::string_view("\x0d\x00\x00\x00\x00", 5));
+}
+
+TEST(ProtoWireEncoder, Fixed64) {
+  absl::Cord data;
+  ProtoWireEncoder encoder("foo.Bar", data);
+  EXPECT_TRUE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 0);
+  EXPECT_OK(encoder.WriteTag(ProtoWireTag(1, ProtoWireType::kFixed64)));
+  EXPECT_OK(encoder.WriteFixed64(0.0));
+  encoder.EnsureFullyEncoded();
+  EXPECT_FALSE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 9);
+  EXPECT_EQ(data, absl::string_view("\x09\x00\x00\x00\x00\x00\x00\x00\x00", 9));
+}
+
+TEST(ProtoWireEncoder, LengthDelimited) {
+  absl::Cord data;
+  ProtoWireEncoder encoder("foo.Bar", data);
+  EXPECT_TRUE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 0);
+  EXPECT_OK(encoder.WriteTag(ProtoWireTag(1, ProtoWireType::kLengthDelimited)));
+  EXPECT_OK(encoder.WriteLengthDelimited(absl::Cord("foo")));
+  encoder.EnsureFullyEncoded();
+  EXPECT_FALSE(encoder.empty());
+  EXPECT_EQ(encoder.size(), 5);
+  EXPECT_EQ(data,
+            "\x0a\x03"
+            "foo");
+}
+
 }  // namespace
 
 }  // namespace cel::internal