From b43987fcc8fb738e7fe36f52c0b40c250b741f10 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 20 May 2021 14:01:31 +1000
Subject: [PATCH] Fix up some crashes around piece availability

---
 peerconn.go | 16 +++++++++++-----
 torrent.go  | 15 ++++++++++-----
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/peerconn.go b/peerconn.go
index 976081ab..8e3f5443 100644
--- a/peerconn.go
+++ b/peerconn.go
@@ -396,7 +396,9 @@ func (p *Peer) close() {
 	p.discardPieceInclination()
 	p._pieceRequestOrder.Clear()
 	p.peerImpl.onClose()
-	p.t.decPeerPieceAvailability(p)
+	if p.t != nil {
+		p.t.decPeerPieceAvailability(p)
+	}
 	for _, f := range p.callbacks.PeerClosed {
 		f(p)
 	}
@@ -1651,19 +1653,23 @@ func (cn *Peer) peerMaxRequests() int {
 	return cn.PeerMaxRequests
 }
 
-// Returns the pieces the peer has claimed to have.
+// Returns the pieces the peer could have based on their claims. If we don't know how many pieces
+// are in the torrent, it could be a very large range the peer has sent HaveAll.
 func (cn *PeerConn) PeerPieces() bitmap.Bitmap {
 	cn.locker().RLock()
 	defer cn.locker().RUnlock()
 	return cn.newPeerPieces()
 }
 
-// Returns a new Bitmap that includes bits for all pieces the peer claims to have.
+// Returns a new Bitmap that includes bits for all pieces the peer could have based on their claims.
 func (cn *Peer) newPeerPieces() bitmap.Bitmap {
 	ret := cn._peerPieces.Copy()
 	if cn.peerSentHaveAll {
-
-		ret.AddRange(0, bitmap.BitRange(cn.t.numPieces()))
+		if cn.t.haveInfo() {
+			ret.AddRange(0, bitmap.BitRange(cn.t.numPieces()))
+		} else {
+			ret.AddRange(0, bitmap.ToEnd)
+		}
 	}
 	return ret
 }
diff --git a/torrent.go b/torrent.go
index 5190b65e..e500a47d 100644
--- a/torrent.go
+++ b/torrent.go
@@ -436,16 +436,18 @@ func (t *Torrent) onSetInfo() {
 		p.onGotInfo(t.info)
 	})
 	for i := range t.pieces {
-		t.updatePieceCompletion(pieceIndex(i))
 		p := &t.pieces[i]
-		if !p.storageCompletionOk {
-			// t.logger.Printf("piece %s completion unknown, queueing check", p)
-			t.queuePieceCheck(pieceIndex(i))
-		}
+		// Need to add availability before updating piece completion, as that may result in conns
+		// being dropped.
 		if p.availability != 0 {
 			panic(p.availability)
 		}
 		p.availability = int64(t.pieceAvailabilityFromPeers(i))
+		t.updatePieceCompletion(pieceIndex(i))
+		if !p.storageCompletionOk {
+			// t.logger.Printf("piece %s completion unknown, queueing check", p)
+			t.queuePieceCheck(pieceIndex(i))
+		}
 	}
 	t.cl.event.Broadcast()
 	t.gotMetainfo.Set()
@@ -1374,6 +1376,9 @@ func (t *Torrent) deletePeerConn(c *PeerConn) (ret bool) {
 }
 
 func (t *Torrent) decPeerPieceAvailability(p *Peer) {
+	if !t.haveInfo() {
+		return
+	}
 	p.newPeerPieces().IterTyped(func(i int) bool {
 		p.t.decPieceAvailability(i)
 		return true
-- 
2.48.1