diff --git a/p2p/host/resource-manager/scope.go b/p2p/host/resource-manager/scope.go index 872f0eb686..2a83c69508 100644 --- a/p2p/host/resource-manager/scope.go +++ b/p2p/host/resource-manager/scope.go @@ -2,6 +2,8 @@ package rcmgr import ( "fmt" + "math" + "math/big" "strings" "sync" @@ -79,14 +81,55 @@ func IsSpan(name string) bool { return strings.Contains(name, ".span-") } +func addInt64WithOverflow(a int64, b int64) (c int64, ok bool) { + c = a + b + return c, (c > a) == (b > 0) +} + +// mulInt64WithOverflow checks for overflow in multiplying two int64s. See +// https://groups.google.com/g/golang-nuts/c/h5oSN5t3Au4/m/KaNQREhZh0QJ +func mulInt64WithOverflow(a, b int64) (c int64, ok bool) { + const mostPositive = 1<<63 - 1 + const mostNegative = -(mostPositive + 1) + c = a * b + if a == 0 || b == 0 || a == 1 || b == 1 { + return c, true + } + if a == mostNegative || b == mostNegative { + return c, false + } + return c, c/b == a +} + // Resources implementation func (rc *resources) checkMemory(rsvp int64, prio uint8) error { - // overflow check; this also has the side effect that we cannot reserve negative memory. - newmem := rc.memory + rsvp + if rsvp < 0 { + return fmt.Errorf("can't reserve negative memory. rsvp=%v", rsvp) + } + limit := rc.limit.GetMemoryLimit() - threshold := (1 + int64(prio)) * limit / 256 + if limit == math.MaxInt64 { + // Special case where we've set max limits. + return nil + } + + newmem, addOk := addInt64WithOverflow(rc.memory, rsvp) + + threshold, mulOk := mulInt64WithOverflow(1+int64(prio), limit) + if !mulOk { + thresholdBig := big.NewInt(limit) + thresholdBig = thresholdBig.Mul(thresholdBig, big.NewInt(1+int64(prio))) + thresholdBig.Rsh(thresholdBig, 8) // Divide 256 + if !thresholdBig.IsInt64() { + // Shouldn't happen since the threshold can only be <= limit + threshold = limit + } + threshold = thresholdBig.Int64() + } else { + threshold = threshold / 256 + } - if newmem > threshold { + if !addOk || newmem > threshold { return &errMemoryLimitExceeded{ current: rc.memory, attempted: rsvp, diff --git a/p2p/host/resource-manager/scope_test.go b/p2p/host/resource-manager/scope_test.go index af192a7bdc..0464cb0c05 100644 --- a/p2p/host/resource-manager/scope_test.go +++ b/p2p/host/resource-manager/scope_test.go @@ -1,9 +1,12 @@ package rcmgr import ( + "math" "testing" + "testing/quick" "github.com/libp2p/go-libp2p/core/network" + "github.com/stretchr/testify/require" ) func checkResources(t *testing.T, rc *resources, st network.ScopeStat) { @@ -29,6 +32,77 @@ func checkResources(t *testing.T, rc *resources, st network.ScopeStat) { } } +func TestCheckMemory(t *testing.T) { + t.Run("overflows", func(t *testing.T) { + rc := resources{limit: &BaseLimit{ + Memory: math.MaxInt64 - 1, + StreamsInbound: 1, + StreamsOutbound: 1, + Streams: 1, + ConnsInbound: 1, + ConnsOutbound: 1, + Conns: 1, + FD: 1, + }} + rc.memory = math.MaxInt64 - 1 + require.Error(t, rc.checkMemory(2, network.ReservationPriorityAlways)) + + rc.memory = 1024 + require.NoError(t, rc.checkMemory(1, network.ReservationPriorityAlways)) + }) + + t.Run("negative mem", func(t *testing.T) { + rc := resources{limit: &BaseLimit{ + Memory: math.MaxInt64, + StreamsInbound: 1, + StreamsOutbound: 1, + Streams: 1, + ConnsInbound: 1, + ConnsOutbound: 1, + Conns: 1, + FD: 1, + }} + rc.memory = math.MaxInt64 + + require.Error(t, rc.checkMemory(-1, network.ReservationPriorityAlways)) + }) + + f := func(limit uint64, res uint64, currentMem uint64, priShift uint8) bool { + limit = (limit % math.MaxInt64) + 1 + if limit < 1024 { + // We set the min to 1KiB + limit = 1024 + } + currentMem = (currentMem % limit) // We can't have reserved more than our limit + res = (res >> 14) // We won't resonably ever have a reservation > 2^50 + rc := resources{limit: &BaseLimit{ + Memory: int64(limit), + StreamsInbound: 1, + StreamsOutbound: 1, + Streams: 1, + ConnsInbound: 1, + ConnsOutbound: 1, + Conns: 1, + FD: 1, + }} + rc.memory = int64(currentMem) + + priShift = (priShift % 9) + // Check different priorties at 2^0, 2^1,...2^8. This lets our math be correct in the check below (and avoid overflows). + pri := uint8((1 << priShift) - 1) + + err := rc.checkMemory(int64(res), pri) + if limit == math.MaxInt64 && err == nil { + // Special case logic + return true + } + + return (err != nil) == (uint64(res)+uint64(rc.memory) > (uint64(limit) >> uint64(8-priShift))) + } + + require.NoError(t, quick.Check(f, nil)) +} + func TestResources(t *testing.T) { rc := resources{limit: &BaseLimit{ Memory: 4096,