refactoring package

headware · Jun 28, 2019 · 137ed5d · 137ed5d
1 parent b734745
commit 137ed5d
Show file tree

Hide file tree

Showing 20 changed files with 528 additions and 166 deletions.
diff --git a/cmd/bt.go b/cmd/bt.go
@@ -1,10 +1,10 @@
 package main
 
 import (
-	"gopeed/down/bt"
+	"gopeed/down/bt/metainfo"
 )
 
 func main() {
-	metaInfo, _ := bt.ParseFromFile("e:/Game.of.Thrones.S08E05.720p.WEB.H264-MEMENTO.torrent")
+	metaInfo, _ := metainfo.ParseFromFile("e:/Game.of.Thrones.S08E05.720p.WEB.H264-MEMENTO.torrent")
 	metaInfo.Tracker()
 }
diff --git a/common/buffer.go b/common/buffer.go
diff --git a/down/bt/client.go b/down/bt/client.go
@@ -2,13 +2,15 @@ package bt
 
 import (
 	"crypto/rand"
+	"gopeed/down/bt/metainfo"
+	"gopeed/down/bt/torrent"
 	"net/url"
 )
 
 type Client struct {
 	PeerID [20]byte
 
-	torrents []Torrent
+	// torrents []torrent.Torrent
 }
 
 func NewClient() *Client {
@@ -31,19 +33,19 @@ func (client *Client) AddTorrent(rawurl string) error {
 	if err != nil {
 		return err
 	}
-	var metaInfo *MetaInfo
+	var metaInfo *metainfo.MetaInfo
 	switch parse.Scheme {
 	case "magnet":
 	// TODO 磁力链接解析
 	default:
-		metaInfo, err = ParseFromFile(rawurl)
+		metaInfo, err = metainfo.ParseFromFile(rawurl)
 		break
 	}
 	if err != nil {
 		return err
 	}
 
-	torrent := Torrent{
+	torrent := torrent.Torrent{
 		client:   client,
 		MetaInfo: metaInfo,
 	}

diff --git a/down/bt/client_test.go b/down/bt/client_test.go
@@ -5,20 +5,17 @@ import (
 )
 
 func TestClient_AddTorrent(t *testing.T) {
-	client := NewClient()
-	metaInfo, err := ParseFromFile("testdata/Game.of.Thrones.S08E05.720p.WEB.H264-MEMENTO.torrent")
+	/*client := NewClient()
+	metaInfo, err := metainfo.ParseFromFile("testdata/Game.of.Thrones.S08E05.720p.WEB.H264-MEMENTO.torrent")
 	if err != nil {
 		panic(err)
 	}
-	torrent := &Torrent{
-		client:   client,
-		MetaInfo: metaInfo,
-	}
+	torrent := &
 	torrent.MetaInfo.Announce = "udp://tracker.opentrackr.org:1337/announce"
 	torrent.MetaInfo.AnnounceList = [][]string{}
 	err = torrent.Tracker()
 	if err != nil {
 		panic(err)
 		return
-	}
+	}*/
 }
diff --git a/down/bt/metainfo.go → down/bt/metainfo/metainfo.go b/down/bt/metainfo.go → down/bt/metainfo/metainfo.go
@@ -1,4 +1,4 @@
-package bt
+package metainfo
 
 import (
 	"crypto/sha1"

diff --git a/down/bt/peer/handshake.go b/down/bt/peer/handshake.go
@@ -0,0 +1,51 @@
+package peer
+
+import (
+	"bytes"
+	"encoding/binary"
+)
+
+const (
+	ProtocolIdentifier       = "BitTorrent protocol"
+	ProtocolIdentifierLength = 0x13
+)
+
+// Handshake of Peer wire protocol
+// Per https://wiki.theory.org/index.php/BitTorrentSpecification#Handshake
+type Handshake struct {
+	Pstrlen  byte
+	Pstr     [19]byte
+	Reserved [8]byte
+	InfoHash [20]byte
+	PeerID   [20]byte
+}
+
+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 {
+	writer := new(bytes.Buffer)
+	err := binary.Write(writer, binary.BigEndian, handshake)
+	if err != nil {
+		panic(err)
+	}
+	return writer.Bytes()
+}
+
+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
+}
diff --git a/down/bt/peer/message.go b/down/bt/peer/message.go
@@ -0,0 +1,77 @@
+package peer
+
+import (
+	"encoding/binary"
+)
+
+type MessageType int
+
+const (
+	Keepalive MessageType = -1
+	Choke     MessageType = iota
+	Unchoke
+	Interested
+	NotInterested
+	Have
+	Bitfield
+	Request
+	Piece
+	Cancel
+)
+
+// 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
+	Payload []byte
+}
+
+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 (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 {
+	bts := []byte(mb)
+	index := i / 8
+	if index >= len(bts) {
+		return false
+	}
+	return bts[index]&(1<<uint(7-i%8)) > 0
+}
diff --git a/down/bt/peer/message_test.go b/down/bt/peer/message_test.go
@@ -0,0 +1,79 @@
+package peer
+
+import (
+	"strconv"
+	"testing"
+)
+
+func TestMsgBitfield_get(t *testing.T) {
+	type args struct {
+		i int
+	}
+	tests := []struct {
+		name string
+		mb   MsgBitfield
+		args args
+		want bool
+	}{
+		{
+			"case-0",
+			MsgBitfield([]byte{convertByte("10000000")}),
+			args{
+				0,
+			},
+			true,
+		},
+		{
+			"case-1",
+			MsgBitfield([]byte{convertByte("01000000")}),
+			args{
+				1,
+			},
+			true,
+		},
+		{
+			"case-2",
+			MsgBitfield([]byte{convertByte("00000000")}),
+			args{
+				2,
+			},
+			false,
+		},
+		{
+			"case-10",
+			MsgBitfield([]byte{convertByte("00000000"), convertByte("00100000")}),
+			args{
+				10,
+			},
+			true,
+		},
+		{
+			"case-11",
+			MsgBitfield([]byte{convertByte("00000000"), convertByte("00100000")}),
+			args{
+				11,
+			},
+			false,
+		},
+		{
+			"case-20",
+			MsgBitfield([]byte{convertByte("00000000"), convertByte("00100000")}),
+			args{
+				20,
+			},
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := tt.mb.get(tt.args.i); got != tt.want {
+				t.Errorf("MsgBitfield.get() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func convertByte(bit string) byte {
+	i, _ := strconv.ParseUint(bit, 2, 8)
+	return byte(i)
+}
diff --git a/down/bt/peer/peer.go b/down/bt/peer/peer.go
@@ -0,0 +1,30 @@
+package peer
+
+import (
+	"encoding/binary"
+	"math/rand"
+	"net"
+	"strconv"
+)
+
+type Peer struct {
+	IP   uint32
+	Port uint16
+}
+
+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))
+}
+
+// 生成Peer ID，规则为前三位固定字母(-GP)+SemVer(xyz)+End(-),后面随机生成
+// 参考：https://wiki.theory.org/index.php/BitTorrentSpecification#peer_id
+func GenPeerID() [20]byte {
+	peerID := [20]byte{'-', 'G', 'P', '0', '0', '1', '-'}
+	_, err := rand.Read(peerID[7:])
+	if err != nil {
+		panic(err)
+	}
+	return peerID
+}
diff --git a/down/bt/peer.go → down/bt/peer/peer.go.bak b/down/bt/peer.go → down/bt/peer/peer.go.bak
@@ -1,11 +1,12 @@
-package bt
+package peer
 
 import (
 	"bufio"
 	"bytes"
 	"encoding/binary"
 	"fmt"
 	"github.com/cenkalti/mse"
+	"gopeed/down/bt/metainfo"
 	"io"
 	"net"
 	"strconv"
@@ -310,6 +311,6 @@ func decodeMessage(buf []byte) *Message {
 	return msg
 }
 
-func doBitfield(peer *Peer, metaInfo *MetaInfo, conn net.Conn) {
+func doBitfield(peer *Peer, metaInfo *metainfo.MetaInfo, conn net.Conn) {
 
 }
diff --git a/down/bt/peer_test.go → down/bt/peer/peer_test.go.bak b/down/bt/peer_test.go → down/bt/peer/peer_test.go.bak
@@ -1,4 +1,4 @@
-package bt
+package peer
 
 import (
 	"fmt"