单个peer的文件下载实现

down/bt/metainfo/metainfo.go

@@ -22,7 +22,9 @@ type MetaInfo struct {
 	UrlList []string `json:"url-list"`
 	Info    *Info    `json:"info"`
 
-	ExtraInfo `json:"-"`
+	infoHash    [20]byte     `json:"-"`
+	totalSize   uint64       `json:"-"`
+	fileDetails []FileDetail `json:"-"`
 }
 
 type Info struct {
@@ -34,13 +36,15 @@ type Info struct {
 }
 
 type File struct {
-	Length uint64 `json:"length"`
-	Path   string `json:"path"`
+	Length uint64   `json:"length"`
+	Path   []string `json:"path"`
 }
 
-type ExtraInfo struct {
-	InfoHash  [20]byte
-	TotalSize uint64
+type FileDetail struct {
+	Length uint64
+	Path   []string
+	Begin  int64
+	End    int64
 }
 
 func ParseFromFile(path string) (*MetaInfo, error) {
@@ -63,7 +67,7 @@ func ParseFromFile(path string) (*MetaInfo, error) {
 		// encode info hash
 		info := btDecode["info"].(map[string]interface{})
 		infoBtEncode := bencode.Encode(info)
-		metaInfo.InfoHash = sha1.Sum(infoBtEncode[:])
+		metaInfo.infoHash = sha1.Sum(infoBtEncode[:])
 		// Split pieces hash
 		pieces := info["pieces"].(string)
 		if pieces != "" {
@@ -75,7 +79,7 @@ func ParseFromFile(path string) (*MetaInfo, error) {
 			}
 		}
 	}
-	metaInfo.TotalSize = getTotalSize(&metaInfo)
+	calcFileSize(&metaInfo)
 	return &metaInfo, nil
 }
 
@@ -83,7 +87,7 @@ func ParseFromFile(path string) (*MetaInfo, error) {
 func (metaInfo *MetaInfo) GetPieceSize(index int) uint64 {
 	// 是否为最后一个分片
 	if index == len(metaInfo.Info.Pieces)-1 {
-		size := metaInfo.TotalSize % metaInfo.Info.PieceLength
+		size := metaInfo.totalSize % metaInfo.Info.PieceLength
 		if size > 0 {
 			return size
 		}
@@ -91,17 +95,40 @@ func (metaInfo *MetaInfo) GetPieceSize(index int) uint64 {
 	return metaInfo.Info.PieceLength
 }
 
-func getTotalSize(metaInfo *MetaInfo) uint64 {
+// 获取所有文件的大小
+func (metaInfo *MetaInfo) GetTotalSize() uint64 {
+	return metaInfo.totalSize
+}
+
+// 获取info hash
+func (metaInfo *MetaInfo) GetInfoHash() [20]byte {
+	return metaInfo.infoHash
+}
+
+// 获取所有文件的开始偏移字节
+func (metaInfo *MetaInfo) GetFileDetails() []FileDetail {
+	return metaInfo.fileDetails
+}
+
+// 计算所有文件的大小
+func calcFileSize(metaInfo *MetaInfo) {
 	if metaInfo.Info != nil {
-		if len(metaInfo.Info.Files) > 0 {
+		fileCount := len(metaInfo.Info.Files)
+		if fileCount > 0 {
+			metaInfo.fileDetails = make([]FileDetail, fileCount)
 			var length uint64
-			for _, f := range metaInfo.Info.Files {
+			for i, f := range metaInfo.Info.Files {
+				metaInfo.fileDetails[i] = FileDetail{
+					Length: f.Length,
+					Path:   f.Path,
+					Begin:  int64(length),
+					End:    int64(length + f.Length),
+				}
 				length += f.Length
 			}
-			return length
+			metaInfo.totalSize = length
 		} else {
-			return metaInfo.Info.Length
+			metaInfo.totalSize = metaInfo.Info.Length
 		}
 	}
-	return 0
 }

down/bt/metainfo/metainfo_test.go

@@ -10,5 +10,6 @@ func TestParseFromFile(t *testing.T) {
 	if err != nil {
 		panic(err)
 	}
-	fmt.Println(metaInfo)
+	fmt.Printf("%+v\n", metaInfo)
+	fmt.Printf("%+v\n", metaInfo.Info.Files)
 }

down/bt/peer/message/message.go

@@ -29,14 +29,14 @@ type Message struct {
 	ID     ID
 }
 
-func (msg Message) Encode() []byte {
+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) {
+func (msg *Message) Decode(buf []byte) {
 	msg.Length = binary.BigEndian.Uint32(buf)
 	msg.ID = ID(buf[4])
 }

down/bt/peer/message/msg_bitfield.go

@@ -2,13 +2,13 @@ package message
 
 // bitfield: <len=0001+X><id=5><bitfield>
 type Bitfield struct {
-	Message
+	*Message
 	payload []byte
 }
 
 func NewBitfield(payload []byte) *Bitfield {
 	return &Bitfield{
-		Message: Message{
+		Message: &Message{
 			Length: uint32(len(payload) + 1),
 			ID:     IdBitfield,
 		},
@@ -21,8 +21,9 @@ func (b *Bitfield) Encode() []byte {
 }
 
 func (b *Bitfield) Decode(buf []byte) Serialize {
+	b.Message = &Message{}
 	b.Message.Decode(buf)
-	copy(b.payload, buf[5:])
+	b.payload = buf[5:]
 	return b
 }
 

down/bt/peer/message/msg_piece.go

@@ -1,23 +1,41 @@
 package message
 
-import "io"
+import "encoding/binary"
 
 // piece: <len=0009+X><id=7><index><begin><block>
 type Piece struct {
-	Message
+	*Message
 	Index uint32
 	Begin uint32
-	Block io.Reader
+	Block []byte
 }
 
-func newPiece(index uint32, begin uint32, reader io.Reader, length uint32) *Piece {
+func NewPiece(index uint32, begin uint32, block []byte) *Piece {
 	return &Piece{
-		Message: Message{
-			Length: 9 + length,
+		Message: &Message{
+			Length: 9 + uint32(len(block)),
 			ID:     IdPiece,
 		},
 		Index: index,
 		Begin: begin,
-		Block: io.LimitReader(reader, int64(length)),
+		// 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
+}

down/bt/peer/message/msg_request.go

@@ -26,9 +26,9 @@ func NewRequest(index uint32, begin uint32, length uint32) *Request {
 
 func (r *Request) Encode() []byte {
 	buf := make([]byte, 12)
-	binary.BigEndian.PutUint32(buf[:4], r.Index)
+	binary.BigEndian.PutUint32(buf[0:4], r.Index)
 	binary.BigEndian.PutUint32(buf[4:8], r.Begin)
-	binary.BigEndian.PutUint32(buf[8:], r.Length)
+	binary.BigEndian.PutUint32(buf[8:12], r.Length)
 	return append(r.Message.Encode(), buf...)
 }
 

down/bt/torrent/peer_state.go

@@ -5,11 +5,15 @@ import (
 	"encoding/binary"
 	"fmt"
 	"github.com/cenkalti/mse"
+	"github.com/monkeyWie/gopeed/down/bt/metainfo"
 	"github.com/monkeyWie/gopeed/down/bt/peer"
 	"github.com/monkeyWie/gopeed/down/bt/peer/message"
 	"io"
 	"math"
 	"net"
+	"os"
+	"path"
+	"time"
 )
 
 type peerState struct {
@@ -32,12 +36,13 @@ type peerState struct {
 // 使用MSE加密来避免运营商对bt流量的封锁，基本上现在市面上BT客户端都默认开启了，不用MSE的话很多Peer拒绝连接
 // see http://wiki.vuze.com/w/Message_Stream_Encryption
 func (ps *peerState) dialMse() error {
-	conn, err := net.Dial("tcp", ps.Address())
+	conn, err := net.DialTimeout("tcp", ps.Address(), time.Second*time.Duration(30))
 	if err != nil {
 		return err
 	}
 	mseConn := mse.WrapConn(conn)
-	_, err = mseConn.HandshakeOutgoing(ps.torrent.MetaInfo.InfoHash[:], mse.PlainText, nil)
+	infoHash := ps.torrent.MetaInfo.GetInfoHash()
+	_, err = mseConn.HandshakeOutgoing(infoHash[:], mse.PlainText, nil)
 	if err != nil {
 		mseConn.Close()
 		return err
@@ -50,7 +55,7 @@ func (ps *peerState) dialMse() error {
 // see https://wiki.theory.org/index.php/BitTorrentSpecification#Handshake
 func (ps *peerState) handshake() (*peer.Handshake, error) {
 	handshakeRes, err := func() (*peer.Handshake, error) {
-		handshakeReq := peer.NewHandshake([8]byte{}, ps.torrent.MetaInfo.InfoHash, ps.torrent.PeerID)
+		handshakeReq := peer.NewHandshake([8]byte{}, ps.torrent.MetaInfo.GetInfoHash(), ps.torrent.PeerID)
 		_, err := ps.conn.Write(handshakeReq.Encode())
 		if err != nil {
 			return nil, err
@@ -104,7 +109,6 @@ func (ps *peerState) download() error {
 			// 	keepalive
 		} else {
 			messageID := message.ID(buf[4])
-			payload := buf[5:]
 			switch messageID {
 			case message.IdChoke:
 				break
@@ -118,12 +122,12 @@ func (ps *peerState) download() error {
 			case message.IdHave:
 				break
 			case message.IdBitfield:
-				ps.handleBitfield(payload)
+				ps.handleBitfield(buf)
 				break
 			case message.IdRequest:
 				break
 			case message.IdPiece:
-				ps.handlePiece()
+				ps.handlePiece(buf)
 				break
 			case message.IdCancel:
 				break
@@ -133,6 +137,10 @@ func (ps *peerState) download() error {
 	return nil
 }
 
+func downloadPiece() {
+
+}
+
 func (ps *peerState) handleUnchoke() {
 	ps.peerChoking = false
 	// 如果客户端对peer感兴趣并且peer没有choked客户端，就可以开始下载了
@@ -148,10 +156,7 @@ func (ps *peerState) handleUnchoke() {
 					blockLength := math.Min(2<<13, float64(pieceLength-lastOffset))
 					ps.write(message.NewRequest(uint32(index), uint32(lastOffset), uint32(blockLength)).Encode())
 					lastOffset += uint64(blockLength)
-					fmt.Println("request piece")
 				}
-				// TODO 测试代码
-				break
 			}
 			return
 		}
@@ -159,8 +164,9 @@ func (ps *peerState) handleUnchoke() {
 	ps.conn.Close()
 }
 
-func (ps *peerState) handleBitfield(payload []byte) {
-	ps.bitfield = message.NewBitfield(payload)
+func (ps *peerState) handleBitfield(buf []byte) {
+	ps.bitfield = &message.Bitfield{}
+	ps.bitfield.Decode(buf)
 	have := ps.getHavePieces(ps.bitfield)
 	if len(have) > 0 {
 		// 表示对该peer感兴趣，并且不choked该peer
@@ -175,8 +181,85 @@ func (ps *peerState) handleBitfield(payload []byte) {
 }
 
 // 处理分片下载响应
-func (ps *peerState) handlePiece() {
-	fmt.Println("one piece")
+func (ps *peerState) handlePiece(buf []byte) {
+	piece := &message.Piece{}
+	piece.Decode(buf)
+	info := ps.torrent.MetaInfo.Info
+	blockLength := int64(len(piece.Block))
+	pieceBegin := int64(piece.Index)*int64(info.PieceLength) + int64(piece.Begin)
+	// 获取对应要的文件
+	var fileBlocks []fileBlock
+	if len(info.Files) == 0 {
+		// 单文件
+		fileBlocks = append(fileBlocks, fileBlock{
+			filepath:   info.Name,
+			fileSeek:   pieceBegin,
+			blockBegin: 0,
+			blockEnd:   blockLength,
+		})
+	} else {
+		// 获取要写入的第一个文件和偏移
+		writeIndex := getWriteFile(pieceBegin, ps.torrent.MetaInfo.GetFileDetails())
+		// block对应种子的偏移
+		blockBegin := pieceBegin
+		// block对应文件的偏移
+		var blockSeek int64 = 0
+		for _, f := range ps.torrent.MetaInfo.GetFileDetails()[writeIndex:] {
+			// 计算文件可剩余写入字节数
+			fileWritable := f.End - blockBegin
+			// 计算block剩余写入字节数
+			blockWritable := blockLength - blockSeek
+			fb := fileBlock{
+				filepath:   f.Path[0],
+				fileSeek:   blockBegin - f.Begin,
+				blockBegin: blockSeek,
+			}
+			if fileWritable >= blockWritable {
+				// 若够block写入直接跳出循环
+				fb.blockEnd = blockSeek + blockWritable
+				fileBlocks = append(fileBlocks, fb)
+				break
+			} else {
+				// 否则计算剩余可写字节数，写入到下一个文件中
+				fb.blockEnd = blockSeek + fileWritable
+				fileBlocks = append(fileBlocks, fb)
+				blockBegin = f.End
+				blockSeek += fileWritable
+			}
+		}
+	}
+	fmt.Println(fileBlocks)
+	for _, f := range fileBlocks {
+		func() {
+			name := path.Join("e:/testbt/download", f.filepath)
+			file, err := os.OpenFile(name, os.O_RDWR|os.O_CREATE, 0644)
+			if err != nil {
+				panic(err)
+			}
+			defer file.Close()
+			_, err = file.WriteAt(piece.Block[f.blockBegin:f.blockEnd], f.fileSeek)
+			if err != nil {
+				fmt.Println(err)
+			}
+		}()
+	}
+}
+
+// 获取piece对应要写入的文件
+func getWriteFile(pieceBegin int64, fds []metainfo.FileDetail) int {
+	for i, f := range fds {
+		if f.Begin <= pieceBegin && f.End > pieceBegin {
+			return i
+		}
+	}
+	return -1
+}
+
+type fileBlock struct {
+	filepath   string
+	fileSeek   int64
+	blockBegin int64
+	blockEnd   int64
 }
 
 // 获取peer能提供需要下载的文件分片

down/bt/torrent/peer_state_test.go

@@ -43,7 +43,7 @@ func buildTorrent() *Torrent {
 	if err != nil {
 		panic(err)
 	}
-	metaInfo.Announce = "udp://tracker.opentrackr.org:1337/announce"
+	metaInfo.Announce = "udp://exodus.desync.com:6969/announce"
 	metaInfo.AnnounceList = [][]string{}
 	return NewTorrent(peer.GenPeerID(), metaInfo)
 }

down/bt/tracker/tracker.go

@@ -204,11 +204,11 @@ func (tracker *Tracker) connect(conn *net.UDPConn, timeout int64) (response *udp
 
 func (tracker *Tracker) announce(conn *net.UDPConn, timeout int64, connectionId uint64) (response *udpAnnounceResponse, err error) {
 	request := newUdpAnnounceRequest(connectionId)
-	request.infoHash = tracker.MetaInfo.InfoHash
+	request.infoHash = tracker.MetaInfo.GetInfoHash()
 	request.peerID = tracker.PeerID
 	request.downloaded = 0
 	request.uploaded = 0
-	request.left = tracker.MetaInfo.TotalSize - request.downloaded
+	request.left = tracker.MetaInfo.GetTotalSize() - request.downloaded
 	request.event = 0
 	request.ip = 0
 	request.key = 0
@@ -282,12 +282,13 @@ func (tracker *Tracker) httpTracker(url *url.URL) (peers []peer.Peer, err error)
 	peerID := tracker.PeerID
 
 	query := url.Query()
-	query.Add("info_hash", string(metaInfo.InfoHash[:]))
+	infoHash := metaInfo.GetInfoHash()
+	query.Add("info_hash", string(infoHash[:]))
 	query.Add("peer_id", string(peerID[:]))
 	query.Add("port", "6882")
 	query.Add("uploaded", "0")
 	query.Add("downloaded", "0")
-	query.Add("left", strconv.FormatInt(int64(metaInfo.TotalSize), 10))
+	query.Add("left", strconv.FormatInt(int64(metaInfo.GetTotalSize()), 10))
 	url.RawQuery = query.Encode()
 	request, err := http.NewRequest("GET", url.String(), nil)
 	if err != nil {