From 8f362d2b1500809cbae2ee08c2b42a6f226a10b8 Mon Sep 17 00:00:00 2001
From: Jeromy <why@ipfs.io>
Date: Thu, 8 Sep 2016 12:21:29 -0700
Subject: [PATCH] dht: protect against a panic in case record on pbmessage is
 nil

License: MIT
Signed-off-by: Jeromy <why@ipfs.io>
---
 routing/dht/dht_test.go | 17 +++++++++++++++--
 routing/dht/handlers.go | 10 +++++++---
 routing/dht/records.go  |  4 ++++
 3 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/routing/dht/dht_test.go b/routing/dht/dht_test.go
index 0abc27ed769..3bd8880e070 100644
--- a/routing/dht/dht_test.go
+++ b/routing/dht/dht_test.go
@@ -11,14 +11,15 @@ import (
 
 	key "github.com/ipfs/go-ipfs/blocks/key"
 	routing "github.com/ipfs/go-ipfs/routing"
+	pb "github.com/ipfs/go-ipfs/routing/dht/pb"
 	record "github.com/ipfs/go-ipfs/routing/record"
 	ci "github.com/ipfs/go-ipfs/thirdparty/testutil/ci"
 	travisci "github.com/ipfs/go-ipfs/thirdparty/testutil/ci/travis"
-	ds "gx/ipfs/QmTxLSvdhwg68WJimdS6icLPhZi28aTp6b7uihC2Yb47Xk/go-datastore"
-	dssync "gx/ipfs/QmTxLSvdhwg68WJimdS6icLPhZi28aTp6b7uihC2Yb47Xk/go-datastore/sync"
 
 	pstore "gx/ipfs/QmQdnfvZQuhdT93LNc5bos52wAmdr3G2p6G8teLJMEN32P/go-libp2p-peerstore"
 	peer "gx/ipfs/QmRBqJF7hb8ZSpRcMwUt8hNhydWcxGEhtk81HKq6oUwKvs/go-libp2p-peer"
+	ds "gx/ipfs/QmTxLSvdhwg68WJimdS6icLPhZi28aTp6b7uihC2Yb47Xk/go-datastore"
+	dssync "gx/ipfs/QmTxLSvdhwg68WJimdS6icLPhZi28aTp6b7uihC2Yb47Xk/go-datastore/sync"
 	netutil "gx/ipfs/QmVCe3SNMjkcPgnpFhZs719dheq6xE7gJwjzV7aWcUM4Ms/go-libp2p/p2p/test/util"
 	ma "gx/ipfs/QmYzDkkgAEmrcNzFCiYo6L1dTX4EAG1gZkbtdbd9trL4vd/go-multiaddr"
 	u "gx/ipfs/QmZNVWh8LLjAavuQ2JXuFmuYH3C11xo988vSgp7UQrTRj1/go-ipfs-util"
@@ -826,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")
+	}
+}
diff --git a/routing/dht/handlers.go b/routing/dht/handlers.go
index feaf9aea2f5..edb3d906057 100644
--- a/routing/dht/handlers.go
+++ b/routing/dht/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()))
 
diff --git a/routing/dht/records.go b/routing/dht/records.go
index d920e9843ef..57146975007 100644
--- a/routing/dht/records.go
+++ b/routing/dht/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 *pb.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