Encode packet header after Quic cloner is sure the packet can be cloned

Summary:
otherwise we keep encoding packet headers and potentially write into
the write buffer without actually generating a packet. It leaves a trail of bad
headers inside the buffer

Reviewed By: mjoras

Differential Revision: D21626733

fbshipit-source-id: 3bdcf6277beccb09b390a590ba2bb0eb8e68e6c1

quic/api/QuicPacketScheduler.cpp

@@ -541,6 +541,12 @@ SchedulingResult CloningScheduler::scheduleFramesForPacket(
     if (opPnSpace != PacketNumberSpace::AppData) {
       continue;
     }
+    size_t prevSize = 0;
+    if (conn_.transportSettings.dataPathType ==
+        DataPathType::ContinuousMemory) {
+      ScopedBufAccessor scopedBufAccessor(conn_.bufAccessor);
+      prevSize = scopedBufAccessor.buf()->length();
+    }
     // Reusing the same builder throughout loop bodies will lead to frames
     // belong to different original packets being written into the same clone
     // packet. So re-create a builder every time.
@@ -560,9 +566,6 @@ SchedulingResult CloningScheduler::scheduleFramesForPacket(
           header,
           getAckState(conn_, builderPnSpace).largestAckedByPeer);
     }
-    internalBuilder->setCipherOverhead(cipherOverhead_);
-    internalBuilder->encodePacketHeader();
-    PacketRebuilder rebuilder(*internalBuilder, conn_);
     // We shouldn't clone Handshake packet.
     if (iter->isHandshake) {
       continue;
@@ -573,13 +576,16 @@ SchedulingResult CloningScheduler::scheduleFramesForPacket(
         conn_.outstandingPacketEvents.count(*iter->associatedEvent) == 0) {
       continue;
     }
-
-    // The writableBytes here is an optimization. If the writableBytes is too
-    // small for this packet. rebuildFromPacket should fail anyway.
+    // I think this only fail if udpSendPacketLen somehow shrinks in the middle
+    // of a connection.
     if (iter->encodedSize > writableBytes + cipherOverhead_) {
       continue;
     }
 
+    internalBuilder->setCipherOverhead(cipherOverhead_);
+    internalBuilder->encodePacketHeader();
+    PacketRebuilder rebuilder(*internalBuilder, conn_);
+
     // TODO: It's possible we write out a packet that's larger than the packet
     // size limit. For example, when the packet sequence number has advanced to
     // a point where we need more bytes to encoded it than that of the original
@@ -593,6 +599,22 @@ SchedulingResult CloningScheduler::scheduleFramesForPacket(
     if (rebuildResult) {
       return SchedulingResult(
           std::move(rebuildResult), std::move(*internalBuilder).buildPacket());
+    } else if (
+        conn_.transportSettings.dataPathType ==
+        DataPathType::ContinuousMemory) {
+      // When we use Inplace packet building and reuse the write buffer, even if
+      // the packet rebuild has failed, there might be some bytes already
+      // written into the buffer and the buffer tail pointer has already moved.
+      // We need to roll back the tail pointer to the position before the packet
+      // building to exclude those bytes. Otherwise these bytes will be sitting
+      // in between legit packets inside the buffer and will either cause errors
+      // further down the write path, or be sent out and then dropped at peer
+      // when peer fail to parse them.
+      internalBuilder.reset();
+      CHECK(conn_.bufAccessor && conn_.bufAccessor->ownsBuffer());
+      ScopedBufAccessor scopedBufAccessor(conn_.bufAccessor);
+      auto& buf = scopedBufAccessor.buf();
+      buf->trimEnd(buf->length() - prevSize);
     }
   }
   return SchedulingResult(folly::none, folly::none);

quic/api/test/QuicPacketSchedulerTest.cpp

@@ -1210,5 +1210,88 @@ TEST_F(QuicPacketSchedulerTest, StreamFrameSchedulerRemoveOne) {
   EXPECT_EQ(*builder.frames_[0].asWriteStreamFrame(), f1);
 }
 
+TEST_F(
+    QuicPacketSchedulerTest,
+    CloningSchedulerWithInplaceBuilderDoNotEncodeHeaderWithoutBuild) {
+  QuicClientConnectionState conn(
+      FizzClientQuicHandshakeContext::Builder().build());
+  conn.transportSettings.dataPathType = DataPathType::ContinuousMemory;
+  SimpleBufAccessor bufAccessor(2000);
+  auto buf = bufAccessor.obtain();
+  EXPECT_EQ(buf->length(), 0);
+  bufAccessor.release(std::move(buf));
+  conn.bufAccessor = &bufAccessor;
+
+  FrameScheduler noopScheduler("frame");
+  ASSERT_FALSE(noopScheduler.hasData());
+  CloningScheduler cloningScheduler(noopScheduler, conn, "Little Hurry", 0);
+  addOutstandingPacket(conn);
+  // There needs to have retransmittable frame for the rebuilder to work
+  conn.outstandingPackets.back().packet.frames.push_back(
+      MaxDataFrame(conn.flowControlState.advertisedMaxOffset));
+  // Lie about the encodedSize to let the Cloner skip it:
+  conn.outstandingPackets.back().encodedSize = kDefaultUDPSendPacketLen * 2;
+  EXPECT_TRUE(cloningScheduler.hasData());
+
+  ASSERT_FALSE(noopScheduler.hasData());
+  ShortHeader header(
+      ProtectionType::KeyPhaseOne,
+      conn.clientConnectionId.value_or(getTestConnectionId()),
+      getNextPacketNum(conn, PacketNumberSpace::AppData));
+  InplaceQuicPacketBuilder builder(
+      bufAccessor,
+      conn.udpSendPacketLen,
+      std::move(header),
+      conn.ackStates.appDataAckState.largestAckedByPeer);
+  auto result = cloningScheduler.scheduleFramesForPacket(
+      std::move(builder), kDefaultUDPSendPacketLen);
+  EXPECT_FALSE(result.packetEvent.has_value());
+
+  // Nothing was written into the buffer:
+  EXPECT_TRUE(bufAccessor.ownsBuffer());
+  buf = bufAccessor.obtain();
+  EXPECT_EQ(buf->length(), 0);
+}
+
+TEST_F(
+    QuicPacketSchedulerTest,
+    CloningSchedulerWithInplaceBuilderRollbackBufWhenFailToRebuild) {
+  QuicClientConnectionState conn(
+      FizzClientQuicHandshakeContext::Builder().build());
+  conn.transportSettings.dataPathType = DataPathType::ContinuousMemory;
+  SimpleBufAccessor bufAccessor(2000);
+  auto buf = bufAccessor.obtain();
+  EXPECT_EQ(buf->length(), 0);
+  bufAccessor.release(std::move(buf));
+  conn.bufAccessor = &bufAccessor;
+
+  FrameScheduler noopScheduler("frame");
+  ASSERT_FALSE(noopScheduler.hasData());
+  CloningScheduler cloningScheduler(noopScheduler, conn, "HotPot", 0);
+  addOutstandingPacket(conn);
+  // Not adding frame to this outstanding packet so that rebuild will fail:
+  ASSERT_TRUE(conn.outstandingPackets.back().packet.frames.empty());
+  EXPECT_TRUE(cloningScheduler.hasData());
+
+  ASSERT_FALSE(noopScheduler.hasData());
+  ShortHeader header(
+      ProtectionType::KeyPhaseOne,
+      conn.clientConnectionId.value_or(getTestConnectionId()),
+      getNextPacketNum(conn, PacketNumberSpace::AppData));
+  InplaceQuicPacketBuilder builder(
+      bufAccessor,
+      conn.udpSendPacketLen,
+      std::move(header),
+      conn.ackStates.appDataAckState.largestAckedByPeer);
+  auto result = cloningScheduler.scheduleFramesForPacket(
+      std::move(builder), kDefaultUDPSendPacketLen);
+  EXPECT_FALSE(result.packetEvent.has_value());
+
+  // Nothing was written into the buffer:
+  EXPECT_TRUE(bufAccessor.ownsBuffer());
+  buf = bufAccessor.obtain();
+  EXPECT_EQ(buf->length(), 0);
+}
+
 } // namespace test
 } // namespace quic