包名重构

protocol/bt/peer/message/message.go

@@ -0,0 +1,96 @@
+package message
+
+import (
+	"encoding/binary"
+)
+
+// message ID
+type ID byte
+
+const (
+	IdChoke ID = iota
+	IdUnchoke
+	IdInterested
+	IdNotinterested
+	IdHave
+	IdBitfield
+	IdRequest
+	IdPiece
+	IdCancel
+)
+
+/*type Message struct {
+	Length  uint32
+	ID      MessageID
+	Payload []byte
+}*/
+type Message struct {
+	Length uint32
+	ID     ID
+}
+
+func (msg *Message) Encode() []byte {
+	buf := make([]byte, 5)
+	binary.BigEndian.PutUint32(buf, msg.Length)
+	buf[4] = byte(msg.ID)
+	return buf
+}
+
+func (msg *Message) Decode(buf []byte) {
+	msg.Length = binary.BigEndian.Uint32(buf)
+	msg.ID = ID(buf[4])
+}
+
+type Serialize interface {
+	Encode() []byte
+	Decode(buf []byte) Serialize
+}
+
+/*func NewMessage(id MessageID, payload []byte) *Message {
+	message := &Message{ID: id, Payload: payload}
+	if id == Keepalive {
+		message.Length = 0
+	} else {
+		message.Length = uint32(len(payload) + 1)
+	}
+	return message
+}
+
+func (msg *Message) Encode() []byte {
+	if msg.ID == Keepalive {
+		return make([]byte, 4)
+	} else {
+		buf := make([]byte, 5+len(msg.Payload))
+		binary.BigEndian.PutUint32(buf, msg.Length)
+		buf[4] = byte(msg.ID)
+		copy(buf[5:], msg.Payload[:])
+		return buf
+	}
+}
+
+func (msg *Message) Decode(buf []byte) {
+	msg.Length = binary.BigEndian.Uint32(buf)
+	if msg.Length == 0 {
+		msg.ID = -1
+	} else {
+		msg.ID = MessageID(buf[4])
+		if msg.Length > 4 {
+			msg.Payload = buf[5:]
+		}
+	}
+}*/
+
+// Message protocol
+// Length prefix| message ID | Payload
+// 4-byte       | 1-byte     | remaining bytes
+// 100          | 5          | 99-bytes
+// https://wiki.theory.org/index.php/BitTorrentSpecification#Messages
+func SplitMessage(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	if !atEOF && len(data) > 4 {
+		length := int(binary.BigEndian.Uint32(data))
+		if len(data)-4 >= length {
+			return length + 4, data[:length+4], nil
+		}
+	}
+	return
+}

protocol/bt/peer/message/msg_bitfield.go

@@ -0,0 +1,60 @@
+package message
+
+// bitfield: <len=0001+X><id=5><bitfield>
+type Bitfield struct {
+	*Message
+	payload []byte
+}
+
+func NewBitfield(payload []byte) *Bitfield {
+	return &Bitfield{
+		Message: &Message{
+			Length: uint32(len(payload) + 1),
+			ID:     IdBitfield,
+		},
+		payload: payload,
+	}
+}
+
+func (b *Bitfield) Encode() []byte {
+	return append(b.Message.Encode(), b.payload...)
+}
+
+func (b *Bitfield) Decode(buf []byte) Serialize {
+	b.Message = &Message{}
+	b.Message.Decode(buf)
+	b.payload = buf[5:]
+	return b
+}
+
+// 某个分片是否下载完成
+func (b *Bitfield) IsComplete(i int) bool {
+	index := i / 8
+	if index >= len(b.payload) {
+		return false
+	}
+	return b.payload[index]&(1<<(7-i%8)) > 0
+}
+
+// 给定一组分片下载状态，计算出当前peer能提供下载的分片下标
+func (b *Bitfield) Have(pieces []bool) []int {
+	arr := make([]int, 0)
+	length := len(pieces) / 8
+	if len(pieces)%8 != 0 {
+		length++
+	}
+
+	for i := 0; i < length; i++ {
+		for j := 0; j < 8; j++ {
+			index := i*8 + j
+			// 如果此分片在本地还未下载，检查peer是否能提供该分片下载
+			if index < len(pieces) &&
+				!pieces[index] &&
+				i < len(b.payload) &&
+				b.payload[i]&(1<<(7-j)) > 0 {
+				arr = append(arr, index)
+			}
+		}
+	}
+	return arr
+}

protocol/bt/peer/message/msg_bitfield_test.go

@@ -0,0 +1,162 @@
+package message
+
+import (
+	"reflect"
+	"testing"
+)
+
+func TestBitfield_IsComplete(t *testing.T) {
+	type fields struct {
+		Message Message
+		payload []byte
+	}
+	type args struct {
+		i int
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   bool
+	}{
+		{
+			"case-0",
+			fields{
+				payload: []byte{0b10000000},
+			},
+			args{
+				0,
+			},
+			true,
+		},
+		{
+			"case-1",
+			fields{
+				payload: []byte{0b01000000},
+			},
+			args{
+				1,
+			},
+			true,
+		},
+		{
+			"case-2",
+			fields{
+				payload: []byte{0b00000000, 0b00100000},
+			},
+			args{
+				2,
+			},
+			false,
+		},
+		{
+			"case-10",
+			fields{
+				payload: []byte{0b00000000, 0b00100000},
+			},
+			args{
+				10,
+			},
+			true,
+		},
+		{
+			"case-11",
+			fields{
+				payload: []byte{0b00000000, 0b00100000},
+			},
+			args{
+				11,
+			},
+			false,
+		},
+		{
+			"case-20",
+			fields{
+				payload: []byte{0b00000000, 0b00100000},
+			},
+			args{
+				20,
+			},
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b := &Bitfield{
+				Message: &tt.fields.Message,
+				payload: tt.fields.payload,
+			}
+			if got := b.IsComplete(tt.args.i); got != tt.want {
+				t.Errorf("Bitfield.IsComplete() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestBitfield_Have(t *testing.T) {
+	type fields struct {
+		Message Message
+		payload []byte
+	}
+	type args struct {
+		pieces []bool
+	}
+	tests := []struct {
+		name   string
+		fields fields
+		args   args
+		want   []int
+	}{
+		{
+			"case-1",
+			fields{
+				payload: []byte{0b11111111, 0b00000000},
+			},
+			args{
+				[]bool{false, false, false, false, false, false, false, false},
+			},
+			[]int{0, 1, 2, 3, 4, 5, 6, 7},
+		},
+		{
+			"case-2",
+			fields{
+				payload: []byte{0b01010101, 0b01000000},
+			},
+			args{
+				[]bool{false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false},
+			},
+			[]int{1, 3, 5, 7, 9},
+		},
+		{
+			"case-3",
+			fields{
+				payload: []byte{0b11010101, 0b01000000},
+			},
+			args{
+				[]bool{false, true, false},
+			},
+			[]int{0},
+		},
+		{
+			"case-4",
+			fields{
+				payload: []byte{0b11010101},
+			},
+			args{
+				[]bool{false, true, false, false, true, false, false, true, false},
+			},
+			[]int{0, 3, 5},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			b := &Bitfield{
+				Message: &tt.fields.Message,
+				payload: tt.fields.payload,
+			}
+			if got := b.Have(tt.args.pieces); !reflect.DeepEqual(got, tt.want) {
+				t.Errorf("Bitfield.Have() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

protocol/bt/peer/message/msg_choke.go

@@ -0,0 +1,24 @@
+package message
+
+// choke: <len=0001><id=0>
+type Choke struct {
+	Message
+}
+
+func NewChoke() *Choke {
+	return &Choke{
+		Message{
+			Length: 1,
+			ID:     IdChoke,
+		},
+	}
+}
+
+func (c *Choke) Encode() []byte {
+	return c.Message.Encode()
+}
+
+func (c *Choke) Decode(buf []byte) Serialize {
+	c.Message.Decode(buf)
+	return c
+}

protocol/bt/peer/message/msg_interested.go

@@ -0,0 +1,24 @@
+package message
+
+// choke: <len=0001><id=0>
+type Interested struct {
+	Message
+}
+
+func NewInterested() *Interested {
+	return &Interested{
+		Message{
+			Length: 1,
+			ID:     IdInterested,
+		},
+	}
+}
+
+func (c *Interested) Encode() []byte {
+	return c.Message.Encode()
+}
+
+func (c *Interested) Decode(buf []byte) Serialize {
+	c.Message.Decode(buf)
+	return c
+}

protocol/bt/peer/message/msg_keepalive.go

@@ -0,0 +1,12 @@
+package message
+
+// 	keep-alive: <len=0000>
+type Keepalive struct{}
+
+func (k *Keepalive) Encode() []byte {
+	return []byte{0, 0, 0, 0}
+}
+
+func (k *Keepalive) Decode(buf []byte) Serialize {
+	return &Keepalive{}
+}

protocol/bt/peer/message/msg_not_interested.go

@@ -0,0 +1,24 @@
+package message
+
+// choke: <len=0001><id=0>
+type NotInterested struct {
+	Message
+}
+
+func NewNotInterested() *NotInterested {
+	return &NotInterested{
+		Message{
+			Length: 1,
+			ID:     IdNotinterested,
+		},
+	}
+}
+
+func (ni *NotInterested) Encode() []byte {
+	return ni.Message.Encode()
+}
+
+func (ni *NotInterested) Decode(buf []byte) Serialize {
+	ni.Message.Decode(buf)
+	return ni
+}

protocol/bt/peer/message/msg_piece.go

@@ -0,0 +1,41 @@
+package message
+
+import "encoding/binary"
+
+// piece: <len=0009+X><id=7><index><begin><block>
+type Piece struct {
+	*Message
+	Index uint32
+	Begin uint32
+	Block []byte
+}
+
+func NewPiece(index uint32, begin uint32, block []byte) *Piece {
+	return &Piece{
+		Message: &Message{
+			Length: 9 + uint32(len(block)),
+			ID:     IdPiece,
+		},
+		Index: index,
+		Begin: begin,
+		// Block: io.LimitReader(reader, int64(length)),
+		Block: block,
+	}
+}
+
+func (p *Piece) Encode() []byte {
+	buf := make([]byte, 8)
+	binary.BigEndian.PutUint32(buf[:4], p.Index)
+	binary.BigEndian.PutUint32(buf[4:8], p.Begin)
+	buf = append(buf, p.Block...)
+	return append(p.Message.Encode(), buf...)
+}
+
+func (p *Piece) Decode(buf []byte) Serialize {
+	p.Message = &Message{}
+	p.Message.Decode(buf)
+	p.Index = binary.BigEndian.Uint32(buf[5:9])
+	p.Begin = binary.BigEndian.Uint32(buf[9:13])
+	p.Block = buf[13:]
+	return p
+}

protocol/bt/peer/message/msg_request.go

@@ -0,0 +1,41 @@
+package message
+
+import (
+	"encoding/binary"
+)
+
+// request: <len=0013><id=6><index><begin><length>
+type Request struct {
+	Message
+	Index  uint32
+	Begin  uint32
+	Length uint32
+}
+
+func NewRequest(index uint32, begin uint32, length uint32) *Request {
+	return &Request{
+		Message: Message{
+			Length: 13,
+			ID:     IdRequest,
+		},
+		Index:  index,
+		Begin:  begin,
+		Length: length,
+	}
+}
+
+func (r *Request) Encode() []byte {
+	buf := make([]byte, 12)
+	binary.BigEndian.PutUint32(buf[0:4], r.Index)
+	binary.BigEndian.PutUint32(buf[4:8], r.Begin)
+	binary.BigEndian.PutUint32(buf[8:12], r.Length)
+	return append(r.Message.Encode(), buf...)
+}
+
+func (r *Request) Decode(buf []byte) Serialize {
+	r.Message.Decode(buf)
+	r.Length = binary.BigEndian.Uint32(buf[5:9])
+	r.Begin = binary.BigEndian.Uint32(buf[9:13])
+	r.Index = binary.BigEndian.Uint32(buf[13:17])
+	return r
+}

protocol/bt/peer/message/msg_unchoke.go

@@ -0,0 +1,24 @@
+package message
+
+// 	unchoke: <len=0001><id=1>
+type Unchoke struct {
+	Message
+}
+
+func NewUnchoke() *Unchoke {
+	return &Unchoke{
+		Message{
+			Length: 1,
+			ID:     IdUnchoke,
+		},
+	}
+}
+
+func (u *Unchoke) Encode() []byte {
+	return u.Message.Encode()
+}
+
+func (u *Unchoke) Decode(buf []byte) Serialize {
+	u.Message.Decode(buf)
+	return u
+}