Refactor out sending address token logic into a stand alone method in QuicSession, and make it an interface in QuicConnectionVisitorInterface. PiperOrigin-RevId: 348531227 Change-Id: I213c82ee797269b42ba46691eaa03c254f035fcd

diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index f272a6a..dd453cb 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h

@@ -219,6 +219,10 @@
   // Consider the client address gets validated (and therefore remove
   // amplification factor) once the |token| gets successfully validated.
   virtual bool ValidateToken(absl::string_view token) const = 0;
+
+  // Called by the server to send another token.
+  // Return false if the crypto stream fail to generate one.
+  virtual void MaybeSendAddressToken() = 0;
 };
 
 // Interface which gets callbacks from the QuicConnection at interesting

diff --git a/quic/core/quic_session.cc b/quic/core/quic_session.cc
index b2d59c5..e0a5e09 100644
--- a/quic/core/quic_session.cc
+++ b/quic/core/quic_session.cc

@@ -1681,21 +1681,28 @@
         connection()->version().HasIetfQuicFrames()) {
       QUIC_RELOADABLE_FLAG_COUNT_N(quic_enable_token_based_address_validation,
                                    1, 2);
-      std::string address_token = GetCryptoStream()->GetAddressToken();
-      if (!address_token.empty()) {
-        const size_t buf_len = address_token.length() + 1;
-        auto buffer = std::make_unique<char[]>(buf_len);
-        QuicDataWriter writer(buf_len, buffer.get());
-        // Add prefix 0 for token sent in NEW_TOKEN frame.
-        writer.WriteUInt8(0);
-        writer.WriteBytes(address_token.data(), address_token.length());
-        control_frame_manager_.WriteOrBufferNewToken(
-            absl::string_view(buffer.get(), buf_len));
-      }
+      MaybeSendAddressToken();
     }
   }
 }
 
+void QuicSession::MaybeSendAddressToken() {
+  DCHECK(perspective_ == Perspective::IS_SERVER &&
+         connection()->version().HasIetfQuicFrames());
+  std::string address_token = GetCryptoStream()->GetAddressToken();
+  if (address_token.empty()) {
+    return;
+  }
+  const size_t buf_len = address_token.length() + 1;
+  auto buffer = std::make_unique<char[]>(buf_len);
+  QuicDataWriter writer(buf_len, buffer.get());
+  // Add prefix 0 for token sent in NEW_TOKEN frame.
+  writer.WriteUInt8(0);
+  writer.WriteBytes(address_token.data(), address_token.length());
+  control_frame_manager_.WriteOrBufferNewToken(
+      absl::string_view(buffer.get(), buf_len));
+}
+
 void QuicSession::DiscardOldDecryptionKey(EncryptionLevel level) {
   if (!connection()->version().KnowsWhichDecrypterToUse()) {
     return;

diff --git a/quic/core/quic_session.h b/quic/core/quic_session.h
index d849743..6067c8d 100644
--- a/quic/core/quic_session.h
+++ b/quic/core/quic_session.h

@@ -153,6 +153,7 @@
   std::unique_ptr<QuicEncrypter> CreateCurrentOneRttEncrypter() override;
   void BeforeConnectionCloseSent() override {}
   bool ValidateToken(absl::string_view token) const override;
+  void MaybeSendAddressToken() override;
 
   // QuicStreamFrameDataProducer
   WriteStreamDataResult WriteStreamData(QuicStreamId id,

diff --git a/quic/test_tools/quic_test_utils.h b/quic/test_tools/quic_test_utils.h
index 50afd2d..428d695 100644
--- a/quic/test_tools/quic_test_utils.h
+++ b/quic/test_tools/quic_test_utils.h

@@ -599,6 +599,7 @@
               (override));
   MOCK_METHOD(void, BeforeConnectionCloseSent, (), (override));
   MOCK_METHOD(bool, ValidateToken, (absl::string_view), (const, override));
+  MOCK_METHOD(void, MaybeSendAddressToken, (), (override));
 };
 
 class MockQuicConnectionHelper : public QuicConnectionHelperInterface {

diff --git a/quic/test_tools/simulator/quic_endpoint.h b/quic/test_tools/simulator/quic_endpoint.h
index fd2fa06..de755f0 100644
--- a/quic/test_tools/simulator/quic_endpoint.h
+++ b/quic/test_tools/simulator/quic_endpoint.h

@@ -106,6 +106,7 @@
   bool ValidateToken(absl::string_view /*token*/) const override {
     return true;
   }
+  void MaybeSendAddressToken() override {}
 
   // End QuicConnectionVisitorInterface implementation.