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 {