Separate bitfield into its own package

bitfield/bitfield.go

@@ -0,0 +1,18 @@
+package bitfield
+
+// A Bitfield represents the pieces that a peer has
+type Bitfield []byte
+
+// HasPiece tells if a bitfield has a particular index set
+func (bf Bitfield) HasPiece(index int) bool {
+	byteIndex := index / 8
+	offset := index % 8
+	return bf[byteIndex]>>(7-offset)&1 != 0
+}
+
+// SetPiece sets a bit in the bitfield
+func (bf Bitfield) SetPiece(index int) {
+	byteIndex := index / 8
+	offset := index % 8
+	bf[byteIndex] |= 1 << (7 - offset)
+}

bitfield/bitfield_test.go

@@ -0,0 +1,44 @@
+package bitfield
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestHasPiece(t *testing.T) {
+	bf := Bitfield{0b01010100, 0b01010100}
+	outputs := []bool{false, true, false, true, false, true, false, false, false, true, false, true, false, true, false, false}
+	for i := 0; i < len(outputs); i++ {
+		assert.Equal(t, outputs[i], bf.HasPiece(i))
+	}
+}
+
+func TestSetPiece(t *testing.T) {
+	tests := []struct {
+		input Bitfield
+		index int
+		outpt Bitfield
+	}{
+		{
+			input: Bitfield{0b01010100, 0b01010100},
+			index: 4,
+			outpt: Bitfield{0b01011100, 0b01010100},
+		},
+		{
+			input: Bitfield{0b01010100, 0b01010100},
+			index: 9,
+			outpt: Bitfield{0b01010100, 0b01010100},
+		},
+		{
+			input: Bitfield{0b01010100, 0b01010100},
+			index: 15,
+			outpt: Bitfield{0b01010100, 0b01010101},
+		},
+	}
+	for _, test := range tests {
+		bf := test.input
+		bf.SetPiece(test.index)
+		assert.Equal(t, test.outpt, bf)
+	}
+}

message/message.go

@@ -30,9 +30,6 @@ type Message struct {
 	Payload []byte
 }
 
-// A Bitfield represents the pieces that a peer has
-type Bitfield []byte
-
 // FormatRequest creates a REQUEST message
 func FormatRequest(index, begin, length int) *Message {
 	payload := make([]byte, 12)
@@ -164,17 +161,3 @@ func (m *Message) String() string {
 	}
 	return fmt.Sprintf("%s [%d]", m.name(), len(m.Payload))
 }
-
-// HasPiece tells if a bitfield has a particular index set
-func (bf Bitfield) HasPiece(index int) bool {
-	byteIndex := index / 8
-	offset := index % 8
-	return bf[byteIndex]>>(7-offset)&1 != 0
-}
-
-// SetPiece sets a bit in the bitfield
-func (bf Bitfield) SetPiece(index int) {
-	byteIndex := index / 8
-	offset := index % 8
-	bf[byteIndex] |= 1 << (7 - offset)
-}

message/message_test.go

@@ -263,40 +263,3 @@ func TestString(t *testing.T) {
 		assert.Equal(t, test.output, s)
 	}
 }
-
-func TestHasPiece(t *testing.T) {
-	bf := Bitfield{0b01010100, 0b01010100}
-	outputs := []bool{false, true, false, true, false, true, false, false, false, true, false, true, false, true, false, false}
-	for i := 0; i < len(outputs); i++ {
-		assert.Equal(t, outputs[i], bf.HasPiece(i))
-	}
-}
-
-func TestSetPiece(t *testing.T) {
-	tests := []struct {
-		input Bitfield
-		index int
-		outpt Bitfield
-	}{
-		{
-			input: Bitfield{0b01010100, 0b01010100},
-			index: 4,
-			outpt: Bitfield{0b01011100, 0b01010100},
-		},
-		{
-			input: Bitfield{0b01010100, 0b01010100},
-			index: 9,
-			outpt: Bitfield{0b01010100, 0b01010100},
-		},
-		{
-			input: Bitfield{0b01010100, 0b01010100},
-			index: 15,
-			outpt: Bitfield{0b01010100, 0b01010101},
-		},
-	}
-	for _, test := range tests {
-		bf := test.input
-		bf.SetPiece(test.index)
-		assert.Equal(t, test.outpt, bf)
-	}
-}

p2p/client.go

@@ -7,6 +7,8 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/veggiedefender/torrent-client/bitfield"
+
 	"github.com/veggiedefender/torrent-client/message"
 
 	"github.com/veggiedefender/torrent-client/handshake"
@@ -18,7 +20,7 @@ type client struct {
 	peerID   [20]byte
 	conn     net.Conn
 	reader   *bufio.Reader
-	bitfield message.Bitfield
+	bitfield bitfield.Bitfield
 	choked   bool
 }
 
@@ -39,7 +41,7 @@ func completeHandshake(conn net.Conn, r *bufio.Reader, infohash, peerID [20]byte
 	return res, nil
 }
 
-func recvBitfield(conn net.Conn, r *bufio.Reader) (message.Bitfield, error) {
+func recvBitfield(conn net.Conn, r *bufio.Reader) (bitfield.Bitfield, error) {
 	conn.SetDeadline(time.Now().Local().Add(5 * time.Second))
 	defer conn.SetDeadline(time.Time{}) // Disable the deadline