From 75cde79d65d86571cd365e8bb51808dd8d10a5ce Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Mon, 1 Dec 2014 03:37:33 -0600
Subject: [PATCH] Improve the interface with client and download strategies

---
 client.go           | 29 ++++++++++++++++++++++++-----
 cmd/torrent/main.go | 45 +++++++++++++++++++++++----------------------
 2 files changed, 47 insertions(+), 27 deletions(-)

diff --git a/client.go b/client.go
index 74141627..2be26443 100644
--- a/client.go
+++ b/client.go
@@ -107,6 +107,12 @@ func (me *Client) PrioritizeDataRegion(ih InfoHash, off, len_ int64) error {
 	if !t.haveInfo() {
 		return errors.New("missing metadata")
 	}
+	i := int(off / int64(t.UsualPieceSize()))
+	me.queueFirstHash(t, i)
+	i++
+	if i < t.NumPieces() {
+		me.queueFirstHash(t, i)
+	}
 	me.downloadStrategy.TorrentPrioritize(t, off, len_)
 	for _, cn := range t.Conns {
 		me.replenishConnRequests(t, cn)
@@ -1430,6 +1436,18 @@ type Torrent struct {
 	*torrent
 }
 
+func (t Torrent) AddPeers(pp []Peer) error {
+	return t.cl.AddPeers(t.torrent.InfoHash, pp)
+}
+
+func (t Torrent) DownloadAll() {
+	t.cl.mu.Lock()
+	for i := 0; i < t.NumPieces(); i++ {
+		t.cl.queueFirstHash(t.torrent, i)
+	}
+	t.cl.mu.Unlock()
+}
+
 func (me Torrent) ReadAt(p []byte, off int64) (n int, err error) {
 	err = me.cl.PrioritizeDataRegion(me.InfoHash, off, int64(len(p)))
 	if err != nil {
@@ -1525,27 +1543,28 @@ func (me *Client) addTorrent(t *torrent) (err error) {
 }
 
 // Adds the torrent to the client.
-func (me *Client) AddTorrent(metaInfo *metainfo.MetaInfo) (err error) {
+func (me *Client) AddTorrent(metaInfo *metainfo.MetaInfo) (t Torrent, err error) {
+	t.cl = me
 	var ih InfoHash
 	CopyExact(&ih, metaInfo.Info.Hash)
-	t, err := newTorrent(ih, metaInfo.AnnounceList, me.halfOpenLimit)
+	t.torrent, err = newTorrent(ih, metaInfo.AnnounceList, me.halfOpenLimit)
 	if err != nil {
 		return
 	}
 	me.mu.Lock()
 	defer me.mu.Unlock()
-	err = me.addTorrent(t)
+	err = me.addTorrent(t.torrent)
 	if err != nil {
 		return
 	}
-	err = me.setMetaData(t, metaInfo.Info.Info, metaInfo.Info.Bytes)
+	err = me.setMetaData(t.torrent, metaInfo.Info.Info, metaInfo.Info.Bytes)
 	if err != nil {
 		return
 	}
 	return
 }
 
-func (me *Client) AddTorrentFromFile(name string) (err error) {
+func (me *Client) AddTorrentFromFile(name string) (t Torrent, err error) {
 	mi, err := metainfo.LoadFromFile(name)
 	if err != nil {
 		err = fmt.Errorf("error loading metainfo from file: %s", err)
diff --git a/cmd/torrent/main.go b/cmd/torrent/main.go
index af65f227..121a7b83 100644
--- a/cmd/torrent/main.go
+++ b/cmd/torrent/main.go
@@ -51,30 +51,27 @@ func main() {
 		return
 	}
 	for _, arg := range flag.Args() {
-		var ih torrent.InfoHash
-		if strings.HasPrefix(arg, "magnet:") {
-			m, err := torrent.ParseMagnetURI(arg)
-			if err != nil {
-				log.Fatalf("error parsing magnet uri: %s", err)
-			}
-			ih = m.InfoHash
-			_, err = client.AddMagnet(arg)
-			if err != nil {
-				log.Fatalf("error adding magnet: %s", err)
-			}
-		} else {
-			metaInfo, err := metainfo.LoadFromFile(arg)
-			if err != nil {
-				log.Fatal(err)
+		t := func() torrent.Torrent {
+			if strings.HasPrefix(arg, "magnet:") {
+				t, err := client.AddMagnet(arg)
+				if err != nil {
+					log.Fatalf("error adding magnet: %s", err)
+				}
+				return t
+			} else {
+				metaInfo, err := metainfo.LoadFromFile(arg)
+				if err != nil {
+					log.Fatal(err)
+				}
+				t, err := client.AddTorrent(metaInfo)
+				if err != nil {
+					log.Fatal(err)
+				}
+				return t
 			}
-			err = client.AddTorrent(metaInfo)
-			if err != nil {
-				log.Fatal(err)
-			}
-			util.CopyExact(&ih, metaInfo.Info.Hash)
-		}
+		}()
 		// client.PrioritizeDataRegion(ih, 0, 999999999)
-		err := client.AddPeers(ih, func() []torrent.Peer {
+		err := t.AddPeers(func() []torrent.Peer {
 			if *testPeer == "" {
 				return nil
 			}
@@ -90,6 +87,10 @@ func main() {
 		if err != nil {
 			log.Fatal(err)
 		}
+		go func() {
+			<-t.GotMetainfo
+			t.DownloadAll()
+		}()
 	}
 	if *seed {
 		select {}
-- 
2.48.1