prefix lookup implementation

dht.go

@@ -148,6 +148,10 @@ type IpfsDHT struct {
 
 	// configuration variables for tests
 	testAddressUpdateProcessing bool
+
+	// length of prefix of keys for provider lookups
+	// if 0, the whole key is used.
+	prefixLength int
 }
 
 // Assert that IPFS assumptions about interfaces aren't broken. These aren't a
@@ -196,6 +200,8 @@ func New(ctx context.Context, h host.Host, options ...Option) (*IpfsDHT, error)
 		return nil, err
 	}
 
+	dht.prefixLength = cfg.PrefixLookupLength
+
 	dht.testAddressUpdateProcessing = cfg.TestAddressUpdateProcessing
 
 	dht.auto = cfg.Mode
@@ -685,13 +691,22 @@ func (dht *IpfsDHT) FindLocal(id peer.ID) peer.AddrInfo {
 
 // nearestPeersToQuery returns the routing tables closest peers.
 func (dht *IpfsDHT) nearestPeersToQuery(pmes *pb.Message, count int) []peer.ID {
+	key := pmes.GetKey()
+
 	if pmes.GetType() == pb.Message_GET_PROVIDERS {
 		// for GET_PROVIDERS messages, the message key is the hashed multihash, so don't hash it again
-		closer := dht.routingTable.NearestPeers(kb.ID(string(pmes.GetKey())), count)
-		return closer
+		if len(key) < 32 {
+			// prefix lookup
+			closer := dht.routingTable.NearestPeersToPrefix(kb.ID(string(key)), count)
+			return closer
+		} else {
+			// normal non-prefixed lookup
+			closer := dht.routingTable.NearestPeers(kb.ID(string(key)), count)
+			return closer
+		}
 	}
 
-	closer := dht.routingTable.NearestPeers(kb.ConvertKey(string(pmes.GetKey())), count)
+	closer := dht.routingTable.NearestPeers(kb.ConvertKey(string(key)), count)
 	return closer
 }
 

dht_options.go

@@ -36,6 +36,14 @@ const DefaultPrefix protocol.ID = "/ipfs"
 
 type Option = dhtcfg.Option
 
+// PrefixLookups configures the DHT to use prefixes of keys for provider lookups.
+func PrefixLookups(prefixLen int) Option {
+	return func(c *dhtcfg.Config) error {
+		c.PrefixLookupLength = prefixLen
+		return nil
+	}
+}
+
 // ProviderStore sets the provider storage manager.
 func ProviderStore(ps providers.ProviderStore) Option {
 	return func(c *dhtcfg.Config) error {

dht_test.go

@@ -1014,6 +1014,55 @@ func TestProvidesAsync(t *testing.T) {
 	}
 }
 
+func TestProvides_PrefixLookup(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	dhts := setupDHTS(t, ctx, 4)
+	defer func() {
+		for i := 0; i < 4; i++ {
+			dhts[i].Close()
+			defer dhts[i].host.Close()
+		}
+	}()
+
+	connect(t, ctx, dhts[0], dhts[1])
+	connect(t, ctx, dhts[1], dhts[2])
+	connect(t, ctx, dhts[1], dhts[3])
+
+	for _, k := range testCaseCids {
+		logger.Debugf("announcing provider for %s", k)
+		if err := dhts[3].Provide(ctx, k, true); err != nil {
+			t.Fatal(err)
+		}
+	}
+
+	time.Sleep(time.Millisecond * 6)
+
+	n := 0
+	for _, c := range testCaseCids {
+		n = (n + 1) % 3
+
+		logger.Debugf("getting providers for %s from %d", c, n)
+		ctxT, cancel := context.WithTimeout(ctx, time.Second)
+		defer cancel()
+		dhts[n].prefixLength = 16 // half the hashed CID for now
+		provchan := dhts[n].FindProvidersAsync(ctxT, c, 1)
+
+		select {
+		case prov := <-provchan:
+			if prov.ID == "" {
+				t.Fatal("Got back nil provider")
+			}
+			if prov.ID != dhts[3].self {
+				t.Fatal("Got back wrong provider")
+			}
+		case <-ctxT.Done():
+			t.Fatal("Did not get a provider back.")
+		}
+	}
+}
+
 func TestLayeredGet(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()

fullrt/dht.go

@@ -95,6 +95,8 @@ type FullRT struct {
 	timeoutPerOp time.Duration
 
 	bulkSendParallelism int
+
+	prefixLength int
 }
 
 // NewFullRT creates a DHT client that tracks the full network. It takes a protocol prefix for the given network,
@@ -185,6 +187,8 @@ func NewFullRT(h host.Host, protocolPrefix protocol.ID, options ...Option) (*Ful
 		crawlerInterval: time.Minute * 60,
 
 		bulkSendParallelism: 20,
+
+		prefixLength: dhtcfg.PrefixLookupLength,
 	}
 
 	rt.wg.Add(1)
@@ -1215,14 +1219,17 @@ func (dht *FullRT) FindProvidersAsync(ctx context.Context, key cid.Cid, count in
 
 	keyMH := key.Hash()
 
-	logger.Debugw("finding providers", "cid", key, "mh", internal.LoggableProviderRecordBytes(keyMH))
 	go dht.findProvidersAsyncRoutine(ctx, keyMH, count, peerOut)
 	return peerOut
 }
 
 func (dht *FullRT) findProvidersAsyncRoutine(ctx context.Context, key multihash.Multihash, count int, peerOut chan peer.AddrInfo) {
 	defer close(peerOut)
 
+	// hash multihash for double-hashing implementation
+	mhHash := internal.Sha256Multihash(key)
+	logger.Debugw("finding providers", "cid", key, "mhHash", mhHash, "mh", internal.LoggableProviderRecordBytes(key))
+
 	findAll := count == 0
 	ps := make(map[peer.ID]struct{})
 	psLock := &sync.Mutex{}
@@ -1242,7 +1249,7 @@ func (dht *FullRT) findProvidersAsyncRoutine(ctx context.Context, key multihash.
 		return len(ps)
 	}
 
-	provs, err := dht.ProviderManager.GetProviders(ctx, key)
+	provs, err := dht.ProviderManager.GetProviders(ctx, mhHash[:])
 	if err != nil {
 		return
 	}
@@ -1278,7 +1285,14 @@ func (dht *FullRT) findProvidersAsyncRoutine(ctx context.Context, key multihash.
 			ID:   p,
 		})
 
-		provs, closest, err := dht.protoMessenger.GetProviders(ctx, p, key)
+		var lookupKey []byte
+		if dht.prefixLength != 0 {
+			lookupKey = mhHash[:dht.prefixLength]
+		} else {
+			lookupKey = mhHash[:]
+		}
+
+		provs, closest, err := dht.protoMessenger.GetProviders(ctx, p, lookupKey)
 		if err != nil {
 			return err
 		}
@@ -1287,12 +1301,27 @@ func (dht *FullRT) findProvidersAsyncRoutine(ctx context.Context, key multihash.
 
 		// Add unique providers from request, up to 'count'
 		for _, prov := range provs {
-			dht.maybeAddAddrs(prov.ID, prov.Addrs, peerstore.TempAddrTTL)
+			// if this is a prefix lookup, the providers might not actually have
+			// the content we're looking for. discard all that don't
+			if dht.prefixLength > 0 {
+				var hasContent bool
+				for _, key := range prov.Keys {
+					if bytes.Equal(key, mhHash[:]) {
+						hasContent = true
+						break
+					}
+				}
+				if !hasContent {
+					continue
+				}
+			}
+
+			dht.maybeAddAddrs(prov.AddrInfo.ID, prov.AddrInfo.Addrs, peerstore.TempAddrTTL)
 			logger.Debugf("got provider: %s", prov)
-			if psTryAdd(prov.ID) {
+			if psTryAdd(prov.AddrInfo.ID) {
 				logger.Debugf("using provider: %s", prov)
 				select {
-				case peerOut <- *prov:
+				case peerOut <- *prov.AddrInfo:
 				case <-ctx.Done():
 					logger.Debug("context timed out sending more providers")
 					return ctx.Err()

go.mod

@@ -17,7 +17,6 @@ require (
 	github.com/ipfs/go-log v1.0.5
 	github.com/jbenet/goprocess v0.1.4
 	github.com/libp2p/go-libp2p v0.22.0
-	github.com/libp2p/go-libp2p-core v0.20.0
 	github.com/libp2p/go-libp2p-kbucket v0.4.7
 	github.com/libp2p/go-libp2p-record v0.2.0
 	github.com/libp2p/go-libp2p-routing-helpers v0.2.3
@@ -71,7 +70,7 @@ require (
 	github.com/libp2p/go-cidranger v1.1.0 // indirect
 	github.com/libp2p/go-flow-metrics v0.1.0 // indirect
 	github.com/libp2p/go-libp2p-asn-util v0.2.0 // indirect
-	github.com/libp2p/go-libp2p-peerstore v0.8.0 // indirect
+	github.com/libp2p/go-libp2p-core v0.20.0 // indirect
 	github.com/libp2p/go-nat v0.1.0 // indirect
 	github.com/libp2p/go-openssl v0.1.0 // indirect
 	github.com/libp2p/go-reuseport v0.2.0 // indirect
@@ -123,3 +122,9 @@ require (
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	lukechampine.com/blake3 v1.1.7 // indirect
 )
+
+replace github.com/libp2p/go-libp2p-kbucket => ../../ChainSafe/go-libp2p-kbucket
+
+// replace github.com/libp2p/go-libp2p-kbucket => github.com/ChainSafe/go-libp2p-kbucket v0.0.0-20220908203920-02442af318da
+// replace github.com/libp2p/go-libp2p-kbucket/keyspace => github.com/ChainSafe/go-libp2p-kbucket/keyspace v0.0.0-20220908203920-02442af318da
+// replace github.com/libp2p/go-libp2p-kbucket/peerdiversity => github.com/ChainSafe/go-libp2p-kbucket/peerdiversity v0.0.0-20220908203920-02442af318da