From bddf5aaa2f737569e5f96de3676c5123b82e60f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felf=C3=B6ldi=20Zsolt?= <zsfelfoldi@gmail.com>
Date: Mon, 23 Nov 2020 10:18:33 +0100
Subject: [PATCH] les/utils: protect against WeightedRandomSelect overflow
 (#21839)

Also fixes a bug in les/flowcontrol that caused the overflow.
---
 les/flowcontrol/control.go   |  4 ++++
 les/utils/weighted_select.go | 11 +++++++++++
 2 files changed, 15 insertions(+)

diff --git a/les/flowcontrol/control.go b/les/flowcontrol/control.go
index 490013677..4f0de8231 100644
--- a/les/flowcontrol/control.go
+++ b/les/flowcontrol/control.go
@@ -19,6 +19,7 @@ package flowcontrol
 
 import (
 	"fmt"
+	"math"
 	"sync"
 	"time"
 
@@ -316,6 +317,9 @@ func (node *ServerNode) CanSend(maxCost uint64) (time.Duration, float64) {
 	node.lock.RLock()
 	defer node.lock.RUnlock()
 
+	if node.params.BufLimit == 0 {
+		return time.Duration(math.MaxInt64), 0
+	}
 	now := node.clock.Now()
 	node.recalcBLE(now)
 	maxCost += uint64(safetyMargin) * node.params.MinRecharge / uint64(fcTimeConst)
diff --git a/les/utils/weighted_select.go b/les/utils/weighted_select.go
index d6db3c0e6..9b4413afb 100644
--- a/les/utils/weighted_select.go
+++ b/les/utils/weighted_select.go
@@ -17,7 +17,10 @@
 package utils
 
 import (
+	"math"
 	"math/rand"
+
+	"github.com/ethereum/go-ethereum/log"
 )
 
 type (
@@ -54,6 +57,14 @@ func (w *WeightedRandomSelect) IsEmpty() bool {
 
 // setWeight sets an item's weight to a specific value (removes it if zero)
 func (w *WeightedRandomSelect) setWeight(item WrsItem, weight uint64) {
+	if weight > math.MaxInt64-w.root.sumWeight {
+		// old weight is still included in sumWeight, remove and check again
+		w.setWeight(item, 0)
+		if weight > math.MaxInt64-w.root.sumWeight {
+			log.Error("WeightedRandomSelect overflow", "sumWeight", w.root.sumWeight, "new weight", weight)
+			weight = math.MaxInt64 - w.root.sumWeight
+		}
+	}
 	idx, ok := w.idx[item]
 	if ok {
 		w.root.setWeight(idx, weight)