i2pd/libi2pd/KadDHT.cpp

373 lines
7.2 KiB
C++
Raw Permalink Normal View History

2023-02-17 06:14:02 +03:00
/*
* Copyright (c) 2023, The PurpleI2P Project
*
* This file is part of Purple i2pd project and licensed under BSD3
*
* See full license text in LICENSE file at top of project tree
*
*/
#include "KadDHT.h"
namespace i2p
{
namespace data
{
DHTNode::DHTNode ():
2023-02-22 03:08:12 +03:00
zero (nullptr), one (nullptr)
2023-02-17 06:14:02 +03:00
{
}
DHTNode::~DHTNode ()
{
if (zero) delete zero;
if (one) delete one;
}
2023-02-22 03:08:12 +03:00
void DHTNode::MoveRouterUp (bool fromOne)
2023-02-17 06:14:02 +03:00
{
DHTNode *& side = fromOne ? one : zero;
if (side)
{
2023-02-22 03:08:12 +03:00
if (router) router = nullptr; // shouldn't happen
router = side->router;
side->router = nullptr;
2023-02-17 06:14:02 +03:00
delete side;
side = nullptr;
}
}
DHTTable::DHTTable ():
m_Size (0)
{
m_Root = new DHTNode;
}
DHTTable::~DHTTable ()
{
delete m_Root;
}
2023-02-22 03:08:12 +03:00
void DHTTable::Clear ()
2023-02-17 06:14:02 +03:00
{
2023-02-22 03:08:12 +03:00
m_Size = 0;
delete m_Root;
m_Root = new DHTNode;
}
void DHTTable::Insert (const std::shared_ptr<RouterInfo>& r)
{
if (!r) return;
return Insert (r, m_Root, 0);
2023-02-17 06:14:02 +03:00
}
2023-02-22 03:08:12 +03:00
void DHTTable::Insert (const std::shared_ptr<RouterInfo>& r, DHTNode * root, int level)
2023-02-17 06:14:02 +03:00
{
2023-02-22 03:08:12 +03:00
if (root->router)
2023-02-17 06:14:02 +03:00
{
2023-02-22 03:08:12 +03:00
if (root->router->GetIdentHash () == r->GetIdentHash ())
2023-02-17 06:14:02 +03:00
{
2023-02-22 03:08:12 +03:00
root->router = r; // replace
return;
2023-02-17 06:14:02 +03:00
}
2023-02-22 03:08:12 +03:00
auto r2 = root->router;
root->router = nullptr; m_Size--;
2023-02-17 06:14:02 +03:00
int bit1, bit2;
do
{
2023-02-22 03:08:12 +03:00
bit1 = r->GetIdentHash ().GetBit (level);
bit2 = r2->GetIdentHash ().GetBit (level);
2023-02-17 06:14:02 +03:00
if (bit1 == bit2)
{
if (bit1)
{
2023-03-24 23:15:47 +03:00
if (root->one) return; // something wrong
2023-02-17 06:14:02 +03:00
root->one = new DHTNode;
root = root->one;
}
else
{
2023-03-24 23:15:47 +03:00
if (root->zero) return; // something wrong
2023-02-17 06:14:02 +03:00
root->zero = new DHTNode;
root = root->zero;
}
level++;
}
}
while (bit1 == bit2);
if (!root->zero)
root->zero = new DHTNode;
if (!root->one)
root->one = new DHTNode;
if (bit1)
{
2023-02-22 03:08:12 +03:00
Insert (r2, root->zero, level + 1);
Insert (r, root->one, level + 1);
2023-02-17 06:14:02 +03:00
}
else
{
2023-02-22 03:08:12 +03:00
Insert (r2, root->one, level + 1);
Insert (r, root->zero, level + 1);
2023-02-17 06:14:02 +03:00
}
}
else
{
if (!root->zero && !root->one)
{
2023-02-22 03:08:12 +03:00
root->router = r; m_Size++;
return;
2023-02-17 06:14:02 +03:00
}
2023-02-22 03:08:12 +03:00
int bit = r->GetIdentHash ().GetBit (level);
2023-02-17 06:14:02 +03:00
if (bit)
{
if (!root->one)
root->one = new DHTNode;
2023-02-22 03:08:12 +03:00
Insert (r, root->one, level + 1);
2023-02-17 06:14:02 +03:00
}
else
{
if (!root->zero)
root->zero = new DHTNode;
2023-02-22 03:08:12 +03:00
Insert (r, root->zero, level + 1);
2023-02-17 06:14:02 +03:00
}
}
}
bool DHTTable::Remove (const IdentHash& h)
{
return Remove (h, m_Root, 0);
}
bool DHTTable::Remove (const IdentHash& h, DHTNode * root, int level)
{
if (root)
{
2023-02-22 03:08:12 +03:00
if (root->router && root->router->GetIdentHash () == h)
2023-02-17 06:14:02 +03:00
{
2023-02-22 03:08:12 +03:00
root->router = nullptr;
2023-02-17 06:14:02 +03:00
m_Size--;
return true;
}
int bit = h.GetBit (level);
if (bit)
{
if (root->one && Remove (h, root->one, level + 1))
{
if (root->one->IsEmpty ())
{
delete root->one;
root->one = nullptr;
2023-02-22 03:08:12 +03:00
if (root->zero && root->zero->router)
root->MoveRouterUp (false);
2023-02-17 06:14:02 +03:00
}
2023-02-22 03:08:12 +03:00
else if (root->one->router && !root->zero)
root->MoveRouterUp (true);
2023-02-17 06:14:02 +03:00
return true;
}
}
else
{
if (root->zero && Remove (h, root->zero, level + 1))
{
if (root->zero->IsEmpty ())
{
delete root->zero;
root->zero = nullptr;
2023-02-22 03:08:12 +03:00
if (root->one && root->one->router)
root->MoveRouterUp (true);
2023-02-17 06:14:02 +03:00
}
2023-02-22 03:08:12 +03:00
else if (root->zero->router && !root->one)
root->MoveRouterUp (false);
2023-02-17 06:14:02 +03:00
return true;
}
}
}
return false;
}
2023-02-22 23:58:20 +03:00
std::shared_ptr<RouterInfo> DHTTable::FindClosest (const IdentHash& h, const Filter& filter) const
2023-02-17 06:14:02 +03:00
{
2023-02-22 03:08:12 +03:00
if (filter) m_Filter = filter;
auto r = FindClosest (h, m_Root, 0);
m_Filter = nullptr;
return r;
2023-02-17 06:14:02 +03:00
}
2023-02-22 23:58:20 +03:00
std::shared_ptr<RouterInfo> DHTTable::FindClosest (const IdentHash& h, DHTNode * root, int level) const
2023-02-17 06:14:02 +03:00
{
2023-02-22 03:08:12 +03:00
bool split = false;
do
{
if (root->router)
return (!m_Filter || m_Filter (root->router)) ? root->router : nullptr;
split = root->zero && root->one;
if (!split)
{
if (root->zero) root = root->zero;
else if (root->one) root = root->one;
else return nullptr;
level++;
}
}
while (!split);
2023-02-17 06:14:02 +03:00
int bit = h.GetBit (level);
if (bit)
{
if (root->one)
2023-02-22 03:08:12 +03:00
{
auto r = FindClosest (h, root->one, level + 1);
if (r) return r;
}
2023-02-17 06:14:02 +03:00
if (root->zero)
2023-02-22 03:08:12 +03:00
{
auto r = FindClosest (h, root->zero, level + 1);
if (r) return r;
}
2023-02-17 06:14:02 +03:00
}
else
{
if (root->zero)
2023-02-22 03:08:12 +03:00
{
auto r = FindClosest (h, root->zero, level + 1);
if (r) return r;
}
2023-02-17 06:14:02 +03:00
if (root->one)
2023-02-22 03:08:12 +03:00
{
auto r = FindClosest (h, root->one, level + 1);
if (r) return r;
}
2023-02-17 06:14:02 +03:00
}
return nullptr;
}
2023-02-19 03:45:31 +03:00
2023-02-22 23:58:20 +03:00
std::vector<std::shared_ptr<RouterInfo> > DHTTable::FindClosest (const IdentHash& h, size_t num, const Filter& filter) const
2023-02-19 03:45:31 +03:00
{
2023-02-22 03:08:12 +03:00
std::vector<std::shared_ptr<RouterInfo> > vec;
2023-02-19 03:45:31 +03:00
if (num > 0)
2023-02-22 03:08:12 +03:00
{
if (filter) m_Filter = filter;
2023-02-19 03:45:31 +03:00
FindClosest (h, num, m_Root, 0, vec);
2023-02-22 03:08:12 +03:00
m_Filter = nullptr;
}
2023-02-19 03:45:31 +03:00
return vec;
}
2023-02-22 23:58:20 +03:00
void DHTTable::FindClosest (const IdentHash& h, size_t num, DHTNode * root, int level, std::vector<std::shared_ptr<RouterInfo> >& hashes) const
2023-02-19 03:45:31 +03:00
{
if (hashes.size () >= num) return;
2023-02-22 03:08:12 +03:00
bool split = false;
do
2023-02-19 03:45:31 +03:00
{
2023-02-22 03:08:12 +03:00
if (root->router)
{
if (!m_Filter || m_Filter (root->router))
hashes.push_back (root->router);
return;
}
split = root->zero && root->one;
if (!split)
{
if (root->zero) root = root->zero;
else if (root->one) root = root->one;
else return;
level++;
}
}
while (!split);
2023-02-19 03:45:31 +03:00
int bit = h.GetBit (level);
if (bit)
{
if (root->one)
FindClosest (h, num, root->one, level + 1, hashes);
if (hashes.size () < num && root->zero)
FindClosest (h, num, root->zero, level + 1, hashes);
}
else
{
if (root->zero)
FindClosest (h, num, root->zero, level + 1, hashes);
if (hashes.size () < num && root->one)
FindClosest (h, num, root->one, level + 1, hashes);
}
}
2023-02-22 03:08:12 +03:00
2023-02-22 23:58:20 +03:00
void DHTTable::Cleanup (const Filter& filter)
2023-02-22 03:08:12 +03:00
{
if (filter)
{
m_Filter = filter;
Cleanup (m_Root);
m_Filter = nullptr;
}
else
Clear ();
}
void DHTTable::Cleanup (DHTNode * root)
{
if (!root) return;
if (root->router)
{
if (!m_Filter || !m_Filter (root->router))
{
m_Size--;
root->router = nullptr;
}
return;
}
if (root->zero)
{
Cleanup (root->zero);
if (root->zero->IsEmpty ())
{
delete root->zero;
root->zero = nullptr;
}
}
if (root->one)
{
Cleanup (root->one);
if (root->one->IsEmpty ())
{
delete root->one;
root->one = nullptr;
if (root->zero && root->zero->router)
root->MoveRouterUp (false);
}
else if (root->one->router && !root->zero)
root->MoveRouterUp (true);
}
}
2023-02-17 06:14:02 +03:00
void DHTTable::Print (std::stringstream& s)
{
Print (s, m_Root, 0);
}
void DHTTable::Print (std::stringstream& s, DHTNode * root, int level)
{
if (!root) return;
s << std::string (level, '-');
2023-02-22 03:08:12 +03:00
if (root->router)
2023-02-17 06:14:02 +03:00
{
if (!root->zero && !root->one)
2023-02-22 03:08:12 +03:00
s << '>' << GetIdentHashAbbreviation (root->router->GetIdentHash ());
2023-02-17 06:14:02 +03:00
else
s << "error";
}
s << std::endl;
if (root->zero)
{
s << std::string (level, '-') << "0" << std::endl;
Print (s, root->zero, level + 1);
}
if (root->one)
{
s << std::string (level, '-') << "1" << std::endl;
Print (s, root->one, level + 1);
}
}
}
}