Add back SizeEnforcedPacketBuilder
Summary: ^
MODIFY
quic/codec/QuicPacketBuilder.cpp
quic/codec/QuicPacketBuilder.cpp
@@ -286,6 +286,94 @@ Buf StatelessResetPacketBuilder::buildPacket() && {
return std::move(data_);
}
+RegularSizeEnforcedPacketBuilder::RegularSizeEnforcedPacketBuilder(
+ Packet packet,
+ uint64_t enforcedSize,
+ uint32_t cipherOverhead)
+ : packet_(std::move(packet.packet)),
+ header_(std::move(packet.header)),
+ body_(std::move(packet.body)),
+ bodyAppender_(body_.get(), kAppenderGrowthSize),
+ enforcedSize_(enforcedSize),
+ cipherOverhead_(cipherOverhead) {}
+
+bool RegularSizeEnforcedPacketBuilder::canBuildPacket() const noexcept {
+ // We only force size of packets with short header, because d6d probes always
+ // have short headers and there's no other situations for this type of builder
+ const ShortHeader* shortHeader = packet_.header.asShort();
+ // We also don't want to send packets longer than kDefaultMaxUDPPayload
+ return shortHeader && enforcedSize_ <= kDefaultMaxUDPPayload &&
+ (body_->computeChainDataLength() + header_->computeChainDataLength() +
+ cipherOverhead_ <
+ enforcedSize_);
+}
+
+PacketBuilderInterface::Packet
+RegularSizeEnforcedPacketBuilder::buildPacket() && {
+ // Store counters on the stack to overhead from function calls
+ size_t extraDataWritten = 0;
+ size_t bodyLength = body_->computeChainDataLength();
+ size_t headerLength = header_->computeChainDataLength();
+ while (extraDataWritten + bodyLength + headerLength + cipherOverhead_ <
+ enforcedSize_) {
+ QuicInteger paddingType(static_cast<uint8_t>(FrameType::PADDING));
+ paddingType.encode([&](auto val) { bodyAppender_.writeBE(val); });
+ extraDataWritten++;
+ }
+ return Packet(std::move(packet_), std::move(header_), std::move(body_));
+}
+
+InplaceSizeEnforcedPacketBuilder::InplaceSizeEnforcedPacketBuilder(
+ BufAccessor& bufAccessor,
+ Packet packet,
+ uint64_t enforcedSize,
+ uint32_t cipherOverhead)
+ : bufAccessor_(bufAccessor),
+ iobuf_(bufAccessor_.obtain()),
+ packet_(std::move(packet.packet)),
+ header_(std::move(packet.header)),
+ body_(std::move(packet.body)),
+ enforcedSize_(enforcedSize),
+ cipherOverhead_(cipherOverhead) {}
+
+bool InplaceSizeEnforcedPacketBuilder::canBuildPacket() const noexcept {
+ const ShortHeader* shortHeader = packet_.header.asShort();
+ size_t encryptedPacketSize =
+ header_->length() + body_->length() + cipherOverhead_;
+ size_t delta = enforcedSize_ - encryptedPacketSize;
+ return shortHeader && enforcedSize_ <= kDefaultMaxUDPPayload &&
+ encryptedPacketSize < enforcedSize_ && iobuf_->tailroom() >= delta;
+}
+
+PacketBuilderInterface::Packet
+InplaceSizeEnforcedPacketBuilder::buildPacket() && {
+ // Create bodyWriter
+ size_t encryptedPacketSize =
+ header_->length() + body_->length() + cipherOverhead_;
+ size_t paddingSize = enforcedSize_ - encryptedPacketSize;
+ BufWriter bodyWriter(*iobuf_, paddingSize);
+
+ // Store counters on the stack to overhead from function calls
+ size_t extraDataWritten = 0;
+ size_t bodyLength = body_->computeChainDataLength();
+ size_t headerLength = header_->computeChainDataLength();
+ while (extraDataWritten + bodyLength + headerLength + cipherOverhead_ <
+ enforcedSize_) {
+ QuicInteger paddingType(static_cast<uint8_t>(FrameType::PADDING));
+ paddingType.encode([&](auto val) { bodyWriter.writeBE(val); });
+ extraDataWritten++;
+ }
+
+ PacketBuilderInterface::Packet builtPacket(
+ std::move(packet_),
+ std::move(header_),
+ folly::IOBuf::wrapBuffer(body_->data(), iobuf_->tail() - body_->data()));
+
+ // Release internal iobuf
+ bufAccessor_.release(std::move(iobuf_));
+ return builtPacket;
+}
+
VersionNegotiationPacketBuilder::VersionNegotiationPacketBuilder(
ConnectionId sourceConnectionId,
ConnectionId destinationConnectionId,

MODIFY
quic/codec/QuicPacketBuilder.h
quic/codec/QuicPacketBuilder.h
@@ -255,6 +255,88 @@ class RegularQuicPacketBuilder final : public PacketBuilderInterface {
folly::Optional<PacketNumEncodingResult> packetNumberEncoding_;
};
+/**
+ * A less involving interface for packet builder, this enables polymorphism
+ * for wrapper-like packet builders (e.g. RegularSizeEnforcedPacketBuilder).
+ */
+class WrapperPacketBuilderInterface {
+ public:
+ using Packet = PacketBuilderInterface::Packet;
+
+ virtual ~WrapperPacketBuilderInterface() = default;
+
+ FOLLY_NODISCARD virtual bool canBuildPacket() const noexcept = 0;
+
+ virtual Packet buildPacket() && = 0;
+};
+
+/**
+ * This builder will enforce the packet size by appending padding frames for
+ * chained memory. This means appending IOBuf at the end of the chain. The
+ * caller should ensure canBuildPacket() returns true before constructing the
+ * builder.
+ */
+class RegularSizeEnforcedPacketBuilder : public WrapperPacketBuilderInterface {
+ public:
+ using Packet = PacketBuilderInterface::Packet;
+
+ explicit RegularSizeEnforcedPacketBuilder(
+ Packet packet,
+ uint64_t enforcedSize,
+ uint32_t cipherOverhead);
+
+ /**
+ * Returns true when packet has short header, and that enforced size >
+ * current packet size + cipher overhead, otherwise false
+ */
+ FOLLY_NODISCARD bool canBuildPacket() const noexcept override;
+
+ Packet buildPacket() && override;
+
+ private:
+ RegularQuicWritePacket packet_;
+ Buf header_;
+ Buf body_;
+ BufAppender bodyAppender_;
+ uint64_t enforcedSize_;
+ uint32_t cipherOverhead_;
+};
+
+/**
+ * This builder will enforce the packet size by appending padding frames for
+ * continuous memory. This means pushing padding frame directly to the current
+ * tail offset. The caller should ensure canBuildPacket() returns true before
+ * constructing the builder.
+ */
+class InplaceSizeEnforcedPacketBuilder : public WrapperPacketBuilderInterface {
+ public:
+ using Packet = PacketBuilderInterface::Packet;
+
+ explicit InplaceSizeEnforcedPacketBuilder(
+ BufAccessor& bufAccessor,
+ Packet packet,
+ uint64_t enforcedSize,
+ uint32_t cipherOverhead);
+
+ /**
+ * Returns true when packet has short header, and that enforced size> current
+ * packet size + cipher oveahead and that iobuf has enough tailroom,
+ * otherwise false
+ */
+ FOLLY_NODISCARD bool canBuildPacket() const noexcept override;
+
+ Packet buildPacket() && override;
+
+ private:
+ BufAccessor& bufAccessor_;
+ Buf iobuf_;
+ RegularQuicWritePacket packet_;
+ Buf header_;
+ Buf body_;
+ uint64_t enforcedSize_;
+ uint32_t cipherOverhead_;
+};
+
class VersionNegotiationPacketBuilder {
public:
explicit VersionNegotiationPacketBuilder(

MODIFY
quic/codec/test/QuicPacketBuilderTest.cpp
quic/codec/test/QuicPacketBuilderTest.cpp
@@ -109,7 +109,7 @@ class QuicPacketBuilderTest : public TestWithParam<TestFlavor> {
folly::assume_unreachable();
}
- private:
+ protected:
std::unique_ptr<BufAccessor> simpleBufAccessor_;
};
@@ -277,6 +277,67 @@ TEST_P(QuicPacketBuilderTest, ShortHeaderRegularPacket) {
EXPECT_EQ(pktNum, decodedHeader.getPacketSequenceNum());
}
+TEST_P(QuicPacketBuilderTest, EnforcePacketSizeWithCipherOverhead) {
+ auto connId = getTestConnectionId();
+ PacketNum pktNum = 222;
+
+ PacketNum largestAckedPacketNum = 0;
+ size_t cipherOverhead = 2;
+ uint64_t enforcedSize = 1400;
+ auto aead = std::make_unique<NiceMock<MockAead>>();
+ auto aead_ = aead.get();
+ EXPECT_CALL(*aead_, _inplaceEncrypt(_, _, _))
+ .WillRepeatedly(Invoke([&](auto& buf, auto, auto) {
+ auto overhead = folly::IOBuf::create(1000);
+ overhead->append(cipherOverhead);
+ auto clone = buf->clone();
+ clone->prependChain(std::move(overhead));
+ return std::move(clone);
+ }));
+
+ auto builder = testBuilderProvider(
+ GetParam(),
+ kDefaultUDPSendPacketLen,
+ ShortHeader(ProtectionType::KeyPhaseZero, connId, pktNum),
+ largestAckedPacketNum,
+ 2000);
+ builder->accountForCipherOverhead(cipherOverhead);
+ builder->encodePacketHeader();
+
+ // write out at least one frame
+ writeFrame(PaddingFrame(), *builder);
+ EXPECT_TRUE(builder->canBuildPacket());
+ auto builtOut = std::move(*builder).buildPacket();
+
+ auto param = GetParam();
+ if (param == TestFlavor::Regular) {
+ EXPECT_EQ(builtOut.body->isManagedOne(), true);
+ RegularSizeEnforcedPacketBuilder sizeEnforcedBuilder(
+ std::move(builtOut), enforcedSize, cipherOverhead);
+ EXPECT_TRUE(sizeEnforcedBuilder.canBuildPacket());
+ auto out = std::move(sizeEnforcedBuilder).buildPacket();
+ EXPECT_EQ(
+ out.header->computeChainDataLength() +
+ out.body->computeChainDataLength(),
+ enforcedSize - cipherOverhead);
+ auto buf = packetToBuf(out, aead_);
+ EXPECT_EQ(buf->computeChainDataLength(), enforcedSize);
+
+ } else {
+ EXPECT_EQ(builtOut.body->isManagedOne(), false);
+ InplaceSizeEnforcedPacketBuilder sizeEnforcedBuilder(
+ *simpleBufAccessor_, std::move(builtOut), enforcedSize, cipherOverhead);
+ EXPECT_TRUE(sizeEnforcedBuilder.canBuildPacket());
+ auto out = std::move(sizeEnforcedBuilder).buildPacket();
+ EXPECT_EQ(
+ out.header->computeChainDataLength() +
+ out.body->computeChainDataLength(),
+ enforcedSize - cipherOverhead);
+ auto buf = packetToBuf(out, aead_);
+ EXPECT_EQ(buf->computeChainDataLength(), enforcedSize);
+ }
+}
+
TEST_P(QuicPacketBuilderTest, ShortHeaderWithNoFrames) {
auto connId = getTestConnectionId();
PacketNum pktNum = 222;

Add back SizeEnforcedPacketBuilder

Summary: ^

Reviewed By: yangchi

Differential Revision: D23432975

fbshipit-source-id: 65892a1bb5146dd12b5f3f89ec689f9d4326d035