diff --git a/libi2pd/Datagram.cpp b/libi2pd/Datagram.cpp index eb83efe7..b69e173a 100644 --- a/libi2pd/Datagram.cpp +++ b/libi2pd/Datagram.cpp @@ -21,7 +21,7 @@ namespace datagram DatagramDestination::DatagramDestination (std::shared_ptr owner, bool gzip, DatagramVersion version): m_Owner (owner), m_DefaultReceiver (nullptr), m_DefaultRawReceiver (nullptr), - m_Gzip (gzip), m_Version (version) + m_DefaultReceiverPort (0), m_DefaultRawReceiverPort (0), m_Gzip (gzip), m_Version (version) { if (m_Gzip) m_Deflator.reset (new i2p::data::GzipDeflator); @@ -353,8 +353,9 @@ namespace datagram void DatagramDestination::SetReceiver (const Receiver& receiver, uint16_t port) { std::lock_guard lock(m_ReceiversMutex); - m_ReceiversByPorts[port] = receiver; - if (!m_DefaultReceiver) { + if (port) m_ReceiversByPorts[port] = receiver; + if (!m_DefaultReceiver) + { m_DefaultReceiver = receiver; m_DefaultReceiverPort = port; } @@ -363,8 +364,9 @@ namespace datagram void DatagramDestination::ResetReceiver (uint16_t port) { std::lock_guard lock(m_ReceiversMutex); - m_ReceiversByPorts.erase (port); - if (m_DefaultReceiverPort == port) { + if (port) m_ReceiversByPorts.erase (port); + if (!port || m_DefaultReceiverPort == port) + { m_DefaultReceiver = nullptr; m_DefaultReceiverPort = 0; } @@ -374,8 +376,9 @@ namespace datagram void DatagramDestination::SetRawReceiver (const RawReceiver& receiver, uint16_t port) { std::lock_guard lock(m_RawReceiversMutex); - m_RawReceiversByPorts[port] = receiver; - if (!m_DefaultRawReceiver) { + if (port) m_RawReceiversByPorts[port] = receiver; + if (!m_DefaultRawReceiver) + { m_DefaultRawReceiver = receiver; m_DefaultRawReceiverPort = port; } @@ -384,8 +387,9 @@ namespace datagram void DatagramDestination::ResetRawReceiver (uint16_t port) { std::lock_guard lock(m_RawReceiversMutex); - m_RawReceiversByPorts.erase (port); - if (m_DefaultRawReceiverPort == port) { + if (port) m_RawReceiversByPorts.erase (port); + if (!port || m_DefaultRawReceiverPort == port) + { m_DefaultRawReceiver = nullptr; m_DefaultRawReceiverPort = 0; } diff --git a/libi2pd/Datagram.h b/libi2pd/Datagram.h index 0a35c1d3..1f7e9ba6 100644 --- a/libi2pd/Datagram.h +++ b/libi2pd/Datagram.h @@ -150,11 +150,11 @@ namespace datagram void HandleDataMessagePayload (uint16_t fromPort, uint16_t toPort, const uint8_t * buf, size_t len, uint8_t protocolType, i2p::garlic::ECIESX25519AEADRatchetSession * from); - void SetReceiver (const Receiver& receiver, uint16_t port); - void ResetReceiver (uint16_t port); + void SetReceiver (const Receiver& receiver, uint16_t port = 0); + void ResetReceiver (uint16_t port = 0); - void SetRawReceiver (const RawReceiver& receiver, uint16_t port); - void ResetRawReceiver (uint16_t port); + void SetRawReceiver (const RawReceiver& receiver, uint16_t port = 0); + void ResetRawReceiver (uint16_t port = 0); std::shared_ptr GetInfoForRemote(const i2p::data::IdentHash & remote); diff --git a/libi2pd_client/UDPTunnel.cpp b/libi2pd_client/UDPTunnel.cpp index 3533dcae..c5fb3cf0 100644 --- a/libi2pd_client/UDPTunnel.cpp +++ b/libi2pd_client/UDPTunnel.cpp @@ -426,13 +426,10 @@ namespace client dgram->SetReceiver (std::bind (&I2PUDPClientTunnel::HandleRecvFromI2P, this, std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4, - std::placeholders::_5, std::placeholders::_6), - RemotePort - ); + std::placeholders::_5, std::placeholders::_6)); + dgram->SetRawReceiver (std::bind (&I2PUDPClientTunnel::HandleRecvFromI2PRaw, this, - std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4), - RemotePort - ); + std::placeholders::_1, std::placeholders::_2, std::placeholders::_3, std::placeholders::_4)); m_LocalDest->Start (); if (m_ResolveThread == nullptr) @@ -443,9 +440,10 @@ namespace client void I2PUDPClientTunnel::Stop () { auto dgram = m_LocalDest->GetDatagramDestination (); - if (dgram) { - dgram->ResetReceiver (RemotePort); - dgram->ResetRawReceiver (RemotePort); + if (dgram) + { + dgram->ResetReceiver (); + dgram->ResetRawReceiver (); } m_cancel_resolve = true;