Quic Padding one more try

Summary:
Move the Initial padding code into main scheduler so that ack packets
in Initial space will also be padded.

Reviewed By: mjoras

Differential Revision: D21090338

fbshipit-source-id: dd92c2a1c4d00c58bf2470e2b070c43881a70187

quic/api/QuicPacketScheduler.cpp

@@ -170,6 +170,18 @@ SchedulingResult FrameScheduler::scheduleFramesForPacket(
     streamFrameScheduler_->writeStreams(wrapper);
   }
 
+  if (builder.hasFramesPending()) {
+    const LongHeader* longHeader = builder.getPacketHeader().asLong();
+    bool initialPacket =
+        longHeader && longHeader->getHeaderType() == LongHeader::Types::Initial;
+    if (initialPacket) {
+      // This is the initial packet, we need to fill er up.
+      while (wrapper.remainingSpaceInPkt() > 0) {
+        writeFrame(PaddingFrame(), builder);
+      }
+    }
+  }
+
   return SchedulingResult(folly::none, std::move(builder).buildPacket());
 }
 
@@ -475,17 +487,6 @@ bool CryptoStreamScheduler::writeCryptoData(PacketBuilderInterface& builder) {
       cryptoDataWritten = true;
     }
   }
-  if (cryptoDataWritten) {
-    const LongHeader* longHeader = builder.getPacketHeader().asLong();
-    bool initialPacket =
-        longHeader && longHeader->getHeaderType() == LongHeader::Types::Initial;
-    if (initialPacket) {
-      // This is the initial packet, we need to fill er up.
-      while (builder.remainingSpaceInPkt() > 0) {
-        writeFrame(PaddingFrame(), builder);
-      }
-    }
-  }
   return cryptoDataWritten;
 }
 

quic/api/test/QuicPacketSchedulerTest.cpp

@@ -106,38 +106,100 @@ TEST_F(QuicPacketSchedulerTest, CryptoPaddingInitialPacket) {
   QuicClientConnectionState conn(
       FizzClientQuicHandshakeContext::Builder().build());
   auto connId = getTestConnectionId();
-  LongHeader longHeader1(
+  LongHeader longHeader(
       LongHeader::Types::Initial,
       getTestConnectionId(1),
       connId,
       getNextPacketNum(conn, PacketNumberSpace::Initial),
       QuicVersion::MVFST);
   increaseNextPacketNum(conn, PacketNumberSpace::Initial);
-  RegularQuicPacketBuilder builder1(
+  RegularQuicPacketBuilder builder(
       conn.udpSendPacketLen,
-      std::move(longHeader1),
+      std::move(longHeader),
       conn.ackStates.initialAckState.largestAckedByPeer);
-  CryptoStreamScheduler scheduler(
-      conn, *getCryptoStream(*conn.cryptoState, EncryptionLevel::Initial));
+  FrameScheduler cryptoOnlyScheduler =
+      std::move(
+          FrameScheduler::Builder(
+              conn,
+              EncryptionLevel::Initial,
+              LongHeader::typeToPacketNumberSpace(LongHeader::Types::Initial),
+              "CryptoOnlyScheduler")
+              .cryptoFrames())
+          .build();
   writeDataToQuicStream(
       conn.cryptoState->initialStream, folly::IOBuf::copyBuffer("chlo"));
-  scheduler.writeCryptoData(builder1);
-  EXPECT_EQ(builder1.remainingSpaceInPkt(), 0);
+  auto result = cryptoOnlyScheduler.scheduleFramesForPacket(
+      std::move(builder), conn.udpSendPacketLen);
+  auto packetLength = result.packet->header->computeChainDataLength() +
+      result.packet->body->computeChainDataLength();
+  EXPECT_EQ(conn.udpSendPacketLen, packetLength);
+}
 
-  LongHeader longHeader2(
-      LongHeader::Types::Handshake,
+TEST_F(QuicPacketSchedulerTest, PaddingInitialPureAcks) {
+  QuicClientConnectionState conn(
+      FizzClientQuicHandshakeContext::Builder().build());
+  auto connId = getTestConnectionId();
+  LongHeader longHeader(
+      LongHeader::Types::Initial,
       connId,
       connId,
-      getNextPacketNum(conn, PacketNumberSpace::Handshake),
+      getNextPacketNum(conn, PacketNumberSpace::Initial),
       QuicVersion::MVFST);
-  RegularQuicPacketBuilder builder2(
+  RegularQuicPacketBuilder builder(
       conn.udpSendPacketLen,
-      std::move(longHeader2),
+      std::move(longHeader),
       conn.ackStates.handshakeAckState.largestAckedByPeer);
-  writeDataToQuicStream(
-      conn.cryptoState->initialStream, folly::IOBuf::copyBuffer("finished"));
-  scheduler.writeCryptoData(builder2);
-  EXPECT_GT(builder2.remainingSpaceInPkt(), 0);
+  conn.ackStates.initialAckState.largestRecvdPacketTime = Clock::now();
+  conn.ackStates.initialAckState.needsToSendAckImmediately = true;
+  conn.ackStates.initialAckState.acks.insert(10);
+  FrameScheduler acksOnlyScheduler =
+      std::move(
+          FrameScheduler::Builder(
+              conn,
+              EncryptionLevel::Initial,
+              LongHeader::typeToPacketNumberSpace(LongHeader::Types::Initial),
+              "AcksOnlyScheduler")
+              .ackFrames())
+          .build();
+  auto result = acksOnlyScheduler.scheduleFramesForPacket(
+      std::move(builder), conn.udpSendPacketLen);
+  auto packetLength = result.packet->header->computeChainDataLength() +
+      result.packet->body->computeChainDataLength();
+  EXPECT_EQ(conn.udpSendPacketLen, packetLength);
+}
+
+TEST_F(QuicPacketSchedulerTest, PaddingUpToWrapperSize) {
+  QuicClientConnectionState conn(
+      FizzClientQuicHandshakeContext::Builder().build());
+  auto connId = getTestConnectionId();
+  size_t cipherOverhead = 30;
+  LongHeader longHeader(
+      LongHeader::Types::Initial,
+      connId,
+      connId,
+      getNextPacketNum(conn, PacketNumberSpace::Initial),
+      QuicVersion::MVFST);
+  RegularQuicPacketBuilder builder(
+      conn.udpSendPacketLen,
+      std::move(longHeader),
+      conn.ackStates.handshakeAckState.largestAckedByPeer);
+  conn.ackStates.initialAckState.largestRecvdPacketTime = Clock::now();
+  conn.ackStates.initialAckState.needsToSendAckImmediately = true;
+  conn.ackStates.initialAckState.acks.insert(10);
+  FrameScheduler acksOnlyScheduler =
+      std::move(
+          FrameScheduler::Builder(
+              conn,
+              EncryptionLevel::Initial,
+              LongHeader::typeToPacketNumberSpace(LongHeader::Types::Initial),
+              "AcksOnlyScheduler")
+              .ackFrames())
+          .build();
+  auto result = acksOnlyScheduler.scheduleFramesForPacket(
+      std::move(builder), conn.udpSendPacketLen - cipherOverhead);
+  auto packetLength = result.packet->header->computeChainDataLength() +
+      result.packet->body->computeChainDataLength();
+  EXPECT_EQ(conn.udpSendPacketLen - cipherOverhead, packetLength);
 }
 
 TEST_F(QuicPacketSchedulerTest, CryptoServerInitialPadded) {
@@ -154,13 +216,56 @@ TEST_F(QuicPacketSchedulerTest, CryptoServerInitialPadded) {
       conn.udpSendPacketLen,
       std::move(longHeader1),
       conn.ackStates.initialAckState.largestAckedByPeer);
-  CryptoStreamScheduler scheduler(
-      conn, *getCryptoStream(*conn.cryptoState, EncryptionLevel::Initial));
+  FrameScheduler scheduler =
+      std::move(
+          FrameScheduler::Builder(
+              conn,
+              EncryptionLevel::Initial,
+              LongHeader::typeToPacketNumberSpace(LongHeader::Types::Initial),
+              "CryptoOnlyScheduler")
+              .cryptoFrames())
+          .build();
   writeDataToQuicStream(
       conn.cryptoState->initialStream, folly::IOBuf::copyBuffer("shlo"));
-  scheduler.writeCryptoData(builder1);
-  EXPECT_EQ(builder1.remainingSpaceInPkt(), 0);
-  nextPacketNum++;
+  auto result = scheduler.scheduleFramesForPacket(
+      std::move(builder1), conn.udpSendPacketLen);
+  auto packetLength = result.packet->header->computeChainDataLength() +
+      result.packet->body->computeChainDataLength();
+  EXPECT_EQ(conn.udpSendPacketLen, packetLength);
+}
+
+TEST_F(QuicPacketSchedulerTest, PadTwoInitialPackets) {
+  QuicServerConnectionState conn;
+  auto connId = getTestConnectionId();
+  PacketNum nextPacketNum = getNextPacketNum(conn, PacketNumberSpace::Initial);
+  LongHeader longHeader1(
+      LongHeader::Types::Initial,
+      getTestConnectionId(1),
+      connId,
+      nextPacketNum,
+      QuicVersion::MVFST);
+  RegularQuicPacketBuilder builder1(
+      conn.udpSendPacketLen,
+      std::move(longHeader1),
+      conn.ackStates.initialAckState.largestAckedByPeer);
+  FrameScheduler scheduler =
+      std::move(
+          FrameScheduler::Builder(
+              conn,
+              EncryptionLevel::Initial,
+              LongHeader::typeToPacketNumberSpace(LongHeader::Types::Initial),
+              "CryptoOnlyScheduler")
+              .cryptoFrames())
+          .build();
+  writeDataToQuicStream(
+      conn.cryptoState->initialStream, folly::IOBuf::copyBuffer("shlo"));
+  auto result = scheduler.scheduleFramesForPacket(
+      std::move(builder1), conn.udpSendPacketLen);
+  auto packetLength = result.packet->header->computeChainDataLength() +
+      result.packet->body->computeChainDataLength();
+  EXPECT_EQ(conn.udpSendPacketLen, packetLength);
+
+  increaseNextPacketNum(conn, PacketNumberSpace::Initial);
   LongHeader longHeader2(
       LongHeader::Types::Initial,
       getTestConnectionId(1),
@@ -172,9 +277,12 @@ TEST_F(QuicPacketSchedulerTest, CryptoServerInitialPadded) {
       std::move(longHeader2),
       conn.ackStates.initialAckState.largestAckedByPeer);
   writeDataToQuicStream(
-      conn.cryptoState->initialStream, folly::IOBuf::copyBuffer("shlo"));
-  scheduler.writeCryptoData(builder2);
-  EXPECT_EQ(builder2.remainingSpaceInPkt(), 0);
+      conn.cryptoState->initialStream, folly::IOBuf::copyBuffer("shlo again"));
+  auto result2 = scheduler.scheduleFramesForPacket(
+      std::move(builder2), conn.udpSendPacketLen);
+  packetLength = result2.packet->header->computeChainDataLength() +
+      result2.packet->body->computeChainDataLength();
+  EXPECT_EQ(conn.udpSendPacketLen, packetLength);
 }
 
 TEST_F(QuicPacketSchedulerTest, CryptoPaddingRetransmissionClientInitial) {
@@ -191,12 +299,22 @@ TEST_F(QuicPacketSchedulerTest, CryptoPaddingRetransmissionClientInitial) {
       conn.udpSendPacketLen,
       std::move(longHeader),
       conn.ackStates.initialAckState.largestAckedByPeer);
-  CryptoStreamScheduler scheduler(
-      conn, *getCryptoStream(*conn.cryptoState, EncryptionLevel::Initial));
+  FrameScheduler scheduler =
+      std::move(
+          FrameScheduler::Builder(
+              conn,
+              EncryptionLevel::Initial,
+              LongHeader::typeToPacketNumberSpace(LongHeader::Types::Initial),
+              "CryptoOnlyScheduler")
+              .cryptoFrames())
+          .build();
   conn.cryptoState->initialStream.lossBuffer.push_back(
       StreamBuffer{folly::IOBuf::copyBuffer("chlo"), 0, false});
-  scheduler.writeCryptoData(builder);
-  EXPECT_EQ(builder.remainingSpaceInPkt(), 0);
+  auto result = scheduler.scheduleFramesForPacket(
+      std::move(builder), conn.udpSendPacketLen);
+  auto packetLength = result.packet->header->computeChainDataLength() +
+      result.packet->body->computeChainDataLength();
+  EXPECT_EQ(conn.udpSendPacketLen, packetLength);
 }
 
 TEST_F(QuicPacketSchedulerTest, CryptoSchedulerOnlySingleLossFits) {
@@ -239,14 +357,25 @@ TEST_F(QuicPacketSchedulerTest, CryptoWritePartialLossBuffer) {
       25,
       std::move(longHeader),
       conn.ackStates.initialAckState.largestAckedByPeer);
-  CryptoStreamScheduler scheduler(
-      conn, *getCryptoStream(*conn.cryptoState, EncryptionLevel::Initial));
+  FrameScheduler cryptoOnlyScheduler =
+      std::move(
+          FrameScheduler::Builder(
+              conn,
+              EncryptionLevel::Initial,
+              LongHeader::typeToPacketNumberSpace(LongHeader::Types::Initial),
+              "CryptoOnlyScheduler")
+              .cryptoFrames())
+          .build();
   conn.cryptoState->initialStream.lossBuffer.push_back(StreamBuffer{
       folly::IOBuf::copyBuffer("return the special duration value max"),
       0,
       false});
-  EXPECT_TRUE(scheduler.writeCryptoData(builder));
-  EXPECT_EQ(builder.remainingSpaceInPkt(), 0);
+  auto result = cryptoOnlyScheduler.scheduleFramesForPacket(
+      std::move(builder), conn.udpSendPacketLen);
+  auto packetLength = result.packet->header->computeChainDataLength() +
+      result.packet->body->computeChainDataLength();
+  EXPECT_LE(packetLength, 25);
+  EXPECT_TRUE(result.packet->packet.frames[0].asWriteCryptoFrame() != nullptr);
   EXPECT_FALSE(conn.cryptoState->initialStream.lossBuffer.empty());
 }
 

quic/codec/QuicPacketBuilder.cpp

@@ -531,4 +531,12 @@ uint32_t InplaceQuicPacketBuilder::getHeaderBytes() const {
       (isLongHeader ? packetNumberEncoding_->length + kMaxPacketLenSize : 0);
 }
 
+bool RegularQuicPacketBuilder::hasFramesPending() const {
+  return !packet_.frames.empty();
+}
+
+bool InplaceQuicPacketBuilder::hasFramesPending() const {
+  return !packet_.frames.empty();
+}
+
 } // namespace quic

quic/codec/QuicPacketBuilder.h

@@ -98,6 +98,8 @@ class PacketBuilderInterface {
    */
   [[nodiscard]] virtual uint32_t getHeaderBytes() const = 0;
 
+  [[nodiscard]] virtual bool hasFramesPending() const = 0;
+
   virtual Packet buildPacket() && = 0;
 };
 
@@ -143,6 +145,8 @@ class InplaceQuicPacketBuilder final : public PacketBuilderInterface {
 
   [[nodiscard]] uint32_t getHeaderBytes() const override;
 
+  [[nodiscard]] bool hasFramesPending() const override;
+
  private:
   folly::IOBuf& iobuf_;
   BufWriter bufWriter_;
@@ -207,6 +211,8 @@ class RegularQuicPacketBuilder final : public PacketBuilderInterface {
 
   void setCipherOverhead(uint8_t overhead) noexcept override;
 
+  [[nodiscard]] bool hasFramesPending() const override;
+
  private:
   void writeHeaderBytes(PacketNum largestAckedPacketNum);
   void encodeLongHeader(
@@ -362,6 +368,10 @@ class PacketBuilderWrapper : public PacketBuilderInterface {
     return builder.getHeaderBytes();
   }
 
+  [[nodiscard]] bool hasFramesPending() const override {
+    return builder.hasFramesPending();
+  }
+
  private:
   PacketBuilderInterface& builder;
   uint32_t diff;

quic/codec/test/Mocks.h

@@ -76,6 +76,7 @@ class MockQuicPacketBuilder : public PacketBuilderInterface {
   GMOCK_METHOD1_(, noexcept, , setCipherOverhead, void(uint8_t));
   GMOCK_METHOD0_(, noexcept, , canBuildPacketNonConst, bool());
   GMOCK_METHOD0_(, const, , getHeaderBytes, uint32_t());
+  GMOCK_METHOD0_(, const, , hasFramesPending, bool());
 
   bool canBuildPacket() const noexcept override {
     return const_cast<MockQuicPacketBuilder&>(*this).canBuildPacketNonConst();