From 3d49a29288cad7ee9e47a92e26fd01f0c9fae4b5 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Mon, 25 Oct 2021 20:52:44 +1100
Subject: [PATCH] Support minimum peer extensions

---
 client.go             | 5 +++++
 cmd/torrent/main.go   | 7 +++++++
 config.go             | 2 ++
 test/issue377_test.go | 3 +++
 testing.go            | 3 +++
 torrent_test.go       | 2 ++
 6 files changed, 22 insertions(+)

diff --git a/client.go b/client.go
index b4d3f17a..6e90fa75 100644
--- a/client.go
+++ b/client.go
@@ -932,6 +932,11 @@ func (cl *Client) runReceivedConn(c *PeerConn) {
 // Client lock must be held before entering this.
 func (cl *Client) runHandshookConn(c *PeerConn, t *Torrent) error {
 	c.setTorrent(t)
+	for i, b := range cl.config.MinPeerExtensions {
+		if c.PeerExtensionBytes[i]&b != b {
+			return fmt.Errorf("peer did not meet minimum peer extensions: %x", c.PeerExtensionBytes)
+		}
+	}
 	if c.PeerID == cl.peerID {
 		if c.outgoing {
 			connsToSelf.Add(1)
diff --git a/cmd/torrent/main.go b/cmd/torrent/main.go
index 06b76f6b..50604cf6 100644
--- a/cmd/torrent/main.go
+++ b/cmd/torrent/main.go
@@ -22,6 +22,7 @@ import (
 	"github.com/anacrolix/missinggo/v2"
 	"github.com/anacrolix/tagflag"
 	"github.com/anacrolix/torrent/bencode"
+	pp "github.com/anacrolix/torrent/peer_protocol"
 	"github.com/anacrolix/torrent/version"
 	"github.com/davecgh/go-spew/spew"
 	"github.com/dustin/go-humanize"
@@ -186,6 +187,9 @@ type DownloadCmd struct {
 	UtpPeers        bool `default:"true"`
 	Webtorrent      bool `default:"true"`
 	DisableWebseeds bool
+	// Don't progress past handshake for peer connections where the peer doesn't offer the fast
+	// extension.
+	RequireFastExtension bool
 
 	Ipv4 bool `default:"true"`
 	Ipv6 bool `default:"true"`
@@ -316,6 +320,9 @@ func downloadErr(flags downloadFlags) error {
 	if flags.Quiet {
 		clientConfig.Logger = log.Discard
 	}
+	if flags.RequireFastExtension {
+		clientConfig.MinPeerExtensions.SetBit(pp.ExtensionBitFast, true)
+	}
 	clientConfig.MaxUnverifiedBytes = flags.MaxUnverifiedBytes.Int64()
 
 	var stop missinggo.SynchronizedEvent
diff --git a/config.go b/config.go
index 14c129cf..3d370d55 100644
--- a/config.go
+++ b/config.go
@@ -146,6 +146,8 @@ type ClientConfig struct {
 	DHTOnQuery func(query *krpc.Msg, source net.Addr) (propagate bool)
 
 	Extensions PeerExtensionBits
+	// Bits that peers must have set to proceed past handshakes.
+	MinPeerExtensions PeerExtensionBits
 
 	DisableWebtorrent bool
 	DisableWebseeds   bool
diff --git a/test/issue377_test.go b/test/issue377_test.go
index 7456e9c4..01caaed6 100644
--- a/test/issue377_test.go
+++ b/test/issue377_test.go
@@ -51,6 +51,9 @@ func testReceiveChunkStorageFailure(t *testing.T, seederFast bool) {
 	defer testutil.ExportStatusWriter(seederClient, "s", t)()
 	leecherClientConfig := torrent.TestingConfig(t)
 	leecherClientConfig.Debug = true
+	// Don't require fast extension, whether the seeder will provide it or not (so we can test mixed
+	// cases).
+	leecherClientConfig.MinPeerExtensions.SetBit(pp.ExtensionBitFast, false)
 	justOneNetwork(leecherClientConfig)
 	leecherClient, err := torrent.NewClient(leecherClientConfig)
 	require.NoError(t, err)
diff --git a/testing.go b/testing.go
index 6a9d892d..d681796d 100644
--- a/testing.go
+++ b/testing.go
@@ -3,6 +3,8 @@ package torrent
 import (
 	"testing"
 	"time"
+
+	pp "github.com/anacrolix/torrent/peer_protocol"
 )
 
 func TestingConfig(t testing.TB) *ClientConfig {
@@ -15,6 +17,7 @@ func TestingConfig(t testing.TB) *ClientConfig {
 	cfg.DisableAcceptRateLimiting = true
 	cfg.ListenPort = 0
 	cfg.KeepAliveTimeout = time.Millisecond
+	cfg.MinPeerExtensions.SetBit(pp.ExtensionBitFast, true)
 	//cfg.Debug = true
 	//cfg.Logger = cfg.Logger.WithText(func(m log.Msg) string {
 	//	t := m.Text()
diff --git a/torrent_test.go b/torrent_test.go
index 02e56e2d..2b96f191 100644
--- a/torrent_test.go
+++ b/torrent_test.go
@@ -163,6 +163,8 @@ func TestPieceHashFailed(t *testing.T) {
 func TestTorrentMetainfoIncompleteMetadata(t *testing.T) {
 	cfg := TestingConfig(t)
 	cfg.Debug = true
+	// Disable this just because we manually initiate a connection without it.
+	cfg.MinPeerExtensions.SetBit(pp.ExtensionBitFast, false)
 	cl, err := NewClient(cfg)
 	require.NoError(t, err)
 	defer cl.Close()
-- 
2.48.1