In QuicSession::OnCanWrite, do not send stream data when PTO fires and handshake is not confirmed.

Protected by FLAGS_quic_reloadable_flag_quic_donot_pto_stream_data_before_handshake_confirmed.

PiperOrigin-RevId: 481377538

quiche/quic/core/quic_connection.cc

diff --git a/quiche/quic/core/quic_connection.cc b/quiche/quic/core/quic_connection.cc
index 5c44172..565419d 100644
--- a/quiche/quic/core/quic_connection.cc
+++ b/quiche/quic/core/quic_connection.cc
@@ -2438,8 +2438,7 @@
 
   if (perspective_ == Perspective::IS_SERVER &&
       version().CanSendCoalescedPackets() && !IsHandshakeConfirmed()) {
-    if (in_on_retransmission_time_out_ &&
-        coalesced_packet_.NumberOfPackets() == 0u) {
+    if (in_probe_time_out_ && coalesced_packet_.NumberOfPackets() == 0u) {
       // PTO fires while handshake is not confirmed. Do not preempt handshake
       // data with stream data.
       QUIC_CODE_COUNT(quic_try_to_send_half_rtt_data_when_pto_fires);
@@ -7088,15 +7087,15 @@
 QuicConnection::ScopedRetransmissionTimeoutIndicator::
     ScopedRetransmissionTimeoutIndicator(QuicConnection* connection)
     : connection_(connection) {
-  QUICHE_DCHECK(!connection_->in_on_retransmission_time_out_)
+  QUICHE_DCHECK(!connection_->in_probe_time_out_)
       << "ScopedRetransmissionTimeoutIndicator is not supposed to be nested";
-  connection_->in_on_retransmission_time_out_ = true;
+  connection_->in_probe_time_out_ = true;
 }
 
 QuicConnection::ScopedRetransmissionTimeoutIndicator::
     ~ScopedRetransmissionTimeoutIndicator() {
-  QUICHE_DCHECK(connection_->in_on_retransmission_time_out_);
-  connection_->in_on_retransmission_time_out_ = false;
+  QUICHE_DCHECK(connection_->in_probe_time_out_);
+  connection_->in_probe_time_out_ = false;
 }
 
 void QuicConnection::RestoreToLastValidatedPath(

quiche/quic/core/quic_connection.h

diff --git a/quiche/quic/core/quic_connection.h b/quiche/quic/core/quic_connection.h
index 36e9bbf..153c142 100644
--- a/quiche/quic/core/quic_connection.h
+++ b/quiche/quic/core/quic_connection.h
@@ -1263,6 +1263,8 @@
     context_.bug_listener.swap(bug_listener);
   }
 
+  bool in_probe_time_out() const { return in_probe_time_out_; }
+
   // Ensures the network blackhole delay is longer than path degrading delay.
   static QuicTime::Delta CalculateNetworkBlackholeDelay(
       QuicTime::Delta blackhole_delay, QuicTime::Delta path_degrading_delay,
@@ -1537,7 +1539,7 @@
     QuicConnection* connection_;
   };
 
-  // A class which sets and clears in_on_retransmission_time_out_ when entering
+  // A class which sets and clears in_probe_time_out_ when entering
   // and exiting OnRetransmissionTimeout, respectively.
   class QUIC_EXPORT_PRIVATE ScopedRetransmissionTimeoutIndicator {
    public:
@@ -2244,7 +2246,7 @@
   bool have_decrypted_first_one_rtt_packet_ = false;
 
   // True if we are currently processing OnRetransmissionTimeout.
-  bool in_on_retransmission_time_out_ = false;
+  bool in_probe_time_out_ = false;
 
   QuicPathValidator path_validator_;
 

quiche/quic/core/quic_flags_list.h

diff --git a/quiche/quic/core/quic_flags_list.h b/quiche/quic/core/quic_flags_list.h
index d54304f..d3079d5 100644
--- a/quiche/quic/core/quic_flags_list.h
+++ b/quiche/quic/core/quic_flags_list.h
@@ -51,6 +51,8 @@
 QUIC_FLAG(quic_reloadable_flag_quic_enable_disable_resumption, true)
 // If true, discard INITIAL packet if the key has been dropped.
 QUIC_FLAG(quic_reloadable_flag_quic_discard_initial_packet_with_key_dropped, true)
+// If true, do not PTO new stream data before handshake confirmed.
+QUIC_FLAG(quic_reloadable_flag_quic_donot_pto_stream_data_before_handshake_confirmed, true)
 // If true, do not issue a new connection ID that has been claimed by another connection.
 QUIC_FLAG(quic_reloadable_flag_quic_check_cid_collision_when_issue_new_cid, true)
 // If true, do not mark stream connection level write blocked if its write side has been closed.

quiche/quic/core/quic_session.cc

diff --git a/quiche/quic/core/quic_session.cc b/quiche/quic/core/quic_session.cc
index b6b0371..b4abf52 100644
--- a/quiche/quic/core/quic_session.cc
+++ b/quiche/quic/core/quic_session.cc
@@ -639,6 +639,14 @@
   if (control_frame_manager_.WillingToWrite()) {
     control_frame_manager_.OnCanWrite();
   }
+  if (GetQuicReloadableFlag(
+          quic_donot_pto_stream_data_before_handshake_confirmed) &&
+      version().UsesTls() && GetHandshakeState() != HANDSHAKE_CONFIRMED &&
+      connection_->in_probe_time_out()) {
+    QUIC_CODE_COUNT(quic_donot_pto_stream_data_before_handshake_confirmed);
+    // Do not PTO stream data before handshake gets confirmed.
+    return;
+  }
   // TODO(b/147146815): this makes all datagrams go before stream data.  We
   // should have a better priority scheme for this.
   if (!datagram_queue_.empty()) {

quiche/quic/core/quic_session_test.cc

diff --git a/quiche/quic/core/quic_session_test.cc b/quiche/quic/core/quic_session_test.cc
index 916ed1b..3e66183 100644
--- a/quiche/quic/core/quic_session_test.cc
+++ b/quiche/quic/core/quic_session_test.cc
@@ -3128,6 +3128,40 @@
   }
 }
 
+TEST_P(QuicSessionTestServer, DonotPtoStreamDataBeforeHandshakeConfirmed) {
+  if (!session_.version().UsesTls()) {
+    return;
+  }
+  EXPECT_NE(HANDSHAKE_CONFIRMED, session_.GetHandshakeState());
+
+  TestCryptoStream* crypto_stream = session_.GetMutableCryptoStream();
+  EXPECT_FALSE(crypto_stream->HasBufferedCryptoFrames());
+  std::string data(1350, 'a');
+  EXPECT_CALL(*connection_, SendCryptoData(ENCRYPTION_INITIAL, 1350, 0))
+      .WillOnce(Return(1000));
+  crypto_stream->WriteCryptoData(ENCRYPTION_INITIAL, data);
+  ASSERT_TRUE(crypto_stream->HasBufferedCryptoFrames());
+
+  TestStream* stream = session_.CreateOutgoingBidirectionalStream();
+
+  session_.MarkConnectionLevelWriteBlocked(stream->id());
+  // Buffered crypto data gets sent.
+  EXPECT_CALL(*connection_, SendCryptoData(ENCRYPTION_INITIAL, _, _))
+      .WillOnce(Return(350));
+  if (GetQuicReloadableFlag(
+          quic_donot_pto_stream_data_before_handshake_confirmed)) {
+    // Verify stream data is not sent on PTO before handshake confirmed.
+    EXPECT_CALL(*stream, OnCanWrite()).Times(0);
+  } else {
+    EXPECT_CALL(*stream, OnCanWrite());
+  }
+
+  // Fire PTO.
+  QuicConnectionPeer::SetInProbeTimeOut(connection_, true);
+  session_.OnCanWrite();
+  EXPECT_FALSE(crypto_stream->HasBufferedCryptoFrames());
+}
+
 // A client test class that can be used when the automatic configuration is not
 // desired.
 class QuicSessionTestClientUnconfigured : public QuicSessionTestBase {

quiche/quic/test_tools/quic_connection_peer.cc

diff --git a/quiche/quic/test_tools/quic_connection_peer.cc b/quiche/quic/test_tools/quic_connection_peer.cc
index eb28e2d..f7c6424 100644
--- a/quiche/quic/test_tools/quic_connection_peer.cc
+++ b/quiche/quic/test_tools/quic_connection_peer.cc
@@ -560,5 +560,11 @@
   connection->FlushCoalescedPacket();
 }
 
+// static
+void QuicConnectionPeer::SetInProbeTimeOut(QuicConnection* connection,
+                                           bool value) {
+  connection->in_probe_time_out_ = value;
+}
+
 }  // namespace test
 }  // namespace quic

quiche/quic/test_tools/quic_connection_peer.h

diff --git a/quiche/quic/test_tools/quic_connection_peer.h b/quiche/quic/test_tools/quic_connection_peer.h
index 169a33c..932ddba 100644
--- a/quiche/quic/test_tools/quic_connection_peer.h
+++ b/quiche/quic/test_tools/quic_connection_peer.h
@@ -226,6 +226,8 @@
   static void FlushCoalescedPacket(QuicConnection* connection);
 
   static QuicAlarm* GetMultiPortProbingAlarm(QuicConnection* connection);
+
+  static void SetInProbeTimeOut(QuicConnection* connection, bool value);
 };
 
 }  // namespace test