diff --git a/SSU.cpp b/SSU.cpp index d033ec36..69b105f8 100644 --- a/SSU.cpp +++ b/SSU.cpp @@ -27,6 +27,12 @@ namespace ssu SSUSession::~SSUSession () { delete m_DHKeysPair; + for (auto it: m_IncomleteMessages) + if (it.second) + { + DeleteI2NPMessage (it.second->msg); + delete it.second; + } } void SSUSession::CreateAESandMacKey (uint8_t * pubKey, uint8_t * aesKey, uint8_t * macKey) @@ -673,9 +679,29 @@ namespace ssu auto it = m_IncomleteMessages.find (msgID); if (it != m_IncomleteMessages.end ()) { - msg = it->second; - memcpy (msg->buf + msg->len, buf, fragmentSize); - msg->len += fragmentSize; + if (fragmentNum == it->second->nextFragmentNum) + { + // expected fragment + msg = it->second->msg; + memcpy (msg->buf + msg->len, buf, fragmentSize); + msg->len += fragmentSize; + it->second->nextFragmentNum++; + } + else if (fragmentNum < it->second->nextFragmentNum) + // duplicate fragment + LogPrint ("Duplicate fragment ", fragmentNum, " of message ", msgID, ". Ignored"); + else + { + // missing fragment + LogPrint ("Missing fragments from ", it->second->nextFragmentNum, " to ", fragmentNum - 1, " of message ", msgID); + //TODO + } + + if (isLast) + { + delete it->second; + m_IncomleteMessages.erase (it); + } } else // TODO: @@ -691,12 +717,10 @@ namespace ssu if (msg) { if (!fragmentNum && !isLast) - m_IncomleteMessages[msgID] = msg; + m_IncomleteMessages[msgID] = new IncompleteMessage (msg); if (isLast) { SendMsgAck (msgID); - if (fragmentNum > 0) - m_IncomleteMessages.erase (msgID); msg->FromSSU (msgID); if (m_State == eSessionStateEstablished) i2p::HandleI2NPMessage (msg); diff --git a/SSU.h b/SSU.h index 3f7ac96b..8ad3fecf 100644 --- a/SSU.h +++ b/SSU.h @@ -124,6 +124,14 @@ namespace ssu void HandleTerminationTimer (const boost::system::error_code& ecode); private: + + struct IncompleteMessage + { + I2NPMessage * msg; + uint8_t nextFragmentNum; + + IncompleteMessage (I2NPMessage * m): msg (m), nextFragmentNum (1) {}; + }; SSUServer& m_Server; boost::asio::ip::udp::endpoint m_RemoteEndpoint; @@ -137,7 +145,7 @@ namespace ssu CryptoPP::CBC_Mode::Encryption m_Encryption; CryptoPP::CBC_Mode::Decryption m_Decryption; uint8_t m_SessionKey[32], m_MacKey[32]; - std::map m_IncomleteMessages; + std::map m_IncomleteMessages; std::list m_DelayedMessages; };