From 2e9ecd5a320ac9963abfe6209b250459013d5f6f Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Mon, 27 Dec 2021 22:19:04 +1100
Subject: [PATCH] Reject peer requests on data read failures

---
 peerconn.go | 31 +++++++++++++++++--------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/peerconn.go b/peerconn.go
index c71581ac..6866d450 100644
--- a/peerconn.go
+++ b/peerconn.go
@@ -527,12 +527,7 @@ func (cn *PeerConn) choke(msg messageWriter) (more bool) {
 	more = msg(pp.Message{
 		Type: pp.Choke,
 	})
-	if cn.fastEnabled() {
-		for r := range cn.peerRequests {
-			// TODO: Don't reject pieces in allowed fast set.
-			cn.reject(r)
-		}
-	} else {
+	if !cn.fastEnabled() {
 		cn.peerRequests = nil
 	}
 	return
@@ -1016,7 +1011,6 @@ func (c *PeerConn) onReadRequest(r Request) error {
 	value := &peerRequestState{}
 	c.peerRequests[r] = value
 	go c.peerRequestDataReader(r, value)
-	// c.tickleWriter()
 	return nil
 }
 
@@ -1058,14 +1052,23 @@ func (c *PeerConn) peerRequestDataReadFailed(err error, r Request) {
 		// here.
 		c.t.updatePieceCompletion(i)
 	}
-	// If we failed to send a chunk, choke the peer to ensure they flush all their requests. We've
-	// probably dropped a piece from storage, but there's no way to communicate this to the peer. If
-	// they ask for it again, we'll kick them to allow us to send them an updated bitfield on the
-	// next connect. TODO: Support rejecting here too.
-	if c.choking {
-		c.logger.WithDefaultLevel(log.Warning).Printf("already choking peer, requests might not be rejected correctly")
+	// We've probably dropped a piece from storage, but there's no way to communicate this to the
+	// peer. If they ask for it again, we kick them allowing us to send them updated piece states if
+	// we reconnect. TODO: Instead, we could just try to update them with Bitfield or HaveNone and
+	// if they kick us for breaking protocol, on reconnect we will be compliant again (at least
+	// initially).
+	if c.fastEnabled() {
+		c.reject(r)
+	} else {
+		if c.choking {
+			// If fast isn't enabled, I think we would have wiped all peer requests when we last
+			// choked, and requests while we're choking would be ignored. It could be possible that
+			// a peer request data read completed concurrently to it being deleted elsewhere.
+			c.logger.WithDefaultLevel(log.Warning).Printf("already choking peer, requests might not be rejected correctly")
+		}
+		// Choking a non-fast peer should cause them to flush all their requests.
+		c.choke(c.write)
 	}
-	c.choke(c.write)
 }
 
 func readPeerRequestData(r Request, c *PeerConn) ([]byte, error) {
-- 
2.44.0