From d6c88d0c4bf60c759cf38d09a0af98e106d70f4e Mon Sep 17 00:00:00 2001 From: liwei Date: Sat, 6 Jul 2019 14:35:19 +0800 Subject: [PATCH] bitfield handle --- down/bt/peer/message.go | 118 +++++++---- down/bt/peer/message_test.go | 63 +++++- down/bt/peer/peer.go.bak | 316 ----------------------------- down/bt/peer/peer_test.go.bak | 41 ---- down/bt/torrent/peer_state.go | 56 ++++- down/bt/torrent/peer_state_test.go | 40 ++-- down/bt/torrent/torrent.go | 5 - 7 files changed, 216 insertions(+), 423 deletions(-) delete mode 100644 down/bt/peer/peer.go.bak delete mode 100644 down/bt/peer/peer_test.go.bak diff --git a/down/bt/peer/message.go b/down/bt/peer/message.go index 0e710297..54706b3c 100644 --- a/down/bt/peer/message.go +++ b/down/bt/peer/message.go @@ -4,32 +4,66 @@ import ( "encoding/binary" ) -type MessageType int +type MessageID byte const ( - Keepalive MessageType = -1 - Choke MessageType = iota - Unchoke - Interested - NotInterested - Have - Bitfield - Request - Piece - Cancel + Keepalive MessageID = -1 + Choke MessageID = 0 + Unchoke MessageID = 1 + Interested MessageID = 2 + NotInterested MessageID = 3 + Have MessageID = 4 + Bitfield MessageID = 5 + Request MessageID = 6 + Piece MessageID = 7 + Cancel MessageID = 8 ) -// 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 type Message struct { Length uint32 - Type int + ID MessageID Payload []byte } +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) + buf = append(buf, 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)) @@ -40,34 +74,10 @@ func SplitMessage(data []byte, atEOF bool) (advance int, token []byte, err error return } -func (msg *Message) Encode() []byte { - if msg.Type == -1 { - return make([]byte, 4) - } else { - buf := make([]byte, 5+len(msg.Payload)) - binary.BigEndian.PutUint32(buf, msg.Length) - buf[4] = byte(msg.Type) - buf = append(buf, msg.Payload[:]...) - return buf - } -} - -func (msg *Message) Decode(buf []byte) { - msg.Length = binary.BigEndian.Uint32(buf) - if msg.Length == 0 { - msg.Type = -1 - } else { - msg.Type = int(buf[4]) - if msg.Length > 4 { - msg.Payload = buf[5:] - } - } -} - type MsgBitfield []byte // 某个分片是否下载完成 -func (mb MsgBitfield) get(i int) bool { +func (mb MsgBitfield) IsComplete(i int) bool { bts := []byte(mb) index := i / 8 if index >= len(bts) { @@ -75,3 +85,27 @@ func (mb MsgBitfield) get(i int) bool { } return bts[index]&(1< 0 } + +// 给定一组分片下载状态，计算出当前peer能提供下载的分片下标 +func (mb MsgBitfield) Have(pieces []bool) []int { + arr := make([]int, 0) + bts := []byte(mb) + 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(bts) && + uint(bts[i])&(1< 0 { + arr = append(arr, index) + } + } + } + return arr +} diff --git a/down/bt/peer/message_test.go b/down/bt/peer/message_test.go index 019e51fa..0dd64576 100644 --- a/down/bt/peer/message_test.go +++ b/down/bt/peer/message_test.go @@ -1,11 +1,13 @@ package peer import ( + "fmt" + "reflect" "strconv" "testing" ) -func TestMsgBitfield_get(t *testing.T) { +func TestMsgBitfield_IsComplete(t *testing.T) { type args struct { i int } @@ -66,13 +68,70 @@ func TestMsgBitfield_get(t *testing.T) { } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - if got := tt.mb.get(tt.args.i); got != tt.want { + if got := tt.mb.IsComplete(tt.args.i); got != tt.want { t.Errorf("MsgBitfield.get() = %v, want %v", got, tt.want) } }) } } +func TestMsgBitfield_Have(t *testing.T) { + type args struct { + pieces []bool + } + tests := []struct { + name string + mb MsgBitfield + args args + want []int + }{ + { + "case-1", + MsgBitfield([]byte{convertByte("11111111"), convertByte("00000000")}), + args{ + []bool{false, false, false, false, false, false, false, false}, + }, + []int{0, 1, 2, 3, 4, 5, 6, 7}, + }, + { + "case-2", + MsgBitfield([]byte{convertByte("01010101"), convertByte("01000000")}), + 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", + MsgBitfield([]byte{convertByte("11010101"), convertByte("01000000")}), + args{ + []bool{false, true, false}, + }, + []int{0}, + }, + { + "case-4", + MsgBitfield([]byte{convertByte("11010101")}), + 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) { + if got := tt.mb.Have(tt.args.pieces); !reflect.DeepEqual(got, tt.want) { + t.Errorf("MsgBitfield.Have() = %v, want %v", got, tt.want) + } + }) + } +} + +func TestMsg(t *testing.T) { + fmt.Println(Choke) + fmt.Println(Bitfield) +} + func convertByte(bit string) byte { i, _ := strconv.ParseUint(bit, 2, 8) return byte(i) diff --git a/down/bt/peer/peer.go.bak b/down/bt/peer/peer.go.bak deleted file mode 100644 index 3142c056..00000000 --- a/down/bt/peer/peer.go.bak +++ /dev/null @@ -1,316 +0,0 @@ -package peer - -import ( - "bufio" - "bytes" - "encoding/binary" - "fmt" - "github.com/cenkalti/mse" - "gopeed/down/bt/metainfo" - "io" - "net" - "strconv" -) - -type MessageType byte - -const ( - ProtocolIdentifier = "BitTorrent protocol" - ProtocolIdentifierLength = 0x13 - - Choke MessageType = iota - Unchoke - Interested - NotInterested - Have - Bitfield - Request - Piece - Cancel -) - -type Peer struct { - IP uint32 - Port uint16 - - Torrent *Torrent - conn net.Conn - - // this client is choking the peer - AmChoking bool - // this client is interested in the peer - AmInterested bool - // peer is choking this client - PeerChoking bool - // peer is interested in this client - PeerInterested bool -} - -func (peer *Peer) Address() string { - bts := make([]byte, 4) - binary.BigEndian.PutUint32(bts, peer.IP) - return net.IP(bts).String() + ":" + strconv.Itoa(int(peer.Port)) -} - -type Handshake struct { - Pstrlen byte - Pstr [19]byte - Reserved [8]byte - InfoHash [20]byte - PeerID [20]byte -} - -type Message struct { - Length uint32 - Type byte - Payload []byte -} - -type MessageBitfield []bool - -func decodeBitfield(payload []byte) MessageBitfield { - bits := make([]bool, len(payload)*8) - for i, b := range payload { - for j := 0; j < 8; j++ { - bits[i*8+j] = b&(1< 0 - } - } - return bits -} - -func newHandshake(reserved [8]byte, infoHash [20]byte, peerID [20]byte) *Handshake { - var arr [ProtocolIdentifierLength]byte - copy(arr[:], ProtocolIdentifier) - return &Handshake{ - Pstrlen: ProtocolIdentifierLength, - Pstr: arr, - Reserved: reserved, - InfoHash: infoHash, - PeerID: peerID, - } -} - -func (handshake *Handshake) encode() ([]byte, error) { - writer := new(bytes.Buffer) - err := binary.Write(writer, binary.BigEndian, handshake) - if err != nil { - return nil, err - } - return writer.Bytes(), nil -} - -func (handshake *Handshake) decode(buf []byte) error { - reader := bytes.NewReader(buf) - err := binary.Read(reader, binary.BigEndian, handshake) - if err != nil { - return err - } - return nil -} - -func dialMse(address string, infoHash [20]byte) (net.Conn, error) { - conn, err := net.Dial("tcp", address) - if err != nil { - return nil, err - } - mseConn := mse.WrapConn(conn) - _, err = mseConn.HandshakeOutgoing(infoHash[:], mse.PlainText, nil) - if err != nil { - return nil, err - } - return mseConn, nil -} - -func (peer *Peer) DoDownload() error { - if peer.Port == 0 { - return fmt.Errorf("error port %d", peer.Port) - } - conn, err := dialMse(peer.Address(), peer.Torrent.MetaInfo.InfoHash) - if err != nil { - return err - } - defer conn.Close() - err = peer.doHandshake(conn) - if err != nil { - return err - } - // handshake success,init status - peer.AmChoking = true - peer.AmInterested = false - peer.PeerChoking = true - peer.PeerInterested = false - - scanner := bufio.NewScanner(conn) - scanner.Split(splitMessage) - for scanner.Scan() { - buf := scanner.Bytes() - message := decodeMessage(buf) - // keep-alive message - if message.Length == 0 { - - } else { - switch MessageType(buf[4]) { - case Choke: - peer.PeerChoking = true - break - case Unchoke: - peer.PeerChoking = false - break - case Interested: - peer.PeerInterested = true - break - case NotInterested: - peer.PeerInterested = false - break - case Have: - break - case Bitfield: - bitfield := decodeBitfield(message.Payload) - fmt.Println(bitfield) - break - case Request: - break - case Piece: - break - case Cancel: - break - } - } - } - - /*var ( - // This client is choking the peer - amChocking = true - // This client is interested in the peer - amInterested = false - // Peer is choking this client - peerChocking = true - // Peer is interested in this client - peerInterested = false - ) - - // Receive message - var cacheBuf []byte - buf := make([]byte, 8192) - for { - scanner := bufio.NewScanner(conn) - n, err := conn.Read(buf) - if n > 0 { - if len(cacheBuf) > 0 { - buf = append(cacheBuf, buf...) - } - msg, err := decodeMessage(buf, n) - if err != nil { - if err == NMD { - cacheBuf = make([]byte, n) - copy(cacheBuf, buf[0:n]) - } - return err - } - - msgLength := binary.BigEndian.Uint32(buf) - // keep-alive message - if msgLength == 0 { - - } else { - switch MessageType(buf[4]) { - case Choke: - peerChocking = true - break - case Unchoke: - peerChocking = false - break - case Interested: - peerInterested = true - break - case NotInterested: - peerInterested = false - break - case Have: - break - case Bitfield: - break - case Request: - break - case Piece: - break - case Cancel: - break - - } - - } - - } - if err != nil { - if err == io.EOF { - return nil - } - return err - } - }*/ - return nil -} - -// Handshake of Peer wire protocol -// Per https://wiki.theory.org/index.php/BitTorrentSpecification#Handshake -func (peer *Peer) doHandshake(conn net.Conn) error { - metaInfo := peer.Torrent.MetaInfo - peerID := peer.Torrent.client.PeerID - - reserved := [8]byte{} - /* reserved[5] = 0x10 - reserved[6] = 0x0 - reserved[7] = 0x5*/ - handshakeReq := newHandshake(reserved, metaInfo.InfoHash, peerID) - buf, err := handshakeReq.encode() - if err != nil { - return err - } - conn.Write(buf) - - var read [68]byte - _, err = io.ReadFull(conn, read[:]) - if err != nil { - return err - } - handshakeRes := new(Handshake) - err = handshakeRes.decode(read[:]) - if err != nil { - return err - } - if handshakeRes.InfoHash != handshakeReq.InfoHash { - return err - } - return nil -} - -// 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 -} - -func decodeMessage(buf []byte) *Message { - msg := &Message{} - msg.Length = binary.BigEndian.Uint32(buf) - // keep-alive message - if msg.Length > 0 { - msg.Type = buf[4] - msg.Payload = buf[5:] - } - return msg -} - -func doBitfield(peer *Peer, metaInfo *metainfo.MetaInfo, conn net.Conn) { - -} diff --git a/down/bt/peer/peer_test.go.bak b/down/bt/peer/peer_test.go.bak deleted file mode 100644 index 2b797c7a..00000000 --- a/down/bt/peer/peer_test.go.bak +++ /dev/null @@ -1,41 +0,0 @@ -package peer - -import ( - "fmt" - "reflect" - "testing" -) - -func Test_decodeBitfield(t *testing.T) { - fmt.Printf("%08b", 0xFF) - type args struct { - payload []byte - } - tests := []struct { - name string - args args - want MessageBitfield - }{ - { - "case-11111111", - args{ - []byte{0xFF}, - }, - []bool{true, true, true, true, true, true, true, true}, - }, - { - "case-10101010", - args{ - []byte{0xAA}, - }, - []bool{true, false, true, false, true, false, true, false}, - }, - } - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - if got := decodeBitfield(tt.args.payload); !reflect.DeepEqual(got, tt.want) { - t.Errorf("decodeBitfield() = %v, want %v", got, tt.want) - } - }) - } -} diff --git a/down/bt/torrent/peer_state.go b/down/bt/torrent/peer_state.go index 565dc864..fc70e8d8 100644 --- a/down/bt/torrent/peer_state.go +++ b/down/bt/torrent/peer_state.go @@ -21,6 +21,8 @@ type peerState struct { peerChoking bool // peer is interested in this client peerInterested bool + + bitfield peer.MsgBitfield } // 使用MSE加密来避免运营商对bt流量的封锁，基本上现在市面上BT客户端都默认开启了，不用MSE的话很多Peer拒绝连接 @@ -85,14 +87,14 @@ func (ps *peerState) download() error { buf := scanner.Bytes() message := &peer.Message{} message.Decode(buf) - switch peer.MessageType(message.Type) { + switch message.ID { case peer.Keepalive: break case peer.Choke: ps.peerChoking = true break case peer.Unchoke: - ps.peerChoking = false + ps.handleUnchoke() break case peer.Interested: ps.peerInterested = true @@ -103,9 +105,7 @@ func (ps *peerState) download() error { case peer.Have: break case peer.Bitfield: - bitfield := peer.MsgBitfield(message.Payload) - // 判断peer能提供需要下载的文件分片 - + ps.handleBitfield(message.Payload) break case peer.Request: break @@ -117,3 +117,49 @@ func (ps *peerState) download() error { } return nil } + +func (ps *peerState) handleUnchoke() { + ps.peerChoking = false + // 如果客户端对peer感兴趣并且peer没有choked客户端，就可以开始下载了 + if ps.amInterested { + have := ps.getHavePieces(ps.bitfield) + if len(have) > 0 { + /*for _, index := range have { + ps.conn.Write((&peer.Message{13, peer.Request,})) + }*/ + return + } + } + ps.conn.Close() +} + +func (ps *peerState) handleBitfield(payload []byte) { + ps.bitfield = payload + have := ps.getHavePieces(ps.bitfield) + if len(have) > 0 { + // 表示对该peer感兴趣，并且不choked该peer + _, err := ps.conn.Write(peer.NewMessage(peer.Interested, []byte{}).Encode()) + if err != nil { + ps.conn.Close() + return + } + ps.amInterested = true + _, err = ps.conn.Write(peer.NewMessage(peer.Unchoke, []byte{}).Encode()) + if err != nil { + ps.conn.Close() + return + } + ps.amChoking = false + } else { + ps.conn.Close() + } +} + +// 获取peer能提供需要下载的文件分片 +func (ps *peerState) getHavePieces(bitfield peer.MsgBitfield) []int { + states := make([]bool, len(ps.torrent.PieceState)) + for i := range states { + states[i] = ps.torrent.PieceState[i].complete + } + return bitfield.Have(states) +} diff --git a/down/bt/torrent/peer_state_test.go b/down/bt/torrent/peer_state_test.go index dc60bd20..f23792ec 100644 --- a/down/bt/torrent/peer_state_test.go +++ b/down/bt/torrent/peer_state_test.go @@ -9,17 +9,8 @@ import ( ) func Test_peerState_handshake(t *testing.T) { - torrent := buildTorrent() - tracker := &tracker.Tracker{ - PeerID: torrent.PeerID, - MetaInfo: torrent.MetaInfo, - } - peers, err := tracker.Tracker() - if err != nil { - panic(err) - } - fmt.Println("Tracker end") - ps := getUsablePeer(torrent, peers) + + ps := getUsablePeer() // 发出握手请求 handshake, err := ps.handshake() @@ -31,6 +22,20 @@ func Test_peerState_handshake(t *testing.T) { fmt.Println(handshake) } +func Test_peerState_download(t *testing.T) { + + ps := getUsablePeer() + + // 发出握手请求 + _, err := ps.handshake() + if err != nil { + panic(err) + } + // 开始下载 + err = ps.download() + +} + func buildTorrent() *Torrent { metaInfo, err := metainfo.ParseFromFile("../testdata/Game.of.Thrones.S08E05.720p.WEB.H264-MEMENTO.torrent") if err != nil { @@ -41,7 +46,18 @@ func buildTorrent() *Torrent { return NewTorrent(peer.GenPeerID(), metaInfo) } -func getUsablePeer(torrent *Torrent, peers []peer.Peer) *peerState { +// 获取一个能连接上的peer +func getUsablePeer() *peerState { + torrent := buildTorrent() + tracker := &tracker.Tracker{ + PeerID: torrent.PeerID, + MetaInfo: torrent.MetaInfo, + } + peers, err := tracker.Tracker() + if err != nil { + panic(err) + } + fmt.Println("Tracker end") ch := make(chan *peerState, 1) for i := range peers { go peerTest(torrent, &peers[i], ch) diff --git a/down/bt/torrent/torrent.go b/down/bt/torrent/torrent.go index d2bfa19a..c03e2917 100644 --- a/down/bt/torrent/torrent.go +++ b/down/bt/torrent/torrent.go @@ -2,7 +2,6 @@ package torrent import ( "gopeed/down/bt/metainfo" - "gopeed/down/bt/peer" ) type Torrent struct { @@ -21,7 +20,3 @@ func NewTorrent(peerID [20]byte, metaInfo *metainfo.MetaInfo) *Torrent { } return torrent } - -// 检查一个peer能否提供需要下载的文件分片 -func CheckPeer(bitfield peer.MsgBitfield) { -}