Use all encryption levels in quic_crypto_stream.cc

ENCRYPTION_HANDSHAKE was missing from some places. This resulted in a bug
where CRYPTO frames sent at the ENCRYPTION_HANDSHAKE level (which only
happens when TLS is used as the crypto handshake protocol - QUIC crypto
sends CRYPTO frames in ENCRYPTION_ZERO_RTT but not ENCRYPTION_HANDSHAKE)
won't get retransmitted if lost.

gfe-relnote: Retransmit lost TLS ENCRYPTION_HANDSHAKE messages in QUIC, gated by ietf draft quic versions.
PiperOrigin-RevId: 305789426
Change-Id: I54db405c394ee0127f39e69e839be67136db61e8
diff --git a/quic/core/quic_crypto_stream.cc b/quic/core/quic_crypto_stream.cc
index 35f33c2..f0c7e06 100644
--- a/quic/core/quic_crypto_stream.cc
+++ b/quic/core/quic_crypto_stream.cc
@@ -235,12 +235,21 @@
   QuicStream::OnStreamDataConsumed(bytes_consumed);
 }
 
+namespace {
+
+constexpr std::array<EncryptionLevel, NUM_ENCRYPTION_LEVELS>
+AllEncryptionLevels() {
+  return {ENCRYPTION_INITIAL, ENCRYPTION_HANDSHAKE, ENCRYPTION_ZERO_RTT,
+          ENCRYPTION_FORWARD_SECURE};
+}
+
+}  // namespace
+
 bool QuicCryptoStream::HasPendingCryptoRetransmission() const {
   if (!QuicVersionUsesCryptoFrames(session()->transport_version())) {
     return false;
   }
-  for (EncryptionLevel level :
-       {ENCRYPTION_INITIAL, ENCRYPTION_ZERO_RTT, ENCRYPTION_FORWARD_SECURE}) {
+  for (EncryptionLevel level : AllEncryptionLevels()) {
     if (substreams_[level].send_buffer.HasPendingRetransmission()) {
       return true;
     }
@@ -251,8 +260,7 @@
 void QuicCryptoStream::WritePendingCryptoRetransmission() {
   QUIC_BUG_IF(!QuicVersionUsesCryptoFrames(session()->transport_version()))
       << "Versions less than 47 don't write CRYPTO frames";
-  for (EncryptionLevel level :
-       {ENCRYPTION_INITIAL, ENCRYPTION_ZERO_RTT, ENCRYPTION_FORWARD_SECURE}) {
+  for (EncryptionLevel level : AllEncryptionLevels()) {
     QuicStreamSendBuffer* send_buffer = &substreams_[level].send_buffer;
     while (send_buffer->HasPendingRetransmission()) {
       auto pending = send_buffer->NextPendingRetransmission();
@@ -353,9 +361,11 @@
   if (!QuicVersionUsesCryptoFrames(session()->transport_version())) {
     return stream_bytes_read();
   }
-  return substreams_[ENCRYPTION_INITIAL].sequencer.NumBytesConsumed() +
-         substreams_[ENCRYPTION_ZERO_RTT].sequencer.NumBytesConsumed() +
-         substreams_[ENCRYPTION_FORWARD_SECURE].sequencer.NumBytesConsumed();
+  uint64_t bytes_read = 0;
+  for (EncryptionLevel level : AllEncryptionLevels()) {
+    bytes_read += substreams_[level].sequencer.NumBytesConsumed();
+  }
+  return bytes_read;
 }
 
 uint64_t QuicCryptoStream::BytesReadOnLevel(EncryptionLevel level) const {
@@ -408,8 +418,7 @@
 void QuicCryptoStream::WriteBufferedCryptoFrames() {
   QUIC_BUG_IF(!QuicVersionUsesCryptoFrames(session()->transport_version()))
       << "Versions less than 47 don't use CRYPTO frames";
-  for (EncryptionLevel level :
-       {ENCRYPTION_INITIAL, ENCRYPTION_ZERO_RTT, ENCRYPTION_FORWARD_SECURE}) {
+  for (EncryptionLevel level : AllEncryptionLevels()) {
     QuicStreamSendBuffer* send_buffer = &substreams_[level].send_buffer;
     const size_t data_length =
         send_buffer->stream_offset() - send_buffer->stream_bytes_written();
@@ -431,8 +440,7 @@
 bool QuicCryptoStream::HasBufferedCryptoFrames() const {
   QUIC_BUG_IF(!QuicVersionUsesCryptoFrames(session()->transport_version()))
       << "Versions less than 47 don't use CRYPTO frames";
-  for (EncryptionLevel level :
-       {ENCRYPTION_INITIAL, ENCRYPTION_ZERO_RTT, ENCRYPTION_FORWARD_SECURE}) {
+  for (EncryptionLevel level : AllEncryptionLevels()) {
     const QuicStreamSendBuffer& send_buffer = substreams_[level].send_buffer;
     DCHECK_GE(send_buffer.stream_offset(), send_buffer.stream_bytes_written());
     if (send_buffer.stream_offset() > send_buffer.stream_bytes_written()) {
@@ -461,8 +469,7 @@
   if (!QuicVersionUsesCryptoFrames(session()->transport_version())) {
     return QuicStream::IsWaitingForAcks();
   }
-  for (EncryptionLevel level :
-       {ENCRYPTION_INITIAL, ENCRYPTION_ZERO_RTT, ENCRYPTION_FORWARD_SECURE}) {
+  for (EncryptionLevel level : AllEncryptionLevels()) {
     if (substreams_[level].send_buffer.stream_bytes_outstanding()) {
       return true;
     }
diff --git a/quic/core/quic_crypto_stream_test.cc b/quic/core/quic_crypto_stream_test.cc
index 436e572..6365bc7 100644
--- a/quic/core/quic_crypto_stream_test.cc
+++ b/quic/core/quic_crypto_stream_test.cc
@@ -253,6 +253,48 @@
   EXPECT_EQ(ENCRYPTION_FORWARD_SECURE, connection_->encryption_level());
 }
 
+// Regression test for handling the missing ENCRYPTION_HANDSHAKE in
+// quic_crypto_stream.cc. This test is essentially the same as
+// RetransmitCryptoDataInCryptoFrames, except it uses ENCRYPTION_HANDSHAKE in
+// place of ENCRYPTION_ZERO_RTT.
+TEST_F(QuicCryptoStreamTest, RetransmitEncryptionHandshakeLevelCryptoFrames) {
+  if (!QuicVersionUsesCryptoFrames(connection_->transport_version())) {
+    return;
+  }
+  EXPECT_CALL(*connection_, SendCryptoData(_, _, _)).Times(0);
+  InSequence s;
+  // Send [0, 1000) in ENCRYPTION_INITIAL.
+  EXPECT_EQ(ENCRYPTION_INITIAL, connection_->encryption_level());
+  std::string data(1000, 'a');
+  EXPECT_CALL(*connection_, SendCryptoData(ENCRYPTION_INITIAL, 1000, 0))
+      .WillOnce(Invoke(connection_,
+                       &MockQuicConnection::QuicConnection_SendCryptoData));
+  stream_->WriteCryptoData(ENCRYPTION_INITIAL, data);
+  // Send [1000, 2000) in ENCRYPTION_HANDSHAKE.
+  connection_->SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
+  std::unique_ptr<NullEncrypter> encrypter =
+      std::make_unique<NullEncrypter>(Perspective::IS_CLIENT);
+  connection_->SetEncrypter(ENCRYPTION_HANDSHAKE, std::move(encrypter));
+  EXPECT_EQ(ENCRYPTION_HANDSHAKE, connection_->encryption_level());
+  EXPECT_CALL(*connection_, SendCryptoData(ENCRYPTION_HANDSHAKE, 1000, 0))
+      .WillOnce(Invoke(connection_,
+                       &MockQuicConnection::QuicConnection_SendCryptoData));
+  stream_->WriteCryptoData(ENCRYPTION_HANDSHAKE, data);
+  connection_->SetDefaultEncryptionLevel(ENCRYPTION_FORWARD_SECURE);
+  EXPECT_EQ(ENCRYPTION_FORWARD_SECURE, connection_->encryption_level());
+
+  // Lost [1000, 1200).
+  QuicCryptoFrame lost_frame(ENCRYPTION_HANDSHAKE, 0, 200);
+  stream_->OnCryptoFrameLost(&lost_frame);
+  EXPECT_TRUE(stream_->HasPendingCryptoRetransmission());
+  // Verify [1000, 1200) is sent.
+  EXPECT_CALL(*connection_, SendCryptoData(ENCRYPTION_HANDSHAKE, 200, 0))
+      .WillOnce(Invoke(connection_,
+                       &MockQuicConnection::QuicConnection_SendCryptoData));
+  stream_->WritePendingCryptoRetransmission();
+  EXPECT_FALSE(stream_->HasPendingCryptoRetransmission());
+}
+
 TEST_F(QuicCryptoStreamTest, NeuterUnencryptedStreamData) {
   if (QuicVersionUsesCryptoFrames(connection_->transport_version())) {
     return;