buffer: add a method for getting only the first slice (envoyproxy#14050)

Signed-off-by: Greg Greenway <[email protected]>

Author: ggreenway

include/envoy/buffer/buffer.h

@@ -160,6 +160,14 @@ class Instance {
   virtual RawSliceVector
   getRawSlices(absl::optional<uint64_t> max_slices = absl::nullopt) const PURE;
 
+  /**
+   * Fetch the valid data pointer and valid data length of the first non-zero-length
+   * slice in the buffer.
+   * @return RawSlice the first non-empty slice in the buffer, or {nullptr, 0} if the buffer
+   * is empty.
+   */
+  virtual RawSlice frontSlice() const PURE;
+
   /**
    * Transfer ownership of the front slice to the caller. Must only be called if the
    * buffer is not empty otherwise the implementation will have undefined behavior.

source/common/buffer/buffer_impl.cc

@@ -201,6 +201,17 @@ RawSliceVector OwnedImpl::getRawSlices(absl::optional<uint64_t> max_slices) cons
   return raw_slices;
 }
 
+RawSlice OwnedImpl::frontSlice() const {
+  // Ignore zero-size slices and return the first slice with data.
+  for (const auto& slice : slices_) {
+    if (slice->dataSize() > 0) {
+      return RawSlice{slice->data(), slice->dataSize()};
+    }
+  }
+
+  return {nullptr, 0};
+}
+
 SliceDataPtr OwnedImpl::extractMutableFrontSlice() {
   RELEASE_ASSERT(length_ > 0, "Extract called on empty buffer");
   // Remove zero byte fragments from the front of the queue to ensure

source/common/buffer/buffer_impl.h

@@ -564,6 +564,7 @@ class OwnedImpl : public LibEventInstance {
   void copyOut(size_t start, uint64_t size, void* data) const override;
   void drain(uint64_t size) override;
   RawSliceVector getRawSlices(absl::optional<uint64_t> max_slices = absl::nullopt) const override;
+  RawSlice frontSlice() const override;
   SliceDataPtr extractMutableFrontSlice() override;
   uint64_t length() const override;
   void* linearize(uint32_t size) override;

source/common/buffer/zero_copy_input_stream_impl.cc

@@ -29,10 +29,9 @@ void ZeroCopyInputStreamImpl::drainLastSlice() {
 bool ZeroCopyInputStreamImpl::Next(const void** data, int* size) {
   drainLastSlice();
 
-  Buffer::RawSliceVector slices = buffer_->getRawSlices(1);
+  Buffer::RawSlice slice = buffer_->frontSlice();
 
-  if (!slices.empty() && slices[0].len_ > 0) {
-    auto& slice = slices[0];
+  if (slice.len_ > 0) {
     *data = slice.mem_;
     *size = slice.len_;
     position_ = slice.len_;

source/extensions/quic_listeners/quiche/active_quic_listener.cc

@@ -110,11 +110,11 @@ void ActiveQuicListener::onDataWorker(Network::UdpRecvData&& data) {
       quic::QuicTime::Delta::FromMicroseconds(std::chrono::duration_cast<std::chrono::microseconds>(
                                                   data.receive_time_.time_since_epoch())
                                                   .count());
-  ASSERT(data.buffer_->getRawSlices().size() == 1);
-  Buffer::RawSliceVector slices = data.buffer_->getRawSlices(/*max_slices=*/1);
+  Buffer::RawSlice slice = data.buffer_->frontSlice();
+  ASSERT(data.buffer_->length() == slice.len_);
   // TODO(danzh): pass in TTL and UDP header.
-  quic::QuicReceivedPacket packet(reinterpret_cast<char*>(slices[0].mem_), slices[0].len_,
-                                  timestamp, /*owns_buffer=*/false, /*ttl=*/0, /*ttl_valid=*/false,
+  quic::QuicReceivedPacket packet(reinterpret_cast<char*>(slice.mem_), slice.len_, timestamp,
+                                  /*owns_buffer=*/false, /*ttl=*/0, /*ttl_valid=*/false,
                                   /*packet_headers=*/nullptr, /*headers_length=*/0,
                                   /*owns_header_buffer*/ false);
   quic_dispatcher_->ProcessPacket(self_address, peer_address, packet);

source/extensions/quic_listeners/quiche/envoy_quic_client_connection.cc

@@ -62,9 +62,9 @@ void EnvoyQuicClientConnection::processPacket(
           std::chrono::duration_cast<std::chrono::microseconds>(receive_time.time_since_epoch())
               .count());
   ASSERT(buffer->getRawSlices().size() == 1);
-  Buffer::RawSliceVector slices = buffer->getRawSlices(/*max_slices=*/1);
-  quic::QuicReceivedPacket packet(reinterpret_cast<char*>(slices[0].mem_), slices[0].len_,
-                                  timestamp, /*owns_buffer=*/false, /*ttl=*/0, /*ttl_valid=*/false,
+  Buffer::RawSlice slice = buffer->frontSlice();
+  quic::QuicReceivedPacket packet(reinterpret_cast<char*>(slice.mem_), slice.len_, timestamp,
+                                  /*owns_buffer=*/false, /*ttl=*/0, /*ttl_valid=*/false,
                                   /*packet_headers=*/nullptr, /*headers_length=*/0,
                                   /*owns_header_buffer*/ false);
   ProcessUdpPacket(envoyIpAddressToQuicSocketAddress(local_address->ip()),

source/extensions/quic_listeners/quiche/platform/quic_mem_slice_impl.cc

@@ -29,15 +29,11 @@ QuicMemSliceImpl::QuicMemSliceImpl(Envoy::Buffer::Instance& buffer, size_t lengt
 }
 
 const char* QuicMemSliceImpl::data() const {
-  Envoy::Buffer::RawSliceVector slices = single_slice_buffer_.getRawSlices(/*max_slices=*/1);
-  ASSERT(slices.size() <= 1);
-  return !slices.empty() ? static_cast<const char*>(slices[0].mem_) : nullptr;
+  return reinterpret_cast<const char*>(single_slice_buffer_.frontSlice().mem_);
 }
 
 size_t QuicMemSliceImpl::firstSliceLength(Envoy::Buffer::Instance& buffer) {
-  Envoy::Buffer::RawSliceVector slices = buffer.getRawSlices(/*max_slices=*/1);
-  ASSERT(slices.size() == 1);
-  return slices[0].len_;
+  return buffer.frontSlice().len_;
 }
 
 } // namespace quic

test/common/buffer/buffer_fuzz.cc

@@ -126,6 +126,8 @@ class StringBuffer : public Buffer::Instance {
     return {{const_cast<char*>(start()), size_}};
   }
 
+  Buffer::RawSlice frontSlice() const override { return {const_cast<char*>(start()), size_}; }
+
   uint64_t length() const override { return size_; }
 
   void* linearize(uint32_t /*size*/) override {

test/common/buffer/owned_impl_test.cc

@@ -1256,6 +1256,13 @@ TEST_F(OwnedImplTest, MoveSmallSliceIntoNotEnoughFreeSpace) {
   TestBufferMove(4096 - 127, 128, 2);
 }
 
+TEST_F(OwnedImplTest, FrontSlice) {
+  Buffer::OwnedImpl buffer;
+  EXPECT_EQ(0, buffer.frontSlice().len_);
+  buffer.add("a");
+  EXPECT_EQ(1, buffer.frontSlice().len_);
+}
+
 } // namespace
 } // namespace Buffer
 } // namespace Envoy

test/extensions/filters/network/postgres_proxy/postgres_decoder_test.cc

@@ -523,6 +523,7 @@ class FakeBuffer : public Buffer::Instance {
   MOCK_METHOD(void, copyOut, (size_t, uint64_t, void*), (const, override));
   MOCK_METHOD(void, drain, (uint64_t), (override));
   MOCK_METHOD(Buffer::RawSliceVector, getRawSlices, (absl::optional<uint64_t>), (const, override));
+  MOCK_METHOD(Buffer::RawSlice, frontSlice, (), (const, override));
   MOCK_METHOD(Buffer::SliceDataPtr, extractMutableFrontSlice, (), (override));
   MOCK_METHOD(uint64_t, length, (), (const, override));
   MOCK_METHOD(void*, linearize, (uint32_t), (override));

test/per_file_coverage.sh

@@ -72,7 +72,7 @@ declare -a KNOWN_LOW_COVERAGE=(
 "source/extensions/wasm_runtime/wavm:0.0" # Noe enabled in coverage build
 "source/extensions/watchdog:69.6" # Death tests within extensions
 "source/extensions/watchdog/profile_action:84.9"
-"source/server:94.6"
+"source/server:94.5"
 "source/server/config_validation:76.6"
 "source/server/admin:95.2"
 )