Only retrieve one value when fetching public key from DHT

records.go

@@ -5,7 +5,6 @@ import (
 	"fmt"
 	"time"
 
-	ctxfrac "github.com/jbenet/go-context/frac"
 	ci "github.com/libp2p/go-libp2p-crypto"
 	peer "github.com/libp2p/go-libp2p-peer"
 	routing "github.com/libp2p/go-libp2p-routing"
@@ -19,6 +18,11 @@ import (
 // it must be rebroadcasted more frequently than once every 'MaxRecordAge'
 const MaxRecordAge = time.Hour * 36
 
+type pubkrs struct {
+	pubk ci.PubKey
+	err  error
+}
+
 func (dht *IpfsDHT) GetPublicKey(ctx context.Context, p peer.ID) (ci.PubKey, error) {
 	log.Debugf("getPublicKey for: %s", p)
 
@@ -34,44 +38,75 @@ func (dht *IpfsDHT) GetPublicKey(ctx context.Context, p peer.ID) (ci.PubKey, err
 		return pk, nil
 	}
 
-	// ok, try the node itself. if they're overwhelmed or slow we can move on.
-	ctxT, cancelFunc := ctxfrac.WithDeadlineFraction(ctx, 0.3)
-	defer cancelFunc()
-	if pk, err := dht.getPublicKeyFromNode(ctx, p); err == nil {
-		err := dht.peerstore.AddPubKey(p, pk)
-		if err != nil {
-			return pk, err
+	// Try getting the public key both directly from the node it identifies
+	// and from the DHT, in parallel
+	resp := make(chan pubkrs, 2)
+	go func() {
+		pubk, err := dht.getPublicKeyFromNode(ctx, p)
+		resp <- pubkrs{pubk, err}
+	}()
+
+	go func() {
+		pubk, err := dht.getPublicKeyFromDHT(ctx, p)
+		resp <- pubkrs{pubk, err}
+	}()
+
+	// Wait for one of the two go routines to return
+	// a public key (or for both to error out)
+	var err error
+	for i := 0; i < 2; i++ {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case r := <-resp:
+			if r.err == nil {
+				// Found the public key
+				err := dht.peerstore.AddPubKey(p, r.pubk)
+				if err != nil {
+					log.Warningf("Failed to add public key to peerstore for %v", p)
+				}
+				return r.pubk, nil
+			}
+			err = r.err
 		}
-		return pk, nil
 	}
 
-	// last ditch effort: let's try the dht.
-	log.Debugf("pk for %s not in peerstore, and peer failed. Trying DHT.", p)
-	pkkey := routing.KeyForPublicKey(p)
+	// Both go routines failed to find a public key
+	return nil, err
+}
 
-	val, err := dht.GetValue(ctxT, pkkey)
+func (dht *IpfsDHT) getPublicKeyFromDHT(ctx context.Context, p peer.ID) (ci.PubKey, error) {
+	// Only retrieve one value, because the public key is immutable
+	// so there's no need to retrieve multiple versions
+	pkkey := routing.KeyForPublicKey(p)
+	vals, err := dht.GetValues(ctx, pkkey, 1)
 	if err != nil {
-		log.Warning("Failed to find requested public key.")
 		return nil, err
 	}
 
-	pk, err = ci.UnmarshalPublicKey(val)
+	if len(vals) == 0 || vals[0].Val == nil {
+		log.Debugf("Could not find public key for %v in DHT", p)
+		return nil, routing.ErrNotFound
+	}
+
+	_, err = dht.Selector.BestRecord(pkkey, [][]byte{vals[0].Val})
 	if err != nil {
-		log.Debugf("Failed to unmarshal public key: %s", err)
+		log.Errorf("Failed to verify public key for %v retrieved from DHT: %s", p, err)
 		return nil, err
 	}
 
-	return pk, dht.peerstore.AddPubKey(p, pk)
+	log.Debugf("DHT got public key for %s from DHT", p)
+	return dht.pubkFromVal(p, vals[0].Val)
 }
 
 func (dht *IpfsDHT) getPublicKeyFromNode(ctx context.Context, p peer.ID) (ci.PubKey, error) {
-
 	// check locally, just in case...
 	pk := dht.peerstore.PubKey(p)
 	if pk != nil {
 		return pk, nil
 	}
 
+	// Get the key from the node itself
 	pkkey := routing.KeyForPublicKey(p)
 	pmes, err := dht.getValueSingle(ctx, p, pkkey)
 	if err != nil {
@@ -81,29 +116,30 @@ func (dht *IpfsDHT) getPublicKeyFromNode(ctx context.Context, p peer.ID) (ci.Pub
 	// node doesn't have key :(
 	record := pmes.GetRecord()
 	if record == nil {
-		return nil, fmt.Errorf("Node not responding with its public key: %s", p)
+		return nil, fmt.Errorf("Node %v not responding with its public key", p)
 	}
 
-	// Success! We were given the value. we don't need to check
-	// validity because a) we can't. b) we know the hash of the
-	// key we're looking for.
 	val := record.GetValue()
-	log.Debug("DHT got a value from other peer")
+	log.Debugf("DHT got public key from node %v itself", p)
+	return dht.pubkFromVal(p, val)
+}
 
-	pk, err = ci.UnmarshalPublicKey(val)
+func (dht *IpfsDHT) pubkFromVal(p peer.ID, val []byte) (ci.PubKey, error) {
+	pubk, err := ci.UnmarshalPublicKey(val)
 	if err != nil {
+		log.Errorf("DHT could not unmarshall public key for %v", p)
 		return nil, err
 	}
 
-	id, err := peer.IDFromPublicKey(pk)
+	// Make sure the public key matches the peer ID
+	id, err := peer.IDFromPublicKey(pubk)
 	if err != nil {
+		log.Errorf("DHT could not extract peer id from public key for %v", p)
 		return nil, err
 	}
 	if id != p {
-		return nil, fmt.Errorf("public key does not match id: %s", p)
+		return nil, fmt.Errorf("public key %v does not match peer %v", id, p)
 	}
 
-	// ok! it's valid. we got it!
-	log.Debugf("DHT got public key from node itself.")
-	return pk, nil
+	return pubk, nil
 }

records_test.go

@@ -4,11 +4,17 @@ import (
 	"context"
 	"crypto/rand"
 	"testing"
+	"time"
 
+	proto "github.com/gogo/protobuf/proto"
+	u "github.com/ipfs/go-ipfs-util"
 	ci "github.com/libp2p/go-libp2p-crypto"
 	peer "github.com/libp2p/go-libp2p-peer"
+	record "github.com/libp2p/go-libp2p-record"
+	routing "github.com/libp2p/go-libp2p-routing"
 )
 
+// Check that GetPublicKey() correctly extracts a public key
 func TestPubkeyExtract(t *testing.T) {
 	_, pk, err := ci.GenerateEd25519Key(rand.Reader)
 	if err != nil {
@@ -32,3 +38,289 @@ func TestPubkeyExtract(t *testing.T) {
 		t.Fatal("got incorrect public key out")
 	}
 }
+
+// Check that GetPublicKey() correctly retrieves a public key from the peerstore
+func TestPubkeyPeerstore(t *testing.T) {
+	ctx := context.Background()
+	dht := setupDHT(ctx, t, false)
+
+	r := u.NewSeededRand(15) // generate deterministic keypair
+	_, pubk, err := ci.GenerateKeyPairWithReader(ci.RSA, 1024, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	id, err := peer.IDFromPublicKey(pubk)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = dht.peerstore.AddPubKey(id, pubk)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rpubk, err := dht.GetPublicKey(context.Background(), id)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !pubk.Equals(rpubk) {
+		t.Fatal("got incorrect public key")
+	}
+}
+
+// Check that GetPublicKey() correctly retrieves a public key directly
+// from the node it identifies
+func TestPubkeyDirectFromNode(t *testing.T) {
+	ctx := context.Background()
+
+	dhtA := setupDHT(ctx, t, false)
+	dhtB := setupDHT(ctx, t, false)
+
+	defer dhtA.Close()
+	defer dhtB.Close()
+	defer dhtA.host.Close()
+	defer dhtB.host.Close()
+
+	connect(t, ctx, dhtA, dhtB)
+
+	pubk, err := dhtA.GetPublicKey(context.Background(), dhtB.self)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	id, err := peer.IDFromPublicKey(pubk)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if id != dhtB.self {
+		t.Fatal("got incorrect public key")
+	}
+}
+
+// Check that GetPublicKey() correctly retrieves a public key
+// from the DHT
+func TestPubkeyFromDHT(t *testing.T) {
+	ctx := context.Background()
+
+	dhtA := setupDHT(ctx, t, false)
+	dhtB := setupDHT(ctx, t, false)
+
+	defer dhtA.Close()
+	defer dhtB.Close()
+	defer dhtA.host.Close()
+	defer dhtB.host.Close()
+
+	connect(t, ctx, dhtA, dhtB)
+
+	r := u.NewSeededRand(15) // generate deterministic keypair
+	_, pubk, err := ci.GenerateKeyPairWithReader(ci.RSA, 1024, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	id, err := peer.IDFromPublicKey(pubk)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pkkey := routing.KeyForPublicKey(id)
+	pkbytes, err := pubk.Bytes()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Store public key on node B
+	err = dhtB.PutValue(ctx, pkkey, pkbytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Retrieve public key on node A
+	rpubk, err := dhtA.GetPublicKey(ctx, id)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !pubk.Equals(rpubk) {
+		t.Fatal("got incorrect public key")
+	}
+}
+
+// Check that GetPublicKey() correctly returns an error when the
+// public key is not available directly from the node or on the DHT
+func TestPubkeyNotFound(t *testing.T) {
+	ctx := context.Background()
+
+	dhtA := setupDHT(ctx, t, false)
+	dhtB := setupDHT(ctx, t, false)
+
+	defer dhtA.Close()
+	defer dhtB.Close()
+	defer dhtA.host.Close()
+	defer dhtB.host.Close()
+
+	connect(t, ctx, dhtA, dhtB)
+
+	r := u.NewSeededRand(15) // generate deterministic keypair
+	_, pubk, err := ci.GenerateKeyPairWithReader(ci.RSA, 1024, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	id, err := peer.IDFromPublicKey(pubk)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Attempt to retrieve public key on node A (should be not found)
+	_, err = dhtA.GetPublicKey(ctx, id)
+	if err == nil {
+		t.Fatal("Expected not found error")
+	}
+}
+
+// Check that GetPublicKey() returns an error when
+// the DHT returns the wrong key
+func TestPubkeyBadKeyFromDHT(t *testing.T) {
+	ctx := context.Background()
+
+	dhtA := setupDHT(ctx, t, false)
+	dhtB := setupDHT(ctx, t, false)
+
+	defer dhtA.Close()
+	defer dhtB.Close()
+	defer dhtA.host.Close()
+	defer dhtB.host.Close()
+
+	connect(t, ctx, dhtA, dhtB)
+
+	r := u.NewSeededRand(15) // generate deterministic keypair
+	_, pubk, err := ci.GenerateKeyPairWithReader(ci.RSA, 1024, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	id, err := peer.IDFromPublicKey(pubk)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pkkey := routing.KeyForPublicKey(id)
+
+	_, wrongpubk, err := ci.GenerateKeyPairWithReader(ci.RSA, 1024, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if pubk == wrongpubk {
+		t.Fatal("Public keys shouldn't match here")
+	}
+	wrongbytes, err := wrongpubk.Bytes()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Store incorrect public key on node B
+	rec := record.MakePutRecord(pkkey, wrongbytes)
+	rec.TimeReceived = proto.String(u.FormatRFC3339(time.Now()))
+	err = dhtB.putLocal(pkkey, rec)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Retrieve public key from node A
+	_, err = dhtA.GetPublicKey(ctx, id)
+	if err == nil {
+		t.Fatal("Expected error because public key is incorrect")
+	}
+}
+
+// Check that GetPublicKey() returns the correct value
+// when the DHT returns the wrong key but the direct
+// connection returns the correct key
+func TestPubkeyBadKeyFromDHTGoodKeyDirect(t *testing.T) {
+	ctx := context.Background()
+
+	dhtA := setupDHT(ctx, t, false)
+	dhtB := setupDHT(ctx, t, false)
+
+	defer dhtA.Close()
+	defer dhtB.Close()
+	defer dhtA.host.Close()
+	defer dhtB.host.Close()
+
+	connect(t, ctx, dhtA, dhtB)
+
+	r := u.NewSeededRand(15) // generate deterministic keypair
+	pkkey := routing.KeyForPublicKey(dhtB.self)
+
+	_, wrongpubk, err := ci.GenerateKeyPairWithReader(ci.RSA, 1024, r)
+	if err != nil {
+		t.Fatal(err)
+	}
+	wrongbytes, err := wrongpubk.Bytes()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Store incorrect public key on node B
+	rec := record.MakePutRecord(pkkey, wrongbytes)
+	rec.TimeReceived = proto.String(u.FormatRFC3339(time.Now()))
+	err = dhtB.putLocal(pkkey, rec)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Retrieve public key from node A
+	pubk, err := dhtA.GetPublicKey(ctx, dhtB.self)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	id, err := peer.IDFromPublicKey(pubk)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// The incorrect public key retrieved from the DHT
+	// should be ignored in favour of the correct public
+	// key retieved from the node directly
+	if id != dhtB.self {
+		t.Fatal("got incorrect public key")
+	}
+}
+
+// Check that GetPublicKey() returns the correct value
+// when both the DHT returns the correct key and the direct
+// connection returns the correct key
+func TestPubkeyGoodKeyFromDHTGoodKeyDirect(t *testing.T) {
+	ctx := context.Background()
+
+	dhtA := setupDHT(ctx, t, false)
+	dhtB := setupDHT(ctx, t, false)
+
+	defer dhtA.Close()
+	defer dhtB.Close()
+	defer dhtA.host.Close()
+	defer dhtB.host.Close()
+
+	connect(t, ctx, dhtA, dhtB)
+
+	pubk := dhtB.peerstore.PubKey(dhtB.self)
+	pkbytes, err := pubk.Bytes()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Store public key on node B
+	pkkey := routing.KeyForPublicKey(dhtB.self)
+	err = dhtB.PutValue(ctx, pkkey, pkbytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Retrieve public key on node A
+	rpubk, err := dhtA.GetPublicKey(ctx, dhtB.self)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if !pubk.Equals(rpubk) {
+		t.Fatal("got incorrect public key")
+	}
+}