diff --git a/les/utils/weighted_select.go b/les/utils/weighted_select.go index d6db3c0e657c..cce06da93003 100644 --- a/les/utils/weighted_select.go +++ b/les/utils/weighted_select.go @@ -83,13 +83,13 @@ func (w *WeightedRandomSelect) Choose() WrsItem { if w.root.sumWeight == 0 { return nil } - val := uint64(rand.Int63n(int64(w.root.sumWeight))) + val := rand.Uint64() % w.root.sumWeight choice, lastWeight := w.root.choose(val) weight := w.wfn(choice) if weight != lastWeight { w.setWeight(choice, weight) } - if weight >= lastWeight || uint64(rand.Int63n(int64(lastWeight))) < weight { + if weight >= lastWeight || rand.Uint64()%lastWeight < weight { return choice } } diff --git a/les/utils/weighted_select_test.go b/les/utils/weighted_select_test.go index 3e1c0ad9873a..0098f8a66573 100644 --- a/les/utils/weighted_select_test.go +++ b/les/utils/weighted_select_test.go @@ -66,3 +66,17 @@ func TestWeightedRandomSelect(t *testing.T) { testFn(100000) testFn(1000000) } + +// TestOOB tests values which doesn't fit in int64 +func TestOOB(t *testing.T) { + s := NewWeightedRandomSelect(func(i interface{}) uint64 { + // Dummy weight function to return a very large weight + return uint64(0xffffffffffffffff) + }) + s.Update(testWrsItem{idx: 0, widx: nil}) + // int64 conversion should make the sumweight negative + if int64(s.root.sumWeight) >= 0 { + t.Fatalf("test is dysfunctional, sumweight not negative: %d", int64(s.root.sumWeight)) + } + s.Choose() +}