From dcfb61fe6587fdc89359f8603c157e94ad8f4245 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 27 Jun 2024 14:33:18 +1000
Subject: [PATCH] Add peer bootstrapping integration test

---
 client.go                         |   1 +
 tests/peers-bootstrapping/main.go | 133 ++++++++++++++++++++++++++++++
 torrent.go                        |   3 +-
 3 files changed, 136 insertions(+), 1 deletion(-)
 create mode 100644 tests/peers-bootstrapping/main.go

diff --git a/client.go b/client.go
index e657564bda..c87defa8b2 100644
--- a/client.go
+++ b/client.go
@@ -238,6 +238,7 @@ func NewClient(cfg *ClientConfig) (cl *Client, err error) {
 	cl.init(cfg)
 	go cl.acceptLimitClearer()
 	cl.initLogger()
+	//cl.logger.Levelf(log.Critical, "test after init")
 	defer func() {
 		if err != nil {
 			cl.Close()
diff --git a/tests/peers-bootstrapping/main.go b/tests/peers-bootstrapping/main.go
new file mode 100644
index 0000000000..14cda31fe9
--- /dev/null
+++ b/tests/peers-bootstrapping/main.go
@@ -0,0 +1,133 @@
+package main
+
+import (
+	"crypto/rand"
+	"fmt"
+	_ "github.com/anacrolix/envpprof"
+	"github.com/anacrolix/log"
+	"github.com/anacrolix/sync"
+	"github.com/anacrolix/torrent"
+	"github.com/anacrolix/torrent/bencode"
+	"github.com/anacrolix/torrent/metainfo"
+	"github.com/dustin/go-humanize"
+	"golang.org/x/exp/slog"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"time"
+)
+
+func assertNil(x any) {
+	if x != nil {
+		panic(x)
+	}
+}
+
+func newClientConfig() *torrent.ClientConfig {
+	cfg := torrent.NewDefaultClientConfig()
+	cfg.ListenPort = 0
+	cfg.NoDHT = true
+	cfg.NoDefaultPortForwarding = true
+	cfg.Seed = true
+	cfg.Debug = false
+	return cfg
+}
+
+func main() {
+	tmpDir, err := os.MkdirTemp("", "peers-bootstrapping")
+	assertNil(err)
+	slog.Info("made temp dir", slog.String("tmpDir", tmpDir))
+	sourceDir := filepath.Join(tmpDir, "source")
+	assertNil(os.Mkdir(sourceDir, 0o700))
+	f, err := os.Create(filepath.Join(sourceDir, "file"))
+	assertNil(err)
+	_, err = io.CopyN(f, rand.Reader, 1<<30)
+	assertNil(err)
+	assertNil(f.Close())
+	var info metainfo.Info
+	err = info.BuildFromFilePath(f.Name())
+	assertNil(err)
+	var mi metainfo.MetaInfo
+	mi.InfoBytes, err = bencode.Marshal(info)
+	assertNil(err)
+	var clients []*torrent.Client
+	var torrents []*torrent.Torrent
+	clientConfig := newClientConfig()
+	clientConfig.DataDir = sourceDir
+	initialClient, err := torrent.NewClient(clientConfig)
+	assertNil(err)
+	clientIndex := 0
+	addClientAndTorrent := func(cl *torrent.Client, t *torrent.Torrent) int {
+		clients = append(clients, cl)
+		torrents = append(torrents, t)
+		ret := clientIndex
+		http.HandleFunc(
+			fmt.Sprintf("/%v", ret),
+			func(w http.ResponseWriter, r *http.Request) {
+				cl.WriteStatus(w)
+			})
+		clientIndex++
+		return ret
+	}
+	initialTorrent, err := initialClient.AddTorrent(&mi)
+	assertNil(err)
+	addClientAndTorrent(initialClient, initialTorrent)
+	//initialTorrent.VerifyData()
+	<-initialTorrent.Complete().On()
+	var allDownloaded sync.WaitGroup
+	var notCompleted sync.Map
+	http.HandleFunc(
+		"/notCompleted",
+		func(w http.ResponseWriter, r *http.Request) {
+			notCompleted.Range(func(key, value any) bool {
+				fmt.Fprintln(w, key)
+				return true
+			})
+		})
+	for range 5 {
+		clientIndex := clientIndex
+		storageDir := filepath.Join(tmpDir, fmt.Sprintf("client%v", clientIndex))
+		clientConfig := newClientConfig()
+		clientConfig.DataDir = storageDir
+		clientConfig.Logger = log.Default.WithValues(slog.Int("clientIndex", clientIndex))
+		//clientConfig.Logger.Levelf(log.Critical, "test")
+		client, err := torrent.NewClient(clientConfig)
+		assertNil(err)
+		t, _ := client.AddTorrentInfoHash(mi.HashInfoBytes())
+		addClientAndTorrent(client, t)
+		allDownloaded.Add(1)
+		notCompleted.Store(clientIndex, nil)
+		go func() {
+			<-t.GotInfo()
+			t.DownloadAll()
+			<-t.Complete().On()
+			notCompleted.Delete(clientIndex)
+			slog.Info("leecher completed", slog.Int("clientIndex", clientIndex))
+			allDownloaded.Done()
+		}()
+		t.AddClientPeer(initialClient)
+	}
+	go func() {
+		for range time.Tick(time.Second) {
+			for _, t := range torrents {
+				for _, cl := range clients {
+					t.AddClientPeer(cl)
+				}
+			}
+		}
+	}()
+	allDownloaded.Wait()
+	slog.Info("all leechers downloaded")
+	for clientIndex, cl := range clients {
+		stats := cl.Stats()
+		written := stats.BytesWritten
+		read := stats.BytesRead
+		fmt.Printf(
+			"client %v wrote %v read %v\n",
+			clientIndex,
+			humanize.Bytes(uint64(written.Int64())),
+			humanize.Bytes(uint64(read.Int64())),
+		)
+	}
+}
diff --git a/torrent.go b/torrent.go
index 440fc891c1..28bda274a2 100644
--- a/torrent.go
+++ b/torrent.go
@@ -520,7 +520,8 @@ func (t *Torrent) setInfo(info *metainfo.Info) error {
 	t.nameMu.Lock()
 	t.info = info
 	t.nameMu.Unlock()
-	t._chunksPerRegularPiece = chunkIndexType((pp.Integer(t.usualPieceSize()) + t.chunkSize - 1) / t.chunkSize)
+	t._chunksPerRegularPiece = chunkIndexType(
+		(pp.Integer(t.usualPieceSize()) + t.chunkSize - 1) / t.chunkSize)
 	t.updateComplete()
 	t.displayName = "" // Save a few bytes lol.
 	t.initFiles()