In quic, consider frame expansion when extra padding is needed to meet the minimum plaintext size needed for header protection. protected by gfe2_reloadable_flag_quic_fix_extra_padding_bytes.

PiperOrigin-RevId: 318497490
Change-Id: If5ade70a1992e6daa7c6d116cb778cb88210981c

quic/core/quic_connection_test.cc

diff --git a/quic/core/quic_connection_test.cc b/quic/core/quic_connection_test.cc
index 7fc45ca..e751f7b 100644
--- a/quic/core/quic_connection_test.cc
+++ b/quic/core/quic_connection_test.cc
@@ -11366,7 +11366,8 @@
   if (!connection_.version().CanSendCoalescedPackets()) {
     return;
   }
-  if (GetQuicReloadableFlag(quic_fix_min_crypto_frame_size)) {
+  if (GetQuicReloadableFlag(quic_fix_extra_padding_bytes) ||
+      GetQuicReloadableFlag(quic_fix_min_crypto_frame_size)) {
     EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(1);
   } else {
     EXPECT_CALL(visitor_, OnHandshakePacketSent()).Times(0);
@@ -11382,7 +11383,8 @@
     SendStreamDataToPeer(2, std::string(1286, 'a'), 0, NO_FIN, nullptr);
     connection_.SetDefaultEncryptionLevel(ENCRYPTION_HANDSHAKE);
     // Try to coalesce a HANDSHAKE packet after 1-RTT packet.
-    if (GetQuicReloadableFlag(quic_fix_min_crypto_frame_size)) {
+    if (GetQuicReloadableFlag(quic_fix_extra_padding_bytes) ||
+        GetQuicReloadableFlag(quic_fix_min_crypto_frame_size)) {
       // Verify soft max packet length gets resumed and handshake packet gets
       // successfully sent.
       connection_.SendCryptoDataWithString("a", 0, ENCRYPTION_HANDSHAKE);

quic/core/quic_packet_creator.cc

diff --git a/quic/core/quic_packet_creator.cc b/quic/core/quic_packet_creator.cc
index a31b179..56c3050 100644
--- a/quic/core/quic_packet_creator.cc
+++ b/quic/core/quic_packet_creator.cc
@@ -432,10 +432,8 @@
   size_t min_frame_size =
       QuicFramer::GetMinCryptoFrameSize(write_length, offset);
   size_t min_plaintext_bytes = min_frame_size;
-  if (fix_min_crypto_frame_size_ && queued_frames_.empty()) {
-    // TODO(fayang): to make this more general, we need to move the checking
-    // when trying to add a frame when GetSerializedFrameLength. Remove soft
-    // limit if current packet needs extra padding and the space is too small.
+  if (!fix_extra_padding_bytes_ && fix_min_crypto_frame_size_ &&
+      queued_frames_.empty()) {
     QUIC_RELOADABLE_FLAG_COUNT(quic_fix_min_crypto_frame_size);
     min_plaintext_bytes =
         std::max(min_frame_size, MinPlaintextPacketSize(framer_->version()));
@@ -1567,6 +1565,43 @@
       EncryptionlevelToLongHeaderType(packet_.encryption_level);
 }
 
+size_t QuicPacketCreator::GetSerializedFrameLength(const QuicFrame& frame) {
+  size_t serialized_frame_length = framer_->GetSerializedFrameLength(
+      frame, BytesFree(), queued_frames_.empty(),
+      /* last_frame_in_packet= */ true, GetPacketNumberLength());
+  if (!fix_extra_padding_bytes_) {
+    return serialized_frame_length;
+  }
+  if (!framer_->version().HasHeaderProtection() ||
+      serialized_frame_length == 0) {
+    return serialized_frame_length;
+  }
+  // Calculate frame bytes and bytes free with this frame added.
+  const size_t frame_bytes = PacketSize() - PacketHeaderSize() +
+                             ExpansionOnNewFrame() + serialized_frame_length;
+  if (frame_bytes >= MinPlaintextPacketSize(framer_->version())) {
+    // No extra bytes is needed.
+    return serialized_frame_length;
+  }
+  QUIC_RELOADABLE_FLAG_COUNT_N(quic_fix_extra_padding_bytes, 1, 3);
+  if (BytesFree() < serialized_frame_length) {
+    QUIC_BUG << ENDPOINT << "Frame does not fit: " << frame;
+    return 0;
+  }
+  // Please note bytes_free does not take |frame|'s expansion into account.
+  size_t bytes_free = BytesFree() - serialized_frame_length;
+  // Extra bytes needed (this is NOT padding needed) should be at least 1
+  // padding + expansion.
+  const size_t extra_bytes_needed =
+      std::max(1 + ExpansionOnNewFrameWithLastFrame(frame),
+               MinPlaintextPacketSize(framer_->version()) - frame_bytes);
+  if (bytes_free < extra_bytes_needed) {
+    // This frame does not fit.
+    return 0;
+  }
+  return serialized_frame_length;
+}
+
 bool QuicPacketCreator::AddFrame(const QuicFrame& frame,
                                  TransmissionType transmission_type) {
   QUIC_DVLOG(1) << ENDPOINT << "Adding frame with transmission type "
@@ -1594,14 +1629,10 @@
     }
   }
 
-  size_t frame_len = framer_->GetSerializedFrameLength(
-      frame, BytesFree(), queued_frames_.empty(),
-      /* last_frame_in_packet= */ true, GetPacketNumberLength());
+  size_t frame_len = GetSerializedFrameLength(frame);
   if (frame_len == 0 && RemoveSoftMaxPacketLength()) {
     // Remove soft max_packet_length and retry.
-    frame_len = framer_->GetSerializedFrameLength(
-        frame, BytesFree(), queued_frames_.empty(),
-        /* last_frame_in_packet= */ true, GetPacketNumberLength());
+    frame_len = GetSerializedFrameLength(frame);
   }
   if (frame_len == 0) {
     QUIC_DVLOG(1) << "Flushing because current open packet is full when adding "
@@ -1655,6 +1686,25 @@
   return true;
 }
 
+void QuicPacketCreator::MaybeAddExtraPaddingForHeaderProtection() {
+  DCHECK(fix_extra_padding_bytes_);
+  if (!framer_->version().HasHeaderProtection() || needs_full_padding_) {
+    return;
+  }
+  const size_t frame_bytes = PacketSize() - PacketHeaderSize();
+  if (frame_bytes >= MinPlaintextPacketSize(framer_->version())) {
+    return;
+  }
+  const QuicByteCount min_header_protection_padding =
+      std::max(1 + ExpansionOnNewFrame(),
+               MinPlaintextPacketSize(framer_->version()) - frame_bytes) -
+      ExpansionOnNewFrame();
+  QUIC_RELOADABLE_FLAG_COUNT_N(quic_fix_extra_padding_bytes, 2, 3);
+  // Update pending_padding_bytes_.
+  pending_padding_bytes_ =
+      std::max(pending_padding_bytes_, min_header_protection_padding);
+}
+
 bool QuicPacketCreator::MaybeCoalesceStreamFrame(const QuicStreamFrame& frame) {
   if (queued_frames_.empty() || queued_frames_.back().type != STREAM_FRAME) {
     return false;
@@ -1736,15 +1786,21 @@
   }
 
   // Header protection requires a minimum plaintext packet size.
+  // TODO(fayang): remove extra_padding_bytes when deprecating
+  // quic_fix_extra_padding_bytes.
   size_t extra_padding_bytes = 0;
-  if (framer_->version().HasHeaderProtection()) {
-    size_t frame_bytes = PacketSize() - PacketHeaderSize();
+  if (fix_extra_padding_bytes_) {
+    MaybeAddExtraPaddingForHeaderProtection();
+  } else {
+    if (framer_->version().HasHeaderProtection()) {
+      size_t frame_bytes = PacketSize() - PacketHeaderSize();
 
-    if (frame_bytes + pending_padding_bytes_ <
-            MinPlaintextPacketSize(framer_->version()) &&
-        !needs_full_padding_) {
-      extra_padding_bytes =
-          MinPlaintextPacketSize(framer_->version()) - frame_bytes;
+      if (frame_bytes + pending_padding_bytes_ <
+              MinPlaintextPacketSize(framer_->version()) &&
+          !needs_full_padding_) {
+        extra_padding_bytes =
+            MinPlaintextPacketSize(framer_->version()) - frame_bytes;
+      }
     }
   }
 

quic/core/quic_packet_creator.h

diff --git a/quic/core/quic_packet_creator.h b/quic/core/quic_packet_creator.h
index ddc5e2d..c579e42 100644
--- a/quic/core/quic_packet_creator.h
+++ b/quic/core/quic_packet_creator.h
@@ -542,6 +542,14 @@
   // Returns true if packet under construction has IETF long header.
   bool HasIetfLongHeader() const;
 
+  // Get serialized frame length. Returns 0 if the frame does not fit into
+  // current packet.
+  size_t GetSerializedFrameLength(const QuicFrame& frame);
+
+  // Add extra padding to pending_padding_bytes_ to meet minimum plaintext
+  // packet size required for header protection.
+  void MaybeAddExtraPaddingForHeaderProtection();
+
   // Does not own these delegates or the framer.
   DelegateInterface* delegate_;
   DebugDelegate* debug_delegate_;
@@ -622,7 +630,11 @@
   // When true, this will override the padding generation code to disable it.
   bool disable_padding_override_ = false;
 
-  bool update_packet_size_ = GetQuicReloadableFlag(quic_update_packet_size);
+  const bool update_packet_size_ =
+      GetQuicReloadableFlag(quic_update_packet_size);
+
+  const bool fix_extra_padding_bytes_ =
+      GetQuicReloadableFlag(quic_fix_extra_padding_bytes);
 };
 
 }  // namespace quic

quic/core/quic_packet_creator_test.cc

diff --git a/quic/core/quic_packet_creator_test.cc b/quic/core/quic_packet_creator_test.cc
index 394b82e..7c4416d 100644
--- a/quic/core/quic_packet_creator_test.cc
+++ b/quic/core/quic_packet_creator_test.cc
@@ -3793,6 +3793,31 @@
   EXPECT_EQ(TestConnectionId(0x33), creator_.GetSourceConnectionId());
 }
 
+// Regresstion test for b/159812345.
+TEST_F(QuicPacketCreatorMultiplePacketsTest, ExtraPaddingNeeded) {
+  if (!framer_.version().HasHeaderProtection()) {
+    return;
+  }
+  delegate_.SetCanWriteAnything();
+
+  EXPECT_CALL(delegate_, OnSerializedPacket(_))
+      .WillOnce(
+          Invoke(this, &QuicPacketCreatorMultiplePacketsTest::SavePacket));
+  MakeIOVector("a", &iov_);
+  creator_.ConsumeData(QuicUtils::GetFirstBidirectionalStreamId(
+                           framer_.transport_version(), Perspective::IS_CLIENT),
+                       &iov_, 1u, iov_.iov_len, 0, FIN);
+  creator_.Flush();
+  ASSERT_FALSE(packets_[0].nonretransmittable_frames.empty());
+  QuicFrame padding = packets_[0].nonretransmittable_frames[0];
+  if (GetQuicReloadableFlag(quic_fix_extra_padding_bytes)) {
+    // Verify stream frame expansion is excluded.
+    padding.padding_frame.num_padding_bytes = 3;
+  } else {
+    padding.padding_frame.num_padding_bytes = 4;
+  }
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic