From 7b89866bbd7e74ea0555954c8de8d238ed8b6298 Mon Sep 17 00:00:00 2001
From: anton <anton.tom@bitbucket.org>
Date: Mon, 28 Sep 2020 14:11:01 +0300
Subject: [PATCH] Handle keep-alive in recvBitfield

---
 client/client.go      | 4 ++++
 client/client_test.go | 5 +++++
 2 files changed, 9 insertions(+)

diff --git a/client/client.go b/client/client.go
index d7b375d..3436443 100644
--- a/client/client.go
+++ b/client/client.go
@@ -52,6 +52,10 @@ func recvBitfield(conn net.Conn) (bitfield.Bitfield, error) {
 	if err != nil {
 		return nil, err
 	}
+	if msg == nil {
+		err := fmt.Errorf("Expected bitfield but got %s", msg)
+		return nil, err
+	}
 	if msg.ID != message.MsgBitfield {
 		err := fmt.Errorf("Expected bitfield but got ID %d", msg.ID)
 		return nil, err
diff --git a/client/client_test.go b/client/client_test.go
index 50c3687..9386b56 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -48,6 +48,11 @@ func TestRecvBitfield(t *testing.T) {
 			output: nil,
 			fails:  true,
 		},
+		"message is keep-alive": {
+			msg:    []byte{0x00, 0x00, 0x00, 0x00},
+			output: nil,
+			fails:  true,
+		},
 	}
 
 	for _, test := range tests {