k-bucket support for proper kad bootstrapping

bucket.go

@@ -1,8 +1,11 @@
+//go:generate go run ./generate
+
 package kbucket
 
 import (
 	"container/list"
 	"sync"
+	"time"
 
 	"github.com/libp2p/go-libp2p-core/peer"
 )
@@ -11,14 +14,32 @@ import (
 type Bucket struct {
 	lk   sync.RWMutex
 	list *list.List
+
+	lastRefreshedAtLk sync.RWMutex
+	lastRefreshedAt   time.Time // the last time we looked up a key in the bucket
 }
 
 func newBucket() *Bucket {
 	b := new(Bucket)
 	b.list = list.New()
+	b.lastRefreshedAt = time.Now()
 	return b
 }
 
+func (b *Bucket) RefreshedAt() time.Time {
+	b.lastRefreshedAtLk.RLock()
+	defer b.lastRefreshedAtLk.RUnlock()
+
+	return b.lastRefreshedAt
+}
+
+func (b *Bucket) ResetRefreshedAt(newTime time.Time) {
+	b.lastRefreshedAtLk.Lock()
+	defer b.lastRefreshedAtLk.Unlock()
+
+	b.lastRefreshedAt = newTime
+}
+
 func (b *Bucket) Peers() []peer.ID {
 	b.lk.RLock()
 	defer b.lk.RUnlock()

generate/main.go

@@ -0,0 +1,70 @@
+package main
+
+import (
+	"crypto/sha256"
+	"encoding/binary"
+	"fmt"
+	"os"
+	"strings"
+
+	mh "github.com/multiformats/go-multihash"
+)
+
+const bits = 16
+const target = 1 << bits
+const idLen = 32 + 2
+
+func main() {
+	pkg := os.Getenv("GOPACKAGE")
+	file := os.Getenv("GOFILE")
+	targetFile := strings.TrimSuffix(file, ".go") + "_prefixmap.go"
+
+	ids := new([target]uint32)
+	found := new([target]bool)
+	count := int32(0)
+
+	out := make([]byte, 32)
+	inp := [idLen]byte{mh.SHA2_256, 32}
+	hasher := sha256.New()
+
+	for i := uint32(0); count < target; i++ {
+		binary.BigEndian.PutUint32(inp[2:], i)
+
+		hasher.Write(inp[:])
+		out = hasher.Sum(out[:0])
+		hasher.Reset()
+
+		prefix := binary.BigEndian.Uint32(out) >> (32 - bits)
+		if !found[prefix] {
+			found[prefix] = true
+			ids[prefix] = i
+			count++
+		}
+	}
+
+	f, err := os.Create(targetFile)
+	if err != nil {
+		panic(err)
+	}
+
+	printf := func(s string, args ...interface{}) {
+		_, err := fmt.Fprintf(f, s, args...)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	printf("package %s\n\n", pkg)
+	printf("// Code generated by generate/generate_map.go DO NOT EDIT\n")
+	printf("var keyPrefixMap = [...]uint32{")
+	for i, j := range ids[:] {
+		if i%16 == 0 {
+			printf("\n\t")
+		} else {
+			printf(" ")
+		}
+		printf("%d,", j)
+	}
+	printf("\n}")
+	f.Close()
+}

go.mod

@@ -6,4 +6,5 @@ require (
 	github.com/libp2p/go-libp2p-core v0.0.1
 	github.com/libp2p/go-libp2p-peerstore v0.1.0
 	github.com/minio/sha256-simd v0.0.0-20190131020904-2d45a736cd16
+	github.com/multiformats/go-multihash v0.0.1
 )

table.go

@@ -2,13 +2,16 @@
 package kbucket
 
 import (
+	"encoding/binary"
 	"errors"
 	"fmt"
+	"math/rand"
 	"sync"
 	"time"
 
 	"github.com/libp2p/go-libp2p-core/peer"
 	"github.com/libp2p/go-libp2p-core/peerstore"
+	mh "github.com/multiformats/go-multihash"
 
 	logging "github.com/ipfs/go-log"
 )
@@ -20,7 +23,6 @@ var ErrPeerRejectedNoCapacity = errors.New("peer rejected; insufficient capacity
 
 // RoutingTable defines the routing table.
 type RoutingTable struct {
-
 	// ID of the local peer
 	local ID
 
@@ -57,6 +59,74 @@ 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 {
+	if bucketID < 0 {
+		panic(fmt.Sprintf("bucketID %d is not non-negative", bucketID))
+	}
+	rt.tabLock.RLock()
+	bucketLen := len(rt.Buckets)
+	rt.tabLock.RUnlock()
+
+	var targetCpl uint
+	if bucketID > (bucketLen - 1) {
+		targetCpl = uint(bucketLen) - 1
+	} else {
+		targetCpl = uint(bucketID)
+	}
+
+	// We can only handle upto 16 bit prefixes
+	if targetCpl > 16 {
+		targetCpl = 16
+	}
+
+	var targetPrefix uint16
+	localPrefix := binary.BigEndian.Uint16(rt.local)
+	if targetCpl < 16 {
+		// For host with ID `L`, an ID `K` belongs to a bucket with ID `B` ONLY IF CommonPrefixLen(L,K) is EXACTLY B.
+		// Hence, to achieve a targetPrefix `T`, we must toggle the (T+1)th bit in L & then copy (T+1) bits from L
+		// to our randomly generated prefix.
+		toggledLocalPrefix := localPrefix ^ (uint16(0x8000) >> targetCpl)
+		randPrefix := uint16(rand.Uint32())
+
+		// Combine the toggled local prefix and the random bits at the correct offset
+		// such that ONLY the first `targetCpl` bits match the local ID.
+		mask := (^uint16(0)) << (16 - (targetCpl + 1))
+		targetPrefix = (toggledLocalPrefix & mask) | (randPrefix & ^mask)
+	} else {
+		targetPrefix = localPrefix
+	}
+
+	// Convert to a known peer ID.
+	key := keyPrefixMap[targetPrefix]
+	id := [34]byte{mh.SHA2_256, 32}
+	binary.BigEndian.PutUint32(id[2:], key)
+	return peer.ID(id[:])
+}
+
+// Returns the bucket for a given ID
+// should NOT modify the peer list on the returned bucket
+func (rt *RoutingTable) BucketForID(id ID) *Bucket {
+	cpl := CommonPrefixLen(id, 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

@@ -48,6 +48,52 @@ func TestBucket(t *testing.T) {
 	}
 }
 
+func TestGenRandPeerID(t *testing.T) {
+	nBuckets := 21
+	local := test.RandPeerIDFatal(t)
+	m := pstore.NewMetrics()
+	rt := NewRoutingTable(1, ConvertPeerID(local), time.Hour, m)
+
+	// create nBuckets
+	for i := 0; i < nBuckets; i++ {
+		for {
+			if p := test.RandPeerIDFatal(t); CommonPrefixLen(ConvertPeerID(local), ConvertPeerID(p)) == i {
+				rt.Update(p)
+				break
+			}
+		}
+	}
+
+	// test bucket for peer
+	peers := rt.ListPeers()
+	for _, p := range peers {
+		b := rt.BucketForID(ConvertPeerID(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 := rt.GenRandPeerID(bucketID)
+
+		// for bucketID upto maxPrefixLen of 16, CPL should be Exactly bucketID
+		if bucketID < 16 {
+			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 {
+			// from bucketID 16 onwards, CPL should be ATLEAST 16
+			if CommonPrefixLen(ConvertPeerID(peerID), rt.local) < 16 {
+				t.Fatalf("cpl should be ATLEAST 16 for bucket %d but got %d, generated peerID is %s", bucketID,
+					CommonPrefixLen(ConvertPeerID(peerID), rt.local), peerID)
+			}
+		}
+
+	}
+}
+
 func TestTableCallbacks(t *testing.T) {
 	local := test.RandPeerIDFatal(t)
 	m := pstore.NewMetrics()

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