diff --git a/dht/DHT Lite.go b/dht/DHT Lite.go index 2a5503e..4e33a4a 100644 --- a/dht/DHT Lite.go +++ b/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 { diff --git a/dht/Hash Table.go b/dht/Hash Table.go index e88a254..551a3f5 100644 --- a/dht/Hash Table.go +++ b/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 } } } diff --git a/dht/Node.go b/dht/Node.go index 2006cf7..ff079f0 100644 --- a/dht/Node.go +++ b/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)