les/utils: UDP rate limiter (#21930)

* les/utils: Limiter

* les/utils: dropped prior weight vs variable cost logic, using fixed weights

* les/utils: always create node selector in addressGroup

* les/utils: renamed request weight to request cost

* les/utils: simplified and improved the DoS penalty mechanism

* les/utils: minor fixes

* les/utils: made selection weight calculation nicer

* les/utils: fixed linter warning

* les/utils: more precise and reliable probabilistic test

* les/utils: fixed linter warning

les/utils/limiter_test.go

+// Copyright 2020 The go-ethereum Authors
+// This file is part of the go-ethereum library.
+//
+// The go-ethereum library is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Lesser General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// The go-ethereum library is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU Lesser General Public License for more details.
+//
+// You should have received a copy of the GNU Lesser General Public License
+// along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
+
+package utils
+
+import (
+	"math/rand"
+	"testing"
+
+	"github.com/ethereum/go-ethereum/p2p/enode"
+)
+
+const (
+	ltTolerance = 0.03
+	ltRounds    = 7
+)
+
+type (
+	ltNode struct {
+		addr, id         int
+		value, exp       float64
+		cost             uint
+		reqRate          float64
+		reqMax, runCount int
+		lastTotalCost    uint
+
+		served, dropped int
+	}
+
+	ltResult struct {
+		node *ltNode
+		ch   chan struct{}
+	}
+
+	limTest struct {
+		limiter            *Limiter
+		results            chan ltResult
+		runCount           int
+		expCost, totalCost uint
+	}
+)
+
+func (lt *limTest) request(n *ltNode) {
+	var (
+		address string
+		id      enode.ID
+	)
+	if n.addr >= 0 {
+		address = string([]byte{byte(n.addr)})
+	} else {
+		var b [32]byte
+		rand.Read(b[:])
+		address = string(b[:])
+	}
+	if n.id >= 0 {
+		id = enode.ID{byte(n.id)}
+	} else {
+		rand.Read(id[:])
+	}
+	lt.runCount++
+	n.runCount++
+	cch := lt.limiter.Add(id, address, n.value, n.cost)
+	go func() {
+		lt.results <- ltResult{n, <-cch}
+	}()
+}
+
+func (lt *limTest) moreRequests(n *ltNode) {
+	maxStart := int(float64(lt.totalCost-n.lastTotalCost) * n.reqRate)
+	if maxStart != 0 {
+		n.lastTotalCost = lt.totalCost
+	}
+	for n.reqMax > n.runCount && maxStart > 0 {
+		lt.request(n)
+		maxStart--
+	}
+}
+
+func (lt *limTest) process() {
+	res := <-lt.results
+	lt.runCount--
+	res.node.runCount--
+	if res.ch != nil {
+		res.node.served++
+		if res.node.exp != 0 {
+			lt.expCost += res.node.cost
+		}
+		lt.totalCost += res.node.cost
+		close(res.ch)
+	} else {
+		res.node.dropped++
+	}
+}
+
+func TestLimiter(t *testing.T) {
+	limTests := [][]*ltNode{
+		{ // one id from an individual address and two ids from a shared address
+			{addr: 0, id: 0, value: 0, cost: 1, reqRate: 1, reqMax: 1, exp: 0.5},
+			{addr: 1, id: 1, value: 0, cost: 1, reqRate: 1, reqMax: 1, exp: 0.25},
+			{addr: 1, id: 2, value: 0, cost: 1, reqRate: 1, reqMax: 1, exp: 0.25},
+		},
+		{ // varying request costs
+			{addr: 0, id: 0, value: 0, cost: 10, reqRate: 0.2, reqMax: 1, exp: 0.5},
+			{addr: 1, id: 1, value: 0, cost: 3, reqRate: 0.5, reqMax: 1, exp: 0.25},
+			{addr: 1, id: 2, value: 0, cost: 1, reqRate: 1, reqMax: 1, exp: 0.25},
+		},
+		{ // different request rate
+			{addr: 0, id: 0, value: 0, cost: 1, reqRate: 2, reqMax: 2, exp: 0.5},
+			{addr: 1, id: 1, value: 0, cost: 1, reqRate: 10, reqMax: 10, exp: 0.25},
+			{addr: 1, id: 2, value: 0, cost: 1, reqRate: 1, reqMax: 1, exp: 0.25},
+		},
+		{ // adding value
+			{addr: 0, id: 0, value: 3, cost: 1, reqRate: 1, reqMax: 1, exp: (0.5 + 0.3) / 2},
+			{addr: 1, id: 1, value: 0, cost: 1, reqRate: 1, reqMax: 1, exp: 0.25 / 2},
+			{addr: 1, id: 2, value: 7, cost: 1, reqRate: 1, reqMax: 1, exp: (0.25 + 0.7) / 2},
+		},
+		{ // DoS attack from a single address with a single id
+			{addr: 0, id: 0, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 1, id: 1, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 2, id: 2, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 3, id: 3, value: 0, cost: 1, reqRate: 10, reqMax: 1000000000, exp: 0},
+		},
+		{ // DoS attack from a single address with different ids
+			{addr: 0, id: 0, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 1, id: 1, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 2, id: 2, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 3, id: -1, value: 0, cost: 1, reqRate: 1, reqMax: 1000000000, exp: 0},
+		},
+		{ // DDoS attack from different addresses with a single id
+			{addr: 0, id: 0, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 1, id: 1, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 2, id: 2, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: -1, id: 3, value: 0, cost: 1, reqRate: 1, reqMax: 1000000000, exp: 0},
+		},
+		{ // DDoS attack from different addresses with different ids
+			{addr: 0, id: 0, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 1, id: 1, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: 2, id: 2, value: 1, cost: 1, reqRate: 1, reqMax: 1, exp: 0.3333},
+			{addr: -1, id: -1, value: 0, cost: 1, reqRate: 1, reqMax: 1000000000, exp: 0},
+		},
+	}
+
+	lt := &limTest{
+		limiter: NewLimiter(100),
+		results: make(chan ltResult),
+	}
+	for _, test := range limTests {
+		lt.expCost, lt.totalCost = 0, 0
+		iterCount := 10000
+		for j := 0; j < ltRounds; j++ {
+			// try to reach expected target range in multiple rounds with increasing iteration counts
+			last := j == ltRounds-1
+			for _, n := range test {
+				lt.request(n)
+			}
+			for i := 0; i < iterCount; i++ {
+				lt.process()
+				for _, n := range test {
+					lt.moreRequests(n)
+				}
+			}
+			for lt.runCount > 0 {
+				lt.process()
+			}
+			if spamRatio := 1 - float64(lt.expCost)/float64(lt.totalCost); spamRatio > 0.5*(1+ltTolerance) {
+				t.Errorf("Spam ratio too high (%f)", spamRatio)
+			}
+			fail, success := false, true
+			for _, n := range test {
+				if n.exp != 0 {
+					if n.dropped > 0 {
+						t.Errorf("Dropped %d requests of non-spam node", n.dropped)
+						fail = true
+					}
+					r := float64(n.served) * float64(n.cost) / float64(lt.expCost)
+					if r < n.exp*(1-ltTolerance) || r > n.exp*(1+ltTolerance) {
+						if last {
+							// print error only if the target is still not reached in the last round
+							t.Errorf("Request ratio (%f) does not match expected value (%f)", r, n.exp)
+						}
+						success = false
+					}
+				}
+			}
+			if fail || success {
+				break
+			}
+			// neither failed nor succeeded; try more iterations to reach probability targets
+			iterCount *= 2
+		}
+	}
+	lt.limiter.Stop()
+}

les/utils/weighted_select.go

@@ -52,17 +52,17 @@ func (w *WeightedRandomSelect) Remove(item WrsItem) {
 
 // IsEmpty returns true if the set is empty
 func (w *WeightedRandomSelect) IsEmpty() bool {
-	return w.root.sumWeight == 0
+	return w.root.sumCost == 0
 }
 
 // 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
+	if weight > math.MaxInt64-w.root.sumCost {
+		// old weight is still included in sumCost, 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
+		if weight > math.MaxInt64-w.root.sumCost {
+			log.Error("WeightedRandomSelect overflow", "sumCost", w.root.sumCost, "new weight", weight)
+			weight = math.MaxInt64 - w.root.sumCost
 		}
 	}
 	idx, ok := w.idx[item]
@@ -75,9 +75,9 @@ func (w *WeightedRandomSelect) setWeight(item WrsItem, weight uint64) {
 		if weight != 0 {
 			if w.root.itemCnt == w.root.maxItems {
 				// add a new level
-				newRoot := &wrsNode{sumWeight: w.root.sumWeight, itemCnt: w.root.itemCnt, level: w.root.level + 1, maxItems: w.root.maxItems * wrsBranches}
+				newRoot := &wrsNode{sumCost: w.root.sumCost, itemCnt: w.root.itemCnt, level: w.root.level + 1, maxItems: w.root.maxItems * wrsBranches}
 				newRoot.items[0] = w.root
-				newRoot.weights[0] = w.root.sumWeight
+				newRoot.weights[0] = w.root.sumCost
 				w.root = newRoot
 			}
 			w.idx[item] = w.root.insert(item, weight)
@@ -91,10 +91,10 @@ func (w *WeightedRandomSelect) setWeight(item WrsItem, weight uint64) {
 // updates its weight and selects another one
 func (w *WeightedRandomSelect) Choose() WrsItem {
 	for {
-		if w.root.sumWeight == 0 {
+		if w.root.sumCost == 0 {
 			return nil
 		}
-		val := uint64(rand.Int63n(int64(w.root.sumWeight)))
+		val := uint64(rand.Int63n(int64(w.root.sumCost)))
 		choice, lastWeight := w.root.choose(val)
 		weight := w.wfn(choice)
 		if weight != lastWeight {
@@ -112,7 +112,7 @@ const wrsBranches = 8 // max number of branches in the wrsNode tree
 type wrsNode struct {
 	items                    [wrsBranches]interface{}
 	weights                  [wrsBranches]uint64
-	sumWeight                uint64
+	sumCost                  uint64
 	level, itemCnt, maxItems int
 }
 
@@ -126,7 +126,7 @@ func (n *wrsNode) insert(item WrsItem, weight uint64) int {
 		}
 	}
 	n.itemCnt++
-	n.sumWeight += weight
+	n.sumCost += weight
 	n.weights[branch] += weight
 	if n.level == 0 {
 		n.items[branch] = item
@@ -150,7 +150,7 @@ func (n *wrsNode) setWeight(idx int, weight uint64) uint64 {
 		oldWeight := n.weights[idx]
 		n.weights[idx] = weight
 		diff := weight - oldWeight
-		n.sumWeight += diff
+		n.sumCost += diff
 		if weight == 0 {
 			n.items[idx] = nil
 			n.itemCnt--
@@ -161,7 +161,7 @@ func (n *wrsNode) setWeight(idx int, weight uint64) uint64 {
 	branch := idx / branchItems
 	diff := n.items[branch].(*wrsNode).setWeight(idx-branch*branchItems, weight)
 	n.weights[branch] += diff
-	n.sumWeight += diff
+	n.sumCost += diff
 	if weight == 0 {
 		n.itemCnt--
 	}