Add unused interface QuicConnection::SendPathChallenge().

This interface is implemented by 2 code paths:
If the PATH_CHALLENGE is sent on the default socket, add the frame as other frames to the current serialize packet in packet creator and flush at its due time. Since the destination address might be different for that packet, use a ScopedPeerAddressContext to flush it if the peer address is different from the one in used at the call site.
If the PATH_CHALLENGE is sent on an alternative socket, send it in a stand alone code path by serializing it into a individual packet and writing it to the wire directly.

PiperOrigin-RevId: 337176402
Change-Id: I90c4b88809209dd59a31a98c58d0de3ecd9c4889
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index cfb6b3e..f3b0eda 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -28,6 +28,7 @@
 #include "net/third_party/quiche/src/quic/core/quic_error_codes.h"
 #include "net/third_party/quiche/src/quic/core/quic_legacy_version_encapsulator.h"
 #include "net/third_party/quiche/src/quic/core/quic_packet_creator.h"
+#include "net/third_party/quiche/src/quic/core/quic_packet_writer.h"
 #include "net/third_party/quiche/src/quic/core/quic_types.h"
 #include "net/third_party/quiche/src/quic/core/quic_utils.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_bug_tracker.h"
@@ -4227,24 +4228,32 @@
       transmitted_connectivity_probe_payload_ = nullptr;
     }
   }
-
   DCHECK_EQ(IsRetransmittable(*probing_packet), NO_RETRANSMITTABLE_DATA);
+  return WritePacketUsingWriter(std::move(probing_packet), probing_writer,
+                                self_address(), peer_address,
+                                /*measure_rtt=*/true);
+}
 
+bool QuicConnection::WritePacketUsingWriter(
+    std::unique_ptr<SerializedPacket> packet,
+    QuicPacketWriter* writer,
+    const QuicSocketAddress& self_address,
+    const QuicSocketAddress& peer_address,
+    bool measure_rtt) {
   const QuicTime packet_send_time = clock_->Now();
   QUIC_DVLOG(2) << ENDPOINT
                 << "Sending path probe packet for server connection ID "
                 << server_connection_id_ << std::endl
-                << quiche::QuicheTextUtils::HexDump(
-                       absl::string_view(probing_packet->encrypted_buffer,
-                                         probing_packet->encrypted_length));
-  WriteResult result = probing_writer->WritePacket(
-      probing_packet->encrypted_buffer, probing_packet->encrypted_length,
-      self_address().host(), peer_address, per_packet_options_);
+                << quiche::QuicheTextUtils::HexDump(absl::string_view(
+                       packet->encrypted_buffer, packet->encrypted_length));
+  WriteResult result = writer->WritePacket(
+      packet->encrypted_buffer, packet->encrypted_length, self_address.host(),
+      peer_address, per_packet_options_);
 
   // If using a batch writer and the probing packet is buffered, flush it.
-  if (probing_writer->IsBatchMode() && result.status == WRITE_STATUS_OK &&
+  if (writer->IsBatchMode() && result.status == WRITE_STATUS_OK &&
       result.bytes_written == 0) {
-    result = probing_writer->Flush();
+    result = writer->Flush();
   }
 
   if (IsWriteError(result.status)) {
@@ -4257,14 +4266,14 @@
 
   if (!sent_packet_manager_.give_sent_packet_to_debug_visitor_after_sent() &&
       debug_visitor_ != nullptr) {
-    debug_visitor_->OnPacketSent(
-        *probing_packet, probing_packet->transmission_type, packet_send_time);
+    debug_visitor_->OnPacketSent(*packet, packet->transmission_type,
+                                 packet_send_time);
   }
 
   // Send in currrent path. Call OnPacketSent regardless of the write result.
-  sent_packet_manager_.OnPacketSent(
-      probing_packet.get(), packet_send_time, probing_packet->transmission_type,
-      NO_RETRANSMITTABLE_DATA, /*measure_rtt=*/true);
+  sent_packet_manager_.OnPacketSent(packet.get(), packet_send_time,
+                                    packet->transmission_type,
+                                    NO_RETRANSMITTABLE_DATA, measure_rtt);
 
   if (sent_packet_manager_.give_sent_packet_to_debug_visitor_after_sent() &&
       debug_visitor_ != nullptr) {
@@ -4274,18 +4283,18 @@
       QUIC_BUG << "Unacked map is empty right after packet is sent";
     } else {
       debug_visitor_->OnPacketSent(
-          probing_packet->packet_number, probing_packet->encrypted_length,
-          probing_packet->has_crypto_handshake,
-          probing_packet->transmission_type, probing_packet->encryption_level,
+          packet->packet_number, packet->encrypted_length,
+          packet->has_crypto_handshake, packet->transmission_type,
+          packet->encryption_level,
           sent_packet_manager_.unacked_packets()
               .rbegin()
               ->retransmittable_frames,
-          probing_packet->nonretransmittable_frames, packet_send_time);
+          packet->nonretransmittable_frames, packet_send_time);
     }
   }
 
   if (IsWriteBlockedStatus(result.status)) {
-    if (probing_writer == writer_) {
+    if (writer == writer_) {
       // Visitor should not be write blocked if the probing writer is not the
       // default packet writer.
       visitor_->OnWriteBlocked();
@@ -5196,6 +5205,26 @@
   return num_rtos_for_blackhole_detection_ > 0;
 }
 
+void QuicConnection::SendPathChallenge(QuicPathFrameBuffer* data_buffer,
+                                       const QuicSocketAddress& self_address,
+                                       const QuicSocketAddress& peer_address,
+                                       QuicPacketWriter* writer) {
+  if (writer == writer_) {
+    // It's on current path, add the PATH_CHALLENGE the same way as other
+    // frames.
+    QuicPacketCreator::ScopedPeerAddressContext context(&packet_creator_,
+                                                        peer_address);
+    packet_creator_.AddPathChallengeFrame(data_buffer);
+    return;
+  }
+  std::unique_ptr<SerializedPacket> probing_packet =
+      packet_creator_.SerializePathChallengeConnectivityProbingPacket(
+          data_buffer);
+  DCHECK_EQ(IsRetransmittable(*probing_packet), NO_RETRANSMITTABLE_DATA);
+  WritePacketUsingWriter(std::move(probing_packet), writer, self_address,
+                         peer_address, /*measure_rtt=*/false);
+}
+
 bool QuicConnection::SendPathResponse(const QuicPathFrameBuffer& data_buffer,
                                       QuicSocketAddress peer_address_to_send) {
   // Send PATH_RESPONSE using the provided peer address. If the creator has been
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index c9f366c..e02b573 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -1079,6 +1079,17 @@
   // false.
   bool MaybeTestLiveness();
 
+  // Send PATH_CHALLENGE using the given path information. If |writer| is the
+  // default writer, PATH_CHALLENGE can be bundled with other frames, and the
+  // containing packet can be buffered if the writer is blocked. Otherwise,
+  // PATH_CHALLENGE will be written in an individual packet and it will be
+  // dropped if write fails. |data_buffer| will be populated with the payload
+  // for future validation.
+  void SendPathChallenge(QuicPathFrameBuffer* data_buffer,
+                         const QuicSocketAddress& self_address,
+                         const QuicSocketAddress& peer_address,
+                         QuicPacketWriter* writer);
+
   bool can_receive_ack_frequency_frame() const {
     return can_receive_ack_frequency_frame_;
   }
@@ -1461,6 +1472,13 @@
   // Send PING at encryption level.
   void SendPingAtLevel(EncryptionLevel level);
 
+  // Write the given packet with |self_address| and |peer_address| using
+  // |writer|.
+  bool WritePacketUsingWriter(std::unique_ptr<SerializedPacket> packet,
+                              QuicPacketWriter* writer,
+                              const QuicSocketAddress& self_address,
+                              const QuicSocketAddress& peer_address,
+                              bool measure_rtt);
   QuicFramer framer_;
 
   // Contents received in the current packet, especially used to identify
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index e40fa67..132bee2 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -18,8 +18,10 @@
 #include "net/third_party/quiche/src/quic/core/crypto/null_encrypter.h"
 #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/frames/quic_connection_close_frame.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_error_codes.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"
@@ -6453,6 +6455,9 @@
 
 TEST_P(QuicConnectionTest, WriterBlockedAfterServerSendsConnectivityProbe) {
   PathProbeTestInit(Perspective::IS_SERVER);
+  if (version().SupportsAntiAmplificationLimit()) {
+    QuicConnectionPeer::SetAddressValidated(&connection_);
+  }
 
   // Block next write so that sending connectivity probe will encounter a
   // blocked write when send a connectivity probe to the peer.
@@ -6463,8 +6468,16 @@
 
   EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, QuicPacketNumber(1), _, _))
       .Times(1);
-  connection_.SendConnectivityProbingPacket(writer_.get(),
-                                            connection_.peer_address());
+  if (connection_.send_path_response() &&
+      VersionHasIetfQuicFrames(GetParam().version.transport_version)) {
+    QuicPathFrameBuffer payload;
+    QuicConnection::ScopedPacketFlusher flusher(&connection_);
+    connection_.SendPathChallenge(&payload, connection_.self_address(),
+                                  connection_.peer_address(), writer_.get());
+  } else {
+    connection_.SendConnectivityProbingPacket(writer_.get(),
+                                              connection_.peer_address());
+  }
 }
 
 TEST_P(QuicConnectionTest, WriterErrorWhenClientSendsConnectivityProbe) {
@@ -11210,6 +11223,178 @@
   EXPECT_TRUE(connection_.connected());
 }
 
+TEST_P(QuicConnectionTest, SendPathChallenge) {
+  if (!VersionHasIetfQuicFrames(connection_.version().transport_version) ||
+      !connection_.send_path_response()) {
+    return;
+  }
+  PathProbeTestInit(Perspective::IS_CLIENT);
+  const QuicSocketAddress kNewSourceAddress(QuicIpAddress::Any6(), 12345);
+  EXPECT_NE(kNewSourceAddress, connection_.self_address());
+  TestPacketWriter new_writer(version(), &clock_, Perspective::IS_CLIENT);
+  QuicPathFrameBuffer payload;
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _))
+      .WillOnce(Invoke([&]() {
+        EXPECT_EQ(1u, new_writer.packets_write_attempts());
+        EXPECT_EQ(1u, new_writer.path_challenge_frames().size());
+        EXPECT_EQ(
+            0, memcmp(payload.data(),
+                      &(new_writer.path_challenge_frames().front().data_buffer),
+                      sizeof(payload)));
+        EXPECT_EQ(1u, new_writer.padding_frames().size());
+        EXPECT_EQ(kNewSourceAddress.host(),
+                  new_writer.last_write_source_address());
+      }));
+  connection_.SendPathChallenge(&payload, kNewSourceAddress,
+                                connection_.peer_address(), &new_writer);
+  EXPECT_EQ(0u, writer_->packets_write_attempts());
+}
+
+TEST_P(QuicConnectionTest, SendPathChallengeUsingBlockedNewSocket) {
+  if (!VersionHasIetfQuicFrames(connection_.version().transport_version) ||
+      !connection_.send_path_response()) {
+    return;
+  }
+  PathProbeTestInit(Perspective::IS_CLIENT);
+  const QuicSocketAddress kNewSourceAddress(QuicIpAddress::Any6(), 12345);
+  EXPECT_NE(kNewSourceAddress, connection_.self_address());
+  TestPacketWriter new_writer(version(), &clock_, Perspective::IS_CLIENT);
+  new_writer.BlockOnNextWrite();
+  QuicPathFrameBuffer payload;
+  EXPECT_CALL(visitor_, OnWriteBlocked()).Times(0);
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _))
+      .WillOnce(Invoke([&]() {
+        // Even though the socket is blocked, the PATH_CHALLENGE should still be
+        // treated as sent.
+        EXPECT_EQ(1u, new_writer.packets_write_attempts());
+        EXPECT_EQ(1u, new_writer.path_challenge_frames().size());
+        EXPECT_EQ(
+            0, memcmp(payload.data(),
+                      &(new_writer.path_challenge_frames().front().data_buffer),
+                      sizeof(payload)));
+        EXPECT_EQ(1u, new_writer.padding_frames().size());
+        EXPECT_EQ(kNewSourceAddress.host(),
+                  new_writer.last_write_source_address());
+      }));
+  connection_.SendPathChallenge(&payload, kNewSourceAddress,
+                                connection_.peer_address(), &new_writer);
+  EXPECT_EQ(0u, writer_->packets_write_attempts());
+
+  new_writer.SetWritable();
+  // Write event on the default socket shouldn't make any difference.
+  connection_.OnCanWrite();
+  EXPECT_EQ(0u, writer_->packets_write_attempts());
+  EXPECT_EQ(1u, new_writer.packets_write_attempts());
+}
+
+TEST_P(QuicConnectionTest, SendPathChallengeWithDefaultSocketBlocked) {
+  if (!VersionHasIetfQuicFrames(connection_.version().transport_version) ||
+      !connection_.send_path_response()) {
+    return;
+  }
+  PathProbeTestInit(Perspective::IS_SERVER);
+  if (version().SupportsAntiAmplificationLimit()) {
+    QuicConnectionPeer::SetAddressValidated(&connection_);
+  }
+  const QuicSocketAddress kNewPeerAddress(QuicIpAddress::Any6(), 12345);
+  writer_->BlockOnNextWrite();
+  QuicPathFrameBuffer payload;
+  // 1st time is after writer returns WRITE_STATUS_BLOCKED. 2nd time is in
+  // ShouldGeneratePacket();
+  EXPECT_CALL(visitor_, OnWriteBlocked()).Times(2u);
+  // This packet isn't sent actually, instead it is buffered in the connection.
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _))
+      .WillOnce(Invoke([&]() {
+        EXPECT_EQ(1u, writer_->path_challenge_frames().size());
+        EXPECT_EQ(
+            0, memcmp(payload.data(),
+                      &(writer_->path_challenge_frames().front().data_buffer),
+                      sizeof(payload)));
+        EXPECT_EQ(1u, writer_->padding_frames().size());
+        EXPECT_EQ(kNewPeerAddress, writer_->last_write_peer_address());
+      }));
+  connection_.SendPathChallenge(&payload, connection_.self_address(),
+                                kNewPeerAddress, writer_.get());
+  EXPECT_EQ(1u, writer_->packets_write_attempts());
+
+  memset(payload.data(), 0, sizeof(payload));
+  // Try again with the new socket blocked from the beginning. The 2nd
+  // PATH_CHALLENGE shouldn't be serialized, but be dropped.
+  connection_.SendPathChallenge(&payload, connection_.self_address(),
+                                kNewPeerAddress, writer_.get());
+  // No more write attempt should be made.
+  EXPECT_EQ(1u, writer_->packets_write_attempts());
+
+  writer_->SetWritable();
+  // OnCanWrite() should actually write out the 1st PATH_CHALLENGE packet
+  // buffered earlier, thus incrementing the write counter.
+  connection_.OnCanWrite();
+  EXPECT_EQ(2u, writer_->packets_write_attempts());
+}
+
+// Tests that write error on the alternate socket should be ignored.
+TEST_P(QuicConnectionTest, SendPathChallengeFailOnNewSocket) {
+  if (!VersionHasIetfQuicFrames(connection_.version().transport_version) ||
+      !connection_.send_path_response()) {
+    return;
+  }
+  PathProbeTestInit(Perspective::IS_CLIENT);
+  const QuicSocketAddress kNewSourceAddress(QuicIpAddress::Any6(), 12345);
+  EXPECT_NE(kNewSourceAddress, connection_.self_address());
+  TestPacketWriter new_writer(version(), &clock_, Perspective::IS_CLIENT);
+  new_writer.SetShouldWriteFail();
+  QuicPathFrameBuffer payload;
+  EXPECT_CALL(visitor_, OnConnectionClosed(_, ConnectionCloseSource::FROM_SELF))
+      .Times(0);
+
+  connection_.SendPathChallenge(&payload, kNewSourceAddress,
+                                connection_.peer_address(), &new_writer);
+  // Regardless of the write error, the PATH_CHALLENGE should still be
+  // treated as sent.
+  EXPECT_EQ(1u, new_writer.packets_write_attempts());
+  EXPECT_EQ(1u, new_writer.path_challenge_frames().size());
+  EXPECT_EQ(0, memcmp(payload.data(),
+                      &(new_writer.path_challenge_frames().front().data_buffer),
+                      sizeof(payload)));
+  EXPECT_EQ(1u, new_writer.padding_frames().size());
+  EXPECT_EQ(kNewSourceAddress.host(), new_writer.last_write_source_address());
+  EXPECT_EQ(0u, writer_->packets_write_attempts());
+  EXPECT_TRUE(connection_.connected());
+}
+
+// Tests that write error while sending PATH_CHALLANGE from the default socket
+// should close the connection.
+TEST_P(QuicConnectionTest, SendPathChallengeFailOnDefaultPath) {
+  if (!VersionHasIetfQuicFrames(connection_.version().transport_version) ||
+      !connection_.send_path_response()) {
+    return;
+  }
+  PathProbeTestInit(Perspective::IS_CLIENT);
+  writer_->SetShouldWriteFail();
+  QuicPathFrameBuffer payload;
+  EXPECT_CALL(visitor_, OnConnectionClosed(_, ConnectionCloseSource::FROM_SELF))
+      .WillOnce(
+          Invoke([](QuicConnectionCloseFrame frame, ConnectionCloseSource) {
+            EXPECT_EQ(QUIC_PACKET_WRITE_ERROR, frame.quic_error_code);
+          }));
+  EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(0u);
+  {
+    // Add a flusher to force flush, otherwise the frames will remain in the
+    // packet creator.
+    QuicConnection::ScopedPacketFlusher flusher(&connection_);
+    connection_.SendPathChallenge(&payload, connection_.self_address(),
+                                  connection_.peer_address(), writer_.get());
+  }
+  EXPECT_EQ(1u, writer_->packets_write_attempts());
+  EXPECT_EQ(1u, writer_->path_challenge_frames().size());
+  EXPECT_EQ(0, memcmp(payload.data(),
+                      &(writer_->path_challenge_frames().front().data_buffer),
+                      sizeof(payload)));
+  EXPECT_EQ(1u, writer_->padding_frames().size());
+  EXPECT_EQ(connection_.peer_address(), writer_->last_write_peer_address());
+  EXPECT_FALSE(connection_.connected());
+}
+
 // Check that if there are two PATH_CHALLENGE frames in the packet, the latter
 // one is ignored.
 TEST_P(QuicConnectionTest, ReceiveMultiplePathChallenge) {
diff --git a/quic/core/quic_packet_creator.cc b/quic/core/quic_packet_creator.cc
index 85960c4..de6f77f 100644
--- a/quic/core/quic_packet_creator.cc
+++ b/quic/core/quic_packet_creator.cc
@@ -2123,11 +2123,39 @@
   packet_.encryption_level = level;
 }
 
+void QuicPacketCreator::AddPathChallengeFrame(QuicPathFrameBuffer* payload) {
+  // Write a PATH_CHALLENGE frame, which has a random 8-byte payload.
+  random_->RandBytes(payload->data(), payload->size());
+  auto path_challenge_frame = new QuicPathChallengeFrame(0, *payload);
+  QuicFrame frame(path_challenge_frame);
+  if (AddPaddedFrameWithRetry(frame)) {
+    return;
+  }
+  // Fail silently if the probing packet cannot be written, path validation
+  // initiator will retry sending automatically.
+  // TODO(danzh) This will consume retry budget, if it causes performance
+  // regression, consider to notify the caller about the sending failure and let
+  // the caller to decide if it worth retrying.
+  QUIC_DVLOG(1) << ENDPOINT << "Can't send PATH_CHALLENGE now";
+  delete path_challenge_frame;
+}
+
 bool QuicPacketCreator::AddPathResponseFrame(
     const QuicPathFrameBuffer& data_buffer) {
   auto path_response =
       new QuicPathResponseFrame(kInvalidControlFrameId, data_buffer);
   QuicFrame frame(path_response);
+  if (AddPaddedFrameWithRetry(frame)) {
+    return true;
+  }
+
+  QUIC_DVLOG(1) << ENDPOINT << "Can't send PATH_RESPONSE now";
+  QUIC_RELOADABLE_FLAG_COUNT_N(quic_send_path_response, 5, 5);
+  delete path_response;
+  return false;
+}
+
+bool QuicPacketCreator::AddPaddedFrameWithRetry(const QuicFrame& frame) {
   if (HasPendingFrames()) {
     if (AddPaddedSavedFrame(frame, NOT_RETRANSMISSION)) {
       // Frame is queued.
@@ -2138,14 +2166,12 @@
   DCHECK(!HasPendingFrames());
   if (!delegate_->ShouldGeneratePacket(NO_RETRANSMITTABLE_DATA,
                                        NOT_HANDSHAKE)) {
-    QUIC_DVLOG(1) << ENDPOINT << "Can't send PATH_RESPONSE now";
-    QUIC_RELOADABLE_FLAG_COUNT_N(quic_send_path_response, 5, 5);
-    delete path_response;
     return false;
   }
   bool success = AddPaddedSavedFrame(frame, NOT_RETRANSMISSION);
   QUIC_BUG_IF(!success);
   return true;
 }
+
 #undef ENDPOINT  // undef for jumbo builds
 }  // namespace quic
diff --git a/quic/core/quic_packet_creator.h b/quic/core/quic_packet_creator.h
index 882b17d..0307566 100644
--- a/quic/core/quic_packet_creator.h
+++ b/quic/core/quic_packet_creator.h
@@ -258,6 +258,11 @@
   // Add PATH_RESPONSE to current packet, flush before or afterwards if needed.
   bool AddPathResponseFrame(const QuicPathFrameBuffer& data_buffer);
 
+  // Add PATH_CHALLENGE to current packet, flush before or afterwards if needed.
+  // This is a best effort adding. It may fail becasue of delegate state, but
+  // it's okay because of path validation retry mechanism.
+  void AddPathChallengeFrame(QuicPathFrameBuffer* payload);
+
   // Returns a dummy packet that is valid but contains no useful information.
   static SerializedPacket NoPacket();
 
@@ -596,6 +601,11 @@
   // Returns true and close connection if it attempts to send unencrypted data.
   bool AttemptingToSendUnencryptedStreamData();
 
+  // Add the given frame to the current packet with full padding. If the current
+  // packet doesn't have enough space, flush once and try again. Return false if
+  // fail to add.
+  bool AddPaddedFrameWithRetry(const QuicFrame& frame);
+
   // Does not own these delegates or the framer.
   DelegateInterface* delegate_;
   DebugDelegate* debug_delegate_;
diff --git a/quic/test_tools/quic_test_utils.cc b/quic/test_tools/quic_test_utils.cc
index 35edea2..8ba7cd4 100644
--- a/quic/test_tools/quic_test_utils.cc
+++ b/quic/test_tools/quic_test_utils.cc
@@ -1392,9 +1392,10 @@
 
 WriteResult TestPacketWriter::WritePacket(const char* buffer,
                                           size_t buf_len,
-                                          const QuicIpAddress& /*self_address*/,
+                                          const QuicIpAddress& self_address,
                                           const QuicSocketAddress& peer_address,
                                           PerPacketOptions* /*options*/) {
+  last_write_source_address_ = self_address;
   last_write_peer_address_ = peer_address;
   // If the buffer is allocated from the pool, return it back to the pool.
   // Note the buffer content doesn't change.
diff --git a/quic/test_tools/quic_test_utils.h b/quic/test_tools/quic_test_utils.h
index f2ad3be..30798c7 100644
--- a/quic/test_tools/quic_test_utils.h
+++ b/quic/test_tools/quic_test_utils.h
@@ -2121,6 +2121,10 @@
 
   SimpleQuicFramer* framer() { return &framer_; }
 
+  const QuicIpAddress& last_write_source_address() const {
+    return last_write_source_address_;
+  }
+
   const QuicSocketAddress& last_write_peer_address() const {
     return last_write_peer_address_;
   }
@@ -2164,7 +2168,8 @@
   QuicHashMap<char*, PacketBuffer*> packet_buffer_pool_index_;
   // Indices in packet_buffer_pool_ that are not allocated.
   std::list<PacketBuffer*> packet_buffer_free_list_;
-  // The peer address passed into WritePacket().
+  // The soruce/peer address passed into WritePacket().
+  QuicIpAddress last_write_source_address_;
   QuicSocketAddress last_write_peer_address_;
 };