torrent/client.go

920 lines
20 KiB
Go

package torrent
import (
"bitbucket.org/anacrolix/go.torrent/peer_protocol"
"bufio"
"container/list"
"crypto"
"crypto/rand"
"encoding"
"errors"
"fmt"
metainfo "github.com/nsf/libtorgo/torrent"
"io"
"launchpad.net/gommap"
"log"
"net"
"os"
"path/filepath"
"sort"
"sync"
"time"
)
const (
PieceHash = crypto.SHA1
maxRequests = 250
chunkSize = 0x4000 // 16KiB
BEP20 = "-GT0000-"
)
type InfoHash [20]byte
type pieceSum [20]byte
func copyHashSum(dst, src []byte) {
if len(dst) != len(src) || copy(dst, src) != len(dst) {
panic("hash sum sizes differ")
}
}
func BytesInfoHash(b []byte) (ih InfoHash) {
if len(b) != len(ih) || copy(ih[:], b) != len(ih) {
panic("bad infohash bytes")
}
return
}
type piece struct {
Hash pieceSum
PendingChunkSpecs map[ChunkSpec]struct{}
Hashing bool
EverHashed bool
}
func (p *piece) Complete() bool {
return len(p.PendingChunkSpecs) == 0 && !p.Hashing && p.EverHashed
}
func lastChunkSpec(pieceLength peer_protocol.Integer) (cs ChunkSpec) {
cs.Begin = (pieceLength - 1) / chunkSize * chunkSize
cs.Length = pieceLength - cs.Begin
return
}
func (t *Torrent) PieceNumPendingBytes(index peer_protocol.Integer) (count peer_protocol.Integer) {
pendingChunks := t.Pieces[index].PendingChunkSpecs
count = peer_protocol.Integer(len(pendingChunks)) * chunkSize
_lastChunkSpec := lastChunkSpec(t.PieceLength(index))
if _lastChunkSpec.Length != chunkSize {
if _, ok := pendingChunks[_lastChunkSpec]; ok {
count += _lastChunkSpec.Length - chunkSize
}
}
return
}
type ChunkSpec struct {
Begin, Length peer_protocol.Integer
}
type Request struct {
Index peer_protocol.Integer
ChunkSpec
}
type Connection struct {
Socket net.Conn
post chan encoding.BinaryMarshaler
write chan []byte
Interested bool
Choked bool
Requests map[Request]struct{}
PeerId [20]byte
PeerInterested bool
PeerChoked bool
PeerRequests map[Request]struct{}
PeerExtensions [8]byte
PeerPieces []bool
}
func (c *Connection) Close() {
c.Socket.Close()
close(c.post)
}
func (c *Connection) PeerHasPiece(index peer_protocol.Integer) bool {
if c.PeerPieces == nil {
return false
}
return c.PeerPieces[index]
}
func (c *Connection) Post(msg encoding.BinaryMarshaler) {
c.post <- msg
}
func (c *Connection) Request(chunk Request) bool {
if len(c.Requests) >= maxRequests {
return false
}
c.SetInterested(true)
if c.PeerChoked {
return false
}
if _, ok := c.Requests[chunk]; !ok {
c.Post(peer_protocol.Message{
Type: peer_protocol.Request,
Index: chunk.Index,
Begin: chunk.Begin,
Length: chunk.Length,
})
}
if c.Requests == nil {
c.Requests = make(map[Request]struct{}, maxRequests)
}
c.Requests[chunk] = struct{}{}
return true
}
func (c *Connection) SetInterested(interested bool) {
if c.Interested == interested {
return
}
c.Post(peer_protocol.Message{
Type: func() peer_protocol.MessageType {
if interested {
return peer_protocol.Interested
} else {
return peer_protocol.NotInterested
}
}(),
})
c.Interested = interested
}
var (
keepAliveBytes [4]byte
)
func (conn *Connection) writer() {
timer := time.NewTimer(0)
for {
if !timer.Reset(time.Minute) {
<-timer.C
}
var b []byte
select {
case <-timer.C:
b = keepAliveBytes[:]
case b = <-conn.write:
if b == nil {
return
}
}
n, err := conn.Socket.Write(b)
if err != nil {
log.Print(err)
break
}
if n != len(b) {
panic("didn't write all bytes")
}
log.Printf("wrote %#v", string(b))
}
}
func (conn *Connection) writeOptimizer() {
pending := list.New()
var nextWrite []byte
defer close(conn.write)
for {
write := conn.write
if pending.Len() == 0 {
write = nil
} else {
var err error
nextWrite, err = pending.Front().Value.(encoding.BinaryMarshaler).MarshalBinary()
if err != nil {
panic(err)
}
}
select {
case msg, ok := <-conn.post:
if !ok {
return
}
pending.PushBack(msg)
case write <- nextWrite:
pending.Remove(pending.Front())
}
}
}
type Torrent struct {
InfoHash InfoHash
Pieces []*piece
Data MMapSpan
MetaInfo *metainfo.MetaInfo
Conns []*Connection
Peers []Peer
Priorities *list.List
}
func (t *Torrent) Close() (err error) {
t.Data.Close()
for _, conn := range t.Conns {
conn.Close()
}
return
}
type pieceByBytesPendingSlice struct {
Pending, Indices []peer_protocol.Integer
}
func (pcs pieceByBytesPendingSlice) Len() int {
return len(pcs.Indices)
}
func (me pieceByBytesPendingSlice) Less(i, j int) bool {
return me.Pending[me.Indices[i]] < me.Pending[me.Indices[j]]
}
func (me pieceByBytesPendingSlice) Swap(i, j int) {
me.Indices[i], me.Indices[j] = me.Indices[j], me.Indices[i]
}
func (t *Torrent) piecesByPendingBytesDesc() (indices []peer_protocol.Integer) {
slice := pieceByBytesPendingSlice{
Pending: make([]peer_protocol.Integer, 0, len(t.Pieces)),
Indices: make([]peer_protocol.Integer, 0, len(t.Pieces)),
}
for i := range t.Pieces {
slice.Pending = append(slice.Pending, t.PieceNumPendingBytes(peer_protocol.Integer(i)))
slice.Indices = append(slice.Indices, peer_protocol.Integer(i))
}
sort.Sort(sort.Reverse(slice))
return slice.Indices
}
func (cl *Client) queuePieceCheck(t *Torrent, pieceIndex peer_protocol.Integer) {
piece := t.Pieces[pieceIndex]
if piece.Hashing {
return
}
piece.Hashing = true
go cl.verifyPiece(t, pieceIndex)
}
func (cl *Client) PrioritizeDataRegion(ih InfoHash, off, len_ int64) {
cl.mu.Lock()
defer cl.mu.Unlock()
t := cl.torrent(ih)
newPriorities := make([]Request, 0, (len_+2*(chunkSize-1))/chunkSize)
for len_ > 0 {
index := peer_protocol.Integer(off / int64(t.PieceLength(0)))
pieceOff := peer_protocol.Integer(off % int64(t.PieceLength(0)))
piece := t.Pieces[index]
if !piece.EverHashed {
cl.queuePieceCheck(t, index)
}
chunk := ChunkSpec{pieceOff / chunkSize * chunkSize, chunkSize}
adv := int64(chunkSize - pieceOff%chunkSize)
off += adv
len_ -= adv
if _, ok := piece.PendingChunkSpecs[chunk]; !ok && !piece.Hashing {
continue
}
newPriorities = append(newPriorities, Request{index, chunk})
}
if len(newPriorities) < 1 {
return
}
log.Print(newPriorities)
if t.Priorities == nil {
t.Priorities = list.New()
}
t.Priorities.PushFront(newPriorities[0])
for _, req := range newPriorities[1:] {
t.Priorities.PushBack(req)
}
for _, cn := range t.Conns {
cl.replenishConnRequests(t, cn)
}
}
func (t *Torrent) WriteChunk(piece int, begin int64, data []byte) (err error) {
_, err = t.Data.WriteAt(data, int64(piece)*t.MetaInfo.PieceLength+begin)
return
}
func (t *Torrent) bitfield() (bf []bool) {
for _, p := range t.Pieces {
bf = append(bf, p.EverHashed && !p.Hashing && len(p.PendingChunkSpecs) == 0)
}
return
}
func (t *Torrent) pieceChunkSpecs(index peer_protocol.Integer) (cs map[ChunkSpec]struct{}) {
cs = make(map[ChunkSpec]struct{}, (t.MetaInfo.PieceLength+chunkSize-1)/chunkSize)
c := ChunkSpec{
Begin: 0,
}
for left := peer_protocol.Integer(t.PieceLength(index)); left > 0; left -= c.Length {
c.Length = left
if c.Length > chunkSize {
c.Length = chunkSize
}
cs[c] = struct{}{}
c.Begin += c.Length
}
return
}
func (t *Torrent) requestHeat() (ret map[Request]int) {
ret = make(map[Request]int)
for _, conn := range t.Conns {
for req, _ := range conn.Requests {
ret[req]++
}
}
return
}
type Peer struct {
Id [20]byte
IP net.IP
Port int
}
func (t *Torrent) PieceLength(piece peer_protocol.Integer) (len_ peer_protocol.Integer) {
if int(piece) == len(t.Pieces)-1 {
len_ = peer_protocol.Integer(t.Data.Size() % t.MetaInfo.PieceLength)
}
if len_ == 0 {
len_ = peer_protocol.Integer(t.MetaInfo.PieceLength)
}
return
}
func (t *Torrent) HashPiece(piece peer_protocol.Integer) (ps pieceSum) {
hash := PieceHash.New()
n, err := t.Data.WriteSectionTo(hash, int64(piece)*t.MetaInfo.PieceLength, t.MetaInfo.PieceLength)
if err != nil {
panic(err)
}
if peer_protocol.Integer(n) != t.PieceLength(piece) {
panic(fmt.Sprintf("hashed wrong number of bytes: expected %d; did %d; piece %d", t.PieceLength(piece), n, piece))
}
copyHashSum(ps[:], hash.Sum(nil))
return
}
type DataSpec struct {
InfoHash
Request
}
type Client struct {
DataDir string
HalfOpenLimit int
PeerId [20]byte
DataReady chan DataSpec
sync.Mutex
mu *sync.Mutex
event sync.Cond
halfOpen int
torrents map[InfoHash]*Torrent
}
var (
ErrDataNotReady = errors.New("data not ready")
)
func (cl *Client) TorrentReadAt(ih InfoHash, off int64, p []byte) (n int, err error) {
cl.mu.Lock()
defer cl.mu.Unlock()
t := cl.torrent(ih)
if t == nil {
err = errors.New("unknown torrent")
return
}
index := peer_protocol.Integer(off / int64(t.PieceLength(0)))
piece := t.Pieces[index]
if !piece.EverHashed {
cl.queuePieceCheck(t, index)
}
if piece.Hashing {
err = ErrDataNotReady
return
}
pieceOff := peer_protocol.Integer(off % int64(t.PieceLength(0)))
high := int(t.PieceLength(index) - pieceOff)
if high < len(p) {
p = p[:high]
}
for cs, _ := range piece.PendingChunkSpecs {
chunkOff := int64(pieceOff) - int64(cs.Begin)
if chunkOff >= int64(t.PieceLength(index)) {
panic(chunkOff)
}
if 0 <= chunkOff && chunkOff < int64(cs.Length) {
// read begins in a pending chunk
err = ErrDataNotReady
return
}
// pending chunk caps available data
if chunkOff < 0 && int64(len(p)) > -chunkOff {
p = p[:-chunkOff]
}
}
return t.Data.ReadAt(p, off)
}
func (c *Client) Start() {
c.mu = &c.Mutex
c.torrents = make(map[InfoHash]*Torrent)
if c.HalfOpenLimit == 0 {
c.HalfOpenLimit = 10
}
o := copy(c.PeerId[:], BEP20)
_, err := rand.Read(c.PeerId[o:])
if err != nil {
panic("error generating peer id")
}
}
func mmapTorrentData(metaInfo *metainfo.MetaInfo, location string) (mms MMapSpan, err error) {
defer func() {
if err != nil {
mms.Close()
mms = nil
}
}()
for _, miFile := range metaInfo.Files {
fileName := filepath.Join(append([]string{location, metaInfo.Name}, miFile.Path...)...)
err = os.MkdirAll(filepath.Dir(fileName), 0777)
if err != nil {
return
}
var file *os.File
file, err = os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, 0666)
if err != nil {
return
}
func() {
defer file.Close()
var fi os.FileInfo
fi, err = file.Stat()
if err != nil {
return
}
if fi.Size() < miFile.Length {
err = file.Truncate(miFile.Length)
if err != nil {
return
}
}
var mMap gommap.MMap
mMap, err = gommap.MapRegion(file.Fd(), 0, miFile.Length, gommap.PROT_READ|gommap.PROT_WRITE, gommap.MAP_SHARED)
if err != nil {
return
}
if int64(len(mMap)) != miFile.Length {
panic("mmap has wrong length")
}
mms = append(mms, MMap{mMap})
}()
if err != nil {
return
}
}
return
}
func (me *Client) torrent(ih InfoHash) *Torrent {
for _, t := range me.torrents {
if t.InfoHash == ih {
return t
}
}
return nil
}
func (me *Client) initiateConn(peer Peer, torrent *Torrent) {
if peer.Id == me.PeerId {
return
}
me.halfOpen++
go func() {
conn, err := net.DialTCP("tcp", nil, &net.TCPAddr{
IP: peer.IP,
Port: peer.Port,
})
me.mu.Lock()
me.halfOpen--
me.openNewConns()
me.mu.Unlock()
if err != nil {
log.Printf("error connecting to peer: %s", err)
return
}
log.Printf("connected to %s", conn.RemoteAddr())
err = me.runConnection(conn, torrent, peer.Id)
if err != nil {
log.Print(err)
}
}()
}
func (me *Torrent) haveAnyPieces() bool {
for _, piece := range me.Pieces {
if piece.Complete() {
return true
}
}
return false
}
func (me *Client) runConnection(sock net.Conn, torrent *Torrent, peerId [20]byte) (err error) {
conn := &Connection{
Socket: sock,
Choked: true,
PeerChoked: true,
write: make(chan []byte),
post: make(chan encoding.BinaryMarshaler),
}
defer conn.Close()
go conn.writer()
go conn.writeOptimizer()
conn.post <- peer_protocol.Bytes(peer_protocol.Protocol)
conn.post <- peer_protocol.Bytes("\x00\x00\x00\x00\x00\x00\x00\x00")
if torrent != nil {
conn.post <- peer_protocol.Bytes(torrent.InfoHash[:])
conn.post <- peer_protocol.Bytes(me.PeerId[:])
}
var b [28]byte
_, err = io.ReadFull(conn.Socket, b[:])
if err != nil {
err = fmt.Errorf("when reading protocol and extensions: %s", err)
return
}
if string(b[:20]) != peer_protocol.Protocol {
err = fmt.Errorf("wrong protocol: %#v", string(b[:20]))
return
}
if 8 != copy(conn.PeerExtensions[:], b[20:]) {
panic("wtf")
}
log.Printf("peer extensions: %#v", string(conn.PeerExtensions[:]))
var infoHash [20]byte
_, err = io.ReadFull(conn.Socket, infoHash[:])
if err != nil {
return fmt.Errorf("reading peer info hash: %s", err)
}
_, err = io.ReadFull(conn.Socket, conn.PeerId[:])
if err != nil {
return fmt.Errorf("reading peer id: %s", err)
}
if torrent == nil {
torrent = me.torrent(infoHash)
if torrent == nil {
return
}
conn.post <- peer_protocol.Bytes(torrent.InfoHash[:])
conn.post <- peer_protocol.Bytes(me.PeerId[:])
}
me.mu.Lock()
me.addConnection(torrent, conn)
if torrent.haveAnyPieces() {
conn.Post(peer_protocol.Message{
Type: peer_protocol.Bitfield,
Bitfield: torrent.bitfield(),
})
}
err = me.connectionLoop(torrent, conn)
if err != nil {
err = fmt.Errorf("during Connection loop: %s", err)
}
me.dropConnection(torrent, conn)
me.mu.Unlock()
return
}
func (me *Client) peerGotPiece(torrent *Torrent, conn *Connection, piece int) {
if conn.PeerPieces == nil {
conn.PeerPieces = make([]bool, len(torrent.Pieces))
}
conn.PeerPieces[piece] = true
if torrent.wantPiece(piece) {
me.replenishConnRequests(torrent, conn)
}
}
func (t *Torrent) wantPiece(index int) bool {
p := t.Pieces[index]
return p.EverHashed && !p.Hashing && len(p.PendingChunkSpecs) != 0
}
func (me *Client) peerUnchoked(torrent *Torrent, conn *Connection) {
me.replenishConnRequests(torrent, conn)
}
func (me *Client) connectionLoop(torrent *Torrent, conn *Connection) error {
decoder := peer_protocol.Decoder{
R: bufio.NewReader(conn.Socket),
MaxLength: 256 * 1024,
}
for {
me.mu.Unlock()
msg := new(peer_protocol.Message)
err := decoder.Decode(msg)
me.mu.Lock()
if err != nil {
return err
}
log.Print(msg.Type)
if msg.Keepalive {
continue
}
switch msg.Type {
case peer_protocol.Choke:
conn.PeerChoked = true
conn.Requests = nil
case peer_protocol.Unchoke:
conn.PeerChoked = false
me.peerUnchoked(torrent, conn)
case peer_protocol.Interested:
conn.PeerInterested = true
case peer_protocol.NotInterested:
conn.PeerInterested = false
case peer_protocol.Have:
me.peerGotPiece(torrent, conn, int(msg.Index))
case peer_protocol.Request:
conn.PeerRequests[Request{
Index: msg.Index,
ChunkSpec: ChunkSpec{msg.Begin, msg.Length},
}] = struct{}{}
case peer_protocol.Bitfield:
if len(msg.Bitfield) < len(torrent.Pieces) {
err = errors.New("received invalid bitfield")
break
}
if conn.PeerPieces != nil {
err = errors.New("received unexpected bitfield")
break
}
conn.PeerPieces = msg.Bitfield[:len(torrent.Pieces)]
for index, has := range conn.PeerPieces {
if has {
me.peerGotPiece(torrent, conn, index)
}
}
case peer_protocol.Piece:
request_ := Request{msg.Index, ChunkSpec{msg.Begin, peer_protocol.Integer(len(msg.Piece))}}
if _, ok := conn.Requests[request_]; !ok {
err = errors.New("unexpected piece")
break
}
delete(conn.Requests, request_)
err = me.downloadedChunk(torrent, msg)
default:
log.Printf("received unknown message type: %#v", msg.Type)
}
if err != nil {
return err
}
me.replenishConnRequests(torrent, conn)
}
}
func (me *Client) dropConnection(torrent *Torrent, conn *Connection) {
conn.Socket.Close()
for i0, c := range torrent.Conns {
if c != conn {
continue
}
i1 := len(torrent.Conns) - 1
if i0 != i1 {
torrent.Conns[i0] = torrent.Conns[i1]
}
torrent.Conns = torrent.Conns[:i1]
return
}
panic("no such Connection")
}
func (me *Client) addConnection(t *Torrent, c *Connection) bool {
for _, c := range t.Conns {
if c.PeerId == c.PeerId {
return false
}
}
t.Conns = append(t.Conns, c)
return true
}
func (me *Client) openNewConns() {
for _, t := range me.torrents {
for len(t.Peers) != 0 {
if me.halfOpen >= me.HalfOpenLimit {
return
}
p := t.Peers[0]
t.Peers = t.Peers[1:]
me.initiateConn(p, t)
}
}
}
func (me *Client) AddPeers(infoHash InfoHash, peers []Peer) error {
me.mu.Lock()
t := me.torrent(infoHash)
if t == nil {
return errors.New("no such torrent")
}
t.Peers = append(t.Peers, peers...)
me.openNewConns()
me.mu.Unlock()
return nil
}
func (me *Client) AddTorrent(metaInfo *metainfo.MetaInfo) error {
torrent := &Torrent{
InfoHash: BytesInfoHash(metaInfo.InfoHash),
MetaInfo: metaInfo,
}
for offset := 0; offset < len(metaInfo.Pieces); offset += PieceHash.Size() {
hash := metaInfo.Pieces[offset : offset+PieceHash.Size()]
if len(hash) != PieceHash.Size() {
return errors.New("bad piece hash in metainfo")
}
piece := &piece{}
copyHashSum(piece.Hash[:], hash)
torrent.Pieces = append(torrent.Pieces, piece)
}
var err error
torrent.Data, err = mmapTorrentData(metaInfo, me.DataDir)
if err != nil {
return err
}
me.mu.Lock()
defer me.mu.Unlock()
if _, ok := me.torrents[torrent.InfoHash]; ok {
return torrent.Close()
}
me.torrents[torrent.InfoHash] = torrent
return nil
}
func (me *Client) WaitAll() {
me.mu.Lock()
for len(me.torrents) != 0 {
me.event.Wait()
}
me.mu.Unlock()
}
func (me *Client) Stop() {
}
func (me *Client) replenishConnRequests(torrent *Torrent, conn *Connection) {
requestHeatMap := torrent.requestHeat()
addRequest := func(req Request) (again bool) {
piece := torrent.Pieces[req.Index]
if piece.Hashing {
return true
}
if piece.Complete() {
return true
}
if requestHeatMap[req] > 0 {
return true
}
return conn.Request(req)
}
if torrent.Priorities != nil {
for e := torrent.Priorities.Front(); e != nil; e = e.Next() {
if !addRequest(e.Value.(Request)) {
return
}
}
}
for _, index := range torrent.piecesByPendingBytesDesc() {
if torrent.PieceNumPendingBytes(index) == torrent.PieceLength(index) {
continue
}
for chunkSpec := range torrent.Pieces[index].PendingChunkSpecs {
if !addRequest(Request{index, chunkSpec}) {
return
}
}
}
conn.SetInterested(false)
}
func (me *Client) downloadedChunk(torrent *Torrent, msg *peer_protocol.Message) (err error) {
request := Request{msg.Index, ChunkSpec{msg.Begin, peer_protocol.Integer(len(msg.Piece))}}
if _, ok := torrent.Pieces[request.Index].PendingChunkSpecs[request.ChunkSpec]; !ok {
log.Printf("got unnecessary chunk: %s", request)
return
}
log.Printf("got chunk %s", request)
err = torrent.WriteChunk(int(msg.Index), int64(msg.Begin), msg.Piece)
if err != nil {
return
}
delete(torrent.Pieces[request.Index].PendingChunkSpecs, request.ChunkSpec)
if len(torrent.Pieces[request.Index].PendingChunkSpecs) == 0 {
me.queuePieceCheck(torrent, request.Index)
return
}
var next *list.Element
for e := torrent.Priorities.Front(); e != nil; e = next {
next = e.Next()
if e.Value.(Request) == request {
torrent.Priorities.Remove(e)
}
}
me.dataReady(DataSpec{torrent.InfoHash, request})
return
}
func (cl *Client) dataReady(ds DataSpec) {
if cl.DataReady == nil {
return
}
go func() {
cl.DataReady <- ds
}()
}
func (me *Client) pieceHashed(t *Torrent, piece peer_protocol.Integer, correct bool) {
p := t.Pieces[piece]
if !p.Hashing {
panic("invalid state")
}
p.Hashing = false
p.EverHashed = true
if correct {
p.PendingChunkSpecs = nil
var next *list.Element
if t.Priorities != nil {
for e := t.Priorities.Front(); e != nil; e = next {
next = e.Next()
if e.Value.(Request).Index == piece {
t.Priorities.Remove(e)
}
}
}
me.dataReady(DataSpec{
t.InfoHash,
Request{
peer_protocol.Integer(piece),
ChunkSpec{0, peer_protocol.Integer(t.PieceLength(piece))},
},
})
} else {
if len(p.PendingChunkSpecs) == 0 {
p.PendingChunkSpecs = t.pieceChunkSpecs(piece)
}
}
for _, conn := range t.Conns {
if correct {
conn.Post(peer_protocol.Message{
Type: peer_protocol.Have,
Index: peer_protocol.Integer(piece),
})
} else {
if conn.PeerHasPiece(piece) {
me.replenishConnRequests(t, conn)
}
}
}
}
func (cl *Client) verifyPiece(t *Torrent, index peer_protocol.Integer) {
sum := t.HashPiece(index)
cl.mu.Lock()
piece := t.Pieces[index]
cl.pieceHashed(t, index, sum == piece.Hash)
piece.Hashing = false
cl.mu.Unlock()
}
func (me *Client) Torrents() (ret []*Torrent) {
me.mu.Lock()
for _, t := range me.torrents {
ret = append(ret, t)
}
me.mu.Unlock()
return
}