Refactor

ahbxl · Sep 12, 2019 · adb7d89 · adb7d89
1 parent 075852b
commit adb7d89
Show file tree

Hide file tree

Showing 19 changed files with 415 additions and 258 deletions.
diff --git a/down/bt/client.go b/down/bt/client.go
@@ -2,8 +2,8 @@ package bt
 
 import (
 	"crypto/rand"
-	"gopeed/down/bt/metainfo"
-	"gopeed/down/bt/torrent"
+	"github.com/monkeyWie/gopeed/down/bt/metainfo"
+	"github.com/monkeyWie/gopeed/down/bt/torrent"
 	"net/url"
 )
 

diff --git a/down/bt/metainfo/metainfo.go b/down/bt/metainfo/metainfo.go
@@ -39,8 +39,8 @@ type File struct {
 }
 
 type ExtraInfo struct {
-	InfoHash [20]byte
-	FileSize uint64
+	InfoHash  [20]byte
+	TotalSize uint64
 }
 
 func ParseFromFile(path string) (*MetaInfo, error) {
@@ -75,10 +75,22 @@ func ParseFromFile(path string) (*MetaInfo, error) {
 			}
 		}
 	}
-	metaInfo.FileSize = getTotalSize(&metaInfo)
+	metaInfo.TotalSize = getTotalSize(&metaInfo)
 	return &metaInfo, nil
 }
 
+// 获取某个分片的大小
+func (metaInfo *MetaInfo) GetPieceSize(index int) uint64 {
+	// 是否为最后一个分片
+	if index == len(metaInfo.Info.Pieces)-1 {
+		size := metaInfo.TotalSize % metaInfo.Info.PieceLength
+		if size > 0 {
+			return size
+		}
+	}
+	return metaInfo.Info.PieceLength
+}
+
 func getTotalSize(metaInfo *MetaInfo) uint64 {
 	if metaInfo.Info != nil {
 		if len(metaInfo.Info.Files) > 0 {

diff --git a/down/bt/peer/message.go → down/bt/peer/message/message.go b/down/bt/peer/message.go → down/bt/peer/message/message.go
@@ -1,31 +1,52 @@
-package peer
+package message
 
 import (
 	"encoding/binary"
 )
 
-type MessageID int
+// message ID
+type ID byte
 
 const (
-	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
+	IdChoke         ID = 0
+	IdUnchoke       ID = 1
+	IdInterested    ID = 2
+	IdNotinterested ID = 3
+	IdHave          ID = 4
+	IdBitfield      ID = 5
+	IdRequest       ID = 6
+	IdPiece         ID = 7
+	IdCancel        ID = 8
 )
 
-type Message struct {
+/*type Message struct {
 	Length  uint32
 	ID      MessageID
 	Payload []byte
+}*/
+type Message struct {
+	Length uint32
+	ID     ID
+}
+
+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) {
+	msg.Length = binary.BigEndian.Uint32(buf)
+	msg.ID = ID(buf[4])
 }
 
-func NewMessage(id MessageID, payload []byte) *Message {
+type Serialize interface {
+	Encode() []byte
+	Decode(buf []byte) Serialize
+}
+
+/*func NewMessage(id MessageID, payload []byte) *Message {
 	message := &Message{ID: id, Payload: payload}
 	if id == Keepalive {
 		message.Length = 0
@@ -57,7 +78,7 @@ func (msg *Message) Decode(buf []byte) {
 			msg.Payload = buf[5:]
 		}
 	}
-}
+}*/
 
 // Message protocol
 // Length prefix| message ID | Payload
@@ -73,39 +94,3 @@ func SplitMessage(data []byte, atEOF bool) (advance int, token []byte, err error
 	}
 	return
 }
-
-type MsgBitfield []byte
-
-// 某个分片是否下载完成
-func (mb MsgBitfield) IsComplete(i int) bool {
-	bts := []byte(mb)
-	index := i / 8
-	if index >= len(bts) {
-		return false
-	}
-	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
-}
diff --git a/down/bt/peer/message/msg_bitfield.go b/down/bt/peer/message/msg_bitfield.go
@@ -0,0 +1,59 @@
+package message
+
+// bitfield: <len=0001+X><id=5><bitfield>
+type Bitfield struct {
+	Message
+	payload []byte
+}
+
+func NewBitfield(payload []byte) *Bitfield {
+	return &Bitfield{
+		Message: Message{
+			Length: uint32(len(payload) + 1),
+			ID:     IdBitfield,
+		},
+		payload: payload,
+	}
+}
+
+func (b *Bitfield) Encode() []byte {
+	return append(b.Message.Encode(), b.payload...)
+}
+
+func (b *Bitfield) Decode(buf []byte) Serialize {
+	b.Message.Decode(buf)
+	copy(b.payload, buf[5:])
+	return b
+}
+
+// 某个分片是否下载完成
+func (b *Bitfield) IsComplete(i int) bool {
+	index := i / 8
+	if index >= len(b.payload) {
+		return false
+	}
+	return b.payload[index]&(1<<(7-i%8)) > 0
+}
+
+// 给定一组分片下载状态，计算出当前peer能提供下载的分片下标
+func (b *Bitfield) Have(pieces []bool) []int {
+	arr := make([]int, 0)
+	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(b.payload) &&
+				b.payload[i]&(1<<(7-j)) > 0 {
+				arr = append(arr, index)
+			}
+		}
+	}
+	return arr
+}
diff --git a/down/bt/peer/message/msg_choke.go b/down/bt/peer/message/msg_choke.go
@@ -0,0 +1,24 @@
+package message
+
+// choke: <len=0001><id=0>
+type Choke struct {
+	Message
+}
+
+func NewChoke() *Choke {
+	return &Choke{
+		Message{
+			Length: 1,
+			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/down/bt/peer/message/msg_interested.go b/down/bt/peer/message/msg_interested.go
@@ -0,0 +1,24 @@
+package message
+
+// choke: <len=0001><id=0>
+type Interested struct {
+	Message
+}
+
+func NewInterested() *Interested {
+	return &Interested{
+		Message{
+			Length: 1,
+			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/down/bt/peer/message/msg_keepalive.go b/down/bt/peer/message/msg_keepalive.go
@@ -0,0 +1,12 @@
+package message
+
+// 	keep-alive: <len=0000>
+type Keepalive struct{}
+
+func (k *Keepalive) Encode() []byte {
+	return []byte{0, 0, 0, 0}
+}
+
+func (k *Keepalive) Decode(buf []byte) Serialize {
+	return &Keepalive{}
+}
diff --git a/down/bt/peer/message/msg_not_interested.go b/down/bt/peer/message/msg_not_interested.go
@@ -0,0 +1,24 @@
+package message
+
+// choke: <len=0001><id=0>
+type NotInterested struct {
+	Message
+}
+
+func NewNotInterested() *NotInterested {
+	return &NotInterested{
+		Message{
+			Length: 1,
+			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/down/bt/peer/message/msg_piece.go b/down/bt/peer/message/msg_piece.go
@@ -0,0 +1,23 @@
+package message
+
+import "io"
+
+// piece: <len=0009+X><id=7><index><begin><block>
+type Piece struct {
+	Message
+	Index uint32
+	Begin uint32
+	Block io.Reader
+}
+
+func newPiece(index uint32, begin uint32, reader io.Reader, length uint32) *Piece {
+	return &Piece{
+		Message: Message{
+			Length: 9 + length,
+			ID:     IdPiece,
+		},
+		Index: index,
+		Begin: begin,
+		Block: io.LimitReader(reader, int64(length)),
+	}
+}
diff --git a/down/bt/peer/message/msg_request.go b/down/bt/peer/message/msg_request.go
@@ -0,0 +1,41 @@
+package message
+
+import (
+	"encoding/binary"
+)
+
+// request: <len=0013><id=6><index><begin><length>
+type Request struct {
+	Message
+	Index  uint32
+	Begin  uint32
+	Length uint32
+}
+
+func NewRequest(index uint32, begin uint32, length uint32) *Request {
+	return &Request{
+		Message: Message{
+			Length: 13,
+			ID:     IdRequest,
+		},
+		Index:  index,
+		Begin:  begin,
+		Length: length,
+	}
+}
+
+func (r *Request) Encode() []byte {
+	buf := make([]byte, 12)
+	binary.BigEndian.PutUint32(buf[:4], r.Index)
+	binary.BigEndian.PutUint32(buf[4:8], r.Begin)
+	binary.BigEndian.PutUint32(buf[8:], r.Length)
+	return append(r.Message.Encode(), 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
+}