Fix threshold calculation (#1722)

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,

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,