/ libi2pd / KadDHT.cpp
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  }