Dropping initial keys at the end of writing to avoid potential missing write keys in the middle of writing.
Protected by FLAGS_quic_reloadable_flag_quic_fix_missing_initial_keys2.
PiperOrigin-RevId: 334170250
Change-Id: I24e118e69c4cba3ddbed57df67340596446fc6d8
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index 642ebc7..ccc2166 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2971,10 +2971,8 @@
packet->nonretransmittable_frames, packet_send_time);
}
}
- if (fix_missing_initial_keys_ &&
- packet->encryption_level == ENCRYPTION_HANDSHAKE) {
- QUIC_RELOADABLE_FLAG_COUNT_N(quic_fix_missing_initial_keys, 1, 2);
- visitor_->OnHandshakePacketSent();
+ if (packet->encryption_level == ENCRYPTION_HANDSHAKE) {
+ handshake_packet_sent_ = true;
}
if (in_flight || !retransmission_alarm_->IsSet()) {
@@ -3833,7 +3831,9 @@
QuicConnection::ScopedPacketFlusher::ScopedPacketFlusher(
QuicConnection* connection)
: connection_(connection),
- flush_and_set_pending_retransmission_alarm_on_delete_(false) {
+ flush_and_set_pending_retransmission_alarm_on_delete_(false),
+ handshake_packet_sent_(connection != nullptr &&
+ connection->handshake_packet_sent_) {
if (connection_ == nullptr) {
return;
}
@@ -3886,6 +3886,13 @@
connection_->FlushCoalescedPacket();
}
connection_->FlushPackets();
+ if (connection_->fix_missing_initial_keys_ && !handshake_packet_sent_ &&
+ connection_->handshake_packet_sent_) {
+ QUIC_RELOADABLE_FLAG_COUNT_N(quic_fix_missing_initial_keys2, 1, 2);
+ // This would cause INITIAL key to be dropped. Drop keys here to avoid
+ // missing the write keys in the middle of writing.
+ connection_->visitor_->OnHandshakePacketSent();
+ }
// Reset transmission type.
connection_->SetTransmissionType(NOT_RETRANSMISSION);
@@ -4675,10 +4682,14 @@
return true;
}
if (fix_missing_initial_keys_) {
- QUIC_RELOADABLE_FLAG_COUNT_N(quic_fix_missing_initial_keys, 2, 2);
- if (!framer_.HasEncrypterOfEncryptionLevel(ENCRYPTION_INITIAL)) {
+ QUIC_RELOADABLE_FLAG_COUNT_N(quic_fix_missing_initial_keys2, 2, 2);
+ if (coalesced_packet_.ContainsPacketOfEncryptionLevel(ENCRYPTION_INITIAL) &&
+ !framer_.HasEncrypterOfEncryptionLevel(ENCRYPTION_INITIAL)) {
// Initial packet will be re-serialized. Neuter it in case initial key has
// been dropped.
+ QUIC_BUG << ENDPOINT
+ << "Coalescer contains initial packet while initial key has "
+ "been dropped.";
coalesced_packet_.NeuterInitialPacket();
}
}
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index e72ffb8..9b4544f 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -821,6 +821,8 @@
// If true, when this flusher goes out of scope, flush connection and set
// retransmission alarm if there is one pending.
bool flush_and_set_pending_retransmission_alarm_on_delete_;
+ // Latched connection's handshake_packet_sent_ on creation of this flusher.
+ const bool handshake_packet_sent_;
};
QuicPacketWriter* writer() { return writer_; }
@@ -1791,8 +1793,11 @@
// Indicate whether coalescing is done.
bool coalescing_done_ = false;
+ // Indicate whether any ENCRYPTION_HANDSHAKE packet has been sent.
+ bool handshake_packet_sent_ = false;
+
const bool fix_missing_initial_keys_ =
- GetQuicReloadableFlag(quic_fix_missing_initial_keys);
+ GetQuicReloadableFlag(quic_fix_missing_initial_keys2);
const bool fix_out_of_order_sending_ =
GetQuicReloadableFlag(quic_fix_out_of_order_sending2);
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index ab32a38..2c791f2 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -10276,7 +10276,9 @@
? QuicPacketNumber(3)
: QuicPacketNumber(4),
_, _));
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ if (!GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ }
connection_.GetRetransmissionAlarm()->Fire();
// Verify 1-RTT packet gets coalesced with handshake retransmission.
EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
@@ -10298,7 +10300,9 @@
OnPacketSent(_, _, handshake_retransmission + 1, _, _));
EXPECT_CALL(*send_algorithm_,
OnPacketSent(_, _, handshake_retransmission, _, _));
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ if (!GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ }
connection_.GetRetransmissionAlarm()->Fire();
// Verify 1-RTT packet gets coalesced with handshake retransmission.
EXPECT_EQ(0x01010101u, writer_->final_bytes_of_last_packet());
@@ -11078,7 +11082,11 @@
connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
std::make_unique<TaggingEncrypter>(0x02));
connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+ if (GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ } else {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+ }
connection_.SendCryptoDataWithString("foo", 0, ENCRYPTION_HANDSHAKE);
// Verify PTO time does not change.
EXPECT_EQ(expected_pto_time,
@@ -11117,8 +11125,11 @@
// Receives packet 1000 in handshake data.
ProcessCryptoPacketAtLevel(1000, ENCRYPTION_HANDSHAKE);
EXPECT_TRUE(connection_.HasPendingAcks());
-
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+ if (GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ } else {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+ }
connection_.SendCryptoDataWithString("foo", 0, ENCRYPTION_HANDSHAKE);
// Receives packet 1001 in handshake data.
@@ -11181,7 +11192,11 @@
connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
std::make_unique<TaggingEncrypter>(0x02));
connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(3);
+ if (GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ } else {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(3);
+ }
// Send HANDSHAKE 2 and 3.
connection_.SendCryptoDataWithString("foo", 0, ENCRYPTION_HANDSHAKE);
connection_.SendCryptoDataWithString("bar", 3, ENCRYPTION_HANDSHAKE);
@@ -11247,7 +11262,11 @@
connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
std::make_unique<TaggingEncrypter>(0x02));
connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+ if (GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ } else {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+ }
std::string handshake_crypto_data(1024, 'a');
connection_.SendCryptoDataWithString(handshake_crypto_data, 0,
ENCRYPTION_HANDSHAKE);
@@ -11319,7 +11338,11 @@
std::make_unique<TaggingEncrypter>(0x02));
connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
// Verify HANDSHAKE packet is coalesced with INITIAL retransmission.
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+ if (GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ } else {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+ }
std::string handshake_crypto_data(1024, 'a');
connection_.SendCryptoDataWithString(handshake_crypto_data, 0,
ENCRYPTION_HANDSHAKE);
@@ -11610,7 +11633,7 @@
SetQuicReloadableFlag(
quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded, true);
if (!connection_.version().CanSendCoalescedPackets() ||
- !GetQuicReloadableFlag(quic_fix_missing_initial_keys)) {
+ !GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
// Cannot set quic_fix_missing_initial_keys in the test since connection_ is
// created since the setup.
return;
@@ -11625,15 +11648,17 @@
connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
std::make_unique<TaggingEncrypter>(0x02));
connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
- {
- QuicConnection::ScopedPacketFlusher flusher(&connection_);
- connection_.GetAckAlarm()->Fire();
- // Verify this ACK packet is on hold.
- EXPECT_EQ(0u, writer_->packets_write_attempts());
-
- // Discard INITIAL key while there is an INITIAL packet in the coalescer.
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).WillOnce(Invoke([this]() {
connection_.RemoveEncrypter(ENCRYPTION_INITIAL);
connection_.NeuterUnencryptedPackets();
+ }));
+ {
+ QuicConnection::ScopedPacketFlusher flusher(&connection_);
+ connection_.SendCryptoDataWithString("foo", 0, ENCRYPTION_HANDSHAKE);
+ // Verify the packet is on hold.
+ EXPECT_EQ(0u, writer_->packets_write_attempts());
+ // Flush pending ACKs.
+ connection_.GetAckAlarm()->Fire();
}
// If not setting
// quic_neuter_initial_packet_in_coalescer_with_initial_key_discarded, there
@@ -12042,7 +12067,9 @@
connection_.GetSendAlarm()->Set(clock_.ApproximateNow());
// Fire ACK alarm.
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ if (!GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ }
connection_.GetAckAlarm()->Fire();
if (GetQuicReloadableFlag(quic_fix_pto_pending_timer_count)) {
// Verify 1-RTT packet is coalesced with handshake packet.
@@ -12057,7 +12084,9 @@
ASSERT_TRUE(connection_.GetRetransmissionAlarm()->IsSet());
if (GetQuicReloadableFlag(quic_fix_pto_pending_timer_count)) {
- EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ if (!GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
+ EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
+ }
} else {
EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(0);
EXPECT_CALL(visitor_, SendPing()).WillOnce(Invoke([this]() {
@@ -12080,12 +12109,11 @@
// Regression test for b/168294218.
TEST_P(QuicConnectionTest, CoalescerHandlesInitialKeyDiscard) {
if (!connection_.version().CanSendCoalescedPackets() ||
- !GetQuicReloadableFlag(quic_fix_missing_initial_keys)) {
+ !GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
return;
}
SetQuicReloadableFlag(quic_discard_initial_packet_with_key_dropped, true);
- // Verify only one HANDSHAKE packet gets sent.
- EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(1);
+ EXPECT_CALL(*send_algorithm_, OnPacketSent(_, _, _, _, _)).Times(2);
EXPECT_CALL(visitor_, OnHandshakePacketSent()).WillOnce(Invoke([this]() {
connection_.RemoveEncrypter(ENCRYPTION_INITIAL);
connection_.NeuterUnencryptedPackets();
@@ -12113,7 +12141,7 @@
// Regresstion test for b/168294218
TEST_P(QuicConnectionTest, ZeroRttRejectionAndMissingInitialKeys) {
if (!connection_.SupportsMultiplePacketNumberSpaces() ||
- !GetQuicReloadableFlag(quic_fix_missing_initial_keys)) {
+ !GetQuicReloadableFlag(quic_fix_missing_initial_keys2)) {
return;
}
// Not defer send in response to packet.
