Added method to get bucket for a given peer

table.go

@@ -133,6 +133,22 @@ func (rt *RoutingTable) GenRandPeerID(bucketID int) (peer.ID, error) {
 	return peer.ID(h[:]), err
 }
 
+// Returns the bucket for a given peer
+// should NOT modify the peer list on the returned bucket
+func (rt *RoutingTable) BucketForPeer(p peer.ID) *Bucket {
+	peerID := ConvertPeerID(p)
+	cpl := CommonPrefixLen(peerID, rt.local)
+
+	rt.tabLock.RLock()
+	defer rt.tabLock.RUnlock()
+	bucketID := cpl
+	if bucketID >= len(rt.Buckets) {
+		bucketID = len(rt.Buckets) - 1
+	}
+
+	return rt.Buckets[bucketID]
+}
+
 // Update adds or moves the given peer to the front of its respective bucket
 func (rt *RoutingTable) Update(p peer.ID) (evicted peer.ID, err error) {
 	peerID := ConvertPeerID(p)

table_test.go

@@ -64,6 +64,16 @@ func TestGenRandPeerID(t *testing.T) {
 		}
 	}
 
+	// test bucket for peer
+	peers := rt.ListPeers()
+	for _, p := range peers {
+		b := rt.BucketForPeer(p)
+		if !b.Has(p) {
+			t.Fatalf("bucket should have peers %s", p.String())
+		}
+	}
+
+	// test generate rand peer ID
 	for bucketID := 0; bucketID < nBuckets; bucketID++ {
 		peerID, err := rt.GenRandPeerID(bucketID)
 		if err != nil || len(peerID) == 0 {
@@ -83,6 +93,7 @@ func TestGenRandPeerID(t *testing.T) {
 					CommonPrefixLen(ConvertPeerID(peerID), rt.local), peerID)
 			}
 		}
+
 	}
 
 	_, err := rt.GenRandPeerID(-1)