Fixed the generation of random peer ID's

generate/main.go

@@ -28,7 +28,8 @@ func main() {
 	hasher := sha256.New()
 
 	for i := uint64(0); count < target; i++ {
-		binary.BigEndian.PutUint64(inp[idLen-8:], i)
+		binary.BigEndian.PutUint64(inp[2:], i)
+
 		hasher.Write(inp[:])
 		out = hasher.Sum(out[:0])
 		hasher.Reset()
@@ -55,7 +56,7 @@ func main() {
 
 	printf("package %s\n\n", pkg)
 	printf("// Code generated by generate/generate_map.go DO NOT EDIT\n")
-	printf("var keyPrefixMap = [...]uint32{")
+	printf("var keyPrefixMap = [...]uint64{")
 	for i, j := range ids[:] {
 		if i%16 == 0 {
 			printf("\n\t")

table.go

@@ -2,9 +2,9 @@
 package kbucket
 
 import (
+	"encoding/binary"
 	"errors"
 	"fmt"
-	"io"
 	"math/rand"
 	"sync"
 	"time"
@@ -20,7 +20,6 @@ var log = logging.Logger("table")
 
 var ErrPeerRejectedHighLatency = errors.New("peer rejected; latency too high")
 var ErrPeerRejectedNoCapacity = errors.New("peer rejected; insufficient capacity")
-var ErrGenRandPeerIDFailed = errors.New("failed to generate random peerID in bucket: exhausted attempts")
 
 // RoutingTable defines the routing table.
 type RoutingTable struct {
@@ -60,13 +59,19 @@ func NewRoutingTable(bucketsize int, localID ID, latency time.Duration, m peerst
 	return rt
 }
 
+// GetAllBuckets is safe to call as rt.Buckets is append-only
+// caller SHOULD NOT modify the returned slice
 func (rt *RoutingTable) GetAllBuckets() []*Bucket {
 	rt.tabLock.RLock()
 	defer rt.tabLock.RUnlock()
 	return rt.Buckets
 }
 
+// GenRandPeerID generates a random peerID in bucket=bucketID
 func (rt *RoutingTable) GenRandPeerID(bucketID int) (peer.ID, error) {
+	if bucketID < 0 {
+		return "", errors.New("bucketID must be non-negative")
+	}
 	rt.tabLock.RLock()
 	bucketLen := len(rt.Buckets)
 	rt.tabLock.RUnlock()
@@ -78,28 +83,54 @@ func (rt *RoutingTable) GenRandPeerID(bucketID int) (peer.ID, error) {
 		targetCpl = bucketID
 	}
 
-	// should give up after a fixed number of attempts so we don't take up too much time/cpu
-	for i := 0; i < 1000; i++ {
-		peerID, err := randPeerID()
-		if err != nil {
-			log.Debugf("failed to generate random peerID in bucket %d, error is %+v", bucketID, err)
-			continue
+	// generate random 16 bits
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	buf := make([]byte, 2)
+	_, err := r.Read(buf)
+	if err != nil {
+		return "", err
+	}
+
+	// replace the first targetCPL bits with those from the hashed local peer ID & toggle the (targetCpl+1)th bit
+	// so that exactly targetCpl bits match
+	numBytes := targetCpl / 8 // number of bytes we need to replace
+	numBits := targetCpl % 8  // number of bits we need to replace after numBytes have been replaced
+
+	// replace the bytes
+	byteIndex := 0
+	for ; byteIndex < numBytes; byteIndex++ {
+		buf[byteIndex] = rt.local[byteIndex]
+	}
+
+	// replace the bits
+	if byteIndex < len(buf) {
+		dstByte := buf[byteIndex]
+		srcByte := rt.local[byteIndex]
+		j := uint(7)
+		for k := 1; k <= numBits; k++ {
+			if isSet(srcByte, j) {
+				dstByte = setBit(dstByte, j)
+			} else {
+				dstByte = clearBit(dstByte, j)
+			}
+			j--
 		}
-		if CommonPrefixLen(ConvertPeerID(peerID), rt.local) == targetCpl {
-			return peerID, err
+
+		// toggle the next bit
+		if isSet(srcByte, j) {
+			dstByte = clearBit(dstByte, j)
+		} else {
+			dstByte = setBit(dstByte, j)
 		}
+		buf[byteIndex] = dstByte
 	}
-	return "", ErrGenRandPeerIDFailed
-}
 
-func randPeerID() (peer.ID, error) {
-	r := rand.New(rand.NewSource(time.Now().UnixNano()))
-	buf := make([]byte, 16)
-	if _, err := io.ReadFull(r, buf); err != nil {
-		return "", err
-	}
-	h, _ := mh.Sum(buf, mh.SHA2_256, -1)
-	return peer.ID(h), nil
+	// get the seed using buf & use it as the hash digest for a SHA2-256 Multihash to get the desired peer ID
+	prefix := binary.BigEndian.Uint16(buf)
+	key := keyPrefixMap[prefix]
+	h := [34]byte{mh.SHA2_256, 32}
+	binary.BigEndian.PutUint64(h[2:], key)
+	return peer.ID(h[:]), err
 }
 
 // Update adds or moves the given peer to the front of its respective bucket

table_test.go

@@ -49,12 +49,13 @@ func TestBucket(t *testing.T) {
 }
 
 func TestGenRandPeerID(t *testing.T) {
+	nBuckets := 16
 	local := test.RandPeerIDFatal(t)
 	m := pstore.NewMetrics()
 	rt := NewRoutingTable(1, ConvertPeerID(local), time.Hour, m)
 
-	// create 3 buckets
-	for i := 0; i < 5; i++ {
+	// create nBuckets
+	for i := 0; i < nBuckets; i++ {
 		for {
 			if p := test.RandPeerIDFatal(t); CommonPrefixLen(ConvertPeerID(local), ConvertPeerID(p)) == i {
 				rt.Update(p)
@@ -63,30 +64,31 @@ func TestGenRandPeerID(t *testing.T) {
 		}
 	}
 
-	for i := 0; i < 5; i++ {
-		peerID, err := rt.GenRandPeerID(i)
+	for bucketID := 0; bucketID < nBuckets; bucketID++ {
+		peerID, err := rt.GenRandPeerID(bucketID)
 		if err != nil || len(peerID) == 0 {
-			t.Fatalf("error %+v & peerID %s for bucket %d", err, peerID, i)
+			t.Fatalf("error %+v & peerID %s for bucket %d", err, peerID, bucketID)
 		}
 
-		var expectedCpl int
-		if i < len(rt.Buckets)-1 {
-			expectedCpl = i
+		// except for the last bucket, CPL should be Exactly bucketID
+		if bucketID < len(rt.Buckets)-1 {
+			if CommonPrefixLen(ConvertPeerID(peerID), rt.local) != bucketID {
+				t.Fatalf("cpl should be %d for bucket %d but got %d, generated peerID is %s", bucketID, bucketID,
+					CommonPrefixLen(ConvertPeerID(peerID), rt.local), peerID)
+			}
 		} else {
-			expectedCpl = len(rt.Buckets)
-		}
-
-		if CommonPrefixLen(ConvertPeerID(peerID), rt.local) != expectedCpl {
-			t.Fatalf("cpl should be %d for bucket %d, generated peerID is %s", expectedCpl, i, peerID)
+			// for the last bucket, CPL should be Atleast the length of the bucket
+			if CommonPrefixLen(ConvertPeerID(peerID), rt.local) < len(rt.Buckets) {
+				t.Fatalf("cpl should be ATLEAST %d for bucket %d but got %d, generated peerID is %s", len(rt.Buckets), bucketID,
+					CommonPrefixLen(ConvertPeerID(peerID), rt.local), peerID)
+			}
 		}
 	}
 
-	// test error case
-	peerID, err := rt.GenRandPeerID(-1)
-	if err != ErrGenRandPeerIDFailed || len(peerID) != 0 {
-		t.Fatalf("should have got err ErrGenRandPeerIDFailed for bucketID -1")
+	_, err := rt.GenRandPeerID(-1)
+	if err == nil {
+		t.Fatalf("should get erorr for bucketID=-1")
 	}
-
 }
 
 func TestTableCallbacks(t *testing.T) {

util.go

@@ -8,7 +8,7 @@ import (
 	ks "github.com/libp2p/go-libp2p-kbucket/keyspace"
 
 	u "github.com/ipfs/go-ipfs-util"
-	sha256 "github.com/minio/sha256-simd"
+	"github.com/minio/sha256-simd"
 )
 
 // Returned if a routing table query returns no results. This is NOT expected
@@ -35,6 +35,22 @@ func xor(a, b ID) ID {
 	return ID(u.XOR(a, b))
 }
 
+func setBit(n byte, pos uint) byte {
+	n |= (1 << pos)
+	return n
+}
+
+func clearBit(n byte, pos uint) byte {
+	mask := byte(^(1 << pos))
+	n &= mask
+	return n
+}
+
+func isSet(n byte, pos uint) bool {
+	val := n & (1 << pos)
+	return (val > 0)
+}
+
 func CommonPrefixLen(a, b ID) int {
 	return ks.ZeroPrefixLen(u.XOR(a, b))
 }

util_test.go

+package kbucket
+
+import (
+	"testing"
+)
+
+func TestIsSet(t *testing.T) {
+	a := byte(2)
+
+	if !isSet(a, 1) {
+		t.Fatal("1st bit should be set")
+	}
+
+	if isSet(a, 0) {
+		t.Fatal("0th bit should not be set")
+	}
+}
+
+func TestSetBit(t *testing.T) {
+	a := byte(1)
+
+	if setBit(a, 1) != 3 {
+		t.Fatal("1st bit should have been set")
+	}
+}
+
+func TestClearBit(t *testing.T) {
+	a := byte(3)
+
+	if clearBit(a, 0) != 2 {
+		t.Fatal("0th bit should have been cleared")
+	}
+}