bitfield handle

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<<uint(7-i%8)) > 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<<uint(7-j)) > 0 {
+				arr = append(arr, index)
+			}
+		}
+	}
+	return arr
+}

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)

down/bt/peer/peer.go.bak

@@ -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<<uint(7-j)) > 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) {
-
-}

down/bt/peer/peer_test.go.bak

@@ -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)
-			}
-		})
-	}
-}

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)
+}

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)

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) {
-}