From 07552a0cd89c85a1b43c448cf71a87a3d0d208a0 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Sun, 18 Feb 2024 12:21:11 +1100
Subject: [PATCH] Add smart ban hash benchmarks and reduce allocations

---
 go.mod               |  1 +
 go.sum               |  1 +
 smartban/smartban.go | 23 +++++++++++++----------
 smartban_test.go     | 39 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 10 deletions(-)
 create mode 100644 smartban_test.go

diff --git a/go.mod b/go.mod
index 68120007..1b0f5816 100644
--- a/go.mod
+++ b/go.mod
@@ -25,6 +25,7 @@ require (
 	github.com/anacrolix/utp v0.1.0
 	github.com/bahlo/generic-list-go v0.2.0
 	github.com/bradfitz/iter v0.0.0-20191230175014-e8f45d346db8
+	github.com/cespare/xxhash v1.1.0
 	github.com/davecgh/go-spew v1.1.1
 	github.com/dustin/go-humanize v1.0.0
 	github.com/edsrzf/mmap-go v1.1.0
diff --git a/go.sum b/go.sum
index ecf85199..c8ce30f9 100644
--- a/go.sum
+++ b/go.sum
@@ -142,6 +142,7 @@ github.com/bradfitz/iter v0.0.0-20191230175014-e8f45d346db8/go.mod h1:spo1JLcs67
 github.com/cenkalti/backoff/v4 v4.1.3 h1:cFAlzYUlVYDysBEH2T5hyJZMh3+5+WCBvSnK6Q8UtC4=
 github.com/cenkalti/backoff/v4 v4.1.3/go.mod h1:scbssz8iZGpm3xbr14ovlUdkxfGXNInqkPWOWmG2CLw=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
+github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
 github.com/cespare/xxhash/v2 v2.1.2 h1:YRXhKfTDauu4ajMg1TPgFO5jnlC2HCbmLXMcTG5cbYE=
diff --git a/smartban/smartban.go b/smartban/smartban.go
index 96e9b759..5a9050eb 100644
--- a/smartban/smartban.go
+++ b/smartban/smartban.go
@@ -1,6 +1,7 @@
 package smartban
 
 import (
+	g "github.com/anacrolix/generics"
 	"sync"
 )
 
@@ -8,7 +9,7 @@ type Cache[Peer, BlockKey, Hash comparable] struct {
 	Hash func([]byte) Hash
 
 	lock   sync.RWMutex
-	blocks map[BlockKey]map[Peer]Hash
+	blocks map[BlockKey][]peerAndHash[Peer, Hash]
 }
 
 type Block[Key any] struct {
@@ -16,8 +17,13 @@ type Block[Key any] struct {
 	Data []byte
 }
 
+type peerAndHash[Peer, Hash any] struct {
+	Peer Peer
+	Hash Hash
+}
+
 func (me *Cache[Peer, BlockKey, Hash]) Init() {
-	me.blocks = make(map[BlockKey]map[Peer]Hash)
+	g.MakeMap(&me.blocks)
 }
 
 func (me *Cache[Peer, BlockKey, Hash]) RecordBlock(peer Peer, key BlockKey, data []byte) {
@@ -25,20 +31,17 @@ func (me *Cache[Peer, BlockKey, Hash]) RecordBlock(peer Peer, key BlockKey, data
 	me.lock.Lock()
 	defer me.lock.Unlock()
 	peers := me.blocks[key]
-	if peers == nil {
-		peers = make(map[Peer]Hash)
-		me.blocks[key] = peers
-	}
-	peers[peer] = hash
+	peers = append(peers, peerAndHash[Peer, Hash]{peer, hash})
+	me.blocks[key] = peers
 }
 
 func (me *Cache[Peer, BlockKey, Hash]) CheckBlock(key BlockKey, data []byte) (bad []Peer) {
 	correct := me.Hash(data)
 	me.lock.RLock()
 	defer me.lock.RUnlock()
-	for peer, hash := range me.blocks[key] {
-		if hash != correct {
-			bad = append(bad, peer)
+	for _, item := range me.blocks[key] {
+		if item.Hash != correct {
+			bad = append(bad, item.Peer)
 		}
 	}
 	return
diff --git a/smartban_test.go b/smartban_test.go
new file mode 100644
index 00000000..2947f521
--- /dev/null
+++ b/smartban_test.go
@@ -0,0 +1,39 @@
+package torrent
+
+import (
+	"crypto/sha1"
+	"github.com/anacrolix/missinggo/v2/iter"
+	"github.com/anacrolix/torrent/smartban"
+	"github.com/cespare/xxhash"
+	"net/netip"
+	"testing"
+)
+
+func benchmarkSmartBanRecordBlock[Sum comparable](b *testing.B, hash func([]byte) Sum) {
+	var cache smartban.Cache[bannableAddr, RequestIndex, Sum]
+	cache.Hash = hash
+	cache.Init()
+	var data [defaultChunkSize]byte
+	var addr netip.Addr
+	b.SetBytes(int64(len(data)))
+	for i := range iter.N(b.N) {
+		cache.RecordBlock(addr, RequestIndex(i), data[:])
+	}
+}
+
+func BenchmarkSmartBanRecordBlock(b *testing.B) {
+	b.Run("xxHash", func(b *testing.B) {
+		var salt [8]byte
+		benchmarkSmartBanRecordBlock(b, func(block []byte) uint64 {
+			h := xxhash.New()
+			// xxHash is not cryptographic, and so we're salting it so attackers can't know a priori
+			// where block data collisions are.
+			h.Write(salt[:])
+			h.Write(block)
+			return h.Sum64()
+		})
+	})
+	b.Run("Sha1", func(b *testing.B) {
+		benchmarkSmartBanRecordBlock(b, sha1.Sum)
+	})
+}
-- 
2.48.1