From 7909084a177861f37921f449858cd4c8bd597ef4 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Tue, 2 Jun 2020 16:18:25 +1000
Subject: [PATCH] Fix webseeds when info isn't available immediately

---
 peerconn.go | 14 +++++++++-----
 torrent.go  | 15 +++++++--------
 web_seed.go |  7 +++++++
 3 files changed, 23 insertions(+), 13 deletions(-)

diff --git a/peerconn.go b/peerconn.go
index 0c328c22..af98660f 100644
--- a/peerconn.go
+++ b/peerconn.go
@@ -18,6 +18,7 @@ import (
 	"github.com/anacrolix/missinggo/v2/bitmap"
 	"github.com/anacrolix/missinggo/v2/prioritybitmap"
 	"github.com/anacrolix/multiless"
+	"github.com/anacrolix/torrent/metainfo"
 	"github.com/pkg/errors"
 
 	"github.com/anacrolix/torrent/bencode"
@@ -44,6 +45,7 @@ type PeerImpl interface {
 	ConnectionFlags() string
 	Close()
 	PostCancel(request)
+	onGotInfo(*metainfo.Info)
 	Drop()
 }
 
@@ -233,13 +235,15 @@ func (cn *peer) completedString() string {
 	return fmt.Sprintf("%d/%d", have, cn.bestPeerNumPieces())
 }
 
-// Correct the PeerPieces slice length. Return false if the existing slice is
-// invalid, such as by receiving badly sized BITFIELD, or invalid HAVE
-// messages.
-func (cn *PeerConn) setNumPieces(num pieceIndex) error {
+func (cn *PeerConn) onGotInfo(info *metainfo.Info) {
+	cn.setNumPieces(info.NumPieces())
+}
+
+// Correct the PeerPieces slice length. Return false if the existing slice is invalid, such as by
+// receiving badly sized BITFIELD, or invalid HAVE messages.
+func (cn *PeerConn) setNumPieces(num pieceIndex) {
 	cn._peerPieces.RemoveRange(bitmap.BitIndex(num), bitmap.ToEnd)
 	cn.peerPiecesChanged()
-	return nil
 }
 
 func eventAgeString(t time.Time) string {
diff --git a/torrent.go b/torrent.go
index 21859af3..a5225518 100644
--- a/torrent.go
+++ b/torrent.go
@@ -406,13 +406,11 @@ func (t *Torrent) setInfo(info *metainfo.Info) error {
 	return nil
 }
 
+// This seems to be all the follow-up tasks after info is set, that can't fail.
 func (t *Torrent) onSetInfo() {
-	for conn := range t.conns {
-		if err := conn.setNumPieces(t.numPieces()); err != nil {
-			t.logger.Printf("closing connection: %s", err)
-			conn.close()
-		}
-	}
+	t.iterPeers(func(p *peer) {
+		p.onGotInfo(t.info)
+	})
 	for i := range t.pieces {
 		t.updatePieceCompletion(pieceIndex(i))
 		p := &t.pieces[i]
@@ -2026,12 +2024,13 @@ func (t *Torrent) addWebSeed(url string) {
 		client: webseed.Client{
 			HttpClient: http.DefaultClient,
 			Url:        url,
-			FileIndex:  t.fileIndex,
-			Info:       t.info,
 		},
 		requests: make(map[request]webseed.Request, maxRequests),
 	}
 	ws.peer.PeerImpl = &ws
+	if t.haveInfo() {
+		ws.onGotInfo(t.info)
+	}
 	t.webSeeds[url] = &ws.peer
 }
 
diff --git a/web_seed.go b/web_seed.go
index 13acc41d..3163ac4d 100644
--- a/web_seed.go
+++ b/web_seed.go
@@ -3,6 +3,8 @@ package torrent
 import (
 	"net/http"
 
+	"github.com/anacrolix/torrent/common"
+	"github.com/anacrolix/torrent/metainfo"
 	pp "github.com/anacrolix/torrent/peer_protocol"
 	"github.com/anacrolix/torrent/segments"
 	"github.com/anacrolix/torrent/webseed"
@@ -31,6 +33,11 @@ type webSeed struct {
 
 var _ PeerImpl = (*webSeed)(nil)
 
+func (ws *webSeed) onGotInfo(info *metainfo.Info) {
+	ws.client.FileIndex = segments.NewIndex(common.LengthIterFromUpvertedFiles(info.UpvertedFiles()))
+	ws.client.Info = info
+}
+
 func (ws *webSeed) PostCancel(r request) {
 	ws.Cancel(r)
 }