Do not QUIC_BUG 0-RTT stream data retransmission before 1-RTT keys arrive.

PiperOrigin-RevId: 329566741
Change-Id: I0a881a90377ddfc3de4063b295e09a8278e38a32

quic/core/http/end_to_end_test.cc

diff --git a/quic/core/http/end_to_end_test.cc b/quic/core/http/end_to_end_test.cc
index ff6d357..33bfee0 100644
--- a/quic/core/http/end_to_end_test.cc
+++ b/quic/core/http/end_to_end_test.cc
@@ -1670,6 +1670,65 @@
   EXPECT_TRUE(client_->client()->EarlyDataAccepted());
 }
 
+// Regression test for b/166836136.
+TEST_P(EndToEndTest, LargePostZeroRTTRejectRequestDuringHandshake) {
+  if (!version_.UsesTls()) {
+    // This test is TLS specific.
+    ASSERT_TRUE(Initialize());
+    return;
+  }
+  // Send a request and then disconnect. This prepares the client to attempt
+  // a 0-RTT handshake for the next request.
+  NiceMock<MockQuicConnectionDebugVisitor> visitor;
+  connection_debug_visitor_ = &visitor;
+  ASSERT_TRUE(Initialize());
+
+  SendSynchronousFooRequestAndCheckResponse();
+  QuicSpdyClientSession* client_session = GetClientSession();
+  ASSERT_TRUE(client_session);
+  EXPECT_FALSE(client_session->EarlyDataAccepted());
+  EXPECT_FALSE(client_session->ReceivedInchoateReject());
+  EXPECT_FALSE(client_->client()->EarlyDataAccepted());
+  EXPECT_FALSE(client_->client()->ReceivedInchoateReject());
+
+  client_->Disconnect();
+
+  client_->Connect();
+  EXPECT_TRUE(client_->client()->WaitForOneRttKeysAvailable());
+  ASSERT_TRUE(client_->client()->connected());
+  EXPECT_EQ(kFooResponseBody, client_->SendSynchronousRequest("/foo"));
+
+  client_session = GetClientSession();
+  ASSERT_TRUE(client_session);
+  EXPECT_TRUE(client_session->EarlyDataAccepted());
+  EXPECT_TRUE(client_->client()->EarlyDataAccepted());
+
+  client_->Disconnect();
+
+  // Restart the server so that the 0-RTT handshake will take 1 RTT.
+  StopServer();
+  server_writer_ = new PacketDroppingTestWriter();
+  StartServer();
+
+  ON_CALL(visitor, OnZeroRttRejected()).WillByDefault(Invoke([this]() {
+    EXPECT_FALSE(GetClientSession()->IsEncryptionEstablished());
+    // Trigger an OnCanWrite() to make sure no unencrypted data will be written.
+    GetClientSession()->OnCanWrite();
+  }));
+
+  // The 0-RTT handshake should fail.
+  client_->Connect();
+  ASSERT_TRUE(client_->client()->WaitForOneRttKeysAvailable());
+  client_->WaitForWriteToFlush();
+  client_->WaitForResponse();
+  ASSERT_TRUE(client_->client()->connected());
+
+  client_session = GetClientSession();
+  ASSERT_TRUE(client_session);
+  EXPECT_FALSE(client_session->EarlyDataAccepted());
+  EXPECT_FALSE(client_->client()->EarlyDataAccepted());
+}
+
 TEST_P(EndToEndTest, RejectWithPacketLoss) {
   // In this test, we intentionally drop the first packet from the
   // server, which corresponds with the initial REJ response from

quic/core/quic_connection.cc

diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index 9e0a493..d89639f 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -2471,10 +2471,10 @@
 }
 
 void QuicConnection::MarkZeroRttPacketsForRetransmission() {
+  sent_packet_manager_.MarkZeroRttPacketsForRetransmission();
   if (debug_visitor_ != nullptr && version().UsesTls()) {
     debug_visitor_->OnZeroRttRejected();
   }
-  sent_packet_manager_.MarkZeroRttPacketsForRetransmission();
 }
 
 void QuicConnection::NeuterUnencryptedPackets() {

quic/core/quic_session.cc

diff --git a/quic/core/quic_session.cc b/quic/core/quic_session.cc
index 8989269..aa7c491 100644
--- a/quic/core/quic_session.cc
+++ b/quic/core/quic_session.cc
@@ -719,8 +719,15 @@
       !QuicUtils::IsCryptoStreamId(transport_version(), id)) {
     // Do not let streams write without encryption. The calling stream will end
     // up write blocked until OnCanWrite is next called.
-    QUIC_BUG << ENDPOINT << "Try to send data of stream " << id
-             << " before encryption is established.";
+    if (was_zero_rtt_rejected_ && !OneRttKeysAvailable()) {
+      DCHECK(version().UsesTls() && perspective() == Perspective::IS_CLIENT);
+      QUIC_BUG_IF(type == NOT_RETRANSMISSION)
+          << ENDPOINT << "Try to send new data on stream " << id
+          << "before 1-RTT keys are available while 0-RTT is rejected.";
+    } else {
+      QUIC_BUG << ENDPOINT << "Try to send data of stream " << id
+               << " before encryption is established.";
+    }
     return QuicConsumedData(0, false);
   }
 

quic/test_tools/quic_test_utils.h

diff --git a/quic/test_tools/quic_test_utils.h b/quic/test_tools/quic_test_utils.h
index aac4a2c..614e958 100644
--- a/quic/test_tools/quic_test_utils.h
+++ b/quic/test_tools/quic_test_utils.h
@@ -1464,6 +1464,8 @@
               OnTransportParametersReceived,
               (const TransportParameters&),
               (override));
+
+  MOCK_METHOD(void, OnZeroRttRejected, (), (override));
 };
 
 class MockReceivedPacketManager : public QuicReceivedPacketManager {