Skip to content

Commit

Permalink
fix: edge case with ~uint64 weight >= platform maxInt
Browse files Browse the repository at this point in the history
Exposed during fuzzing, possible  edge (very edge!) case where when
using a Chooser[any, ~uint64], if a single Choice weight exceeds system
maxInt, then when it gets converted to system int type it no longer
triggers the the check for exceeding maxInt in the running total,
resulting in a Chooser that can panic on Pick.
  • Loading branch information
mroth committed Jul 27, 2023
1 parent 647d2a6 commit bedcf02
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 5 deletions.
2 changes: 2 additions & 0 deletions testdata/fuzz/FuzzNewChooser/a547669aeb7ca0ca
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
go test fuzz v1
[]byte("0000000\x8f00000000")
14 changes: 10 additions & 4 deletions weightedrand.go
Original file line number Diff line number Diff line change
Expand Up @@ -48,11 +48,16 @@ func NewChooser[T any, W integer](choices ...Choice[T, W]) (*Chooser[T, W], erro
totals := make([]int, len(choices))
runningTotal := 0
for i, c := range choices {
weight := int(c.Weight)
if weight < 0 {
if c.Weight < 0 {
continue // ignore negative weights, can never be picked
}

// case of single ~uint64 or similar value that exceeds maxInt on its own
if uint64(c.Weight) >= maxInt {
return nil, errWeightOverflow
}

weight := int(c.Weight) // convert weight to int for internal counter usage
if (maxInt - runningTotal) <= weight {
return nil, errWeightOverflow
}
Expand All @@ -68,8 +73,9 @@ func NewChooser[T any, W integer](choices ...Choice[T, W]) (*Chooser[T, W], erro
}

const (
intSize = 32 << (^uint(0) >> 63) // cf. strconv.IntSize
maxInt = 1<<(intSize-1) - 1
intSize = 32 << (^uint(0) >> 63) // cf. strconv.IntSize
maxInt = 1<<(intSize-1) - 1
maxUint64 = 1<<64 - 1
)

// Possible errors returned by NewChooser, preventing the creation of a Chooser
Expand Down
34 changes: 33 additions & 1 deletion weightedrand_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -81,10 +81,42 @@ func TestNewChooser(t *testing.T) {
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
_, err := NewChooser(tt.cs...)
c, err := NewChooser(tt.cs...)
if err != tt.wantErr {
t.Errorf("NewChooser() error = %v, wantErr %v", err, tt.wantErr)
}

if err == nil { // run a few Picks to make sure there are no panics
for i := 0; i < 10; i++ {
_ = c.Pick()
}
}
})
}

u64tests := []struct {
name string
cs []Choice[rune, uint64]
wantErr error
}{
{
name: "weight overflow from single uint64 exceeding system maxInt",
cs: []Choice[rune, uint64]{{Item: 'a', Weight: maxInt + 1}},
wantErr: errWeightOverflow,
},
}
for _, tt := range u64tests {
t.Run(tt.name, func(t *testing.T) {
c, err := NewChooser(tt.cs...)
if err != tt.wantErr {
t.Errorf("NewChooser() error = %v, wantErr %v", err, tt.wantErr)
}

if err == nil { // run a few Picks to make sure there are no panics
for i := 0; i < 10; i++ {
_ = c.Pick()
}
}
})
}
}
Expand Down

0 comments on commit bedcf02

Please sign in to comment.