From 09477db2bc0535c93a09442c6b5216497b3bc7db Mon Sep 17 00:00:00 2001
From: Bruno Dal Bo <brunodalbo@google.com>
Date: Sat, 18 Sep 2021 05:57:28 -0700
Subject: Avoid ambient clock on ICMP Rate Limiter
PiperOrigin-RevId: 397496920
---
pkg/tcpip/stack/icmp_rate_limit.go | 39 ++++++++++++++++++++++++++++++++++----
pkg/tcpip/stack/stack.go | 2 +-
2 files changed, 36 insertions(+), 5 deletions(-)
(limited to 'pkg/tcpip/stack')
diff --git a/pkg/tcpip/stack/icmp_rate_limit.go b/pkg/tcpip/stack/icmp_rate_limit.go
index 3a20839da..99e5d2df7 100644
--- a/pkg/tcpip/stack/icmp_rate_limit.go
+++ b/pkg/tcpip/stack/icmp_rate_limit.go
@@ -16,6 +16,7 @@ package stack
import (
"golang.org/x/time/rate"
+ "gvisor.dev/gvisor/pkg/tcpip"
)
const (
@@ -31,11 +32,41 @@ const (
// ICMPRateLimiter is a global rate limiter that controls the generation of
// ICMP messages generated by the stack.
type ICMPRateLimiter struct {
- *rate.Limiter
+ limiter *rate.Limiter
+ clock tcpip.Clock
}
// NewICMPRateLimiter returns a global rate limiter for controlling the rate
-// at which ICMP messages are generated by the stack.
-func NewICMPRateLimiter() *ICMPRateLimiter {
- return &ICMPRateLimiter{Limiter: rate.NewLimiter(icmpLimit, icmpBurst)}
+// at which ICMP messages are generated by the stack. The returned limiter
+// does not apply limits to any ICMP types by default.
+func NewICMPRateLimiter(clock tcpip.Clock) *ICMPRateLimiter {
+ return &ICMPRateLimiter{
+ clock: clock,
+ limiter: rate.NewLimiter(icmpLimit, icmpBurst),
+ }
+}
+
+// SetLimit sets a new Limit for the limiter.
+func (l *ICMPRateLimiter) SetLimit(limit rate.Limit) {
+ l.limiter.SetLimitAt(l.clock.Now(), limit)
+}
+
+// Limit returns the maximum overall event rate.
+func (l *ICMPRateLimiter) Limit() rate.Limit {
+ return l.limiter.Limit()
+}
+
+// SetBurst sets a new burst size for the limiter.
+func (l *ICMPRateLimiter) SetBurst(burst int) {
+ l.limiter.SetBurstAt(l.clock.Now(), burst)
+}
+
+// Burst returns the maximum burst size.
+func (l *ICMPRateLimiter) Burst() int {
+ return l.limiter.Burst()
+}
+
+// Allow reports whether one ICMP message may be sent now.
+func (l *ICMPRateLimiter) Allow() bool {
+ return l.limiter.AllowN(l.clock.Now(), 1)
}
diff --git a/pkg/tcpip/stack/stack.go b/pkg/tcpip/stack/stack.go
index 98867a828..428350f31 100644
--- a/pkg/tcpip/stack/stack.go
+++ b/pkg/tcpip/stack/stack.go
@@ -375,7 +375,7 @@ func New(opts Options) *Stack {
stats: opts.Stats.FillIn(),
handleLocal: opts.HandleLocal,
tables: opts.IPTables,
- icmpRateLimiter: NewICMPRateLimiter(),
+ icmpRateLimiter: NewICMPRateLimiter(clock),
seed: seed,
nudConfigs: opts.NUDConfigs,
uniqueIDGenerator: opts.UniqueID,
--
cgit v1.2.3