From e0c2265cb02a65fa45fbb0b441e1a6d077723bdf Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Sun, 4 Feb 2018 19:10:25 +1100
Subject: [PATCH] Support sending HaveAll

---
 client.go          | 23 +++++++++++++++--------
 connection.go      | 32 ++++++++++++++++----------------
 connection_test.go |  6 +++++-
 torrent.go         | 16 +++++++++++-----
 4 files changed, 47 insertions(+), 30 deletions(-)

diff --git a/client.go b/client.go
index d2dac0fa..81a7936b 100644
--- a/client.go
+++ b/client.go
@@ -893,13 +893,20 @@ func (cl *Client) sendInitialMessages(conn *connection, torrent *Torrent) {
 			}(),
 		})
 	}
-	if torrent.haveAnyPieces() {
-		conn.Bitfield(torrent.bitfield())
-	} else if cl.extensionBytes.SupportsFast() && conn.PeerExtensionBytes.SupportsFast() {
-		conn.Post(pp.Message{
-			Type: pp.HaveNone,
-		})
-	}
+	func() {
+		if conn.fastEnabled() {
+			if torrent.haveAllPieces() {
+				conn.Post(pp.Message{Type: pp.HaveAll})
+				conn.sentHaves.AddRange(0, conn.t.NumPieces())
+				return
+			} else if !torrent.haveAnyPieces() {
+				conn.Post(pp.Message{Type: pp.HaveNone})
+				conn.sentHaves.Clear()
+				return
+			}
+		}
+		conn.PostBitfield()
+	}()
 	if conn.PeerExtensionBytes.SupportsDHT() && cl.extensionBytes.SupportsDHT() && cl.dHT != nil {
 		conn.Post(pp.Message{
 			Type: pp.Port,
@@ -1101,7 +1108,7 @@ func (cl *Client) allTorrentsCompleted() bool {
 		if !t.haveInfo() {
 			return false
 		}
-		if t.numPiecesCompleted() != t.numPieces() {
+		if !t.haveAllPieces() {
 			return false
 		}
 	}
diff --git a/connection.go b/connection.go
index ada84b8e..7e475659 100644
--- a/connection.go
+++ b/connection.go
@@ -66,7 +66,7 @@ type connection struct {
 	// Indexed by metadata piece, set to true if posted and pending a
 	// response.
 	metadataRequests []bool
-	sentHaves        []bool
+	sentHaves        bitmap.Bitmap
 
 	// Stuff controlled by the remote peer.
 	PeerID             PeerID
@@ -141,7 +141,7 @@ func (cn *connection) completedString() string {
 // invalid, such as by receiving badly sized BITFIELD, or invalid HAVE
 // messages.
 func (cn *connection) setNumPieces(num int) error {
-	cn.peerPieces.RemoveRange(num, -1)
+	cn.peerPieces.RemoveRange(num, bitmap.ToEnd)
 	cn.peerPiecesChanged()
 	return nil
 }
@@ -468,31 +468,28 @@ func (cn *connection) writer(keepAliveTimeout time.Duration) {
 }
 
 func (cn *connection) Have(piece int) {
-	for piece >= len(cn.sentHaves) {
-		cn.sentHaves = append(cn.sentHaves, false)
-	}
-	if cn.sentHaves[piece] {
+	if cn.sentHaves.Get(piece) {
 		return
 	}
 	cn.Post(pp.Message{
 		Type:  pp.Have,
 		Index: pp.Integer(piece),
 	})
-	cn.sentHaves[piece] = true
+	cn.sentHaves.Add(piece)
 }
 
-func (cn *connection) Bitfield(haves []bool) {
-	if cn.sentHaves != nil {
+func (cn *connection) PostBitfield() {
+	if cn.sentHaves.Len() != 0 {
 		panic("bitfield must be first have-related message sent")
 	}
+	if !cn.t.haveAnyPieces() {
+		return
+	}
 	cn.Post(pp.Message{
 		Type:     pp.Bitfield,
-		Bitfield: haves,
+		Bitfield: cn.t.bitfield(),
 	})
-	// Make a copy of haves, as that's read when the message is marshalled
-	// without the lock. Also it obviously shouldn't change in the Msg due to
-	// changes in .sentHaves.
-	cn.sentHaves = append([]bool(nil), haves...)
+	cn.sentHaves = cn.t.completedPieces.Copy()
 }
 
 // Determines interest and requests to send to a connected peer.
@@ -553,8 +550,11 @@ func iterBitmapsDistinct(skip bitmap.Bitmap, bms ...bitmap.Bitmap) iter.Func {
 
 func (cn *connection) unbiasedPieceRequestOrder() iter.Func {
 	now, readahead := cn.t.readerPiecePriorities()
-	// Pieces to skip include pieces the peer doesn't have
-	skip := bitmap.Flip(cn.peerPieces, 0, cn.t.numPieces())
+	var skip bitmap.Bitmap
+	if !cn.peerSentHaveAll {
+		// Pieces to skip include pieces the peer doesn't have
+		skip = bitmap.Flip(cn.peerPieces, 0, cn.t.numPieces())
+	}
 	// And pieces that we already have.
 	skip.Union(cn.t.completedPieces)
 	// Return an iterator over the different priority classes, minus the skip
diff --git a/connection_test.go b/connection_test.go
index 4d4c7465..6b4ccd6b 100644
--- a/connection_test.go
+++ b/connection_test.go
@@ -24,11 +24,15 @@ func TestSendBitfieldThenHave(t *testing.T) {
 	cl.initLogger()
 	c := cl.newConnection(nil)
 	c.setTorrent(cl.newTorrent(metainfo.Hash{}, nil))
+	c.t.setInfo(&metainfo.Info{
+		Pieces: make([]byte, metainfo.HashSize*3),
+	})
 	c.r = r
 	c.w = w
 	go c.writer(time.Minute)
 	c.mu().Lock()
-	c.Bitfield([]bool{false, true, false})
+	c.t.completedPieces.Add(1)
+	c.PostBitfield( /*[]bool{false, true, false}*/ )
 	c.mu().Unlock()
 	c.mu().Lock()
 	c.Have(2)
diff --git a/torrent.go b/torrent.go
index 12bc48f5..fca91798 100644
--- a/torrent.go
+++ b/torrent.go
@@ -735,6 +735,10 @@ type Peer struct {
 }
 
 func (t *Torrent) pieceLength(piece int) pp.Integer {
+	if t.info.PieceLength == 0 {
+		// There will be no variance amongst pieces. Only pain.
+		return 0
+	}
 	if piece == t.numPieces()-1 {
 		ret := pp.Integer(*t.length % t.info.PieceLength)
 		if ret != 0 {
@@ -762,12 +766,14 @@ func (t *Torrent) hashPiece(piece int) (ret metainfo.Hash) {
 }
 
 func (t *Torrent) haveAnyPieces() bool {
-	for i := range t.pieces {
-		if t.pieceComplete(i) {
-			return true
-		}
+	return t.completedPieces.Len() != 0
+}
+
+func (t *Torrent) haveAllPieces() bool {
+	if !t.haveInfo() {
+		return false
 	}
-	return false
+	return t.completedPieces.Len() == t.numPieces()
 }
 
 func (t *Torrent) havePiece(index int) bool {