diff --git a/quic/codec/Decode.cpp b/quic/codec/Decode.cpp index d7ea878e2..06c5fd40b 100644 --- a/quic/codec/Decode.cpp +++ b/quic/codec/Decode.cpp @@ -840,22 +840,20 @@ QuicFrame parseFrame( throw QuicTransportException( "Invalid frame-type field", TransportErrorCode::FRAME_ENCODING_ERROR); } - queue.trimStart(frameTypeInt->second); + queue.trimStart(cursor - queue.front()); bool consumedQueue = false; bool error = false; SCOPE_EXIT { if (consumedQueue || error) { return; } - if (!queue.empty()) { - queue.trimStart(cursor - queue.front()); - } + queue.trimStart(cursor - queue.front()); }; - // trimStart() may have free'd the IOBuf; - folly::IOBuf emptyBuf{}; - cursor.reset(queue.empty() ? &emptyBuf : queue.front()); + cursor.reset(queue.front()); FrameType frameType = static_cast(frameTypeInt->first); - try { + try + + { switch (frameType) { case FrameType::PADDING: return QuicFrame(decodePaddingFrame(cursor)); diff --git a/quic/common/BufUtil.cpp b/quic/common/BufUtil.cpp index 4289b219e..602b97e02 100644 --- a/quic/common/BufUtil.cpp +++ b/quic/common/BufUtil.cpp @@ -59,9 +59,45 @@ Buf BufQueue::splitAtMost(size_t len) { } size_t BufQueue::trimStartAtMost(size_t amount) { - auto oldChainLength = chainLength_; - splitAtMost(amount); - return oldChainLength - chainLength_; + auto original = amount; + folly::IOBuf* current = chain_.get(); + if (current == nullptr || amount == 0) { + return 0; + } + while (amount > 0) { + if (current->length() >= amount) { + current->trimStart(amount); + amount = 0; + break; + } + amount -= current->length(); + current = current->next(); + if (current == chain_.get()) { + break; + } + } + auto prev = current->prev(); + /** We are potentially in 2 states here, + * 1. we found the entire amount + * 2. or we did not. + * If we did not find the entire amount, then current == + * chain and we can remove the entire chain. + * If we did, then we can split from the chain head to the previous buffer and + * then keep the current buffer. + */ + if (prev != current && current != chain_.get()) { + auto chain = chain_.release(); + current->separateChain(chain, prev); + chain_ = std::unique_ptr(current); + } else if (amount > 0) { + DCHECK_EQ(current, chain_.get()); + chain_ = nullptr; + } + size_t trimmed = original - amount; + DCHECK_GE(chainLength_, trimmed); + chainLength_ -= trimmed; + DCHECK(chainLength_ == 0 || !chain_->empty()); + return trimmed; } // TODO replace users with trimStartAtMost diff --git a/quic/common/test/BufUtilTest.cpp b/quic/common/test/BufUtilTest.cpp index 51dbfc508..fde4c5a60 100644 --- a/quic/common/test/BufUtilTest.cpp +++ b/quic/common/test/BufUtilTest.cpp @@ -187,15 +187,6 @@ TEST(BufQueue, TrimStartAtMost) { checkConsistency(queue); } -TEST(BufQueue, TrimStartOneByte) { - BufQueue queue; - queue.append(IOBuf::copyBuffer(SCL("H"))); - checkConsistency(queue); - queue.trimStart(1); - checkConsistency(queue); - EXPECT_TRUE(queue.front() == nullptr); -} - TEST(BufQueue, CloneBufNull) { BufQueue queue; auto buf = queue.clone();