From d6d411f48c1a1e503aea4ef8b5b0249861d54b0e Mon Sep 17 00:00:00 2001 From: liwei Date: Thu, 23 Apr 2020 09:56:17 +0800 Subject: [PATCH] refactor: message entity --- protocol/bt/metainfo/metainfo_test.go | 1 + protocol/bt/peer/message/message.go | 83 +++++++++++++++---- protocol/bt/peer/message/msg_bitfield.go | 54 ++++++------ protocol/bt/peer/message/msg_bitfield_test.go | 12 ++- protocol/bt/peer/message/msg_choke.go | 16 +--- protocol/bt/peer/message/msg_interested.go | 18 +--- protocol/bt/peer/message/msg_keepalive.go | 12 ++- .../bt/peer/message/msg_not_interested.go | 18 +--- protocol/bt/peer/message/msg_piece.go | 45 +++++----- protocol/bt/peer/message/msg_request.go | 32 +++---- protocol/bt/peer/message/msg_unchoke.go | 16 +--- protocol/bt/torrent/peer_conn.go | 11 ++- 12 files changed, 167 insertions(+), 151 deletions(-) diff --git a/protocol/bt/metainfo/metainfo_test.go b/protocol/bt/metainfo/metainfo_test.go index 2d8b99f0..aa1ab2ba 100644 --- a/protocol/bt/metainfo/metainfo_test.go +++ b/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[:])) } diff --git a/protocol/bt/peer/message/message.go b/protocol/bt/peer/message/message.go index b798f87f..0fc50fab 100644 --- a/protocol/bt/peer/message/message.go +++ b/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 } diff --git a/protocol/bt/peer/message/msg_bitfield.go b/protocol/bt/peer/message/msg_bitfield.go index ee4b1199..54ab50c1 100644 --- a/protocol/bt/peer/message/msg_bitfield.go +++ b/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: 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() } diff --git a/protocol/bt/peer/message/msg_bitfield_test.go b/protocol/bt/peer/message/msg_bitfield_test.go index b9ff7479..7e0e8c48 100644 --- a/protocol/bt/peer/message/msg_bitfield_test.go +++ b/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) diff --git a/protocol/bt/peer/message/msg_choke.go b/protocol/bt/peer/message/msg_choke.go index ed069aa1..e7446914 100644 --- a/protocol/bt/peer/message/msg_choke.go +++ b/protocol/bt/peer/message/msg_choke.go @@ -2,23 +2,13 @@ package message // choke: 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 -} diff --git a/protocol/bt/peer/message/msg_interested.go b/protocol/bt/peer/message/msg_interested.go index e14109b0..ab17bec7 100644 --- a/protocol/bt/peer/message/msg_interested.go +++ b/protocol/bt/peer/message/msg_interested.go @@ -1,24 +1,14 @@ package message -// choke: +// choke: 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 -} diff --git a/protocol/bt/peer/message/msg_keepalive.go b/protocol/bt/peer/message/msg_keepalive.go index 31c8157d..14633c53 100644 --- a/protocol/bt/peer/message/msg_keepalive.go +++ b/protocol/bt/peer/message/msg_keepalive.go @@ -1,12 +1,16 @@ package message // keep-alive: -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) { + } diff --git a/protocol/bt/peer/message/msg_not_interested.go b/protocol/bt/peer/message/msg_not_interested.go index 27d1182a..cbb9eab4 100644 --- a/protocol/bt/peer/message/msg_not_interested.go +++ b/protocol/bt/peer/message/msg_not_interested.go @@ -1,24 +1,14 @@ package message -// choke: +// choke: 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 -} diff --git a/protocol/bt/peer/message/msg_piece.go b/protocol/bt/peer/message/msg_piece.go index 695dee08..182b088d 100644 --- a/protocol/bt/peer/message/msg_piece.go +++ b/protocol/bt/peer/message/msg_piece.go @@ -4,38 +4,39 @@ import "encoding/binary" // piece: 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:] } diff --git a/protocol/bt/peer/message/msg_request.go b/protocol/bt/peer/message/msg_request.go index fd9e3a8e..88165921 100644 --- a/protocol/bt/peer/message/msg_request.go +++ b/protocol/bt/peer/message/msg_request.go @@ -6,36 +6,38 @@ import ( // request: 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]) } diff --git a/protocol/bt/peer/message/msg_unchoke.go b/protocol/bt/peer/message/msg_unchoke.go index cfc0d41a..2fd6e994 100644 --- a/protocol/bt/peer/message/msg_unchoke.go +++ b/protocol/bt/peer/message/msg_unchoke.go @@ -2,23 +2,13 @@ package message // unchoke: 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 -} diff --git a/protocol/bt/torrent/peer_conn.go b/protocol/bt/torrent/peer_conn.go index 11015016..b0349d59 100644 --- a/protocol/bt/torrent/peer_conn.go +++ b/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()