Change quicpacketwriter::getnextwritelocation to return a release function used to free a writer-allocated buffer. use this release function to avoid buffer leak if a writer buffer is allocated but writer->writepacket is not called. protected by --gfe2_reloadable_flag_quic_avoid_leak_writer_buffer.

This is needed to fix a buffer leak when the new PigeonWriter is used.

PiperOrigin-RevId: 312369772
Change-Id: I0d12327b58989ec69e401b03f7b9ebb4d8ce22fd
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index 1f1081c..3519bbc 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2404,6 +2404,14 @@
       buffered_packets_.emplace_back(*packet, self_address(), peer_address());
       break;
     case SEND_TO_WRITER:
+      // At this point, packet->release_encrypted_buffer is either nullptr,
+      // meaning |packet->encrypted_buffer| is a stack buffer, or not-nullptr,
+      /// meaning it's a writer-allocated buffer. Note that connectivity probing
+      // packets do not use this function, so setting release_encrypted_buffer
+      // to nullptr will not cause probing packets to be leaked.
+      //
+      // writer_->WritePacket transfers buffer ownership back to the writer.
+      packet->release_encrypted_buffer = nullptr;
       result = writer_->WritePacket(packet->encrypted_buffer, encrypted_length,
                                     self_address().host(), peer_address(),
                                     per_packet_options_);
@@ -2647,11 +2655,11 @@
   }
 }
 
-char* QuicConnection::GetPacketBuffer() {
+QuicPacketBuffer QuicConnection::GetPacketBuffer() {
   if (version().CanSendCoalescedPackets() && !IsHandshakeConfirmed()) {
     // Do not use writer's packet buffer for coalesced packets which may contain
     // multiple QUIC packets.
-    return nullptr;
+    return {nullptr, nullptr};
   }
   return writer_->GetNextWriteLocation(self_address().host(), peer_address());
 }
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index 3f4329c..87c049b 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -596,7 +596,7 @@
   bool ShouldGeneratePacket(HasRetransmittableData retransmittable,
                             IsHandshake handshake) override;
   const QuicFrames MaybeBundleAckOpportunistically() override;
-  char* GetPacketBuffer() override;
+  QuicPacketBuffer GetPacketBuffer() override;
   void OnSerializedPacket(SerializedPacket packet) override;
   void OnUnrecoverableError(QuicErrorCode error,
                             const std::string& error_details) override;
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index c53b730..fe94fcf 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -18,6 +18,7 @@
 #include "net/third_party/quiche/src/quic/core/crypto/quic_decrypter.h"
 #include "net/third_party/quiche/src/quic/core/crypto/quic_encrypter.h"
 #include "net/third_party/quiche/src/quic/core/quic_connection_id.h"
+#include "net/third_party/quiche/src/quic/core/quic_constants.h"
 #include "net/third_party/quiche/src/quic/core/quic_packets.h"
 #include "net/third_party/quiche/src/quic/core/quic_simple_buffer_allocator.h"
 #include "net/third_party/quiche/src/quic/core/quic_types.h"
@@ -337,6 +338,11 @@
 };
 
 class TestPacketWriter : public QuicPacketWriter {
+  struct PacketBuffer {
+    QUIC_CACHELINE_ALIGNED char buffer[1500];
+    bool in_use = false;
+  };
+
  public:
   TestPacketWriter(ParsedQuicVersion version, MockClock* clock)
       : version_(version),
@@ -345,16 +351,40 @@
     QuicFramerPeer::SetLastSerializedServerConnectionId(framer_.framer(),
                                                         TestConnectionId());
     framer_.framer()->SetInitialObfuscators(TestConnectionId());
+
+    for (int i = 0; i < 128; ++i) {
+      PacketBuffer* p = new PacketBuffer();
+      packet_buffer_pool_.push_back(p);
+      packet_buffer_pool_index_[p->buffer] = p;
+      packet_buffer_free_list_.push_back(p);
+    }
   }
   TestPacketWriter(const TestPacketWriter&) = delete;
   TestPacketWriter& operator=(const TestPacketWriter&) = delete;
 
+  ~TestPacketWriter() override {
+    EXPECT_EQ(packet_buffer_pool_.size(), packet_buffer_free_list_.size())
+        << packet_buffer_pool_.size() - packet_buffer_free_list_.size()
+        << " out of " << packet_buffer_pool_.size()
+        << " packet buffers have been leaked.";
+    for (auto p : packet_buffer_pool_) {
+      delete p;
+    }
+  }
+
   // QuicPacketWriter interface
   WriteResult WritePacket(const char* buffer,
                           size_t buf_len,
                           const QuicIpAddress& /*self_address*/,
                           const QuicSocketAddress& /*peer_address*/,
                           PerPacketOptions* /*options*/) override {
+    // If the buffer is allocated from the pool, return it back to the pool.
+    // Note the buffer content doesn't change.
+    if (packet_buffer_pool_index_.find(const_cast<char*>(buffer)) !=
+        packet_buffer_pool_index_.end()) {
+      FreePacketBuffer(buffer);
+    }
+
     QuicEncryptedPacket packet(buffer, buf_len);
     ++packets_write_attempts_;
 
@@ -441,10 +471,15 @@
 
   bool IsBatchMode() const override { return is_batch_mode_; }
 
-  char* GetNextWriteLocation(
+  QuicPacketBuffer GetNextWriteLocation(
       const QuicIpAddress& /*self_address*/,
       const QuicSocketAddress& /*peer_address*/) override {
-    return nullptr;
+    if (GetQuicReloadableFlag(quic_avoid_leak_writer_buffer)) {
+      return {AllocPacketBuffer(),
+              [this](const char* p) { FreePacketBuffer(p); }};
+    }
+    // Do not use writer buffer for serializing packets.
+    return {nullptr, nullptr};
   }
 
   WriteResult Flush() override {
@@ -592,6 +627,23 @@
   SimpleQuicFramer* framer() { return &framer_; }
 
  private:
+  char* AllocPacketBuffer() {
+    PacketBuffer* p = packet_buffer_free_list_.front();
+    EXPECT_FALSE(p->in_use);
+    p->in_use = true;
+    packet_buffer_free_list_.pop_front();
+    return p->buffer;
+  }
+
+  void FreePacketBuffer(const char* buffer) {
+    auto iter = packet_buffer_pool_index_.find(const_cast<char*>(buffer));
+    ASSERT_TRUE(iter != packet_buffer_pool_index_.end());
+    PacketBuffer* p = iter->second;
+    ASSERT_TRUE(p->in_use);
+    p->in_use = false;
+    packet_buffer_free_list_.push_back(p);
+  }
+
   ParsedQuicVersion version_;
   SimpleQuicFramer framer_;
   size_t last_packet_size_ = 0;
@@ -620,6 +672,12 @@
   QuicTime::Delta write_pause_time_delta_ = QuicTime::Delta::Zero();
   QuicByteCount max_packet_size_ = kMaxOutgoingPacketSize;
   bool supports_release_time_ = false;
+  // Used to verify writer-allocated packet buffers are properly released.
+  std::vector<PacketBuffer*> packet_buffer_pool_;
+  // Buffer address => Address of the owning PacketBuffer.
+  QuicHashMap<char*, PacketBuffer*> packet_buffer_pool_index_;
+  // Indices in packet_buffer_pool_ that are not allocated.
+  std::list<PacketBuffer*> packet_buffer_free_list_;
 };
 
 class TestConnection : public QuicConnection {
diff --git a/quic/core/quic_default_packet_writer.cc b/quic/core/quic_default_packet_writer.cc
index 4222005..b34bed4 100644
--- a/quic/core/quic_default_packet_writer.cc
+++ b/quic/core/quic_default_packet_writer.cc
@@ -54,10 +54,10 @@
   return false;
 }
 
-char* QuicDefaultPacketWriter::GetNextWriteLocation(
+QuicPacketBuffer QuicDefaultPacketWriter::GetNextWriteLocation(
     const QuicIpAddress& /*self_address*/,
     const QuicSocketAddress& /*peer_address*/) {
-  return nullptr;
+  return {nullptr, nullptr};
 }
 
 WriteResult QuicDefaultPacketWriter::Flush() {
diff --git a/quic/core/quic_default_packet_writer.h b/quic/core/quic_default_packet_writer.h
index 36d4d8a..5388cae 100644
--- a/quic/core/quic_default_packet_writer.h
+++ b/quic/core/quic_default_packet_writer.h
@@ -35,8 +35,9 @@
       const QuicSocketAddress& peer_address) const override;
   bool SupportsReleaseTime() const override;
   bool IsBatchMode() const override;
-  char* GetNextWriteLocation(const QuicIpAddress& self_address,
-                             const QuicSocketAddress& peer_address) override;
+  QuicPacketBuffer GetNextWriteLocation(
+      const QuicIpAddress& self_address,
+      const QuicSocketAddress& peer_address) override;
   WriteResult Flush() override;
 
   void set_fd(int fd) { fd_ = fd; }
diff --git a/quic/core/quic_dispatcher.cc b/quic/core/quic_dispatcher.cc
index d3f01b4..ab436f4 100644
--- a/quic/core/quic_dispatcher.cc
+++ b/quic/core/quic_dispatcher.cc
@@ -68,9 +68,9 @@
                                 serialized_packet.encrypted_length, true));
   }
 
-  char* GetPacketBuffer() override {
+  QuicPacketBuffer GetPacketBuffer() override {
     // Let QuicPacketCreator to serialize packets on stack buffer.
-    return nullptr;
+    return {nullptr, nullptr};
   }
 
   void OnUnrecoverableError(QuicErrorCode /*error*/,
diff --git a/quic/core/quic_packet_creator.cc b/quic/core/quic_packet_creator.cc
index 4db87e2..a91d264 100644
--- a/quic/core/quic_packet_creator.cc
+++ b/quic/core/quic_packet_creator.cc
@@ -439,12 +439,23 @@
   }
 
   QUIC_CACHELINE_ALIGNED char stack_buffer[kMaxOutgoingPacketSize];
-  char* serialized_packet_buffer = delegate_->GetPacketBuffer();
-  if (serialized_packet_buffer == nullptr) {
-    serialized_packet_buffer = stack_buffer;
+  QuicOwnedPacketBuffer external_buffer(delegate_->GetPacketBuffer());
+
+  if (!avoid_leak_writer_buffer_ && external_buffer.release_buffer != nullptr) {
+    // This is not a flag count because it is incremented when flag is false.
+    QUIC_CODE_COUNT(quic_avoid_leak_writer_buffer_flag_false_noop_1);
+
+    // Setting it to nullptr to keep the behavior unchanged when flag is false.
+    external_buffer.release_buffer = nullptr;
   }
 
-  SerializePacket(serialized_packet_buffer, kMaxOutgoingPacketSize);
+  if (external_buffer.buffer == nullptr) {
+    external_buffer.buffer = stack_buffer;
+    external_buffer.release_buffer = nullptr;
+  }
+
+  DCHECK_EQ(nullptr, packet_.encrypted_buffer);
+  SerializePacket(std::move(external_buffer), kMaxOutgoingPacketSize);
   OnSerializedPacket();
 }
 
@@ -471,6 +482,12 @@
   packet_.transmission_type = NOT_RETRANSMISSION;
   packet_.encrypted_buffer = nullptr;
   packet_.encrypted_length = 0;
+  if (avoid_leak_writer_buffer_) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_avoid_leak_writer_buffer, 2, 3);
+    QUIC_BUG_IF(packet_.release_encrypted_buffer != nullptr)
+        << "packet_.release_encrypted_buffer should be empty";
+    packet_.release_encrypted_buffer = nullptr;
+  }
   DCHECK(packet_.retransmittable_frames.empty());
   DCHECK(packet_.nonretransmittable_frames.empty());
   packet_.largest_acked.Clear();
@@ -514,7 +531,7 @@
       return 0;
     }
   }
-  SerializePacket(buffer, buffer_len);
+  SerializePacket(QuicOwnedPacketBuffer(buffer, nullptr), buffer_len);
   const size_t encrypted_length = packet_.encrypted_length;
   // Clear frames in packet_. No need to DeleteFrames since frames are owned by
   // initial_packet.
@@ -538,11 +555,23 @@
   FillPacketHeader(&header);
 
   QUIC_CACHELINE_ALIGNED char stack_buffer[kMaxOutgoingPacketSize];
-  char* encrypted_buffer = delegate_->GetPacketBuffer();
-  if (encrypted_buffer == nullptr) {
-    encrypted_buffer = stack_buffer;
+  QuicOwnedPacketBuffer packet_buffer(delegate_->GetPacketBuffer());
+
+  if (!avoid_leak_writer_buffer_ && packet_buffer.release_buffer != nullptr) {
+    // This is not a flag count because it is incremented when flag is false.
+    QUIC_CODE_COUNT(quic_avoid_leak_writer_buffer_flag_false_noop_2);
+
+    // Setting it to nullptr to keep the behavior unchanged when flag is false.
+    packet_buffer.release_buffer = nullptr;
   }
 
+  if (packet_buffer.buffer == nullptr) {
+    packet_buffer.buffer = stack_buffer;
+    packet_buffer.release_buffer = nullptr;
+  }
+
+  char* encrypted_buffer = packet_buffer.buffer;
+
   QuicDataWriter writer(kMaxOutgoingPacketSize, encrypted_buffer);
   size_t length_field_offset = 0;
   if (!framer_->AppendPacketHeader(header, &writer, &length_field_offset)) {
@@ -623,6 +652,14 @@
   packet_size_ = 0;
   packet_.encrypted_buffer = encrypted_buffer;
   packet_.encrypted_length = encrypted_length;
+  if (avoid_leak_writer_buffer_) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_avoid_leak_writer_buffer, 3, 3);
+    packet_.release_encrypted_buffer = std::move(packet_buffer).release_buffer;
+  } else {
+    // If flag --quic_avoid_leak_writer_buffer is false, the release function
+    // should be empty.
+    DCHECK(packet_buffer.release_buffer == nullptr);
+  }
   packet_.retransmittable_frames.push_back(QuicFrame(frame));
   OnSerializedPacket();
 }
@@ -691,7 +728,7 @@
   return false;
 }
 
-void QuicPacketCreator::SerializePacket(char* encrypted_buffer,
+void QuicPacketCreator::SerializePacket(QuicOwnedPacketBuffer encrypted_buffer,
                                         size_t encrypted_buffer_len) {
   DCHECK_LT(0u, encrypted_buffer_len);
   QUIC_BUG_IF(queued_frames_.empty() && pending_padding_bytes_ == 0)
@@ -719,7 +756,7 @@
   // Use the packet_size_ instead of the buffer size to ensure smaller
   // packet sizes are properly used.
   size_t length =
-      framer_->BuildDataPacket(header, queued_frames_, encrypted_buffer,
+      framer_->BuildDataPacket(header, queued_frames_, encrypted_buffer.buffer,
                                packet_size_, packet_.encryption_level);
   if (length == 0) {
     QUIC_BUG << "Failed to serialize " << QuicFramesToString(queued_frames_)
@@ -746,7 +783,7 @@
   const size_t encrypted_length = framer_->EncryptInPlace(
       packet_.encryption_level, packet_.packet_number,
       GetStartOfEncryptedData(framer_->transport_version(), header), length,
-      encrypted_buffer_len, encrypted_buffer);
+      encrypted_buffer_len, encrypted_buffer.buffer);
   if (encrypted_length == 0) {
     QUIC_BUG << "Failed to encrypt packet number " << packet_.packet_number;
     return;
@@ -754,8 +791,17 @@
 
   packet_size_ = 0;
   queued_frames_.clear();
-  packet_.encrypted_buffer = encrypted_buffer;
+  packet_.encrypted_buffer = encrypted_buffer.buffer;
   packet_.encrypted_length = encrypted_length;
+  if (avoid_leak_writer_buffer_) {
+    QUIC_RELOADABLE_FLAG_COUNT_N(quic_avoid_leak_writer_buffer, 1, 3);
+    packet_.release_encrypted_buffer =
+        std::move(encrypted_buffer).release_buffer;
+  } else {
+    // If flag --quic_avoid_leak_writer_buffer is false, the release function
+    // should be empty.
+    DCHECK(encrypted_buffer.release_buffer == nullptr);
+  }
 }
 
 std::unique_ptr<QuicEncryptedPacket>
diff --git a/quic/core/quic_packet_creator.h b/quic/core/quic_packet_creator.h
index 7695587..c25d8f6 100644
--- a/quic/core/quic_packet_creator.h
+++ b/quic/core/quic_packet_creator.h
@@ -41,9 +41,9 @@
    public:
     virtual ~DelegateInterface() {}
     // Get a buffer of kMaxOutgoingPacketSize bytes to serialize the next
-    // packet. If return nullptr, QuicPacketCreator will serialize on a stack
-    // buffer.
-    virtual char* GetPacketBuffer() = 0;
+    // packet. If the return value's buffer is nullptr, QuicPacketCreator will
+    // serialize on a stack buffer.
+    virtual QuicPacketBuffer GetPacketBuffer() = 0;
     // Called when a packet is serialized. Delegate take the ownership of
     // |serialized_packet|.
     virtual void OnSerializedPacket(SerializedPacket serialized_packet) = 0;
@@ -459,7 +459,8 @@
   // retransmitted to packet_.retransmittable_frames. All frames must fit into
   // a single packet.
   // Fails if |encrypted_buffer_len| isn't long enough for the encrypted packet.
-  void SerializePacket(char* encrypted_buffer, size_t encrypted_buffer_len);
+  void SerializePacket(QuicOwnedPacketBuffer encrypted_buffer,
+                       size_t encrypted_buffer_len);
 
   // Called after a new SerialiedPacket is created to call the delegate's
   // OnSerializedPacket and reset state.
@@ -597,6 +598,9 @@
   // accept. There is no limit for QUIC_CRYPTO connections, but QUIC+TLS
   // negotiates this during the handshake.
   QuicByteCount max_datagram_frame_size_;
+
+  const bool avoid_leak_writer_buffer_ =
+      GetQuicReloadableFlag(quic_avoid_leak_writer_buffer);
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_packet_creator_test.cc b/quic/core/quic_packet_creator_test.cc
index 19267c9..ca59097 100644
--- a/quic/core/quic_packet_creator_test.cc
+++ b/quic/core/quic_packet_creator_test.cc
@@ -162,7 +162,8 @@
                        connection_id_.length()),
         data_("foo"),
         creator_(connection_id_, &client_framer_, &delegate_, &producer_) {
-    EXPECT_CALL(delegate_, GetPacketBuffer()).WillRepeatedly(Return(nullptr));
+    EXPECT_CALL(delegate_, GetPacketBuffer())
+        .WillRepeatedly(Return(QuicPacketBuffer()));
     creator_.SetEncrypter(ENCRYPTION_INITIAL, std::make_unique<NullEncrypter>(
                                                   Perspective::IS_CLIENT));
     creator_.SetEncrypter(ENCRYPTION_HANDSHAKE, std::make_unique<NullEncrypter>(
@@ -1725,9 +1726,10 @@
 
 TEST_P(QuicPacketCreatorTest, FlushWithExternalBuffer) {
   creator_.set_encryption_level(ENCRYPTION_FORWARD_SECURE);
-  char external_buffer[kMaxOutgoingPacketSize];
-  char* expected_buffer = external_buffer;
-  EXPECT_CALL(delegate_, GetPacketBuffer()).WillOnce(Return(expected_buffer));
+  char* buffer = new char[kMaxOutgoingPacketSize];
+  QuicPacketBuffer external_buffer = {buffer,
+                                      [](const char* p) { delete[] p; }};
+  EXPECT_CALL(delegate_, GetPacketBuffer()).WillOnce(Return(external_buffer));
 
   QuicFrame frame;
   MakeIOVector("test", &iov_);
@@ -1738,10 +1740,13 @@
       /*needs_full_padding=*/true, NOT_RETRANSMISSION, &frame));
 
   EXPECT_CALL(delegate_, OnSerializedPacket(_))
-      .WillOnce(Invoke([expected_buffer](SerializedPacket serialized_packet) {
-        EXPECT_EQ(expected_buffer, serialized_packet.encrypted_buffer);
+      .WillOnce(Invoke([&external_buffer](SerializedPacket serialized_packet) {
+        EXPECT_EQ(external_buffer.buffer, serialized_packet.encrypted_buffer);
       }));
   creator_.FlushCurrentPacket();
+  if (!GetQuicReloadableFlag(quic_avoid_leak_writer_buffer)) {
+    delete[] buffer;
+  }
 }
 
 // Test for error found in
@@ -2311,7 +2316,7 @@
               MaybeBundleAckOpportunistically,
               (),
               (override));
-  MOCK_METHOD(char*, GetPacketBuffer, (), (override));
+  MOCK_METHOD(QuicPacketBuffer, GetPacketBuffer, (), (override));
   MOCK_METHOD(void, OnSerializedPacket, (SerializedPacket), (override));
   MOCK_METHOD(void,
               OnUnrecoverableError,
@@ -2466,7 +2471,8 @@
                  &delegate_,
                  &producer_),
         ack_frame_(InitAckFrame(1)) {
-    EXPECT_CALL(delegate_, GetPacketBuffer()).WillRepeatedly(Return(nullptr));
+    EXPECT_CALL(delegate_, GetPacketBuffer())
+        .WillRepeatedly(Return(QuicPacketBuffer()));
     creator_.SetEncrypter(
         ENCRYPTION_FORWARD_SECURE,
         std::make_unique<NullEncrypter>(Perspective::IS_CLIENT));
diff --git a/quic/core/quic_packet_writer.h b/quic/core/quic_packet_writer.h
index ab29e15..6efdb60 100644
--- a/quic/core/quic_packet_writer.h
+++ b/quic/core/quic_packet_writer.h
@@ -125,17 +125,22 @@
   // True=Batch mode. False=PassThrough mode.
   virtual bool IsBatchMode() const = 0;
 
-  // PassThrough mode: Return nullptr.
+  // PassThrough mode: Return {nullptr, nullptr}
   //
   // Batch mode:
-  // Return the starting address for the next packet's data. A minimum of
+  // Return the QuicPacketBuffer for the next packet. A minimum of
   // kMaxOutgoingPacketSize is guaranteed to be available from the returned
-  // address. If the internal buffer does not have enough space, nullptr is
-  // returned. All arguments should be identical to the follow-up call to
-  // |WritePacket|, they are here to allow advanced packet memory management in
-  // packet writers, e.g. one packet buffer pool per |peer_address|.
-  virtual char* GetNextWriteLocation(const QuicIpAddress& self_address,
-                                     const QuicSocketAddress& peer_address) = 0;
+  // address. If the internal buffer does not have enough space,
+  // {nullptr, nullptr} is returned. All arguments should be identical to the
+  // follow-up call to |WritePacket|, they are here to allow advanced packet
+  // memory management in packet writers, e.g. one packet buffer pool per
+  // |peer_address|.
+  //
+  // If QuicPacketBuffer.release_buffer is !nullptr, it should be called iff
+  // the caller does not call WritePacket for the returned buffer.
+  virtual QuicPacketBuffer GetNextWriteLocation(
+      const QuicIpAddress& self_address,
+      const QuicSocketAddress& peer_address) = 0;
 
   // PassThrough mode: Return WriteResult(WRITE_STATUS_OK, 0).
   //
diff --git a/quic/core/quic_packet_writer_wrapper.cc b/quic/core/quic_packet_writer_wrapper.cc
index f1f25d8..d2b54d1 100644
--- a/quic/core/quic_packet_writer_wrapper.cc
+++ b/quic/core/quic_packet_writer_wrapper.cc
@@ -45,7 +45,7 @@
   return writer_->IsBatchMode();
 }
 
-char* QuicPacketWriterWrapper::GetNextWriteLocation(
+QuicPacketBuffer QuicPacketWriterWrapper::GetNextWriteLocation(
     const QuicIpAddress& self_address,
     const QuicSocketAddress& peer_address) {
   return writer_->GetNextWriteLocation(self_address, peer_address);
diff --git a/quic/core/quic_packet_writer_wrapper.h b/quic/core/quic_packet_writer_wrapper.h
index 0b331ae..afd3605 100644
--- a/quic/core/quic_packet_writer_wrapper.h
+++ b/quic/core/quic_packet_writer_wrapper.h
@@ -35,8 +35,9 @@
       const QuicSocketAddress& peer_address) const override;
   bool SupportsReleaseTime() const override;
   bool IsBatchMode() const override;
-  char* GetNextWriteLocation(const QuicIpAddress& self_address,
-                             const QuicSocketAddress& peer_address) override;
+  QuicPacketBuffer GetNextWriteLocation(
+      const QuicIpAddress& self_address,
+      const QuicSocketAddress& peer_address) override;
   WriteResult Flush() override;
 
   // Takes ownership of |writer|.
diff --git a/quic/core/quic_packets.h b/quic/core/quic_packets.h
index 04522e8..3e37a20 100644
--- a/quic/core/quic_packets.h
+++ b/quic/core/quic_packets.h
@@ -379,6 +379,8 @@
   SerializedPacket(SerializedPacket&& other);
   ~SerializedPacket();
 
+  // TODO(wub): replace |encrypted_buffer|+|release_encrypted_buffer| by a
+  // QuicOwnedPacketBuffer.
   // Not owned if |release_encrypted_buffer| is nullptr. Otherwise it is
   // released by |release_encrypted_buffer| on destruction.
   const char* encrypted_buffer;
diff --git a/quic/core/quic_types.h b/quic/core/quic_types.h
index f2919ef..e00ebb5 100644
--- a/quic/core/quic_types.h
+++ b/quic/core/quic_types.h
@@ -734,6 +734,45 @@
   bool allow_burst;
 };
 
+// QuicPacketBuffer bundles a buffer and a function that releases it. Note
+// it does not assume ownership of buffer, i.e. it doesn't release the buffer on
+// destruction.
+struct QUIC_NO_EXPORT QuicPacketBuffer {
+  QuicPacketBuffer() = default;
+
+  QuicPacketBuffer(char* buffer,
+                   std::function<void(const char*)> release_buffer)
+      : buffer(buffer), release_buffer(std::move(release_buffer)) {}
+
+  char* buffer = nullptr;
+  std::function<void(const char*)> release_buffer;
+};
+
+// QuicOwnedPacketBuffer is a QuicPacketBuffer that assumes buffer ownership.
+struct QUIC_NO_EXPORT QuicOwnedPacketBuffer : public QuicPacketBuffer {
+  QuicOwnedPacketBuffer(const QuicOwnedPacketBuffer&) = delete;
+  QuicOwnedPacketBuffer& operator=(const QuicOwnedPacketBuffer&) = delete;
+
+  QuicOwnedPacketBuffer(char* buffer,
+                        std::function<void(const char*)> release_buffer)
+      : QuicPacketBuffer(buffer, std::move(release_buffer)) {}
+
+  QuicOwnedPacketBuffer(QuicOwnedPacketBuffer&& owned_buffer)
+      : QuicPacketBuffer(std::move(owned_buffer)) {
+    // |owned_buffer| does not own a buffer any more.
+    owned_buffer.buffer = nullptr;
+  }
+
+  explicit QuicOwnedPacketBuffer(QuicPacketBuffer&& packet_buffer)
+      : QuicPacketBuffer(std::move(packet_buffer)) {}
+
+  ~QuicOwnedPacketBuffer() {
+    if (release_buffer != nullptr && buffer != nullptr) {
+      release_buffer(buffer);
+    }
+  }
+};
+
 }  // namespace quic
 
 #endif  // QUICHE_QUIC_CORE_QUIC_TYPES_H_