Use ScopedEncryptionLevelContext to send data from QuicConnection and QuicSession.

This CL makes more widespread use of ScopedEncryptionLevelContext to explicitly set the default encryption level on a QuicConnection for all data sent, apart from ACK-only packets in gQUIC.

For QuicSession, the sending data includes control frames, messages, stream data and crypto data.
For QuicConnection, the sending data includes PING, CONNECTION_CLOSE and ACK frame (for multiple packet number spaces).

Protected by FLAGS_quic_reloadable_flag_quic_use_encryption_level_context.

PiperOrigin-RevId: 339942006
Change-Id: Ibd8b56cafded3c713c30fe5981c8b5d0aeb30afc
diff --git a/quic/core/quic_connection.cc b/quic/core/quic_connection.cc
index db6facf..a07314d 100644
--- a/quic/core/quic_connection.cc
+++ b/quic/core/quic_connection.cc
@@ -352,7 +352,10 @@
       support_handshake_done_(version().HasHandshakeDone()),
       encrypted_control_frames_(
           GetQuicReloadableFlag(quic_encrypted_control_frames) &&
-          packet_creator_.let_connection_handle_pings()) {
+          packet_creator_.let_connection_handle_pings()),
+      use_encryption_level_context_(
+          encrypted_control_frames_ &&
+          GetQuicReloadableFlag(quic_use_encryption_level_context)) {
   QUIC_BUG_IF(!start_peer_migration_earlier_ && send_path_response_);
   if (GetQuicReloadableFlag(quic_connection_set_initial_self_address)) {
     DCHECK(perspective_ == Perspective::IS_CLIENT ||
@@ -362,6 +365,9 @@
   if (enable_aead_limits_) {
     QUIC_RELOADABLE_FLAG_COUNT(quic_enable_aead_limits);
   }
+  if (use_encryption_level_context_) {
+    QUIC_RELOADABLE_FLAG_COUNT(quic_use_encryption_level_context);
+  }
   QUIC_DLOG(INFO) << ENDPOINT << "Created connection with server connection ID "
                   << server_connection_id
                   << " and version: " << ParsedQuicVersionToString(version());
@@ -3481,7 +3487,9 @@
     return;
   }
   if (packet_creator_.let_connection_handle_pings()) {
-    SendPingAtLevel(encryption_level_);
+    SendPingAtLevel(use_encryption_level_context_
+                        ? framer().GetEncryptionLevelToSendApplicationData()
+                        : encryption_level_);
   } else {
     visitor_->SendPing();
   }
@@ -3907,7 +3915,12 @@
                                                       peer_address());
   if (!SupportsMultiplePacketNumberSpaces()) {
     QUIC_DLOG(INFO) << ENDPOINT << "Sending connection close packet.";
-    SetDefaultEncryptionLevel(GetConnectionCloseEncryptionLevel());
+    if (!use_encryption_level_context_) {
+      SetDefaultEncryptionLevel(GetConnectionCloseEncryptionLevel());
+    }
+    ScopedEncryptionLevelContext context(
+        use_encryption_level_context_ ? this : nullptr,
+        GetConnectionCloseEncryptionLevel());
     if (version().CanSendCoalescedPackets()) {
       coalesced_packet_.Clear();
     }
@@ -3950,7 +3963,11 @@
     }
     QUIC_DLOG(INFO) << ENDPOINT
                     << "Sending connection close packet at level: " << level;
-    SetDefaultEncryptionLevel(level);
+    if (!use_encryption_level_context_) {
+      SetDefaultEncryptionLevel(level);
+    }
+    ScopedEncryptionLevelContext context(
+        use_encryption_level_context_ ? this : nullptr, level);
     // Bundle an ACK of the corresponding packet number space for debugging
     // purpose.
     if (error != QUIC_PACKET_WRITE_ERROR &&
@@ -3978,7 +3995,9 @@
   // Since the connection is closing, if the connection close packets were not
   // sent, then they should be discarded.
   ClearQueuedPackets();
-  SetDefaultEncryptionLevel(current_encryption_level);
+  if (!use_encryption_level_context_) {
+    SetDefaultEncryptionLevel(current_encryption_level);
+  }
 }
 
 void QuicConnection::TearDownLocalConnectionState(
@@ -4926,7 +4945,12 @@
                   << PacketNumberSpaceToString(
                          static_cast<PacketNumberSpace>(i));
     // Switch to the appropriate encryption level.
-    SetDefaultEncryptionLevel(
+    if (!use_encryption_level_context_) {
+      SetDefaultEncryptionLevel(
+          QuicUtils::GetEncryptionLevel(static_cast<PacketNumberSpace>(i)));
+    }
+    ScopedEncryptionLevelContext context(
+        use_encryption_level_context_ ? this : nullptr,
         QuicUtils::GetEncryptionLevel(static_cast<PacketNumberSpace>(i)));
     QuicFrames frames;
     frames.push_back(uber_received_packet_manager_.GetUpdatedAckFrame(
@@ -4942,8 +4966,10 @@
     }
     ResetAckStates();
   }
-  // Restores encryption level.
-  SetDefaultEncryptionLevel(current_encryption_level);
+  if (!use_encryption_level_context_) {
+    // Restores encryption level.
+    SetDefaultEncryptionLevel(current_encryption_level);
+  }
 
   const QuicTime timeout =
       uber_received_packet_manager_.GetEarliestAckTimeout();
diff --git a/quic/core/quic_connection.h b/quic/core/quic_connection.h
index 681a64a..ee477a9 100644
--- a/quic/core/quic_connection.h
+++ b/quic/core/quic_connection.h
@@ -1129,6 +1129,10 @@
 
   bool encrypted_control_frames() const { return encrypted_control_frames_; }
 
+  bool use_encryption_level_context() const {
+    return use_encryption_level_context_;
+  }
+
  protected:
   // Calls cancel() on all the alarms owned by this connection.
   void CancelAllAlarms();
@@ -1929,7 +1933,9 @@
   const bool check_keys_before_writing_ =
       GetQuicReloadableFlag(quic_check_keys_before_writing);
 
-  bool encrypted_control_frames_;
+  const bool encrypted_control_frames_;
+
+  const bool use_encryption_level_context_;
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_datagram_queue_test.cc b/quic/core/quic_datagram_queue_test.cc
index 1fd451c..0e6e3de 100644
--- a/quic/core/quic_datagram_queue_test.cc
+++ b/quic/core/quic_datagram_queue_test.cc
@@ -6,6 +6,7 @@
 
 #include "absl/strings/string_view.h"
 #include "absl/types/optional.h"
+#include "net/third_party/quiche/src/quic/core/crypto/null_encrypter.h"
 #include "net/third_party/quiche/src/quic/core/quic_buffer_allocator.h"
 #include "net/third_party/quiche/src/quic/core/quic_time.h"
 #include "net/third_party/quiche/src/quic/core/quic_types.h"
@@ -38,6 +39,9 @@
         session_(connection_),
         queue_(&session_) {
     session_.SetCryptoStream(new EstablishedCryptoStream(&session_));
+    connection_->SetEncrypter(
+        ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<NullEncrypter>(connection_->perspective()));
   }
 
   QuicMemSlice CreateMemSlice(absl::string_view data) {
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index c5c4df8..d753d05 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -73,6 +73,7 @@
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_testonly_default_false, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_testonly_default_true, true)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_unified_iw_options, false)
+QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_use_encryption_level_context, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_use_fast_huffman_encoder, true)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_quic_use_write_or_buffer_data_at_level, false)
 QUIC_FLAG(FLAGS_quic_reloadable_flag_send_quic_fallback_server_config_on_leto_error, false)
diff --git a/quic/core/quic_session.cc b/quic/core/quic_session.cc
index bee5b10..983f862 100644
--- a/quic/core/quic_session.cc
+++ b/quic/core/quic_session.cc
@@ -749,9 +749,14 @@
 
   SetTransmissionType(type);
   const auto current_level = connection()->encryption_level();
-  if (level.has_value()) {
-    connection()->SetDefaultEncryptionLevel(level.value());
+  if (!use_encryption_level_context()) {
+    if (level.has_value()) {
+      connection()->SetDefaultEncryptionLevel(level.value());
+    }
   }
+  QuicConnection::ScopedEncryptionLevelContext context(
+      use_encryption_level_context() ? connection() : nullptr,
+      use_encryption_level_context() ? level.value() : NUM_ENCRYPTION_LEVELS);
 
   QuicConsumedData data =
       connection_->SendStreamData(id, write_length, offset, state);
@@ -761,8 +766,11 @@
   }
 
   // Restore the encryption level.
-  if (level.has_value()) {
-    connection()->SetDefaultEncryptionLevel(current_level);
+  if (!use_encryption_level_context()) {
+    // Restore the encryption level.
+    if (level.has_value()) {
+      connection()->SetDefaultEncryptionLevel(current_level);
+    }
   }
 
   return data;
@@ -786,11 +794,17 @@
   }
   SetTransmissionType(type);
   const auto current_level = connection()->encryption_level();
-  connection_->SetDefaultEncryptionLevel(level);
+  if (!use_encryption_level_context()) {
+    connection_->SetDefaultEncryptionLevel(level);
+  }
+  QuicConnection::ScopedEncryptionLevelContext context(
+      use_encryption_level_context() ? connection() : nullptr, level);
   const auto bytes_consumed =
       connection_->SendCryptoData(level, write_length, offset);
-  // Restores encryption level.
-  connection_->SetDefaultEncryptionLevel(current_level);
+  if (!use_encryption_level_context()) {
+    // Restores encryption level.
+    connection_->SetDefaultEncryptionLevel(current_level);
+  }
   return bytes_consumed;
 }
 
@@ -814,6 +828,10 @@
     }
   }
   SetTransmissionType(type);
+  QuicConnection::ScopedEncryptionLevelContext context(
+      use_encryption_level_context() ? connection() : nullptr,
+      use_encryption_level_context() ? GetEncryptionLevelToSendApplicationData()
+                                     : NUM_ENCRYPTION_LEVELS);
   return connection_->SendControlFrame(frame);
 }
 
@@ -2373,6 +2391,10 @@
   if (!IsEncryptionEstablished()) {
     return {MESSAGE_STATUS_ENCRYPTION_NOT_ESTABLISHED, 0};
   }
+  QuicConnection::ScopedEncryptionLevelContext context(
+      use_encryption_level_context() ? connection() : nullptr,
+      use_encryption_level_context() ? GetEncryptionLevelToSendApplicationData()
+                                     : NUM_ENCRYPTION_LEVELS);
   MessageStatus result =
       connection_->SendMessage(last_message_id_ + 1, message, flush);
   if (result == MESSAGE_STATUS_SUCCESS) {
diff --git a/quic/core/quic_session.h b/quic/core/quic_session.h
index cb0efbe..775c717 100644
--- a/quic/core/quic_session.h
+++ b/quic/core/quic_session.h
@@ -517,6 +517,11 @@
     return use_write_or_buffer_data_at_level_;
   }
 
+  bool use_encryption_level_context() const {
+    return connection_->use_encryption_level_context() &&
+           use_write_or_buffer_data_at_level_;
+  }
+
  protected:
   using StreamMap = QuicHashMap<QuicStreamId, std::unique_ptr<QuicStream>>;