Merge pull request #7 from libp2p/fix/dht-handle-panic

dht: protect against a panic in case record on pbmessage is nil

dht_test.go

@@ -9,6 +9,8 @@ import (
 	"testing"
 	"time"
 
+	pb "github.com/libp2p/go-libp2p-kad-dht/pb"
+
 	ds "github.com/ipfs/go-datastore"
 	dssync "github.com/ipfs/go-datastore/sync"
 	u "github.com/ipfs/go-ipfs-util"
@@ -825,3 +827,15 @@ func TestConnectCollision(t *testing.T) {
 		dhtB.host.Close()
 	}
 }
+
+func TestBadProtoMessages(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	d := setupDHT(ctx, t)
+
+	nilrec := new(pb.Message)
+	if _, err := d.handlePutValue(ctx, "testpeer", nilrec); err == nil {
+		t.Fatal("should have errored on nil record")
+	}
+}

handlers.go

@@ -150,13 +150,17 @@ func (dht *IpfsDHT) handlePutValue(ctx context.Context, p peer.ID, pmes *pb.Mess
 	defer log.EventBegin(ctx, "handlePutValue", p).Done()
 	dskey := key.Key(pmes.GetKey()).DsKey()
 
-	if err := dht.verifyRecordLocally(pmes.GetRecord()); err != nil {
+	rec := pmes.GetRecord()
+	if rec == nil {
+		log.Infof("Got nil record from: %s", p.Pretty())
+		return nil, errors.New("nil record")
+	}
+
+	if err := dht.verifyRecordLocally(rec); err != nil {
 		log.Warningf("Bad dht record in PUT from: %s. %s", key.Key(pmes.GetRecord().GetAuthor()), err)
 		return nil, err
 	}
 
-	rec := pmes.GetRecord()
-
 	// record the time we receive every record
 	rec.TimeReceived = proto.String(u.FormatRFC3339(time.Now()))
 

records.go

@@ -107,6 +107,10 @@ func (dht *IpfsDHT) getPublicKeyFromNode(ctx context.Context, p peer.ID) (ci.Pub
 // verifyRecordLocally attempts to verify a record. if we do not have the public
 // key, we fail. we do not search the dht.
 func (dht *IpfsDHT) verifyRecordLocally(r *recpb.Record) error {
+	if r == nil {
+		log.Error("nil record passed into verifyRecordLocally")
+		return fmt.Errorf("nil record")
+	}
 
 	if len(r.Signature) > 0 {
 		// First, validate the signature