From b0cb27110be7c74cc9155be3851bd15213a31e27 Mon Sep 17 00:00:00 2001
From: Yang Chi <yangchi@fb.com>
Date: Thu, 7 May 2020 10:52:01 -0700
Subject: [PATCH] Set ciphere overhead in the packet builder used by Quic
 cloner

Summary: as title

Reviewed By: mjoras

Differential Revision: D21383829

fbshipit-source-id: 181085f1f1208ce45579e40c0a0cab28e0bc4aed
---
 quic/api/QuicPacketScheduler.cpp          | 14 ++++--
 quic/api/test/QuicPacketSchedulerTest.cpp | 58 +++++++++++++++++++++++
 2 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/quic/api/QuicPacketScheduler.cpp b/quic/api/QuicPacketScheduler.cpp
index 9f6403205..ad2ffd27e 100644
--- a/quic/api/QuicPacketScheduler.cpp
+++ b/quic/api/QuicPacketScheduler.cpp
@@ -549,17 +549,18 @@ SchedulingResult CloningScheduler::scheduleFramesForPacket(
     std::unique_ptr<PacketBuilderInterface> internalBuilder;
     if (conn_.transportSettings.dataPathType == DataPathType::ChainedMemory) {
       internalBuilder = std::make_unique<RegularQuicPacketBuilder>(
-          conn_.udpSendPacketLen,
+          conn_.udpSendPacketLen - cipherOverhead_,
           header,
           getAckState(conn_, builderPnSpace).largestAckedByPeer);
     } else {
       CHECK(conn_.bufAccessor && conn_.bufAccessor->ownsBuffer());
       internalBuilder = std::make_unique<InplaceQuicPacketBuilder>(
           *conn_.bufAccessor,
-          conn_.udpSendPacketLen,
+          conn_.udpSendPacketLen - cipherOverhead_,
           header,
           getAckState(conn_, builderPnSpace).largestAckedByPeer);
     }
+    internalBuilder->setCipherOverhead(cipherOverhead_);
     internalBuilder->encodePacketHeader();
     PacketRebuilder rebuilder(*internalBuilder, conn_);
     // We shouldn't clone Handshake packet.
@@ -575,11 +576,18 @@ SchedulingResult CloningScheduler::scheduleFramesForPacket(
 
     // The writableBytes here is an optimization. If the writableBytes is too
     // small for this packet. rebuildFromPacket should fail anyway.
-    // TODO: This isn't the ideal way to solve the wrong writableBytes problem.
     if (iter->encodedSize > writableBytes + cipherOverhead_) {
       continue;
     }
 
+    // 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
+    // packet. In that case, if the original packet is already at the packet
+    // size limit, we will generate a packet larger than the limit. We can
+    // either ignore the problem, hoping the packet will be able to travel the
+    // network just fine; Or we can throw away the built packet and send a ping.
+
     // Rebuilder will write the rest of frames
     auto rebuildResult = rebuilder.rebuildFromPacket(*iter);
     if (rebuildResult) {
diff --git a/quic/api/test/QuicPacketSchedulerTest.cpp b/quic/api/test/QuicPacketSchedulerTest.cpp
index ec4ad0dd0..20f27d587 100644
--- a/quic/api/test/QuicPacketSchedulerTest.cpp
+++ b/quic/api/test/QuicPacketSchedulerTest.cpp
@@ -772,6 +772,64 @@ TEST_F(QuicPacketSchedulerTest, CloningSchedulerWithInplaceBuilderFullPacket) {
   EXPECT_EQ(buf->length(), conn.udpSendPacketLen);
 }
 
+TEST_F(QuicPacketSchedulerTest, CloneLargerThanOriginalPacket) {
+  QuicClientConnectionState conn(
+      FizzClientQuicHandshakeContext::Builder().build());
+  conn.udpSendPacketLen = 1000;
+  conn.streamManager->setMaxLocalBidirectionalStreams(10);
+  conn.flowControlState.peerAdvertisedMaxOffset = 100000;
+  conn.flowControlState.peerAdvertisedInitialMaxStreamOffsetBidiRemote = 100000;
+  auto stream = conn.streamManager->createNextBidirectionalStream().value();
+  auto inputData = buildRandomInputData(conn.udpSendPacketLen * 10);
+  writeDataToQuicStream(*stream, inputData->clone(), false);
+  FrameScheduler scheduler = std::move(FrameScheduler::Builder(
+                                           conn,
+                                           EncryptionLevel::AppData,
+                                           PacketNumberSpace::AppData,
+                                           "streamScheduler")
+                                           .streamFrames())
+                                 .build();
+  auto cipherOverhead = 16;
+  PacketNum packetNum = 0;
+  ShortHeader header(
+      ProtectionType::KeyPhaseOne,
+      conn.clientConnectionId.value_or(getTestConnectionId()),
+      packetNum);
+  RegularQuicPacketBuilder builder(
+      conn.udpSendPacketLen,
+      std::move(header),
+      conn.ackStates.appDataAckState.largestAckedByPeer);
+  auto packetResult = scheduler.scheduleFramesForPacket(
+      std::move(builder), conn.udpSendPacketLen - cipherOverhead);
+  auto encodedSize = packetResult.packet->body->computeChainDataLength() +
+      packetResult.packet->header->computeChainDataLength() + cipherOverhead;
+  EXPECT_EQ(encodedSize, conn.udpSendPacketLen);
+  updateConnection(
+      conn,
+      folly::none,
+      packetResult.packet->packet,
+      Clock::now(),
+      encodedSize);
+
+  // make packetNum too larger to be encoded into the same size:
+  packetNum += 0xFF;
+  ShortHeader cloneHeader(
+      ProtectionType::KeyPhaseOne,
+      conn.clientConnectionId.value_or(getTestConnectionId()),
+      packetNum);
+  RegularQuicPacketBuilder throwawayBuilder(
+      conn.udpSendPacketLen,
+      std::move(cloneHeader),
+      conn.ackStates.appDataAckState.largestAckedByPeer);
+  FrameScheduler noopScheduler("noopScheduler");
+  CloningScheduler cloningScheduler(
+      noopScheduler, conn, "CopyCat", cipherOverhead);
+  auto cloneResult = cloningScheduler.scheduleFramesForPacket(
+      std::move(throwawayBuilder), kDefaultUDPSendPacketLen);
+  EXPECT_FALSE(cloneResult.packet.hasValue());
+  EXPECT_FALSE(cloneResult.packetEvent.hasValue());
+}
+
 class AckSchedulingTest : public TestWithParam<PacketNumberSpace> {};
 
 TEST_F(QuicPacketSchedulerTest, AckStateHasAcksToSchedule) {