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.