feat(bt): support disabling DHT and PEX (#1555)

internal/protocol/bt/config.go

@@ -9,4 +9,8 @@ type config struct {
 	SeedRatio float64 `json:"seedRatio"`
 	// SeedTime is the time in seconds to seed after downloading is complete.
 	SeedTime int64 `json:"seedTime"`
+	// DisableDHT and DisablePEX turn off DHT and PEX peer discovery. Like ListenPort,
+	// they are read when the shared client is built.
+	DisableDHT bool `json:"disableDht"`
+	DisablePEX bool `json:"disablePex"`
 }

internal/protocol/bt/fetcher.go

@@ -77,6 +77,8 @@ func (f *Fetcher) initClient() (err error) {
 	cfg.Bep20 = fmt.Sprintf("-GP%s-", parseBep20())
 	cfg.ExtendedHandshakeClientVersion = fmt.Sprintf("Gopeed %s", base.Version)
 	cfg.ListenPort = f.config.ListenPort
+	cfg.NoDHT = f.config.DisableDHT
+	cfg.DisablePEX = f.config.DisablePEX
 	cfg.HTTPProxy = f.ctl.GetProxy(f.meta.Req.Proxy)
 	dnsResolver := &DnsCacheResolver{RefreshTimeout: 5 * time.Minute}
 	cfg.TrackerDialContext = dnsResolver.DialContext
@@ -839,6 +841,8 @@ func (fm *FetcherManager) DefaultConfig() any {
 		SeedKeep:   false,
 		SeedRatio:  1.0,
 		SeedTime:   120 * 60,
+		DisableDHT: false,
+		DisablePEX: false,
 	}
 }
 

internal/protocol/bt/fetcher_test.go

@@ -366,3 +366,52 @@ func TestFetcher_Patch(t *testing.T) {
 		}
 	})
 }
+
+func TestFetcher_PeerDiscoveryConfig(t *testing.T) {
+	tests := []struct {
+		name           string
+		stored         any
+		wantDisableDHT bool
+		wantDisablePEX bool
+	}{
+		{"default config", new(FetcherManager).DefaultConfig(), false, false},
+		{"stored config without the fields", map[string]any{"listenPort": 0, "seedRatio": 1.0}, false, false},
+		{"both disabled", map[string]any{"disableDht": true, "disablePex": true}, true, true},
+		{"only DHT disabled", map[string]any{"disableDht": true}, true, false},
+		{"only PEX disabled", map[string]any{"disablePex": true}, false, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			f := buildFetcherWithConfig(tt.stored).(*Fetcher)
+			if f.config.DisableDHT != tt.wantDisableDHT || f.config.DisablePEX != tt.wantDisablePEX {
+				t.Errorf("config got DisableDHT = %v, DisablePEX = %v, want %v, %v", f.config.DisableDHT, f.config.DisablePEX, tt.wantDisableDHT, tt.wantDisablePEX)
+			}
+		})
+	}
+}
+
+func TestFetcher_PeerDiscoveryDisabledOnClient(t *testing.T) {
+	if err := closeClient(); err != nil {
+		t.Fatal(err)
+	}
+	f := buildFetcherWithConfig(map[string]any{"disableDht": true, "disablePex": true})
+	if err := f.Resolve(&base.Request{URL: "./testdata/test.torrent"}, nil); err != nil {
+		t.Fatal(err)
+	}
+	if !cfg.NoDHT || !cfg.DisablePEX {
+		t.Errorf("client config got NoDHT = %v, DisablePEX = %v, want true", cfg.NoDHT, cfg.DisablePEX)
+	}
+	if n := len(client.DhtServers()); n != 0 {
+		t.Errorf("client got %d DHT servers, want 0", n)
+	}
+}
+
+func buildFetcherWithConfig(stored any) fetcher.Fetcher {
+	fetcher := new(FetcherManager).Build()
+	newController := controller.NewController()
+	newController.GetConfig = func(v any) {
+		json.Unmarshal([]byte(test.ToJson(stored)), v)
+	}
+	fetcher.Setup(newController)
+	return fetcher
+}