diff --git a/quic/api/QuicPacketScheduler.cpp b/quic/api/QuicPacketScheduler.cpp index eff4ff1fa..b1cc5f762 100644 --- a/quic/api/QuicPacketScheduler.cpp +++ b/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; } diff --git a/quic/api/test/QuicPacketSchedulerTest.cpp b/quic/api/test/QuicPacketSchedulerTest.cpp index 508432807..6515ffb79 100644 --- a/quic/api/test/QuicPacketSchedulerTest.cpp +++ b/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()); } diff --git a/quic/codec/QuicPacketBuilder.cpp b/quic/codec/QuicPacketBuilder.cpp index 14826c940..a7675f936 100644 --- a/quic/codec/QuicPacketBuilder.cpp +++ b/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 diff --git a/quic/codec/QuicPacketBuilder.h b/quic/codec/QuicPacketBuilder.h index 72ed0f6a0..5a4c2be9b 100644 --- a/quic/codec/QuicPacketBuilder.h +++ b/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; diff --git a/quic/codec/test/Mocks.h b/quic/codec/test/Mocks.h index 8ce24dfc0..ec0a3ff63 100644 --- a/quic/codec/test/Mocks.h +++ b/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(*this).canBuildPacketNonConst();