In quic, let client bundle handshake data (containing client finished) with handshake ack. protected by existing gfe2_reloadable_flag_quic_bundle_crypto_data_with_initial_ack.

PiperOrigin-RevId: 317663970
Change-Id: Ibe9b10b676899e7496e46810533447fedaf1b3a8
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index 27bb4c4..b5c496f 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -4237,15 +4237,16 @@
   return ENCRYPTION_INITIAL;
 }
 
-void QuicConnection::MaybeBundleCryptoDataWithInitialAck() {
+void QuicConnection::MaybeBundleCryptoDataWithAckOfSpace(
+    PacketNumberSpace space) {
   DCHECK(SupportsMultiplePacketNumberSpaces());
-  const QuicTime initial_ack_timeout =
-      uber_received_packet_manager_.GetAckTimeout(INITIAL_DATA);
-  if (!initial_ack_timeout.IsInitialized() ||
-      (initial_ack_timeout > clock_->ApproximateNow() &&
-       initial_ack_timeout >
-           uber_received_packet_manager_.GetEarliestAckTimeout())) {
-    // Not going to send initial ACK.
+  QUIC_BUG_IF(space == APPLICATION_DATA);
+  const QuicTime ack_timeout =
+      uber_received_packet_manager_.GetAckTimeout(space);
+  if (!ack_timeout.IsInitialized() ||
+      (ack_timeout > clock_->ApproximateNow() &&
+       ack_timeout > uber_received_packet_manager_.GetEarliestAckTimeout())) {
+    // No pending ACK of space.
     return;
   }
   if (coalesced_packet_.length() > 0) {
@@ -4253,9 +4254,8 @@
     // packets.
     return;
   }
-  // Initial ACK will be padded to full anyway, so try to bundle INITIAL crypto
-  // data.
-  sent_packet_manager_.RetransmitInitialDataIfAny();
+
+  sent_packet_manager_.RetransmitDataOfSpaceIfAny(space);
 }
 
 void QuicConnection::SendAllPendingAcks() {
@@ -4265,9 +4265,13 @@
   QuicTime earliest_ack_timeout =
       uber_received_packet_manager_.GetEarliestAckTimeout();
   QUIC_BUG_IF(!earliest_ack_timeout.IsInitialized());
-  if (GetQuicReloadableFlag(quic_bundle_crypto_data_with_initial_ack) &&
-      perspective() == Perspective::IS_SERVER) {
-    MaybeBundleCryptoDataWithInitialAck();
+  if (GetQuicReloadableFlag(quic_bundle_crypto_data_with_initial_ack)) {
+    // On the server side, sends INITIAL data with INITIAL ACK. On the client
+    // side, sends HANDSHAKE data (containing client Finished) with HANDSHAKE
+    // ACK.
+    PacketNumberSpace space =
+        perspective() == Perspective::IS_SERVER ? INITIAL_DATA : HANDSHAKE_DATA;
+    MaybeBundleCryptoDataWithAckOfSpace(space);
     earliest_ack_timeout =
         uber_received_packet_manager_.GetEarliestAckTimeout();
     if (!earliest_ack_timeout.IsInitialized()) {
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index 6cf01f0..bc30a70 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -1318,8 +1318,9 @@
   bool ValidateConfigConnectionIds(const QuicConfig& config);
   bool ValidateConfigConnectionIdsOld(const QuicConfig& config);
 
-  // Called when ACK alarm goes off. Try to bundle INITIAL data with the ACK.
-  void MaybeBundleCryptoDataWithInitialAck();
+  // Called when ACK alarm goes off. Try to bundle crypto data with the ACK of
+  // |space|.
+  void MaybeBundleCryptoDataWithAckOfSpace(PacketNumberSpace space);
 
   // Returns true if an undecryptable packet of |decryption_level| should be
   // buffered (such that connection can try to decrypt it later).
diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 44a1b69..affd89f 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -1373,9 +1373,8 @@
                                     level);
   }
 
-  size_t ProcessCryptoPacketAtLevel(uint64_t number,
-                                    EncryptionLevel /*level*/) {
-    QuicPacketHeader header = ConstructPacketHeader(number, ENCRYPTION_INITIAL);
+  size_t ProcessCryptoPacketAtLevel(uint64_t number, EncryptionLevel level) {
+    QuicPacketHeader header = ConstructPacketHeader(number, level);
     QuicFrames frames;
     if (QuicVersionUsesCryptoFrames(connection_.transport_version())) {
       frames.push_back(QuicFrame(&crypto_frame_));
@@ -1385,10 +1384,10 @@
     frames.push_back(QuicFrame(QuicPaddingFrame(-1)));
     std::unique_ptr<QuicPacket> packet = ConstructPacket(header, frames);
     char buffer[kMaxOutgoingPacketSize];
-    peer_creator_.set_encryption_level(ENCRYPTION_INITIAL);
-    size_t encrypted_length = peer_framer_.EncryptPayload(
-        ENCRYPTION_INITIAL, QuicPacketNumber(number), *packet, buffer,
-        kMaxOutgoingPacketSize);
+    peer_creator_.set_encryption_level(level);
+    size_t encrypted_length =
+        peer_framer_.EncryptPayload(level, QuicPacketNumber(number), *packet,
+                                    buffer, kMaxOutgoingPacketSize);
     connection_.ProcessUdpPacket(
         kSelfAddress, kPeerAddress,
         QuicReceivedPacket(buffer, encrypted_length, clock_.Now(), false));
@@ -11257,7 +11256,7 @@
   EXPECT_EQ(0u, QuicConnectionPeer::NumUndecryptablePackets(&connection_));
 }
 
-TEST_P(QuicConnectionTest, BundleInitialDataWithInitialAck) {
+TEST_P(QuicConnectionTest, ServerBundlesInitialDataWithInitialAck) {
   if (!connection_.SupportsMultiplePacketNumberSpaces()) {
     return;
   }
@@ -11308,6 +11307,44 @@
   }
 }
 
+TEST_P(QuicConnectionTest, ClientBundlesHandshakeDataWithHandshakeAck) {
+  if (!connection_.SupportsMultiplePacketNumberSpaces()) {
+    return;
+  }
+  EXPECT_EQ(Perspective::IS_CLIENT, connection_.perspective());
+  if (QuicVersionUsesCryptoFrames(connection_.transport_version())) {
+    EXPECT_CALL(visitor_, OnCryptoFrame(_)).Times(AnyNumber());
+  }
+  EXPECT_CALL(visitor_, OnStreamFrame(_)).Times(AnyNumber());
+  use_tagging_decrypter();
+  connection_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                           std::make_unique<TaggingEncrypter>(0x02));
+  connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
+  SetDecrypter(ENCRYPTION_HANDSHAKE,
+               std::make_unique<StrictTaggingDecrypter>(0x02));
+  peer_framer_.SetEncrypter(ENCRYPTION_HANDSHAKE,
+                            std::make_unique<TaggingEncrypter>(0x02));
+  // Receives packet 1000 in handshake data.
+  ProcessCryptoPacketAtLevel(1000, ENCRYPTION_HANDSHAKE);
+  EXPECT_TRUE(connection_.HasPendingAcks());
+
+  EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(2);
+  connection_.SendCryptoDataWithString("foo", 0, ENCRYPTION_HANDSHAKE);
+
+  // Receives packet 1001 in handshake data.
+  ProcessCryptoPacketAtLevel(1001, ENCRYPTION_HANDSHAKE);
+  EXPECT_TRUE(connection_.HasPendingAcks());
+  // Receives packet 1002 in handshake data.
+  ProcessCryptoPacketAtLevel(1002, ENCRYPTION_HANDSHAKE);
+  EXPECT_FALSE(writer_->ack_frames().empty());
+  if (GetQuicReloadableFlag(quic_bundle_crypto_data_with_initial_ack)) {
+    // Verify CRYPTO frame is bundled with HANDSHAKE ACK.
+    EXPECT_FALSE(writer_->crypto_frames().empty());
+  } else {
+    EXPECT_TRUE(writer_->crypto_frames().empty());
+  }
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/quic_sent_packet_manager.cc b/quic/core/quic_sent_packet_manager.cc
index eeeddf2..86628ef 100644
--- a/quic/core/quic_sent_packet_manager.cc
+++ b/quic/core/quic_sent_packet_manager.cc
@@ -916,11 +916,11 @@
   pto_exponential_backoff_start_point_ = exponential_backoff_start_point;
 }
 
-void QuicSentPacketManager::RetransmitInitialDataIfAny() {
+void QuicSentPacketManager::RetransmitDataOfSpaceIfAny(
+    PacketNumberSpace space) {
   DCHECK(supports_multiple_packet_number_spaces());
-  if (!unacked_packets_.GetLastInFlightPacketSentTime(INITIAL_DATA)
-           .IsInitialized()) {
-    // No in flight initial data.
+  if (!unacked_packets_.GetLastInFlightPacketSentTime(space).IsInitialized()) {
+    // No in flight data of space.
     return;
   }
   QuicPacketNumber packet_number = unacked_packets_.GetLeastUnacked();
@@ -928,8 +928,7 @@
        it != unacked_packets_.end(); ++it, ++packet_number) {
     if (it->state == OUTSTANDING &&
         unacked_packets_.HasRetransmittableFrames(*it) &&
-        unacked_packets_.GetPacketNumberSpace(it->encryption_level) ==
-            INITIAL_DATA) {
+        unacked_packets_.GetPacketNumberSpace(it->encryption_level) == space) {
       DCHECK(it->in_flight);
       MarkForRetransmission(packet_number, PTO_RETRANSMISSION);
       return;
diff --git a/quic/core/quic_sent_packet_manager.h b/quic/core/quic_sent_packet_manager.h
index ca43d3d..72307b3 100644
--- a/quic/core/quic_sent_packet_manager.h
+++ b/quic/core/quic_sent_packet_manager.h
@@ -394,8 +394,8 @@
   void StartExponentialBackoffAfterNthPto(
       size_t exponential_backoff_start_point);
 
-  // Called to retransmit in flight INITIAL packet if any.
-  void RetransmitInitialDataIfAny();
+  // Called to retransmit in flight packet of |space| if any.
+  void RetransmitDataOfSpaceIfAny(PacketNumberSpace space);
 
   bool supports_multiple_packet_number_spaces() const {
     return unacked_packets_.supports_multiple_packet_number_spaces();
diff --git a/quic/core/quic_sent_packet_manager_test.cc b/quic/core/quic_sent_packet_manager_test.cc
index 33b2470..23aec1e 100644
--- a/quic/core/quic_sent_packet_manager_test.cc
+++ b/quic/core/quic_sent_packet_manager_test.cc
@@ -4079,7 +4079,7 @@
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
         RetransmitDataPacket(4, type, ENCRYPTION_INITIAL);
       })));
-  manager_.RetransmitInitialDataIfAny();
+  manager_.RetransmitDataOfSpaceIfAny(INITIAL_DATA);
   // Verify PTO is re-armed based on packet 2.
   EXPECT_EQ(packet2_sent_time + expected_pto_delay,
             manager_.GetRetransmissionTime());
@@ -4090,7 +4090,7 @@
       .WillOnce(WithArgs<1>(Invoke([this](TransmissionType type) {
         RetransmitDataPacket(5, type, ENCRYPTION_INITIAL);
       })));
-  manager_.RetransmitInitialDataIfAny();
+  manager_.RetransmitDataOfSpaceIfAny(INITIAL_DATA);
   // Verify PTO does not change.
   EXPECT_EQ(packet2_sent_time + expected_pto_delay,
             manager_.GetRetransmissionTime());