From 598879be0b6c2d295b2d1fa98a3e58de39d1a548 Mon Sep 17 00:00:00 2001
From: Cenk Alti <cenkalti@gmail.com>
Date: Thu, 26 Sep 2019 15:56:00 +0300
Subject: [PATCH] fix fast set generation

---
 internal/fast/fast.go      |  2 +-
 internal/fast/fast_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/internal/fast/fast.go b/internal/fast/fast.go
index bdded0eb..93b2ffed 100644
--- a/internal/fast/fast.go
+++ b/internal/fast/fast.go
@@ -28,7 +28,7 @@ func GenerateFastSet(k int, numPieces uint32, infoHash [20]byte, ip net.IP) []ui
 		return false
 	}
 	h := sha1.New() // nolint: gosec
-	for len(a) < k {
+	for j := 0; j < k && len(a) < k; j++ {
 		_, _ = h.Write(x)
 		x = h.Sum(x[:0])
 		h.Reset()
diff --git a/internal/fast/fast_test.go b/internal/fast/fast_test.go
index ee474c4c..79e50ad7 100644
--- a/internal/fast/fast_test.go
+++ b/internal/fast/fast_test.go
@@ -26,4 +26,28 @@ func TestGenerateFastSet(t *testing.T) {
 		t.Log(a)
 		t.FailNow()
 	}
+
+	a = GenerateFastSet(1, 1313, ih, net.IPv4(80, 4, 4, 200))
+	expected = []uint32{1059}
+	if !reflect.DeepEqual(a, expected) {
+		t.Log(expected)
+		t.Log(a)
+		t.FailNow()
+	}
+
+	a = GenerateFastSet(10, 1, ih, net.IPv4(80, 4, 4, 200))
+	expected = []uint32{0}
+	if !reflect.DeepEqual(a, expected) {
+		t.Log(expected)
+		t.Log(a)
+		t.FailNow()
+	}
+
+	a = GenerateFastSet(4, 2, ih, net.IPv4(80, 4, 4, 200))
+	expected = []uint32{1, 0}
+	if !reflect.DeepEqual(a, expected) {
+		t.Log(expected)
+		t.Log(a)
+		t.FailNow()
+	}
 }