torrent/client.go

package torrent

import (
	"crypto"
	"errors"
	metainfo "github.com/nsf/libtorgo/torrent"
	"io"
	"launchpad.net/gommap"
	"os"
	"path/filepath"
)

const (
	PieceHash = crypto.SHA1
)

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) {
		panic("bad infohash bytes")
	}
	return
}

type pieceState uint8

const (
	pieceStateUnknown = iota
	pieceStateComplete
	pieceStateIncomplete
)

type piece struct {
	State pieceState
	Hash  pieceSum
}

type torrent struct {
	InfoHash infoHash
	Pieces   []piece
	Data     MMapSpan
	MetaInfo *metainfo.MetaInfo
}

func (t torrent) PieceSize(piece int) (size int64) {
	if piece == len(t.Pieces)-1 {
		size = t.Data.Size() % t.MetaInfo.PieceLength
	}
	if size == 0 {
		size = t.MetaInfo.PieceLength
	}
	return
}

func (t torrent) PieceReader(piece int) io.Reader {
	return io.NewSectionReader(t.Data, int64(piece)*t.MetaInfo.PieceLength, t.MetaInfo.PieceLength)
}

func (t torrent) HashPiece(piece int) (ps pieceSum) {
	hash := PieceHash.New()
	n, err := io.Copy(hash, t.PieceReader(piece))
	if err != nil {
		panic(err)
	}
	if n != t.PieceSize(piece) {
		panic("hashed wrong number of bytes")
	}
	copyHashSum(ps[:], hash.Sum(nil))
	return
}

type client struct {
	DataDir string

	noTorrents      chan struct{}
	addTorrent      chan *torrent
	torrents        map[infoHash]*torrent
	torrentFinished chan infoHash
	actorTask       chan func()
}

func NewClient(dataDir string) *client {
	c := &client{
		DataDir: dataDir,

		noTorrents:      make(chan struct{}),
		addTorrent:      make(chan *torrent),
		torrents:        make(map[infoHash]*torrent),
		torrentFinished: make(chan infoHash),
		actorTask:       make(chan func()),
	}
	go c.run()
	return c
}

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), 0666)
		if err != nil {
			return
		}
		var file *os.File
		file, err = os.OpenFile(fileName, os.O_CREATE|os.O_RDWR, 0666)
		if err != nil {
			return
		}
		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})
	}
	return
}

func (me *client) AddTorrent(metaInfo *metainfo.MetaInfo) error {
	torrent := &torrent{
		InfoHash: BytesInfoHash(metaInfo.InfoHash),
	}
	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
	}
	torrent.MetaInfo = metaInfo
	me.addTorrent <- torrent
	return nil
}

func (me *client) WaitAll() {
	<-me.noTorrents
}

func (me *client) Close() {
}

func (me *client) withContext(f func()) {
	me.actorTask <- f
}

func (me *client) pieceHashed(ih infoHash, piece int, correct bool) {
	torrent := me.torrents[ih]
	torrent.Pieces[piece].State = func() pieceState {
		if correct {
			return pieceStateComplete
		} else {
			return pieceStateIncomplete
		}
	}()
	for _, piece := range torrent.Pieces {
		if piece.State == pieceStateUnknown {
			return
		}
	}
	go func() {
		me.torrentFinished <- ih
	}()
}

func (me *client) run() {
	for {
		noTorrents := me.noTorrents
		if len(me.torrents) != 0 {
			noTorrents = nil
		}
		select {
		case noTorrents <- struct{}{}:
		case torrent := <-me.addTorrent:
			if _, ok := me.torrents[torrent.InfoHash]; ok {
				break
			}
			me.torrents[torrent.InfoHash] = torrent
			for i := range torrent.Pieces {
				go func(piece int) {
					sum := torrent.HashPiece(piece)
					me.withContext(func() {
						me.pieceHashed(torrent.InfoHash, piece, sum == torrent.Pieces[piece].Hash)
					})
				}(i)
			}
		case infoHash := <-me.torrentFinished:
			delete(me.torrents, infoHash)
		}
	}
}