diff --git a/bitfield/bitfield.go b/bitfield/bitfield.go new file mode 100644 index 0000000..264a995 --- /dev/null +++ b/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) +} diff --git a/bitfield/bitfield_test.go b/bitfield/bitfield_test.go new file mode 100644 index 0000000..31afe17 --- /dev/null +++ b/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) + } +} diff --git a/message/message.go b/message/message.go index 5de1cbc..ac0f0f3 100644 --- a/message/message.go +++ b/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) -} diff --git a/message/message_test.go b/message/message_test.go index cbc7fd0..c65adef 100644 --- a/message/message_test.go +++ b/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) - } -} diff --git a/p2p/client.go b/p2p/client.go index e4aac6c..0914a50 100644 --- a/p2p/client.go +++ b/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