diff --git a/libi2pd/SSU2Session.cpp b/libi2pd/SSU2Session.cpp index b3c539bc..0f53bf7b 100644 --- a/libi2pd/SSU2Session.cpp +++ b/libi2pd/SSU2Session.cpp @@ -35,48 +35,45 @@ namespace transport bool SSU2IncompleteMessage::ConcatOutOfSequenceFragments () { bool isLast = false; - if (nextFragmentNum == 1) - { - if (secondFragment) + while (outOfSequenceFragments) + { + if (outOfSequenceFragments->fragmentNum == nextFragmentNum) { - AttachNextFragment (secondFragment->buf, secondFragment->len); - isLast = secondFragment->isLast; - secondFragment = nullptr; + AttachNextFragment (outOfSequenceFragments->buf, outOfSequenceFragments->len); + isLast = outOfSequenceFragments->isLast; + if (isLast) + outOfSequenceFragments = nullptr; + else + outOfSequenceFragments = outOfSequenceFragments->next; } else - return false; - if (isLast) return true; - } - // might be more - if (outOfSequenceFragments) - { - for (auto it = outOfSequenceFragments->begin (); it != outOfSequenceFragments->end ();) - if (it->first == nextFragmentNum) - { - AttachNextFragment (it->second->buf, it->second->len); - isLast = it->second->isLast; - it = outOfSequenceFragments->erase (it); - } - else - break; + break; } return isLast; } - void SSU2IncompleteMessage::AddOutOfSequenceFragment (int fragmentNum, - std::shared_ptr fragment) + void SSU2IncompleteMessage::AddOutOfSequenceFragment (std::shared_ptr fragment) { - if (!fragmentNum || !fragment) return; // fragment 0 nun allowed - if (fragmentNum < nextFragmentNum) return; // already processed - if (fragmentNum == 1) - { - if (!secondFragment) secondFragment = fragment; - } + if (!fragment || !fragment->fragmentNum) return; // fragment 0 not allowed + if (fragment->fragmentNum < nextFragmentNum) return; // already processed + if (!outOfSequenceFragments) + outOfSequenceFragments = fragment; else { - if (!outOfSequenceFragments) - outOfSequenceFragments.reset (new std::map >()); - outOfSequenceFragments->emplace (fragmentNum, fragment); + auto frag = outOfSequenceFragments; + std::shared_ptr prev; + do + { + if (fragment->fragmentNum < frag->fragmentNum) break; // found + if (fragment->fragmentNum == frag->fragmentNum) return; // duplicate + prev = frag; frag = frag->next; + } + while (frag); + fragment->next = frag; + if (prev) + prev->next = fragment; + else + outOfSequenceFragments = fragment; } lastFragmentInsertTime = i2p::util::GetSecondsSinceEpoch (); } @@ -1839,8 +1836,9 @@ namespace transport auto fragment = m_Server.GetFragmentsPool ().AcquireShared (); memcpy (fragment->buf, buf + 5, len -5); fragment->len = len - 5; + fragment->fragmentNum = fragmentNum; fragment->isLast = isLast; - it->second->AddOutOfSequenceFragment (fragmentNum, fragment); + it->second->AddOutOfSequenceFragment (fragment); } void SSU2Session::HandleRelayRequest (const uint8_t * buf, size_t len) diff --git a/libi2pd/SSU2Session.h b/libi2pd/SSU2Session.h index 97469ed4..08a4b7cf 100644 --- a/libi2pd/SSU2Session.h +++ b/libi2pd/SSU2Session.h @@ -171,18 +171,19 @@ namespace transport { uint8_t buf[SSU2_MAX_PACKET_SIZE]; size_t len; + int fragmentNum; bool isLast; + std::shared_ptr next; }; std::shared_ptr msg; int nextFragmentNum; uint32_t lastFragmentInsertTime; // in seconds - std::shared_ptr secondFragment; // fragment #1 - std::unique_ptr > > outOfSequenceFragments; // fragments #2 and more + std::shared_ptr outOfSequenceFragments; // #1 and more void AttachNextFragment (const uint8_t * fragment, size_t fragmentSize); bool ConcatOutOfSequenceFragments (); // true if message complete - void AddOutOfSequenceFragment (int fragmentNum, std::shared_ptr fragment); + void AddOutOfSequenceFragment (std::shared_ptr fragment); }; struct SSU2SentPacket