Fix QuicSession::WillingAndAbleToWrite to check write keys.

Protected by FLAGS_quic_reloadable_flag_quic_fix_willing_and_able_to_write2.

PiperOrigin-RevId: 339286884
Change-Id: Ifb649dd85bade3a86864b146b20a61ffb56d3b77

quic/core/http/quic_spdy_session_test.cc

diff --git a/quic/core/http/quic_spdy_session_test.cc b/quic/core/http/quic_spdy_session_test.cc
index 84ea617..d95fa6d 100644
--- a/quic/core/http/quic_spdy_session_test.cc
+++ b/quic/core/http/quic_spdy_session_test.cc
@@ -704,6 +704,7 @@
 }
 
 TEST_P(QuicSpdySessionTestServer, OnCanWrite) {
+  CompleteHandshake();
   session_.set_writev_consumes_all_data(true);
   TestStream* stream2 = session_.CreateOutgoingBidirectionalStream();
   TestStream* stream4 = session_.CreateOutgoingBidirectionalStream();
@@ -861,6 +862,7 @@
 }
 
 TEST_P(QuicSpdySessionTestServer, OnCanWriteCongestionControlBlocks) {
+  CompleteHandshake();
   session_.set_writev_consumes_all_data(true);
   InSequence s;
 
@@ -908,6 +910,7 @@
 }
 
 TEST_P(QuicSpdySessionTestServer, OnCanWriteWriterBlocks) {
+  CompleteHandshake();
   // Drive congestion control manually in order to ensure that
   // application-limited signaling is handled correctly.
   MockSendAlgorithm* send_algorithm = new StrictMock<MockSendAlgorithm>;
@@ -1009,6 +1012,7 @@
 
 TEST_P(QuicSpdySessionTestServer,
        OnCanWriteLimitsNumWritesIfFlowControlBlocked) {
+  CompleteHandshake();
   // Drive congestion control manually in order to ensure that
   // application-limited signaling is handled correctly.
   MockSendAlgorithm* send_algorithm = new StrictMock<MockSendAlgorithm>;
@@ -2031,6 +2035,7 @@
 }
 
 TEST_P(QuicSpdySessionTestServer, OnStreamFrameLost) {
+  CompleteHandshake();
   InSequence s;
 
   // Drive congestion control manually.

quic/core/quic_session.cc

diff --git a/quic/core/quic_session.cc b/quic/core/quic_session.cc
index a8d9e6f..5079b46 100644
--- a/quic/core/quic_session.cc
+++ b/quic/core/quic_session.cc
@@ -649,9 +649,16 @@
   // 3) If the crypto or headers streams are blocked, or
   // 4) connection is not flow control blocked and there are write blocked
   // streams.
-  if (QuicVersionUsesCryptoFrames(transport_version()) &&
-      HasPendingHandshake()) {
-    return true;
+  if (QuicVersionUsesCryptoFrames(transport_version())) {
+    if (HasPendingHandshake()) {
+      return true;
+    }
+    if (GetQuicReloadableFlag(quic_fix_willing_and_able_to_write2)) {
+      QUIC_RELOADABLE_FLAG_COUNT(quic_fix_willing_and_able_to_write2);
+      if (!IsEncryptionEstablished()) {
+        return false;
+      }
+    }
   }
   if (control_frame_manager_.WillingToWrite() ||
       !streams_with_pending_retransmission_.empty()) {

quic/core/quic_session_test.cc

diff --git a/quic/core/quic_session_test.cc b/quic/core/quic_session_test.cc
index 46ce528..738ee0b 100644
--- a/quic/core/quic_session_test.cc
+++ b/quic/core/quic_session_test.cc
@@ -104,6 +104,9 @@
           session()->config()->ProcessPeerHello(msg, CLIENT, &error_details);
     }
     EXPECT_THAT(error, IsQuicNoError());
+    session()->OnNewEncryptionKeyAvailable(
+        ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<NullEncrypter>(session()->perspective()));
     session()->OnConfigNegotiated();
     if (session()->connection()->version().handshake_protocol ==
         PROTOCOL_TLS1_3) {
@@ -487,6 +490,16 @@
     closed_streams_.insert(id);
   }
 
+  void CompleteHandshake() {
+    CryptoHandshakeMessage msg;
+    if (connection_->version().HasHandshakeDone() &&
+        connection_->perspective() == Perspective::IS_SERVER) {
+      EXPECT_CALL(*connection_, SendControlFrame(_))
+          .WillOnce(Invoke(&ClearControlFrame));
+    }
+    session_.GetMutableCryptoStream()->OnHandshakeMessage(msg);
+  }
+
   QuicTransportVersion transport_version() const {
     return connection_->transport_version();
   }
@@ -630,12 +643,7 @@
 
 TEST_P(QuicSessionTestServer, OneRttKeysAvailable) {
   EXPECT_FALSE(session_.OneRttKeysAvailable());
-  CryptoHandshakeMessage message;
-  if (connection_->version().HasHandshakeDone()) {
-    EXPECT_CALL(*connection_, SendControlFrame(_));
-  }
-  connection_->SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
-  session_.GetMutableCryptoStream()->OnHandshakeMessage(message);
+  CompleteHandshake();
   EXPECT_TRUE(session_.OneRttKeysAvailable());
 }
 
@@ -928,6 +936,7 @@
 }
 
 TEST_P(QuicSessionTestServer, OnCanWrite) {
+  CompleteHandshake();
   session_.set_writev_consumes_all_data(true);
   TestStream* stream2 = session_.CreateOutgoingBidirectionalStream();
   TestStream* stream4 = session_.CreateOutgoingBidirectionalStream();
@@ -1152,15 +1161,9 @@
 
 TEST_P(QuicSessionTestServer, OnCanWriteBundlesStreams) {
   // Encryption needs to be established before data can be sent.
-  if (connection_->version().HasHandshakeDone()) {
-    EXPECT_CALL(*connection_, SendControlFrame(_))
-        .WillRepeatedly(Invoke(&ClearControlFrame));
-  }
-  CryptoHandshakeMessage msg;
+  CompleteHandshake();
   MockPacketWriter* writer = static_cast<MockPacketWriter*>(
       QuicConnectionPeer::GetWriter(session_.connection()));
-  connection_->SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
-  session_.GetMutableCryptoStream()->OnHandshakeMessage(msg);
 
   // Drive congestion control manually.
   MockSendAlgorithm* send_algorithm = new StrictMock<MockSendAlgorithm>;
@@ -1199,6 +1202,7 @@
 }
 
 TEST_P(QuicSessionTestServer, OnCanWriteCongestionControlBlocks) {
+  CompleteHandshake();
   session_.set_writev_consumes_all_data(true);
   InSequence s;
 
@@ -1246,6 +1250,7 @@
 }
 
 TEST_P(QuicSessionTestServer, OnCanWriteWriterBlocks) {
+  CompleteHandshake();
   // Drive congestion control manually in order to ensure that
   // application-limited signaling is handled correctly.
   MockSendAlgorithm* send_algorithm = new StrictMock<MockSendAlgorithm>;
@@ -2373,6 +2378,7 @@
 }
 
 TEST_P(QuicSessionTestServer, OnStreamFrameLost) {
+  CompleteHandshake();
   InSequence s;
 
   // Drive congestion control manually.