dht.GetClosestContacts introduce custom filter function

dht/DHT Lite.go

@@ -83,8 +83,8 @@ func (dht *DHT) RemoveNode(ID []byte) {
 }
 
 // GetClosestContacts returns the closes contacts in the hash table
-func (dht *DHT) GetClosestContacts(count int, target []byte, ignoredNodes ...[]byte) []*Node {
-	closest := dht.ht.getClosestContacts(count, target, ignoredNodes...)
+func (dht *DHT) GetClosestContacts(count int, target []byte, filterFunc NodeFilterFunc, ignoredNodes ...[]byte) []*Node {
+	closest := dht.ht.getClosestContacts(count, target, filterFunc, ignoredNodes...)
 	return closest.Nodes
 }
 
@@ -118,7 +118,7 @@ func (dht *DHT) iterate(action int, target []byte, data []byte) (value []byte, c
 		return nil, nil, errors.New("unknown iterate type")
 	}
 
-	sl := dht.ht.getClosestContacts(dht.alpha, target)
+	sl := dht.ht.getClosestContacts(dht.alpha, target, nil)
 
 	// We keep a reference to the closestNode. If after performing a search we do not find a closer node, we stop searching.
 	if len(sl.Nodes) == 0 {

dht/Hash Table.go

@@ -86,9 +86,11 @@ func (ht *hashTable) doesNodeExistInBucket(bucket int, node []byte) bool {
 	return false
 }
 
-func (ht *hashTable) getClosestContacts(num int, target []byte, ignoredNodes ...[]byte) *shortList {
+// getClosestContacts returns the closest nodes to the target. filterFunc is optional and allows the caller to filter the nodes.
+func (ht *hashTable) getClosestContacts(num int, target []byte, filterFunc NodeFilterFunc, ignoredNodes ...[]byte) *shortList {
 	ht.mutex.Lock()
 	defer ht.mutex.Unlock()
+
 	// First we need to build the list of adjacent indices to our target in order
 	index := ht.getBucketIndexFromDifferingBit(target)
 	indexList := []int{index}
@@ -109,19 +111,23 @@ func (ht *hashTable) getClosestContacts(num int, target []byte, ignoredNodes ...
 	for leftToAdd > 0 && len(indexList) > 0 {
 		index, indexList = indexList[0], indexList[1:]
 		bucketContacts := len(ht.RoutingTable[index])
+	bucketLoop:
 		for i := 0; i < bucketContacts; i++ {
-			ignored := false
 			for j := 0; j < len(ignoredNodes); j++ {
 				if bytes.Compare(ht.RoutingTable[index][i].ID, ignoredNodes[j]) == 0 {
-					ignored = true
+					continue bucketLoop
 				}
 			}
-			if !ignored {
-				sl.AppendUniqueNodes(ht.RoutingTable[index][i])
-				leftToAdd--
-				if leftToAdd == 0 {
-					break
-				}
+
+			// Use the filter function if set. It allows the caller to only accept certain nodes.
+			if filterFunc != nil && !filterFunc(ht.RoutingTable[index][i]) {
+				continue
+			}
+
+			sl.AppendUniqueNodes(ht.RoutingTable[index][i])
+			leftToAdd--
+			if leftToAdd == 0 {
+				break
 			}
 		}
 	}

dht/Node.go

@@ -67,16 +67,14 @@ func (n *shortList) RemoveNode(ID []byte) {
 }
 
 func (n *shortList) AppendUniqueNodes(nodes ...*Node) {
+nodesLoop:
 	for _, vv := range nodes {
-		exists := false
 		for _, v := range n.Nodes {
 			if bytes.Compare(v.ID, vv.ID) == 0 {
-				exists = true
+				continue nodesLoop
 			}
 		}
-		if !exists {
-			n.Nodes = append(n.Nodes, vv)
-		}
+		n.Nodes = append(n.Nodes, vv)
 	}
 }
 
@@ -134,3 +132,6 @@ type NodeMessage struct {
 	Closest  []*Node
 	Error    error
 }
+
+// NodeFilterFunc is called to filter nodes based on the callers choice
+type NodeFilterFunc func(node *Node) (accept bool)