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.
