refactor: message entity

protocol/bt/metainfo/metainfo_test.go

@@ -16,6 +16,7 @@ func TestParseFromFile(t *testing.T) {
 	fmt.Printf("%+v\n", metaInfo.AnnounceList)
 	fmt.Printf("%+v\n", metaInfo.Info.Files)
 	fmt.Printf("%+v\n", metaInfo.Info.PieceLength)
+	fmt.Printf("%+v\n", len(metaInfo.Info.Pieces))
 	fmt.Printf("%+v\n", hex.EncodeToString(metaInfo.infoHash[:]))
 }
 

protocol/bt/peer/message/message.go

@@ -11,7 +11,7 @@ const (
 	IdChoke ID = iota
 	IdUnchoke
 	IdInterested
-	IdNotinterested
+	IdNotInterested
 	IdHave
 	IdBitfield
 	IdRequest
@@ -24,34 +24,81 @@ const (
 	ID      MessageID
 	Payload []byte
 }*/
-type Message struct {
-	Length uint32
-	ID     ID
+
+type base struct {
+	id ID
 }
 
-func (msg *Message) Encode() []byte {
-	buf := make([]byte, 5)
-	binary.BigEndian.PutUint32(buf, msg.Length)
-	buf[4] = byte(msg.ID)
+func (base *base) ID() ID {
+	return base.id
+}
+
+func (base *base) Encode() []byte {
+	return encode(base, nil)
+}
+
+func (base *base) Decode(body []byte) {
+
+}
+
+type Message interface {
+	ID() ID
+	Encode() []byte
+	Decode(body []byte)
+}
+
+type TestMsg struct {
+	*base
+}
+
+func (base *TestMsg) Encode() []byte {
+	return []byte{1, 2, 3}
+}
+
+func (base *TestMsg) Decode(body []byte) {
+}
+
+func encode(base *base, body []byte) []byte {
+	buf := make([]byte, 5+len(body))
+	binary.BigEndian.PutUint32(buf, uint32(len(buf)))
+	buf[4] = byte(base.id)
+	copy(buf[4:], body)
 	return buf
 }
 
-func (msg *Message) Decode(buf []byte) {
+/*func Decode(buf []byte) Message {
+	head := make([]byte, 5)
+	body := message.Encode()
+	binary.BigEndian.PutUint32(head, uint32(len(head)+len(body)))
+	head[4] = byte(message.ID())
+	return append(head, body...)
+}
+
+func (msg *Message) ToBytes() []byte {
+	buf := make([]byte, 5)
+	binary.BigEndian.PutUint32(buf, msg.Length)
+	buf[4] = byte(msg.ID)
+	if msg.Payload != nil {
+		return append(buf, msg.Payload.Encode()...)
+	} else {
+		return buf
+	}
+}
+
+func (msg *Message) FormBytes(buf []byte) {
 	msg.Length = binary.BigEndian.Uint32(buf)
 	msg.ID = ID(buf[4])
-}
+	if msg.Payload != nil {
+		msg.Payload.Decode(buf[4:])
+	}
+}*/
 
-type Serialize interface {
-	Encode() []byte
-	Decode(buf []byte) Serialize
-}
-
-/*func NewMessage(id MessageID, payload []byte) *Message {
-	message := &Message{ID: id, Payload: payload}
+/*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)
+		message.Length = uint32(len(Payload) + 1)
 	}
 	return message
 }

protocol/bt/peer/message/msg_bitfield.go

@@ -1,58 +1,62 @@
 package message
 
-import "github.com/RoaringBitmap/roaring"
+import (
+	"github.com/RoaringBitmap/roaring"
+	"math"
+)
 
 // bitfield: <len=0001+X><id=5><bitfield>
 type Bitfield struct {
-	*Message
-	payload *roaring.Bitmap
+	*base
+	pieceCount int
+	bitmap     *roaring.Bitmap
 }
 
-func NewBitfield(payload *roaring.Bitmap) *Bitfield {
+func NewBitfield() *Bitfield {
 	return &Bitfield{
-		Message: &Message{
-			// TODO 长度处理
-			Length: 1,
-			ID:     IdBitfield,
+		base: &base{
+			id: IdBitfield,
 		},
-		payload: payload,
 	}
 }
 
+func BuildBitfield(pieceCount int, bitmap *roaring.Bitmap) *Bitfield {
+	bitfield := NewBitfield()
+	bitfield.pieceCount = pieceCount
+	bitfield.bitmap = bitmap
+	return bitfield
+}
+
 func (b *Bitfield) Encode() []byte {
-	buf := b.Message.Encode()
-	for i := 0; i < int(b.Length)-len(buf); i++ {
+	buf := make([]byte, int(math.Ceil(float64(b.pieceCount)/8)))
+	for i := 0; i < len(buf); i++ {
 		temp := byte(0)
 		for j := 0; j < 8; j++ {
 			// 如果piece存在,位数标识为1
-			if b.payload.ContainsInt(i*8 + j) {
+			if b.bitmap.ContainsInt(i*8 + j) {
 				temp = temp | (1 << (7 - j))
 			}
 		}
-		buf = append(buf, temp)
+		buf[i] = temp
 	}
-	return buf
+	return encode(b.base, buf)
 }
 
-func (b *Bitfield) Decode(buf []byte) Serialize {
-	b.Message = &Message{}
-	b.Message.Decode(buf)
-	temp := buf[5:]
-	b.payload = roaring.New()
+func (b *Bitfield) Decode(body []byte) {
+	b.bitmap = roaring.New()
 	// 初始化bitmap
-	for i := 0; i < len(temp); i++ {
+	for i := 0; i < len(body); i++ {
 		for j := 0; j < 8; j++ {
-			if temp[i]&(1<<(7-j)) > 0 {
-				b.payload.AddInt(i*8 + j)
+			if body[i]&(1<<(7-j)) > 0 {
+				b.bitmap.AddInt(i*8 + j)
 			}
 		}
 	}
-	return b
 }
 
 // 某个分片是否下载完成
 func (b *Bitfield) IsComplete(i int) bool {
-	return b.payload.ContainsInt(i)
+	return b.bitmap.ContainsInt(i)
 }
 
 // 给定一组分片下载状态，计算出当前peer能提供下载的分片下标，如果已经拥有则不用提供
@@ -60,6 +64,6 @@ func (b *Bitfield) IsComplete(i int) bool {
 // had & has = 0b10010000 & 0b01111111 = 0b00010000
 // had & has ^ has = 0b00010000 ^ 0b01111111 = 0b01101111
 func (b *Bitfield) Provide(had *roaring.Bitmap) []uint32 {
-	xor := roaring.Xor(roaring.And(had, b.payload), b.payload)
+	xor := roaring.Xor(roaring.And(had, b.bitmap), b.bitmap)
 	return xor.ToArray()
 }

protocol/bt/peer/message/msg_bitfield_test.go

@@ -8,7 +8,6 @@ import (
 
 func TestBitfield_IsComplete(t *testing.T) {
 	type fields struct {
-		Message Message
 		payload *roaring.Bitmap
 	}
 	type args struct {
@@ -44,8 +43,8 @@ func TestBitfield_IsComplete(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			b := &Bitfield{
-				Message: &tt.fields.Message,
-				payload: tt.fields.payload,
+				pieceCount: 100,
+				bitmap:     tt.fields.payload,
 			}
 			if got := b.IsComplete(tt.args.i); got != tt.want {
 				t.Errorf("Bitfield.IsComplete() = %v, want %v", got, tt.want)
@@ -54,9 +53,8 @@ func TestBitfield_IsComplete(t *testing.T) {
 	}
 }
 
-func TestBitfield_Have(t *testing.T) {
+func TestBitfield_Provide(t *testing.T) {
 	type fields struct {
-		Message Message
 		payload *roaring.Bitmap
 	}
 	type args struct {
@@ -102,8 +100,8 @@ func TestBitfield_Have(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			b := &Bitfield{
-				Message: &tt.fields.Message,
-				payload: tt.fields.payload,
+				pieceCount: 100,
+				bitmap:     tt.fields.payload,
 			}
 			if got := b.Provide(tt.args.had); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("Bitfield.Provide() = %v, want %v", got, tt.want)

protocol/bt/peer/message/msg_choke.go

@@ -2,23 +2,13 @@ package message
 
 // choke: <len=0001><id=0>
 type Choke struct {
-	Message
+	*base
 }
 
 func NewChoke() *Choke {
 	return &Choke{
-		Message{
-			Length: 1,
-			ID:     IdChoke,
+		base: &base{
+			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

@@ -1,24 +1,14 @@
 package message
 
-// choke: <len=0001><id=0>
+// choke: <len=0001><id=2>
 type Interested struct {
-	Message
+	*base
 }
 
 func NewInterested() *Interested {
 	return &Interested{
-		Message{
-			Length: 1,
-			ID:     IdInterested,
+		base: &base{
+			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

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

protocol/bt/peer/message/msg_not_interested.go

@@ -1,24 +1,14 @@
 package message
 
-// choke: <len=0001><id=0>
+// choke: <len=0001><id=3>
 type NotInterested struct {
-	Message
+	*base
 }
 
 func NewNotInterested() *NotInterested {
 	return &NotInterested{
-		Message{
-			Length: 1,
-			ID:     IdNotinterested,
+		base: &base{
+			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

@@ -4,38 +4,39 @@ import "encoding/binary"
 
 // piece: <len=0009+X><id=7><index><begin><block>
 type Piece struct {
-	*Message
+	*base
 	Index uint32
 	Begin uint32
 	Block []byte
 }
 
-func NewPiece(index uint32, begin uint32, block []byte) *Piece {
+func NewPiece() *Piece {
 	return &Piece{
-		Message: &Message{
-			Length: 9 + uint32(len(block)),
-			ID:     IdPiece,
+		base: &base{
+			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 BuildPiece(index uint32, begin uint32, block []byte) *Piece {
+	piece := NewPiece()
+	piece.Index = index
+	piece.Begin = begin
+	piece.Block = block
+	return piece
+
 }
 
-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
+func (p *Piece) Encode() []byte {
+	buf := make([]byte, 8+len(p.Block))
+	binary.BigEndian.PutUint32(buf[0:4], p.Index)
+	binary.BigEndian.PutUint32(buf[4:8], p.Begin)
+	copy(buf[8:], p.Block)
+	return encode(p.base, buf)
+}
+
+func (p *Piece) Decode(buf []byte) {
+	p.Index = binary.BigEndian.Uint32(buf[0:4])
+	p.Begin = binary.BigEndian.Uint32(buf[4:8])
+	p.Block = buf[8:]
 }

protocol/bt/peer/message/msg_request.go

@@ -6,36 +6,38 @@ import (
 
 // request: <len=0013><id=6><index><begin><length>
 type Request struct {
-	Message
+	*base
 	Index  uint32
 	Begin  uint32
 	Length uint32
 }
 
-func NewRequest(index uint32, begin uint32, length uint32) *Request {
+func NewRequest() *Request {
 	return &Request{
-		Message: Message{
-			Length: 13,
-			ID:     IdRequest,
+		base: &base{
+			id: IdRequest,
 		},
-		Index:  index,
-		Begin:  begin,
-		Length: length,
 	}
 }
 
+func BuildRequest(index uint32, begin uint32, length uint32) *Request {
+	request := NewRequest()
+	request.Index = index
+	request.Begin = begin
+	request.Length = length
+	return request
+}
+
 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...)
+	return encode(r.base, 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
+func (r *Request) Decode(buf []byte) {
+	r.Length = binary.BigEndian.Uint32(buf[0:4])
+	r.Begin = binary.BigEndian.Uint32(buf[4:8])
+	r.Index = binary.BigEndian.Uint32(buf[8:12])
 }

protocol/bt/peer/message/msg_unchoke.go

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

protocol/bt/torrent/peer_conn.go

@@ -155,8 +155,7 @@ func (pc *peerConn) ready() error {
 			if length == 0 {
 				// 	keepalive
 			} else {
-				messageID := message.ID(buf[4])
-				switch messageID {
+				switch message.ID(buf[4]) {
 				case message.IdChoke:
 					break
 				case message.IdUnchoke:
@@ -164,7 +163,7 @@ func (pc *peerConn) ready() error {
 					break
 				case message.IdInterested:
 					break
-				case message.IdNotinterested:
+				case message.IdNotInterested:
 					break
 				case message.IdHave:
 					break
@@ -245,7 +244,7 @@ func (pc *peerConn) downloadPiece(index int) (err error) {
 			return
 		}
 		// 发起request，对方会响应piece
-		_, err = pc.conn.Write(message.NewRequest(uint32(index), uint32(offset), blockLength).Encode())
+		_, err = pc.conn.Write(message.BuildRequest(uint32(index), uint32(offset), blockLength).Encode())
 		if err != nil {
 			break
 		}
@@ -279,7 +278,7 @@ func (pc *peerConn) handleUnchoke(readyCh chan<- bool) {
 }
 
 func (pc *peerConn) handleBitfield(buf []byte) {
-	pc.bitfield = &message.Bitfield{}
+	pc.bitfield = message.NewBitfield()
 	pc.bitfield.Decode(buf)
 	have := pc.getHavePieces(pc.bitfield)
 	if len(have) > 0 {
@@ -296,7 +295,7 @@ func (pc *peerConn) handleBitfield(buf []byte) {
 
 // 处理下载响应，每次接收到响应直接将block写入到对应文件中
 func (pc *peerConn) handlePiece(buf []byte) {
-	piece := &message.Piece{}
+	piece := message.NewPiece()
 	piece.Decode(buf)
 	info := pc.torrent.MetaInfo.Info
 	fds := pc.torrent.MetaInfo.GetFileDetails()