/* * 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 (): zero (nullptr), one (nullptr) { } DHTNode::~DHTNode () { if (zero) delete zero; if (one) delete one; } void DHTNode::MoveRouterUp (bool fromOne) { DHTNode *& side = fromOne ? one : zero; if (side) { if (router) router = nullptr; // shouldn't happen router = side->router; side->router = nullptr; delete side; side = nullptr; } } DHTTable::DHTTable (): m_Size (0) { m_Root = new DHTNode; } DHTTable::~DHTTable () { delete m_Root; } void DHTTable::Clear () { 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); } void DHTTable::Insert (const std::shared_ptr<RouterInfo>& r, DHTNode * root, int level) { if (root->router) { if (root->router->GetIdentHash () == r->GetIdentHash ()) { root->router = r; // replace return; } auto r2 = root->router; root->router = nullptr; m_Size--; int bit1, bit2; do { bit1 = r->GetIdentHash ().GetBit (level); bit2 = r2->GetIdentHash ().GetBit (level); if (bit1 == bit2) { if (bit1) { if (root->one) return; // something wrong root->one = new DHTNode; root = root->one; } else { if (root->zero) return; // something wrong 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) { Insert (r2, root->zero, level + 1); Insert (r, root->one, level + 1); } else { Insert (r2, root->one, level + 1); Insert (r, root->zero, level + 1); } } else { if (!root->zero && !root->one) { root->router = r; m_Size++; return; } int bit = r->GetIdentHash ().GetBit (level); if (bit) { if (!root->one) root->one = new DHTNode; Insert (r, root->one, level + 1); } else { if (!root->zero) root->zero = new DHTNode; Insert (r, root->zero, level + 1); } } } bool DHTTable::Remove (const IdentHash& h) { return Remove (h, m_Root, 0); } bool DHTTable::Remove (const IdentHash& h, DHTNode * root, int level) { if (root) { if (root->router && root->router->GetIdentHash () == h) { root->router = nullptr; 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; if (root->zero && root->zero->router) root->MoveRouterUp (false); } else if (root->one->router && !root->zero) root->MoveRouterUp (true); return true; } } else { if (root->zero && Remove (h, root->zero, level + 1)) { if (root->zero->IsEmpty ()) { delete root->zero; root->zero = nullptr; if (root->one && root->one->router) root->MoveRouterUp (true); } else if (root->zero->router && !root->one) root->MoveRouterUp (false); return true; } } } return false; } std::shared_ptr<RouterInfo> DHTTable::FindClosest (const IdentHash& h, const Filter& filter) const { if (filter) m_Filter = filter; auto r = FindClosest (h, m_Root, 0); m_Filter = nullptr; return r; } std::shared_ptr<RouterInfo> DHTTable::FindClosest (const IdentHash& h, DHTNode * root, int level) const { 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); int bit = h.GetBit (level); if (bit) { if (root->one) { auto r = FindClosest (h, root->one, level + 1); if (r) return r; } if (root->zero) { auto r = FindClosest (h, root->zero, level + 1); if (r) return r; } } else { if (root->zero) { auto r = FindClosest (h, root->zero, level + 1); if (r) return r; } if (root->one) { auto r = FindClosest (h, root->one, level + 1); if (r) return r; } } return nullptr; } std::vector<std::shared_ptr<RouterInfo> > DHTTable::FindClosest (const IdentHash& h, size_t num, const Filter& filter) const { std::vector<std::shared_ptr<RouterInfo> > vec; if (num > 0) { if (filter) m_Filter = filter; FindClosest (h, num, m_Root, 0, vec); m_Filter = nullptr; } return vec; } void DHTTable::FindClosest (const IdentHash& h, size_t num, DHTNode * root, int level, std::vector<std::shared_ptr<RouterInfo> >& hashes) const { if (hashes.size () >= num) return; bool split = false; do { 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); 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); } } void DHTTable::Cleanup (const Filter& filter) { 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); } } 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, '-'); if (root->router) { if (!root->zero && !root->one) s << '>' << GetIdentHashAbbreviation (root->router->GetIdentHash ()); 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); } } } }