Use WriteOrBufferDataAtLevel to send crypto data. Such that existing WriteOrBufferData is used to send application data.

Change SendHandshakeMessage to take encryption level instead of using current default encryption level.

Protected by FLAGS_quic_reloadable_flag_quic_use_write_or_buffer_data_at_level.

PiperOrigin-RevId: 339927045
Change-Id: Ifa17962e5f70e6d4ed2d5aa8627150872f5dd69d
diff --git a/quic/core/http/quic_headers_stream_test.cc b/quic/core/http/quic_headers_stream_test.cc
index 937b5e8..44ed080 100644
--- a/quic/core/http/quic_headers_stream_test.cc
+++ b/quic/core/http/quic_headers_stream_test.cc
@@ -12,6 +12,7 @@
 #include <vector>
 
 #include "absl/strings/string_view.h"
+#include "net/third_party/quiche/src/quic/core/crypto/null_encrypter.h"
 #include "net/third_party/quiche/src/quic/core/http/spdy_utils.h"
 #include "net/third_party/quiche/src/quic/core/quic_data_writer.h"
 #include "net/third_party/quiche/src/quic/core/quic_utils.h"
@@ -226,6 +227,9 @@
     QuicSpdySessionPeer::SetMaxInboundHeaderListSize(&session_, 256 * 1024);
     EXPECT_CALL(session_, OnCongestionWindowChange(_)).Times(AnyNumber());
     session_.Initialize();
+    connection_->SetEncrypter(
+        quic::ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<quic::NullEncrypter>(connection_->perspective()));
     headers_stream_ = QuicSpdySessionPeer::GetHeadersStream(&session_);
     headers_[":status"] = "200 Ok";
     headers_["content-length"] = "11";
diff --git a/quic/core/http/quic_send_control_stream_test.cc b/quic/core/http/quic_send_control_stream_test.cc
index 53d0d6d..4f1b461 100644
--- a/quic/core/http/quic_send_control_stream_test.cc
+++ b/quic/core/http/quic_send_control_stream_test.cc
@@ -8,6 +8,7 @@
 
 #include "absl/strings/escaping.h"
 #include "absl/strings/string_view.h"
+#include "net/third_party/quiche/src/quic/core/crypto/null_encrypter.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_flags.h"
 #include "net/third_party/quiche/src/quic/test_tools/quic_config_peer.h"
 #include "net/third_party/quiche/src/quic/test_tools/quic_spdy_session_peer.h"
@@ -82,6 +83,9 @@
   void Initialize() {
     EXPECT_CALL(session_, OnCongestionWindowChange(_)).Times(AnyNumber());
     session_.Initialize();
+    connection_->SetEncrypter(
+        ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<NullEncrypter>(connection_->perspective()));
     send_control_stream_ = QuicSpdySessionPeer::GetSendControlStream(&session_);
     QuicConfigPeer::SetReceivedInitialSessionFlowControlWindow(
         session_.config(), kMinimumFlowControlSendWindow);
diff --git a/quic/core/http/quic_spdy_client_session_test.cc b/quic/core/http/quic_spdy_client_session_test.cc
index c8bae78..5444c9d 100644
--- a/quic/core/http/quic_spdy_client_session_test.cc
+++ b/quic/core/http/quic_spdy_client_session_test.cc
@@ -121,6 +121,9 @@
         QuicServerId(kServerHostname, kPort, false),
         client_crypto_config_.get(), &push_promise_index_);
     session_->Initialize();
+    connection_->SetEncrypter(
+        ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<NullEncrypter>(connection_->perspective()));
     crypto_stream_ = static_cast<QuicCryptoClientStream*>(
         session_->GetMutableCryptoStream());
     push_promise_[":path"] = "/bar";
@@ -281,7 +284,7 @@
   char data[] = "hello world";
   EXPECT_QUIC_BUG(
       session_->WritevData(stream->id(), ABSL_ARRAYSIZE(data), 0, NO_FIN,
-                           NOT_RETRANSMISSION, absl::nullopt),
+                           NOT_RETRANSMISSION, ENCRYPTION_INITIAL),
       "Client: Try to send data of stream");
 }
 
diff --git a/quic/core/http/quic_spdy_server_stream_base_test.cc b/quic/core/http/quic_spdy_server_stream_base_test.cc
index 64184a7..4505cd0 100644
--- a/quic/core/http/quic_spdy_server_stream_base_test.cc
+++ b/quic/core/http/quic_spdy_server_stream_base_test.cc
@@ -4,6 +4,7 @@
 
 #include "net/third_party/quiche/src/quic/core/http/quic_spdy_server_stream_base.h"
 
+#include "net/third_party/quiche/src/quic/core/crypto/null_encrypter.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_ptr_util.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_test.h"
 #include "net/third_party/quiche/src/quic/test_tools/quic_spdy_session_peer.h"
@@ -33,6 +34,9 @@
                                         &alarm_factory_,
                                         Perspective::IS_SERVER)) {
     session_.Initialize();
+    session_.connection()->SetEncrypter(
+        ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<NullEncrypter>(session_.perspective()));
     stream_ =
         new TestQuicSpdyServerStream(GetNthClientInitiatedBidirectionalStreamId(
                                          session_.transport_version(), 0),
diff --git a/quic/core/http/quic_spdy_session_test.cc b/quic/core/http/quic_spdy_session_test.cc
index 714c21d..0d27288 100644
--- a/quic/core/http/quic_spdy_session_test.cc
+++ b/quic/core/http/quic_spdy_session_test.cc
@@ -335,7 +335,8 @@
     MakeIOVector("not empty", &iov);
     QuicStreamPeer::SendBuffer(stream).SaveStreamData(&iov, 1, 0, 9);
     QuicConsumedData consumed =
-        WritevData(stream->id(), 9, 0, FIN, NOT_RETRANSMISSION, absl::nullopt);
+        WritevData(stream->id(), 9, 0, FIN, NOT_RETRANSMISSION,
+                   GetEncryptionLevelToSendApplicationData());
     QuicStreamPeer::SendBuffer(stream).OnStreamDataConsumed(
         consumed.bytes_consumed);
     return consumed;
@@ -344,7 +345,7 @@
   QuicConsumedData SendLargeFakeData(QuicStream* stream, int bytes) {
     DCHECK(writev_consumes_all_data_);
     return WritevData(stream->id(), bytes, 0, FIN, NOT_RETRANSMISSION,
-                      absl::nullopt);
+                      GetEncryptionLevelToSendApplicationData());
   }
 
   using QuicSession::closed_streams;
@@ -1469,7 +1470,7 @@
     QuicConfig config;
     CryptoHandshakeMessage crypto_message;
     config.ToHandshakeMessage(&crypto_message, transport_version());
-    crypto_stream->SendHandshakeMessage(crypto_message);
+    crypto_stream->SendHandshakeMessage(crypto_message, ENCRYPTION_INITIAL);
     char buf[1000];
     QuicDataWriter writer(1000, buf, quiche::NETWORK_BYTE_ORDER);
     crypto_stream->WriteStreamData(offset, crypto_message.size(), &writer);
diff --git a/quic/core/qpack/qpack_send_stream_test.cc b/quic/core/qpack/qpack_send_stream_test.cc
index f99b168..80c8986 100644
--- a/quic/core/qpack/qpack_send_stream_test.cc
+++ b/quic/core/qpack/qpack_send_stream_test.cc
@@ -5,6 +5,7 @@
 #include "net/third_party/quiche/src/quic/core/qpack/qpack_send_stream.h"
 
 #include "absl/strings/string_view.h"
+#include "net/third_party/quiche/src/quic/core/crypto/null_encrypter.h"
 #include "net/third_party/quiche/src/quic/core/http/http_constants.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_test.h"
 #include "net/third_party/quiche/src/quic/test_tools/quic_config_peer.h"
@@ -69,6 +70,9 @@
         session_(connection_) {
     EXPECT_CALL(session_, OnCongestionWindowChange(_)).Times(AnyNumber());
     session_.Initialize();
+    connection_->SetEncrypter(
+        ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<NullEncrypter>(connection_->perspective()));
     if (connection_->version().SupportsAntiAmplificationLimit()) {
       QuicConnectionPeer::SetAddressValidated(connection_);
     }
diff --git a/quic/core/quic_crypto_client_handshaker.cc b/quic/core/quic_crypto_client_handshaker.cc
index f744153..74f0724 100644
--- a/quic/core/quic_crypto_client_handshaker.cc
+++ b/quic/core/quic_crypto_client_handshaker.cc
@@ -347,7 +347,7 @@
     chlo_hash_ = CryptoUtils::HashHandshakeMessage(out, Perspective::IS_CLIENT);
     session()->connection()->set_fully_pad_crypto_handshake_packets(
         crypto_config_->pad_inchoate_hello());
-    SendHandshakeMessage(out);
+    SendHandshakeMessage(out, ENCRYPTION_INITIAL);
     return;
   }
 
@@ -374,7 +374,7 @@
   next_state_ = STATE_RECV_SHLO;
   session()->connection()->set_fully_pad_crypto_handshake_packets(
       crypto_config_->pad_full_hello());
-  SendHandshakeMessage(out);
+  SendHandshakeMessage(out, ENCRYPTION_INITIAL);
   // Be prepared to decrypt with the new server write key.
   delegate_->OnNewEncryptionKeyAvailable(
       ENCRYPTION_ZERO_RTT,
diff --git a/quic/core/quic_crypto_handshaker.cc b/quic/core/quic_crypto_handshaker.cc
index d608ead..8adf83b 100644
--- a/quic/core/quic_crypto_handshaker.cc
+++ b/quic/core/quic_crypto_handshaker.cc
@@ -20,13 +20,16 @@
 QuicCryptoHandshaker::~QuicCryptoHandshaker() {}
 
 void QuicCryptoHandshaker::SendHandshakeMessage(
-    const CryptoHandshakeMessage& message) {
+    const CryptoHandshakeMessage& message,
+    EncryptionLevel level) {
   QUIC_DVLOG(1) << ENDPOINT << "Sending " << message.DebugString();
   session()->NeuterUnencryptedData();
   session()->OnCryptoHandshakeMessageSent(message);
   last_sent_handshake_message_tag_ = message.tag();
   const QuicData& data = message.GetSerialized();
-  stream_->WriteCryptoData(session_->connection()->encryption_level(),
+  stream_->WriteCryptoData(session_->use_write_or_buffer_data_at_level()
+                               ? level
+                               : session_->connection()->encryption_level(),
                            data.AsStringPiece());
 }
 
diff --git a/quic/core/quic_crypto_handshaker.h b/quic/core/quic_crypto_handshaker.h
index e5d8d51..b971580 100644
--- a/quic/core/quic_crypto_handshaker.h
+++ b/quic/core/quic_crypto_handshaker.h
@@ -21,7 +21,8 @@
 
   // Sends |message| to the peer.
   // TODO(wtc): return a success/failure status.
-  void SendHandshakeMessage(const CryptoHandshakeMessage& message);
+  void SendHandshakeMessage(const CryptoHandshakeMessage& message,
+                            EncryptionLevel level);
 
   void OnError(CryptoFramer* framer) override;
   void OnHandshakeMessage(const CryptoHandshakeMessage& message) override;
diff --git a/quic/core/quic_crypto_server_stream.cc b/quic/core/quic_crypto_server_stream.cc
index 4c21e55..2bad8ec 100644
--- a/quic/core/quic_crypto_server_stream.cc
+++ b/quic/core/quic_crypto_server_stream.cc
@@ -173,7 +173,8 @@
   if (reply->tag() != kSHLO) {
     session()->connection()->set_fully_pad_crypto_handshake_packets(
         crypto_config_->pad_rej());
-    SendHandshakeMessage(*reply);
+    // Send REJ in plaintext.
+    SendHandshakeMessage(*reply, ENCRYPTION_INITIAL);
     return;
   }
 
@@ -213,7 +214,8 @@
 
   session()->connection()->set_fully_pad_crypto_handshake_packets(
       crypto_config_->pad_shlo());
-  SendHandshakeMessage(*reply);
+  // Send SHLO in ENCRYPTION_ZERO_RTT.
+  SendHandshakeMessage(*reply, ENCRYPTION_ZERO_RTT);
   delegate_->OnNewEncryptionKeyAvailable(
       ENCRYPTION_FORWARD_SECURE,
       std::move(crypto_negotiated_params_->forward_secure_crypters.encrypter));
@@ -284,12 +286,15 @@
 
   QUIC_DVLOG(1) << "Server: Sending server config update: "
                 << message.DebugString();
-  if (!QuicVersionUsesCryptoFrames(transport_version())) {
+
+  if (!session()->use_write_or_buffer_data_at_level() &&
+      !QuicVersionUsesCryptoFrames(transport_version())) {
     const QuicData& data = message.GetSerialized();
     WriteOrBufferData(absl::string_view(data.data(), data.length()), false,
                       nullptr);
   } else {
-    SendHandshakeMessage(message);
+    // Send server config update in ENCRYPTION_FORWARD_SECURE.
+    SendHandshakeMessage(message, ENCRYPTION_FORWARD_SECURE);
   }
 
   ++num_server_config_update_messages_sent_;
diff --git a/quic/core/quic_crypto_stream.cc b/quic/core/quic_crypto_stream.cc
index 688b585..5cdae35 100644
--- a/quic/core/quic_crypto_stream.cc
+++ b/quic/core/quic_crypto_stream.cc
@@ -142,6 +142,11 @@
 void QuicCryptoStream::WriteCryptoData(EncryptionLevel level,
                                        absl::string_view data) {
   if (!QuicVersionUsesCryptoFrames(session()->transport_version())) {
+    if (session()->use_write_or_buffer_data_at_level()) {
+      WriteOrBufferDataAtLevel(data, /*fin=*/false, level,
+                               /*ack_listener=*/nullptr);
+      return;
+    }
     // The QUIC crypto handshake takes care of setting the appropriate
     // encryption level before writing data. Since that is the only handshake
     // supported in versions less than 47, |level| can be ignored here.
@@ -307,7 +312,7 @@
                                             QuicByteCount data_length,
                                             bool /*fin*/,
                                             TransmissionType type) {
-  DCHECK_EQ(HANDSHAKE_RETRANSMISSION, type);
+  DCHECK(type == HANDSHAKE_RETRANSMISSION || type == PTO_RETRANSMISSION);
   QuicIntervalSet<QuicStreamOffset> retransmission(offset,
                                                    offset + data_length);
   // Determine the encryption level to send data. This only needs to be once as
@@ -340,7 +345,7 @@
     QuicByteCount retransmission_length,
     EncryptionLevel encryption_level,
     TransmissionType type) {
-  DCHECK_EQ(HANDSHAKE_RETRANSMISSION, type);
+  DCHECK(type == HANDSHAKE_RETRANSMISSION || type == PTO_RETRANSMISSION);
   const auto consumed = stream_delegate()->WritevData(
       id(), retransmission_length, retransmission_offset, NO_FIN, type,
       encryption_level);
diff --git a/quic/core/quic_flags_list.h b/quic/core/quic_flags_list.h
index de8091e..c5c4df8 100644
--- a/quic/core/quic_flags_list.h
+++ b/quic/core/quic_flags_list.h
@@ -74,6 +74,7 @@
 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_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)
 QUIC_FLAG(FLAGS_quic_restart_flag_dont_fetch_quic_private_keys_from_leto, false)
 QUIC_FLAG(FLAGS_quic_restart_flag_quic_enable_zero_rtt_for_tls_v2, true)
diff --git a/quic/core/quic_framer.cc b/quic/core/quic_framer.cc
index 58bf909..587be13 100644
--- a/quic/core/quic_framer.cc
+++ b/quic/core/quic_framer.cc
@@ -2065,6 +2065,20 @@
   return false;
 }
 
+EncryptionLevel QuicFramer::GetEncryptionLevelToSendApplicationData() const {
+  if (!HasAnEncrypterForSpace(APPLICATION_DATA)) {
+    QUIC_BUG
+        << "Tried to get encryption level to send application data with no "
+           "encrypter available.";
+    return NUM_ENCRYPTION_LEVELS;
+  }
+  if (HasEncrypterOfEncryptionLevel(ENCRYPTION_FORWARD_SECURE)) {
+    return ENCRYPTION_FORWARD_SECURE;
+  }
+  DCHECK(HasEncrypterOfEncryptionLevel(ENCRYPTION_ZERO_RTT));
+  return ENCRYPTION_ZERO_RTT;
+}
+
 bool QuicFramer::AppendPacketHeader(const QuicPacketHeader& header,
                                     QuicDataWriter* writer,
                                     size_t* length_field_offset) {
diff --git a/quic/core/quic_framer.h b/quic/core/quic_framer.h
index 86c8882..e6936e0 100644
--- a/quic/core/quic_framer.h
+++ b/quic/core/quic_framer.h
@@ -621,6 +621,10 @@
   // Returns true if an encrypter of |space| is available.
   bool HasAnEncrypterForSpace(PacketNumberSpace space) const;
 
+  // Returns the encryption level to send application data. This should be only
+  // called with available encrypter for application data.
+  EncryptionLevel GetEncryptionLevelToSendApplicationData() const;
+
   void set_validate_flags(bool value) { validate_flags_ = value; }
 
   Perspective perspective() const { return perspective_; }
diff --git a/quic/core/quic_session.cc b/quic/core/quic_session.cc
index 5f1c5ef..bee5b10 100644
--- a/quic/core/quic_session.cc
+++ b/quic/core/quic_session.cc
@@ -730,6 +730,7 @@
     absl::optional<EncryptionLevel> level) {
   DCHECK(connection_->connected())
       << ENDPOINT << "Try to write stream data when connection is closed.";
+  DCHECK(!use_write_or_buffer_data_at_level_ || level.has_value());
   if (!IsEncryptionEstablished() &&
       !QuicUtils::IsCryptoStreamId(transport_version(), id)) {
     // Do not let streams write without encryption. The calling stream will end
@@ -2502,5 +2503,9 @@
   }
 }
 
+EncryptionLevel QuicSession::GetEncryptionLevelToSendApplicationData() const {
+  return connection_->framer().GetEncryptionLevelToSendApplicationData();
+}
+
 #undef ENDPOINT  // undef for jumbo builds
 }  // namespace quic
diff --git a/quic/core/quic_session.h b/quic/core/quic_session.h
index d162d59..cb0efbe 100644
--- a/quic/core/quic_session.h
+++ b/quic/core/quic_session.h
@@ -493,6 +493,9 @@
   virtual void MaybeSendStopSendingFrame(QuicStreamId id,
                                          QuicRstStreamErrorCode error);
 
+  // Returns the encryption level to send application data.
+  EncryptionLevel GetEncryptionLevelToSendApplicationData() const;
+
   const absl::optional<std::string> user_agent_id() const {
     return user_agent_id_;
   }
@@ -510,6 +513,10 @@
     return liveness_testing_in_progress_;
   }
 
+  bool use_write_or_buffer_data_at_level() const {
+    return use_write_or_buffer_data_at_level_;
+  }
+
  protected:
   using StreamMap = QuicHashMap<QuicStreamId, std::unique_ptr<QuicStream>>;
 
@@ -849,6 +856,9 @@
   bool liveness_testing_in_progress_;
 
   const bool split_up_send_rst_ = GetQuicReloadableFlag(quic_split_up_send_rst);
+
+  const bool use_write_or_buffer_data_at_level_ =
+      GetQuicReloadableFlag(quic_use_write_or_buffer_data_at_level);
 };
 
 }  // namespace quic
diff --git a/quic/core/quic_session_test.cc b/quic/core/quic_session_test.cc
index 129c937..0622508 100644
--- a/quic/core/quic_session_test.cc
+++ b/quic/core/quic_session_test.cc
@@ -350,7 +350,8 @@
     MakeIOVector("not empty", &iov);
     QuicStreamPeer::SendBuffer(stream).SaveStreamData(&iov, 1, 0, 9);
     QuicConsumedData consumed =
-        WritevData(stream->id(), 9, 0, FIN, NOT_RETRANSMISSION, absl::nullopt);
+        WritevData(stream->id(), 9, 0, FIN, NOT_RETRANSMISSION,
+                   GetEncryptionLevelToSendApplicationData());
     QuicStreamPeer::SendBuffer(stream).OnStreamDataConsumed(
         consumed.bytes_consumed);
     return consumed;
@@ -367,7 +368,7 @@
   QuicConsumedData SendLargeFakeData(QuicStream* stream, int bytes) {
     DCHECK(writev_consumes_all_data_);
     return WritevData(stream->id(), bytes, 0, FIN, NOT_RETRANSMISSION,
-                      absl::nullopt);
+                      GetEncryptionLevelToSendApplicationData());
   }
 
   bool UsesPendingStreams() const override { return uses_pending_streams_; }
@@ -1683,7 +1684,7 @@
     QuicConfig config;
     CryptoHandshakeMessage crypto_message;
     config.ToHandshakeMessage(&crypto_message, transport_version());
-    crypto_stream->SendHandshakeMessage(crypto_message);
+    crypto_stream->SendHandshakeMessage(crypto_message, ENCRYPTION_INITIAL);
     char buf[1000];
     QuicDataWriter writer(1000, buf, quiche::NETWORK_BYTE_ORDER);
     crypto_stream->WriteStreamData(offset, crypto_message.size(), &writer);
diff --git a/quic/core/quic_stream.cc b/quic/core/quic_stream.cc
index e9daac7..7b11146 100644
--- a/quic/core/quic_stream.cc
+++ b/quic/core/quic_stream.cc
@@ -639,6 +639,23 @@
     absl::string_view data,
     bool fin,
     QuicReferenceCountedPointer<QuicAckListenerInterface> ack_listener) {
+  if (session()->use_write_or_buffer_data_at_level()) {
+    QUIC_BUG_IF(QuicUtils::IsCryptoStreamId(transport_version(), id_))
+        << ENDPOINT
+        << "WriteOrBufferData is used to send application data, use "
+           "WriteOrBufferDataAtLevel to send crypto data.";
+    return WriteOrBufferDataAtLevel(
+        data, fin, session()->GetEncryptionLevelToSendApplicationData(),
+        ack_listener);
+  }
+  return WriteOrBufferDataInner(data, fin, absl::nullopt, ack_listener);
+}
+
+void QuicStream::WriteOrBufferDataInner(
+    absl::string_view data,
+    bool fin,
+    absl::optional<EncryptionLevel> level,
+    QuicReferenceCountedPointer<QuicAckListenerInterface> ack_listener) {
   if (data.empty() && !fin) {
     QUIC_BUG << "data.empty() && !fin";
     return;
@@ -678,10 +695,20 @@
   }
   if (!had_buffered_data && (HasBufferedData() || fin_buffered_)) {
     // Write data if there is no buffered data before.
-    WriteBufferedData();
+    WriteBufferedData(level);
   }
 }
 
+void QuicStream::WriteOrBufferDataAtLevel(
+    absl::string_view data,
+    bool fin,
+    EncryptionLevel level,
+    QuicReferenceCountedPointer<QuicAckListenerInterface> ack_listener) {
+  DCHECK(session()->use_write_or_buffer_data_at_level());
+  QUIC_RELOADABLE_FLAG_COUNT(quic_use_write_or_buffer_data_at_level);
+  return WriteOrBufferDataInner(data, fin, level, ack_listener);
+}
+
 void QuicStream::OnCanWrite() {
   if (HasDeadlinePassed()) {
     OnDeadlinePassed();
@@ -701,7 +728,11 @@
     return;
   }
   if (HasBufferedData() || (fin_buffered_ && !fin_sent_)) {
-    WriteBufferedData();
+    absl::optional<EncryptionLevel> send_level = absl::nullopt;
+    if (session()->use_write_or_buffer_data_at_level()) {
+      send_level = session()->GetEncryptionLevelToSendApplicationData();
+    }
+    WriteBufferedData(send_level);
   }
   if (!fin_buffered_ && !fin_sent_ && CanWriteNewData()) {
     // Notify upper layer to write new data when buffered data size is below
@@ -779,7 +810,11 @@
 
   if (!had_buffered_data && (HasBufferedData() || fin_buffered_)) {
     // Write data if there is no buffered data before.
-    WriteBufferedData();
+    absl::optional<EncryptionLevel> send_level = absl::nullopt;
+    if (session()->use_write_or_buffer_data_at_level()) {
+      send_level = session()->GetEncryptionLevelToSendApplicationData();
+    }
+    WriteBufferedData(send_level);
   }
 
   return consumed_data;
@@ -1126,6 +1161,10 @@
   if (retransmission.Empty() && !retransmit_fin) {
     return true;
   }
+  absl::optional<EncryptionLevel> send_level = absl::nullopt;
+  if (session()->use_write_or_buffer_data_at_level()) {
+    send_level = session()->GetEncryptionLevelToSendApplicationData();
+  }
   QuicConsumedData consumed(0, false);
   for (const auto& interval : retransmission) {
     QuicStreamOffset retransmission_offset = interval.min();
@@ -1135,7 +1174,7 @@
                            stream_bytes_written());
     consumed = stream_delegate_->WritevData(
         id_, retransmission_length, retransmission_offset,
-        can_bundle_fin ? FIN : NO_FIN, type, absl::nullopt);
+        can_bundle_fin ? FIN : NO_FIN, type, send_level);
     QUIC_DVLOG(1) << ENDPOINT << "stream " << id_
                   << " is forced to retransmit stream data ["
                   << retransmission_offset << ", "
@@ -1157,7 +1196,7 @@
     QUIC_DVLOG(1) << ENDPOINT << "stream " << id_
                   << " retransmits fin only frame.";
     consumed = stream_delegate_->WritevData(id_, 0, stream_bytes_written(), FIN,
-                                            type, absl::nullopt);
+                                            type, send_level);
     if (!consumed.fin_consumed) {
       return false;
     }
@@ -1183,7 +1222,7 @@
   return send_buffer_.WriteStreamData(offset, data_length, writer);
 }
 
-void QuicStream::WriteBufferedData() {
+void QuicStream::WriteBufferedData(absl::optional<EncryptionLevel> level) {
   DCHECK(!write_side_closed_ && (HasBufferedData() || fin_buffered_));
 
   if (session_->ShouldYield(id())) {
@@ -1235,7 +1274,7 @@
   }
   QuicConsumedData consumed_data =
       stream_delegate_->WritevData(id(), write_length, stream_bytes_written(),
-                                   state, NOT_RETRANSMISSION, absl::nullopt);
+                                   state, NOT_RETRANSMISSION, level);
 
   OnStreamDataConsumed(consumed_data.bytes_consumed);
 
@@ -1306,12 +1345,15 @@
 void QuicStream::WritePendingRetransmission() {
   while (HasPendingRetransmission()) {
     QuicConsumedData consumed(0, false);
+    absl::optional<EncryptionLevel> send_level = absl::nullopt;
+    if (session()->use_write_or_buffer_data_at_level()) {
+      send_level = session()->GetEncryptionLevelToSendApplicationData();
+    }
     if (!send_buffer_.HasPendingRetransmission()) {
       QUIC_DVLOG(1) << ENDPOINT << "stream " << id_
                     << " retransmits fin only frame.";
-      consumed =
-          stream_delegate_->WritevData(id_, 0, stream_bytes_written(), FIN,
-                                       LOSS_RETRANSMISSION, absl::nullopt);
+      consumed = stream_delegate_->WritevData(
+          id_, 0, stream_bytes_written(), FIN, LOSS_RETRANSMISSION, send_level);
       fin_lost_ = !consumed.fin_consumed;
       if (fin_lost_) {
         // Connection is write blocked.
@@ -1326,7 +1368,7 @@
           (pending.offset + pending.length == stream_bytes_written());
       consumed = stream_delegate_->WritevData(
           id_, pending.length, pending.offset, can_bundle_fin ? FIN : NO_FIN,
-          LOSS_RETRANSMISSION, absl::nullopt);
+          LOSS_RETRANSMISSION, send_level);
       QUIC_DVLOG(1) << ENDPOINT << "stream " << id_
                     << " tries to retransmit stream data [" << pending.offset
                     << ", " << pending.offset + pending.length
diff --git a/quic/core/quic_stream.h b/quic/core/quic_stream.h
index 467cca9..275e93a 100644
--- a/quic/core/quic_stream.h
+++ b/quic/core/quic_stream.h
@@ -275,15 +275,23 @@
   // stop sending stream-level flow-control updates when this end sends FIN.
   virtual void StopReading();
 
-  // Sends as much of 'data' to the connection as the connection will consume,
-  // and then buffers any remaining data in queued_data_.
-  // If fin is true: if it is immediately passed on to the session,
-  // write_side_closed() becomes true, otherwise fin_buffered_ becomes true.
+  // Sends as much of |data| to the connection on the application encryption
+  // level as the connection will consume, and then buffers any remaining data
+  // in the send buffer. If fin is true: if it is immediately passed on to the
+  // session, write_side_closed() becomes true, otherwise fin_buffered_ becomes
+  // true.
   void WriteOrBufferData(
       absl::string_view data,
       bool fin,
       QuicReferenceCountedPointer<QuicAckListenerInterface> ack_listener);
 
+  // Sends |data| to connection with specified |level|.
+  void WriteOrBufferDataAtLevel(
+      absl::string_view data,
+      bool fin,
+      EncryptionLevel level,
+      QuicReferenceCountedPointer<QuicAckListenerInterface> ack_listener);
+
   // Adds random padding after the fin is consumed for this stream.
   void AddRandomPaddingAfterFin();
 
@@ -455,9 +463,10 @@
   // controller, marks this stream as connection-level write blocked.
   void MaybeSendBlocked();
 
-  // Write buffered data in send buffer. TODO(fayang): Consider combine
-  // WriteOrBufferData, Writev and WriteBufferedData.
-  void WriteBufferedData();
+  // Write buffered data in send buffer.
+  // TODO(fayang): Change absl::optional<EncryptionLevel> to EncryptionLevel
+  // when deprecating quic_use_write_or_buffer_data_at_level.
+  void WriteBufferedData(absl::optional<EncryptionLevel> level);
 
   // Close the read side of the stream.  May cause the stream to be closed.
   void CloseReadSide();
@@ -465,6 +474,14 @@
   // Called when bytes are sent to the peer.
   void AddBytesSent(QuicByteCount bytes);
 
+  // TODO(fayang): Inline this function when deprecating
+  // quic_use_write_or_buffer_data_at_level.
+  void WriteOrBufferDataInner(
+      absl::string_view data,
+      bool fin,
+      absl::optional<EncryptionLevel> level,
+      QuicReferenceCountedPointer<QuicAckListenerInterface> ack_listener);
+
   // Returns true if deadline_ has passed.
   bool HasDeadlinePassed() const;
 
diff --git a/quic/core/quic_stream_test.cc b/quic/core/quic_stream_test.cc
index 563a7fd..d241dac 100644
--- a/quic/core/quic_stream_test.cc
+++ b/quic/core/quic_stream_test.cc
@@ -11,6 +11,7 @@
 #include "absl/base/macros.h"
 #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/frames/quic_rst_stream_frame.h"
 #include "net/third_party/quiche/src/quic/core/quic_connection.h"
 #include "net/third_party/quiche/src/quic/core/quic_constants.h"
@@ -95,7 +96,9 @@
     connection_->AdvanceTime(QuicTime::Delta::FromSeconds(1));
     session_ = std::make_unique<StrictMock<MockQuicSession>>(connection_);
     session_->Initialize();
-
+    connection_->SetEncrypter(
+        ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<NullEncrypter>(connection_->perspective()));
     QuicConfigPeer::SetReceivedInitialSessionFlowControlWindow(
         session_->config(), kMinimumFlowControlSendWindow);
     QuicConfigPeer::SetReceivedInitialMaxStreamDataBytesUnidirectional(
diff --git a/quic/core/stream_delegate_interface.h b/quic/core/stream_delegate_interface.h
index 7c36c91..dc93299 100644
--- a/quic/core/stream_delegate_interface.h
+++ b/quic/core/stream_delegate_interface.h
@@ -26,6 +26,8 @@
   // Called when the stream needs to write data. If |level| is present, the data
   // will be written at the specified |level|. The data will be written
   // at specified transmission |type|.
+  // TODO(fayang): Change absl::optional<EncryptionLevel> to EncryptionLevel
+  // when deprecating quic_use_write_or_buffer_data_at_level.
   virtual QuicConsumedData WritevData(
       QuicStreamId id,
       size_t write_length,
diff --git a/quic/quic_transport/quic_transport_stream_test.cc b/quic/quic_transport/quic_transport_stream_test.cc
index e6c244b..82cd914 100644
--- a/quic/quic_transport/quic_transport_stream_test.cc
+++ b/quic/quic_transport/quic_transport_stream_test.cc
@@ -7,6 +7,7 @@
 #include <memory>
 
 #include "absl/strings/string_view.h"
+#include "net/third_party/quiche/src/quic/core/crypto/null_encrypter.h"
 #include "net/third_party/quiche/src/quic/core/frames/quic_window_update_frame.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_expect_bug.h"
 #include "net/third_party/quiche/src/quic/platform/api/quic_ptr_util.h"
@@ -42,7 +43,9 @@
         session_(connection_) {
     QuicEnableVersion(DefaultVersionForQuicTransport());
     session_.Initialize();
-
+    connection_->SetEncrypter(
+        ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<NullEncrypter>(connection_->perspective()));
     stream_ = new QuicTransportStream(0, &session_, &interface_);
     session_.ActivateStream(QuicWrapUnique(stream_));
 
diff --git a/quic/tools/quic_simple_server_stream_test.cc b/quic/tools/quic_simple_server_stream_test.cc
index 90d1313..4b21cca 100644
--- a/quic/tools/quic_simple_server_stream_test.cc
+++ b/quic/tools/quic_simple_server_stream_test.cc
@@ -11,6 +11,7 @@
 #include "absl/base/macros.h"
 #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/http/http_encoder.h"
 #include "net/third_party/quiche/src/quic/core/http/spdy_utils.h"
 #include "net/third_party/quiche/src/quic/core/quic_error_codes.h"
@@ -251,6 +252,9 @@
     session_.config()->SetInitialSessionFlowControlWindowToSend(
         kInitialSessionFlowControlWindowForTest);
     session_.Initialize();
+    connection_->SetEncrypter(
+        quic::ENCRYPTION_FORWARD_SECURE,
+        std::make_unique<quic::NullEncrypter>(connection_->perspective()));
     if (connection_->version().SupportsAntiAmplificationLimit()) {
       QuicConnectionPeer::SetAddressValidated(connection_);
     }