From 29238ff23f1aa0b9d72b760a942b8a41e494f031 Mon Sep 17 00:00:00 2001 From: Matt Joiner Date: Wed, 18 Mar 2015 18:32:31 +1100 Subject: [PATCH] Change the interface to add torrents to a Client --- client.go | 172 ++++++++++++++++++++++++++++--------------------- client_test.go | 41 +++++++----- torrent.go | 6 +- 3 files changed, 128 insertions(+), 91 deletions(-) diff --git a/client.go b/client.go index 2746aedb..7ff99428 100644 --- a/client.go +++ b/client.go @@ -1335,7 +1335,7 @@ func (cl *Client) completedMetadata(t *torrent) { } // TODO(anacrolix): If this fails, I think something harsher should be // done. - err = cl.setMetaData(t, info, t.MetaData) + err = cl.setMetaData(t, &info, t.MetaData) if err != nil { log.Printf("error setting metadata: %s", err) t.invalidateMetadata() @@ -1844,7 +1844,7 @@ func (cl *Client) setStorage(t *torrent, td data.Data) (err error) { type TorrentDataOpener func(*metainfo.Info) data.Data -func (cl *Client) setMetaData(t *torrent, md metainfo.Info, bytes []byte) (err error) { +func (cl *Client) setMetaData(t *torrent, md *metainfo.Info, bytes []byte) (err error) { err = t.setMetadata(md, bytes, &cl.mu) if err != nil { return @@ -1854,13 +1854,14 @@ func (cl *Client) setMetaData(t *torrent, md metainfo.Info, bytes []byte) (err e log.Printf("error saving torrent file for %s: %s", t, err) } } + cl.event.Broadcast() if strings.Contains(strings.ToLower(md.Name), "porn") { cl.dropTorrent(t.InfoHash) err = errors.New("no porn plx") return } close(t.gotMetainfo) - td := cl.torrentDataOpener(&md) + td := cl.torrentDataOpener(md) err = cl.setStorage(t, td) return } @@ -2172,28 +2173,89 @@ func (cl *Client) torrentCacheMetaInfo(ih InfoHash) (mi *metainfo.MetaInfo, err return } -func (cl *Client) AddMagnet(uri string) (T Torrent, err error) { +// For adding new torrents to a client. +type TorrentSpec struct { + Trackers [][]string + InfoHash InfoHash + Info *metainfo.InfoEx + DisplayName string +} + +func TorrentSpecFromMagnetURI(uri string) (spec *TorrentSpec, err error) { m, err := ParseMagnetURI(uri) if err != nil { return } - mi, err := cl.torrentCacheMetaInfo(m.InfoHash) - if err != nil { - log.Printf("error getting cached metainfo for %x: %s", m.InfoHash[:], err) - } else if mi != nil { - _, err = cl.AddTorrent(mi) - if err != nil { - return - } + spec = &TorrentSpec{ + Trackers: [][]string{m.Trackers}, + DisplayName: m.DisplayName, } + CopyExact(&spec.InfoHash, &m.InfoHash) + return +} + +func TorrentSpecFromMetaInfo(mi *metainfo.MetaInfo) (spec *TorrentSpec) { + spec = &TorrentSpec{ + Trackers: mi.AnnounceList, + Info: &mi.Info, + } + CopyExact(&spec.InfoHash, &mi.Info.Hash) + return +} + +func (cl *Client) AddTorrentSpec(spec *TorrentSpec) (T Torrent, new bool, err error) { + T.cl = cl cl.mu.Lock() defer cl.mu.Unlock() - T, err = cl.addOrMergeTorrent(m.InfoHash, [][]string{m.Trackers}) + + t, ok := cl.torrents[spec.InfoHash] + if ok { + T.torrent = t + return + } + + new = true + + if _, ok := cl.bannedTorrents[spec.InfoHash]; ok { + err = errors.New("banned torrent") + return + } + + t, err = newTorrent(spec.InfoHash) if err != nil { return } - if m.DisplayName != "" { - T.DisplayName = m.DisplayName + if spec.DisplayName != "" { + t.DisplayName = spec.DisplayName + } + if spec.Info != nil { + err = cl.setMetaData(t, &spec.Info.Info, spec.Info.Bytes) + } else { + var mi *metainfo.MetaInfo + mi, err = cl.torrentCacheMetaInfo(spec.InfoHash) + if err != nil { + log.Printf("error getting cached metainfo: %s", err) + } else if mi != nil { + t.addTrackers(mi.AnnounceList) + err = cl.setMetaData(t, &mi.Info.Info, mi.Info.Bytes) + } + } + if err != nil { + return + } + + cl.torrents[spec.InfoHash] = t + T.torrent = t + + T.torrent.pruneTimer = time.AfterFunc(0, func() { + cl.pruneConnectionsUnlocked(T.torrent) + }) + t.addTrackers(spec.Trackers) + if !cl.disableTrackers { + go cl.announceTorrentTrackers(T.torrent) + } + if cl.dHT != nil { + go cl.announceTorrentDHT(T.torrent, true) } return } @@ -2241,63 +2303,6 @@ func (me *Client) dropTorrent(infoHash InfoHash) (err error) { return } -func (me *Client) addOrMergeTorrent(ih InfoHash, announceList [][]string) (T Torrent, err error) { - if _, ok := me.bannedTorrents[ih]; ok { - err = errors.New("banned torrent") - return - } - T.cl = me - var ok bool - T.torrent, ok = me.torrents[ih] - if ok { - T.torrent.addTrackers(announceList) - } else { - T.torrent, err = newTorrent(ih, announceList, me.halfOpenLimit) - if err != nil { - return - } - me.torrents[ih] = T.torrent - if !me.disableTrackers { - go me.announceTorrentTrackers(T.torrent) - } - if me.dHT != nil { - go me.announceTorrentDHT(T.torrent, true) - } - T.torrent.pruneTimer = time.AfterFunc(0, func() { - me.pruneConnectionsUnlocked(T.torrent) - }) - } - return -} - -// Adds a torrent to the client. -func (me *Client) AddTorrent(metaInfo *metainfo.MetaInfo) (t Torrent, err error) { - var ih InfoHash - CopyExact(&ih, metaInfo.Info.Hash) - me.mu.Lock() - defer me.mu.Unlock() - t, err = me.addOrMergeTorrent(ih, metaInfo.AnnounceList) - if err != nil { - return - } - if !t.torrent.haveInfo() { - err = me.setMetaData(t.torrent, metaInfo.Info.Info, metaInfo.Info.Bytes) - if err != nil { - return - } - } - return -} - -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) - return - } - return me.AddTorrent(mi) -} - // Returns true when peers are required, or false if the torrent is closing. func (cl *Client) waitWantPeers(t *torrent) bool { cl.mu.Lock() @@ -2711,3 +2716,26 @@ func (me *Client) Torrents() (ret []Torrent) { me.mu.Unlock() return } + +func (me *Client) AddMagnet(uri string) (T Torrent, err error) { + spec, err := TorrentSpecFromMagnetURI(uri) + if err != nil { + return + } + T, _, err = me.AddTorrentSpec(spec) + return +} + +func (me *Client) AddTorrent(mi *metainfo.MetaInfo) (T Torrent, err error) { + T, _, err = me.AddTorrentSpec(TorrentSpecFromMetaInfo(mi)) + return +} + +func (me *Client) AddTorrentFromFile(filename string) (T Torrent, err error) { + mi, err := metainfo.LoadFromFile(filename) + if err != nil { + return + } + T, _, err = me.AddTorrentSpec(TorrentSpecFromMetaInfo(mi)) + return +} diff --git a/client_test.go b/client_test.go index c7445c00..8d0caf06 100644 --- a/client_test.go +++ b/client_test.go @@ -21,8 +21,12 @@ import ( "github.com/anacrolix/libtorgo/bencode" ) +func init() { + log.SetFlags(log.LstdFlags | log.Lshortfile) +} + var TestingConfig = Config{ - ListenAddr: ":0", + ListenAddr: "localhost:0", NoDHT: true, DisableTrackers: true, NoDefaultBlocklist: true, @@ -30,10 +34,7 @@ var TestingConfig = Config{ } func TestClientDefault(t *testing.T) { - cl, err := NewClient(&Config{ - NoDefaultBlocklist: true, - ListenAddr: ":0", - }) + cl, err := NewClient(&TestingConfig) if err != nil { t.Fatal(err) } @@ -48,10 +49,13 @@ func TestAddDropTorrent(t *testing.T) { defer cl.Close() dir, mi := testutil.GreetingTestTorrent() defer os.RemoveAll(dir) - tt, err := cl.AddTorrent(mi) + tt, new, err := cl.AddTorrentSpec(TorrentSpecFromMetaInfo(mi)) if err != nil { t.Fatal(err) } + if !new { + t.FailNow() + } tt.Drop() } @@ -79,11 +83,11 @@ func TestTorrentInitialState(t *testing.T) { tor, err := newTorrent(func() (ih InfoHash) { util.CopyExact(ih[:], mi.Info.Hash) return - }(), nil, 0) + }()) if err != nil { t.Fatal(err) } - err = tor.setMetadata(mi.Info.Info, mi.Info.Bytes, nil) + err = tor.setMetadata(&mi.Info.Info, mi.Info.Bytes, nil) if err != nil { t.Fatal(err) } @@ -215,9 +219,7 @@ func TestUTPRawConn(t *testing.T) { func TestTwoClientsArbitraryPorts(t *testing.T) { for i := 0; i < 2; i++ { - cl, err := NewClient(&Config{ - ListenAddr: ":0", - }) + cl, err := NewClient(&TestingConfig) if err != nil { t.Fatal(err) } @@ -228,10 +230,17 @@ func TestTwoClientsArbitraryPorts(t *testing.T) { func TestAddDropManyTorrents(t *testing.T) { cl, _ := NewClient(&TestingConfig) defer cl.Close() - var m Magnet for i := range iter.N(1000) { - binary.PutVarint(m.InfoHash[:], int64(i)) - cl.AddMagnet(m.String()) + var spec TorrentSpec + binary.PutVarint(spec.InfoHash[:], int64(i)) + tt, new, err := cl.AddTorrentSpec(&spec) + if err != nil { + t.Error(err) + } + if !new { + t.FailNow() + } + defer tt.Drop() } } @@ -245,7 +254,7 @@ func TestClientTransfer(t *testing.T) { t.Fatal(err) } defer seeder.Close() - seeder.AddTorrent(mi) + seeder.AddTorrentSpec(TorrentSpecFromMetaInfo(mi)) leecherDataDir, err := ioutil.TempDir("", "") if err != nil { t.Fatal(err) @@ -257,7 +266,7 @@ func TestClientTransfer(t *testing.T) { cfg.TorrentDataOpener = blob.NewStore(leecherDataDir).OpenTorrent leecher, _ := NewClient(&cfg) defer leecher.Close() - leecherGreeting, _ := leecher.AddTorrent(mi) + leecherGreeting, _, _ := leecher.AddTorrentSpec(TorrentSpecFromMetaInfo(mi)) leecherGreeting.AddPeers([]Peer{ Peer{ IP: util.AddrIP(seeder.ListenAddr()), diff --git a/torrent.go b/torrent.go index d9b4b0f5..7bd88ddd 100644 --- a/torrent.go +++ b/torrent.go @@ -291,15 +291,15 @@ func infoPieceHashes(info *metainfo.Info) (ret []string) { } // Called when metadata for a torrent becomes available. -func (t *torrent) setMetadata(md metainfo.Info, infoBytes []byte, eventLocker sync.Locker) (err error) { - t.Info = &md +func (t *torrent) setMetadata(md *metainfo.Info, infoBytes []byte, eventLocker sync.Locker) (err error) { + t.Info = md t.length = 0 for _, f := range t.Info.UpvertedFiles() { t.length += f.Length } t.MetaData = infoBytes t.metadataHave = nil - for _, hash := range infoPieceHashes(&md) { + for _, hash := range infoPieceHashes(md) { piece := &piece{} piece.Event.L = eventLocker util.CopyExact(piece.Hash[:], hash)