From a99ead4bf4eba5f3875ada0932f5fc58e89cfca4 Mon Sep 17 00:00:00 2001
From: Esun Kim <veblush@google.com>
Date: Thu, 27 May 2021 11:18:18 -0700
Subject: [PATCH] Added new method to ByteBuffer & Slice (#26014)

---
 include/grpcpp/impl/codegen/byte_buffer.h |  7 +++++
 include/grpcpp/impl/codegen/slice.h       | 11 ++++++++
 src/cpp/util/byte_buffer_cc.cc            | 31 +++++++++++++++++++++++
 test/cpp/util/byte_buffer_test.cc         | 29 +++++++++++++++++++++
 test/cpp/util/slice_test.cc               |  6 +++++
 5 files changed, 84 insertions(+)

diff --git a/include/grpcpp/impl/codegen/byte_buffer.h b/include/grpcpp/impl/codegen/byte_buffer.h
index 2c015f2266f..6959fa29baf 100644
--- a/include/grpcpp/impl/codegen/byte_buffer.h
+++ b/include/grpcpp/impl/codegen/byte_buffer.h
@@ -114,6 +114,13 @@ class ByteBuffer final {
     return *this;
   }
 
+  // If this ByteBuffer's representation is a single flat slice, returns a
+  // slice referencing that array.
+  Status TrySingleSlice(Slice* slice) const;
+
+  /// Dump (read) the buffer contents into \a slics.
+  Status DumpToSingleSlice(Slice* slice) const;
+
   /// Dump (read) the buffer contents into \a slices.
   Status Dump(std::vector<Slice>* slices) const;
 
diff --git a/include/grpcpp/impl/codegen/slice.h b/include/grpcpp/impl/codegen/slice.h
index ea550259e80..d84b842c843 100644
--- a/include/grpcpp/impl/codegen/slice.h
+++ b/include/grpcpp/impl/codegen/slice.h
@@ -74,6 +74,11 @@ class Slice final {
   Slice(const Slice& other)
       : slice_(g_core_codegen_interface->grpc_slice_ref(other.slice_)) {}
 
+  /// Move constructor, steals a reference.
+  Slice(Slice&& other) noexcept : slice_(other.slice_) {
+    other.slice_ = g_core_codegen_interface->grpc_empty_slice();
+  }
+
   /// Assignment, reference count is unchanged.
   Slice& operator=(Slice other) {
     std::swap(slice_, other.slice_);
@@ -107,6 +112,12 @@ class Slice final {
   /// Raw pointer to the end (one byte \em past the last element) of the slice.
   const uint8_t* end() const { return GRPC_SLICE_END_PTR(slice_); }
 
+  /// Returns a substring of the `slice` as another slice.
+  Slice sub(size_t begin, size_t end) const {
+    return Slice(g_core_codegen_interface->grpc_slice_sub(slice_, begin, end),
+                 STEAL_REF);
+  }
+
   /// Raw C slice. Caller needs to call grpc_slice_unref when done.
   grpc_slice c_slice() const {
     return g_core_codegen_interface->grpc_slice_ref(slice_);
diff --git a/src/cpp/util/byte_buffer_cc.cc b/src/cpp/util/byte_buffer_cc.cc
index fb705906455..5c6f22bd7c6 100644
--- a/src/cpp/util/byte_buffer_cc.cc
+++ b/src/cpp/util/byte_buffer_cc.cc
@@ -25,6 +25,37 @@ namespace grpc {
 
 static internal::GrpcLibraryInitializer g_gli_initializer;
 
+Status ByteBuffer::TrySingleSlice(Slice* slice) const {
+  if (!buffer_) {
+    return Status(StatusCode::FAILED_PRECONDITION, "Buffer not initialized");
+  }
+  if ((buffer_->type == GRPC_BB_RAW) &&
+      (buffer_->data.raw.compression == GRPC_COMPRESS_NONE) &&
+      (buffer_->data.raw.slice_buffer.count == 1)) {
+    grpc_slice internal_slice = buffer_->data.raw.slice_buffer.slices[0];
+    *slice = Slice(internal_slice, Slice::ADD_REF);
+    return Status::OK;
+  } else {
+    return Status(StatusCode::FAILED_PRECONDITION,
+                  "Buffer isn't made up of a single uncompressed slice.");
+  }
+}
+
+Status ByteBuffer::DumpToSingleSlice(Slice* slice) const {
+  if (!buffer_) {
+    return Status(StatusCode::FAILED_PRECONDITION, "Buffer not initialized");
+  }
+  grpc_byte_buffer_reader reader;
+  if (!grpc_byte_buffer_reader_init(&reader, buffer_)) {
+    return Status(StatusCode::INTERNAL,
+                  "Couldn't initialize byte buffer reader");
+  }
+  grpc_slice s = grpc_byte_buffer_reader_readall(&reader);
+  *slice = Slice(s, Slice::STEAL_REF);
+  grpc_byte_buffer_reader_destroy(&reader);
+  return Status::OK;
+}
+
 Status ByteBuffer::Dump(std::vector<Slice>* slices) const {
   slices->clear();
   if (!buffer_) {
diff --git a/test/cpp/util/byte_buffer_test.cc b/test/cpp/util/byte_buffer_test.cc
index c63f351a8f0..5d6c1837e4d 100644
--- a/test/cpp/util/byte_buffer_test.cc
+++ b/test/cpp/util/byte_buffer_test.cc
@@ -123,6 +123,35 @@ TEST_F(ByteBufferTest, SerializationMakesCopy) {
   EXPECT_TRUE(send_buffer.Valid());
 }
 
+TEST_F(ByteBufferTest, TrySingleSliceWithSingleSlice) {
+  std::vector<Slice> slices;
+  slices.emplace_back(kContent1);
+  ByteBuffer buffer(&slices[0], 1);
+  Slice slice;
+  EXPECT_TRUE(buffer.TrySingleSlice(&slice).ok());
+  EXPECT_EQ(slice.size(), slices[0].size());
+  EXPECT_EQ(memcmp(slice.begin(), slices[0].begin(), slice.size()), 0);
+}
+
+TEST_F(ByteBufferTest, TrySingleSliceWithMultipleSlices) {
+  std::vector<Slice> slices;
+  slices.emplace_back(kContent1);
+  slices.emplace_back(kContent2);
+  ByteBuffer buffer(&slices[0], 2);
+  Slice slice;
+  EXPECT_FALSE(buffer.TrySingleSlice(&slice).ok());
+}
+
+TEST_F(ByteBufferTest, DumpToSingleSlice) {
+  std::vector<Slice> slices;
+  slices.emplace_back(kContent1);
+  slices.emplace_back(kContent2);
+  ByteBuffer buffer(&slices[0], 2);
+  Slice slice;
+  EXPECT_TRUE(buffer.DumpToSingleSlice(&slice).ok());
+  EXPECT_EQ(strlen(kContent1) + strlen(kContent2), slice.size());
+}
+
 }  // namespace
 }  // namespace grpc
 
diff --git a/test/cpp/util/slice_test.cc b/test/cpp/util/slice_test.cc
index aa27a0ea97a..25041de4041 100644
--- a/test/cpp/util/slice_test.cc
+++ b/test/cpp/util/slice_test.cc
@@ -124,6 +124,12 @@ TEST_F(SliceTest, Add) {
   CheckSlice(spp, kContent);
 }
 
+TEST_F(SliceTest, Sub) {
+  Slice spp("0123456789");
+  Slice sub = spp.sub(1, 9);
+  CheckSlice(sub, "12345678");
+}
+
 TEST_F(SliceTest, Cslice) {
   grpc_slice s = grpc_slice_from_copied_string(kContent);
   Slice spp(s, Slice::STEAL_REF);