KadDHT.cpp
1 /* 2 * Copyright (c) 2023, The PurpleI2P Project 3 * 4 * This file is part of Purple i2pd project and licensed under BSD3 5 * 6 * See full license text in LICENSE file at top of project tree 7 * 8 */ 9 10 #include "KadDHT.h" 11 12 namespace i2p 13 { 14 namespace data 15 { 16 DHTNode::DHTNode (): 17 zero (nullptr), one (nullptr) 18 { 19 } 20 21 DHTNode::~DHTNode () 22 { 23 if (zero) delete zero; 24 if (one) delete one; 25 } 26 27 void DHTNode::MoveRouterUp (bool fromOne) 28 { 29 DHTNode *& side = fromOne ? one : zero; 30 if (side) 31 { 32 if (router) router = nullptr; // shouldn't happen 33 router = side->router; 34 side->router = nullptr; 35 delete side; 36 side = nullptr; 37 } 38 } 39 40 DHTTable::DHTTable (): 41 m_Size (0) 42 { 43 m_Root = new DHTNode; 44 } 45 46 DHTTable::~DHTTable () 47 { 48 delete m_Root; 49 } 50 51 void DHTTable::Clear () 52 { 53 m_Size = 0; 54 delete m_Root; 55 m_Root = new DHTNode; 56 } 57 58 void DHTTable::Insert (const std::shared_ptr<RouterInfo>& r) 59 { 60 if (!r) return; 61 return Insert (r, m_Root, 0); 62 } 63 64 void DHTTable::Insert (const std::shared_ptr<RouterInfo>& r, DHTNode * root, int level) 65 { 66 if (root->router) 67 { 68 if (root->router->GetIdentHash () == r->GetIdentHash ()) 69 { 70 root->router = r; // replace 71 return; 72 } 73 auto r2 = root->router; 74 root->router = nullptr; m_Size--; 75 int bit1, bit2; 76 do 77 { 78 bit1 = r->GetIdentHash ().GetBit (level); 79 bit2 = r2->GetIdentHash ().GetBit (level); 80 if (bit1 == bit2) 81 { 82 if (bit1) 83 { 84 if (root->one) return; // something wrong 85 root->one = new DHTNode; 86 root = root->one; 87 } 88 else 89 { 90 if (root->zero) return; // something wrong 91 root->zero = new DHTNode; 92 root = root->zero; 93 } 94 level++; 95 } 96 } 97 while (bit1 == bit2); 98 99 if (!root->zero) 100 root->zero = new DHTNode; 101 if (!root->one) 102 root->one = new DHTNode; 103 if (bit1) 104 { 105 Insert (r2, root->zero, level + 1); 106 Insert (r, root->one, level + 1); 107 } 108 else 109 { 110 Insert (r2, root->one, level + 1); 111 Insert (r, root->zero, level + 1); 112 } 113 } 114 else 115 { 116 if (!root->zero && !root->one) 117 { 118 root->router = r; m_Size++; 119 return; 120 } 121 int bit = r->GetIdentHash ().GetBit (level); 122 if (bit) 123 { 124 if (!root->one) 125 root->one = new DHTNode; 126 Insert (r, root->one, level + 1); 127 } 128 else 129 { 130 if (!root->zero) 131 root->zero = new DHTNode; 132 Insert (r, root->zero, level + 1); 133 } 134 } 135 } 136 137 bool DHTTable::Remove (const IdentHash& h) 138 { 139 return Remove (h, m_Root, 0); 140 } 141 142 bool DHTTable::Remove (const IdentHash& h, DHTNode * root, int level) 143 { 144 if (root) 145 { 146 if (root->router && root->router->GetIdentHash () == h) 147 { 148 root->router = nullptr; 149 m_Size--; 150 return true; 151 } 152 int bit = h.GetBit (level); 153 if (bit) 154 { 155 if (root->one && Remove (h, root->one, level + 1)) 156 { 157 if (root->one->IsEmpty ()) 158 { 159 delete root->one; 160 root->one = nullptr; 161 if (root->zero && root->zero->router) 162 root->MoveRouterUp (false); 163 } 164 else if (root->one->router && !root->zero) 165 root->MoveRouterUp (true); 166 return true; 167 } 168 } 169 else 170 { 171 if (root->zero && Remove (h, root->zero, level + 1)) 172 { 173 if (root->zero->IsEmpty ()) 174 { 175 delete root->zero; 176 root->zero = nullptr; 177 if (root->one && root->one->router) 178 root->MoveRouterUp (true); 179 } 180 else if (root->zero->router && !root->one) 181 root->MoveRouterUp (false); 182 return true; 183 } 184 } 185 } 186 return false; 187 } 188 189 std::shared_ptr<RouterInfo> DHTTable::FindClosest (const IdentHash& h, const Filter& filter) const 190 { 191 if (filter) m_Filter = filter; 192 auto r = FindClosest (h, m_Root, 0); 193 m_Filter = nullptr; 194 return r; 195 } 196 197 std::shared_ptr<RouterInfo> DHTTable::FindClosest (const IdentHash& h, DHTNode * root, int level) const 198 { 199 bool split = false; 200 do 201 { 202 if (root->router) 203 return (!m_Filter || m_Filter (root->router)) ? root->router : nullptr; 204 split = root->zero && root->one; 205 if (!split) 206 { 207 if (root->zero) root = root->zero; 208 else if (root->one) root = root->one; 209 else return nullptr; 210 level++; 211 } 212 } 213 while (!split); 214 int bit = h.GetBit (level); 215 if (bit) 216 { 217 if (root->one) 218 { 219 auto r = FindClosest (h, root->one, level + 1); 220 if (r) return r; 221 } 222 if (root->zero) 223 { 224 auto r = FindClosest (h, root->zero, level + 1); 225 if (r) return r; 226 } 227 } 228 else 229 { 230 if (root->zero) 231 { 232 auto r = FindClosest (h, root->zero, level + 1); 233 if (r) return r; 234 } 235 if (root->one) 236 { 237 auto r = FindClosest (h, root->one, level + 1); 238 if (r) return r; 239 } 240 } 241 return nullptr; 242 } 243 244 std::vector<std::shared_ptr<RouterInfo> > DHTTable::FindClosest (const IdentHash& h, size_t num, const Filter& filter) const 245 { 246 std::vector<std::shared_ptr<RouterInfo> > vec; 247 if (num > 0) 248 { 249 if (filter) m_Filter = filter; 250 FindClosest (h, num, m_Root, 0, vec); 251 m_Filter = nullptr; 252 } 253 return vec; 254 } 255 256 void DHTTable::FindClosest (const IdentHash& h, size_t num, DHTNode * root, int level, std::vector<std::shared_ptr<RouterInfo> >& hashes) const 257 { 258 if (hashes.size () >= num) return; 259 bool split = false; 260 do 261 { 262 if (root->router) 263 { 264 if (!m_Filter || m_Filter (root->router)) 265 hashes.push_back (root->router); 266 return; 267 } 268 split = root->zero && root->one; 269 if (!split) 270 { 271 if (root->zero) root = root->zero; 272 else if (root->one) root = root->one; 273 else return; 274 level++; 275 } 276 } 277 while (!split); 278 int bit = h.GetBit (level); 279 if (bit) 280 { 281 if (root->one) 282 FindClosest (h, num, root->one, level + 1, hashes); 283 if (hashes.size () < num && root->zero) 284 FindClosest (h, num, root->zero, level + 1, hashes); 285 } 286 else 287 { 288 if (root->zero) 289 FindClosest (h, num, root->zero, level + 1, hashes); 290 if (hashes.size () < num && root->one) 291 FindClosest (h, num, root->one, level + 1, hashes); 292 } 293 } 294 295 void DHTTable::Cleanup (const Filter& filter) 296 { 297 if (filter) 298 { 299 m_Filter = filter; 300 Cleanup (m_Root); 301 m_Filter = nullptr; 302 } 303 else 304 Clear (); 305 } 306 307 void DHTTable::Cleanup (DHTNode * root) 308 { 309 if (!root) return; 310 if (root->router) 311 { 312 if (!m_Filter || !m_Filter (root->router)) 313 { 314 m_Size--; 315 root->router = nullptr; 316 } 317 return; 318 } 319 if (root->zero) 320 { 321 Cleanup (root->zero); 322 if (root->zero->IsEmpty ()) 323 { 324 delete root->zero; 325 root->zero = nullptr; 326 } 327 } 328 if (root->one) 329 { 330 Cleanup (root->one); 331 if (root->one->IsEmpty ()) 332 { 333 delete root->one; 334 root->one = nullptr; 335 if (root->zero && root->zero->router) 336 root->MoveRouterUp (false); 337 } 338 else if (root->one->router && !root->zero) 339 root->MoveRouterUp (true); 340 } 341 } 342 343 void DHTTable::Print (std::stringstream& s) 344 { 345 Print (s, m_Root, 0); 346 } 347 348 void DHTTable::Print (std::stringstream& s, DHTNode * root, int level) 349 { 350 if (!root) return; 351 s << std::string (level, '-'); 352 if (root->router) 353 { 354 if (!root->zero && !root->one) 355 s << '>' << GetIdentHashAbbreviation (root->router->GetIdentHash ()); 356 else 357 s << "error"; 358 } 359 s << std::endl; 360 if (root->zero) 361 { 362 s << std::string (level, '-') << "0" << std::endl; 363 Print (s, root->zero, level + 1); 364 } 365 if (root->one) 366 { 367 s << std::string (level, '-') << "1" << std::endl; 368 Print (s, root->one, level + 1); 369 } 370 } 371 } 372 }