fix fast set generation

cenkalti · Sep 26, 2019 · 598879b · 598879b
1 parent da66340
commit 598879b
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git 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
@@ -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()
+ }
 }