derp/derpserver: implement hierarchical token bucket rate limiting
By adding a server-global parent bucket. Per-client rate limiting is subject to the parent bucket if global rate limiting is enabled. This implementation is experimental, and all related APIs should be considered unstable. Updates tailscale/corp#40291 Signed-off-by: Jordan Whited <jordan@tailscale.com>
This commit is contained in:
committed by
Jordan Whited
parent
5eb0b4be31
commit
d8190e0de5
@@ -181,7 +181,11 @@ type Server struct {
|
|||||||
verifyClientsURL string
|
verifyClientsURL string
|
||||||
verifyClientsURLFailOpen bool
|
verifyClientsURLFailOpen bool
|
||||||
|
|
||||||
mu syncs.Mutex
|
perClientSendQueueDepth int // Sets the client send queue depth for the server.
|
||||||
|
tcpWriteTimeout time.Duration
|
||||||
|
clock tstime.Clock
|
||||||
|
|
||||||
|
mu syncs.Mutex // guards the following fields
|
||||||
closed bool
|
closed bool
|
||||||
netConns map[derp.Conn]chan struct{} // chan is closed when conn closes
|
netConns map[derp.Conn]chan struct{} // chan is closed when conn closes
|
||||||
clients map[key.NodePublic]*clientSet
|
clients map[key.NodePublic]*clientSet
|
||||||
@@ -197,25 +201,10 @@ type Server struct {
|
|||||||
// is gone from the region, we notify all of these watchers,
|
// is gone from the region, we notify all of these watchers,
|
||||||
// calling their funcs in a new goroutine.
|
// calling their funcs in a new goroutine.
|
||||||
peerGoneWatchers map[key.NodePublic]set.HandleSet[func(key.NodePublic)]
|
peerGoneWatchers map[key.NodePublic]set.HandleSet[func(key.NodePublic)]
|
||||||
|
|
||||||
// maps from netip.AddrPort to a client's public key
|
// maps from netip.AddrPort to a client's public key
|
||||||
keyOfAddr map[netip.AddrPort]key.NodePublic
|
keyOfAddr map[netip.AddrPort]key.NodePublic
|
||||||
|
rateConfig RateConfig // server-global and per-client DERP frame rate limiting config
|
||||||
// Sets the client send queue depth for the server.
|
recvLim *xrate.Limiter // server-global DERP frame receive limiter
|
||||||
perClientSendQueueDepth int
|
|
||||||
|
|
||||||
tcpWriteTimeout time.Duration
|
|
||||||
|
|
||||||
// perClientRecvBytesPerSec is the rate limit for receiving data from
|
|
||||||
// a single client connection, in bytes per second. 0 means unlimited.
|
|
||||||
// Mesh peers are exempt from this limit.
|
|
||||||
perClientRecvBytesPerSec uint64
|
|
||||||
// perClientRecvBurst is the burst size in bytes for the per-client
|
|
||||||
// receive rate limiter. Always at least [derp.MaxPacketSize] when
|
|
||||||
// set via [Server.UpdatePerClientRateLimit].
|
|
||||||
perClientRecvBurst uint64
|
|
||||||
|
|
||||||
clock tstime.Clock
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// clientSet represents 1 or more *sclients.
|
// clientSet represents 1 or more *sclients.
|
||||||
@@ -523,16 +512,29 @@ func (s *Server) SetTCPWriteTimeout(d time.Duration) {
|
|||||||
// our token bucket calls for.
|
// our token bucket calls for.
|
||||||
const minRateLimitTokenBucketSize = derp.MaxPacketSize + derp.KeyLen
|
const minRateLimitTokenBucketSize = derp.MaxPacketSize + derp.KeyLen
|
||||||
|
|
||||||
// RateConfig is a JSON-serializable configuration for per-client rate limits.
|
// RateConfig is a JSON-serializable configuration for rate limits. Values are
|
||||||
// Values are in bytes.
|
// in bytes.
|
||||||
type RateConfig struct {
|
type RateConfig struct {
|
||||||
// PerClientRateLimitBytesPerSec represents the per-client
|
// PerClientRateLimitBytesPerSec represents the per-client
|
||||||
// rate limit in bytes per second. A zero value disables rate-limiting.
|
// rate limit in bytes per second. A zero value disables rate limiting.
|
||||||
PerClientRateLimitBytesPerSec uint64 `json:",omitzero"`
|
PerClientRateLimitBytesPerSec uint64 `json:",omitzero"`
|
||||||
// PerClientRateBurstBytes represents the per-client token bucket depth,
|
// PerClientRateBurstBytes represents the per-client token bucket depth,
|
||||||
// or burst, in bytes. Any value lower than [minRateLimitTokenBucketSize]
|
// or burst, in bytes. Any value lower than [minRateLimitTokenBucketSize]
|
||||||
// will be increased to [minRateLimitTokenBucketSize] before application.
|
// will be increased to [minRateLimitTokenBucketSize] before application. Only
|
||||||
|
// relevant if PerClientRateLimitBytesPerSec is nonzero.
|
||||||
PerClientRateBurstBytes uint64 `json:",omitzero"`
|
PerClientRateBurstBytes uint64 `json:",omitzero"`
|
||||||
|
// GlobalRateLimitBytesPerSec represents the global rate limit in bytes per
|
||||||
|
// second. A zero value disables global rate limiting, but per-client (PerClient...)
|
||||||
|
// configuration may still apply. Only relevant if PerClientRateLimitBytesPerSec
|
||||||
|
// is nonzero. If GlobalRateLimitBytesPerSec is nonzero and less than
|
||||||
|
// PerClientRateLimitBytesPerSec, then GlobalRateLimitBytesPerSec will be set
|
||||||
|
// equal to PerClientRateLimitBytesPerSec before application.
|
||||||
|
GlobalRateLimitBytesPerSec uint64 `json:",omitzero"`
|
||||||
|
// GlobalRateBurstBytes represents the global token bucket depth, or burst,
|
||||||
|
// in bytes. Any value lower than [minRateLimitTokenBucketSize] will be increased to
|
||||||
|
// [minRateLimitTokenBucketSize] before application. Only relevant if
|
||||||
|
// PerClientRateLimitBytesPerSec and GlobalRateLimitBytesPerSec are nonzero.
|
||||||
|
GlobalRateBurstBytes uint64 `json:",omitzero"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// LoadRateConfig reads and JSON-unmarshals a [RateConfig] from the file at path.
|
// LoadRateConfig reads and JSON-unmarshals a [RateConfig] from the file at path.
|
||||||
@@ -542,41 +544,56 @@ func LoadRateConfig(path string) (RateConfig, error) {
|
|||||||
}
|
}
|
||||||
b, err := os.ReadFile(path)
|
b, err := os.ReadFile(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return RateConfig{}, fmt.Errorf("reading rate config: %w", err)
|
return RateConfig{}, fmt.Errorf("error reading rate config: %w", err)
|
||||||
}
|
}
|
||||||
var rc RateConfig
|
var rc RateConfig
|
||||||
if err := json.Unmarshal(b, &rc); err != nil {
|
if err := json.Unmarshal(b, &rc); err != nil {
|
||||||
return RateConfig{}, fmt.Errorf("parsing rate config: %w", err)
|
return RateConfig{}, fmt.Errorf("error parsing rate config: %w", err)
|
||||||
}
|
}
|
||||||
return rc, nil
|
return rc, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// LoadAndApplyRateConfig reads a [RateConfig] from the file at path and
|
// LoadAndApplyRateConfig reads a [RateConfig] from the file at path and
|
||||||
// applies it to the server via [Server.UpdatePerClientRateLimit].
|
// applies it to the server via [Server.UpdateRateLimits].
|
||||||
func (s *Server) LoadAndApplyRateConfig(path string) error {
|
func (s *Server) LoadAndApplyRateConfig(path string) error {
|
||||||
rc, err := LoadRateConfig(path)
|
rc, err := LoadRateConfig(path)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
s.UpdatePerClientRateLimit(rc.PerClientRateLimitBytesPerSec, rc.PerClientRateBurstBytes)
|
applied := s.UpdateRateLimits(rc)
|
||||||
s.logf("rate config applied: rate=%d bytes/sec, burst=%d bytes", rc.PerClientRateLimitBytesPerSec, rc.PerClientRateBurstBytes)
|
s.logf("rate config applied: global-rate=%d bytes/sec global-burst=%d bytes client-rate=%d bytes/sec, client-burst=%d bytes",
|
||||||
|
applied.GlobalRateLimitBytesPerSec, applied.GlobalRateBurstBytes, applied.PerClientRateLimitBytesPerSec, applied.PerClientRateBurstBytes)
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
// UpdatePerClientRateLimit sets the per-client receive rate limit in bytes per
|
// UpdateRateLimits sets the receive rate limits, updating all existing client
|
||||||
// second and the burst size in bytes, updating all existing client connections.
|
// connections. It returns the applied config, which may differ from rc. If the
|
||||||
// The applied burst will be at least [minRateLimitTokenBucketSize]. If bytesPerSec is
|
// per-client rate limit is 0, rate limiting is disabled. Mesh peers are always
|
||||||
// 0, rate limiting is disabled. Mesh peers are always exempt from rate limiting.
|
// exempt from rate limiting.
|
||||||
func (s *Server) UpdatePerClientRateLimit(bytesPerSec, burst uint64) {
|
func (s *Server) UpdateRateLimits(rc RateConfig) (applied RateConfig) {
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
defer s.mu.Unlock()
|
||||||
s.perClientRecvBytesPerSec = bytesPerSec
|
if rc.PerClientRateLimitBytesPerSec == 0 {
|
||||||
s.perClientRecvBurst = max(burst, minRateLimitTokenBucketSize)
|
// if per-client is disabled, all rate limiting is disabled
|
||||||
|
rc = RateConfig{}
|
||||||
|
}
|
||||||
|
if rc.PerClientRateLimitBytesPerSec != 0 {
|
||||||
|
rc.PerClientRateBurstBytes = max(rc.PerClientRateBurstBytes, minRateLimitTokenBucketSize)
|
||||||
|
}
|
||||||
|
if rc.GlobalRateLimitBytesPerSec != 0 {
|
||||||
|
rc.GlobalRateLimitBytesPerSec = max(rc.GlobalRateLimitBytesPerSec, rc.PerClientRateLimitBytesPerSec)
|
||||||
|
rc.GlobalRateBurstBytes = max(rc.GlobalRateBurstBytes, minRateLimitTokenBucketSize)
|
||||||
|
s.recvLim = xrate.NewLimiter(xrate.Limit(rc.GlobalRateLimitBytesPerSec), int(rc.GlobalRateBurstBytes))
|
||||||
|
} else {
|
||||||
|
s.recvLim = nil
|
||||||
|
}
|
||||||
|
s.rateConfig = rc
|
||||||
for _, cs := range s.clients {
|
for _, cs := range s.clients {
|
||||||
cs.ForeachClient(func(c *sclient) {
|
cs.ForeachClient(func(c *sclient) {
|
||||||
c.setRateLimit(s.perClientRecvBytesPerSec, s.perClientRecvBurst)
|
c.setRateLimit(rc.PerClientRateLimitBytesPerSec, rc.PerClientRateBurstBytes, s.recvLim)
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
return rc
|
||||||
}
|
}
|
||||||
|
|
||||||
// HasMeshKey reports whether the server is configured with a mesh key.
|
// HasMeshKey reports whether the server is configured with a mesh key.
|
||||||
@@ -741,7 +758,7 @@ func (s *Server) registerClient(c *sclient) {
|
|||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
defer s.mu.Unlock()
|
defer s.mu.Unlock()
|
||||||
|
|
||||||
c.setRateLimit(s.perClientRecvBytesPerSec, s.perClientRecvBurst)
|
c.setRateLimit(s.rateConfig.PerClientRateLimitBytesPerSec, s.rateConfig.PerClientRateBurstBytes, s.recvLim)
|
||||||
|
|
||||||
cs, ok := s.clients[c.key]
|
cs, ok := s.clients[c.key]
|
||||||
if !ok {
|
if !ok {
|
||||||
@@ -1100,7 +1117,7 @@ func (c *sclient) run(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
return fmt.Errorf("client %s: readFrameHeader: %w", c.key.ShortString(), err)
|
return fmt.Errorf("client %s: readFrameHeader: %w", c.key.ShortString(), err)
|
||||||
}
|
}
|
||||||
// Rate limit by DERP frame length (fl), which excludes TLS protocol and
|
// Rate-limit by DERP frame length (fl), which excludes TLS protocol and
|
||||||
// DERP frame length field overheads.
|
// DERP frame length field overheads.
|
||||||
// Note: meshed clients are exempt from rate limits.
|
// Note: meshed clients are exempt from rate limits.
|
||||||
if err := c.rateLimit(int(fl)); err != nil {
|
if err := c.rateLimit(int(fl)); err != nil {
|
||||||
@@ -1316,22 +1333,18 @@ func (c *sclient) handleFrameSendPacket(_ derp.FrameType, fl uint32) error {
|
|||||||
return c.sendPkt(dst, p)
|
return c.sendPkt(dst, p)
|
||||||
}
|
}
|
||||||
|
|
||||||
// setRateLimit updates the per-client receive rate limiter.
|
// setRateLimit updates the receive rate limiter. When bytesPerSec is 0 or the
|
||||||
// When bytesPerSec is 0 or the client is a mesh peer, the limiter is
|
// client is a mesh peer, the limiter is set to nil so that [sclient.rateLimit] is a no-op.
|
||||||
// set to nil so that [sclient.rateLimit] is a no-op.
|
func (c *sclient) setRateLimit(bytesPerSec, burst uint64, parent *xrate.Limiter) {
|
||||||
func (c *sclient) setRateLimit(bytesPerSec uint64, burst uint64) {
|
|
||||||
if bytesPerSec == 0 || c.canMesh {
|
if bytesPerSec == 0 || c.canMesh {
|
||||||
c.recvLim.Store(nil)
|
c.recvLim.Store(nil)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if lim := c.recvLim.Load(); lim != nil {
|
child := xrate.NewLimiter(xrate.Limit(bytesPerSec), int(burst))
|
||||||
// Update in place. SetBurst before SetLimit to avoid a transient
|
lim := &parentChildTokenBuckets{
|
||||||
// state where a new higher rate exceeds the old lower burst.
|
parent: parent,
|
||||||
lim.SetBurst(int(burst))
|
child: child,
|
||||||
lim.SetLimit(xrate.Limit(bytesPerSec))
|
|
||||||
return
|
|
||||||
}
|
}
|
||||||
lim := xrate.NewLimiter(xrate.Limit(bytesPerSec), int(burst))
|
|
||||||
c.recvLim.Store(lim)
|
c.recvLim.Store(lim)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -1350,12 +1363,18 @@ func (c *sclient) rateLimit(n int) error {
|
|||||||
// [minRateLimitTokenBucketSize].
|
// [minRateLimitTokenBucketSize].
|
||||||
//
|
//
|
||||||
// While we could call WaitN multiple times and/or more precisely for
|
// While we could call WaitN multiple times and/or more precisely for
|
||||||
// lim.Burst(), it's better to return early as a larger DERP frame is:
|
// lim.Burst(), it's better to return early as a larger DERP frame:
|
||||||
// 1. unexpected
|
// 1. is unexpected
|
||||||
// 2. only partially read off the socket (bufio)
|
// 2. is only partially read off the socket (bufio)
|
||||||
// 3. would cause the connection to close shortly after rate limiting, anyway.
|
// 3. would cause the connection to close shortly after rate limiting, anyway.
|
||||||
clampedN := min(n, minRateLimitTokenBucketSize)
|
clampedN := min(n, minRateLimitTokenBucketSize)
|
||||||
return lim.WaitN(c.ctx, clampedN)
|
err := lim.child.WaitN(c.ctx, clampedN)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
if lim.parent != nil {
|
||||||
|
return lim.parent.WaitN(c.ctx, clampedN)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
@@ -1792,15 +1811,22 @@ type sclient struct {
|
|||||||
// through us with a peer we have no record of.
|
// through us with a peer we have no record of.
|
||||||
peerGoneLim *rate.Limiter
|
peerGoneLim *rate.Limiter
|
||||||
|
|
||||||
// recvLim is the per-connection receive rate limiter. When rate
|
// recvLim is the receive rate limiter. When rate limiting is enabled for a
|
||||||
// limiting is enabled for a non-mesh client, it points to an
|
// non-mesh client, it points to a [parentChildTokenBuckets]. When rate limiting
|
||||||
// [xrate.Limiter]. When rate limiting is disabled or the client is a
|
// is disabled or the client is a mesh peer, it is nil and [sclient.rateLimit]
|
||||||
// mesh peer, it is nil and [sclient.rateLimit] is a no-op.
|
// is a no-op. Updated atomically by [sclient.setRateLimit] so that
|
||||||
// Updated atomically by [sclient.setRateLimitLocked] so that
|
// [sclient.rateLimit] can load it without holding [Server.mu].
|
||||||
// [sclient.rateLimit] can load it without holding Server.mu.
|
recvLim atomic.Pointer[parentChildTokenBuckets]
|
||||||
// TODO(mikeodr): update to use mono time, requires updates
|
}
|
||||||
// to tstime/rate.Limiter
|
|
||||||
recvLim atomic.Pointer[xrate.Limiter]
|
// parentChildTokenBuckets contains a parent and child token bucket for the
|
||||||
|
// purpose of applying in a hierarchical topology.
|
||||||
|
//
|
||||||
|
// TODO: consider porting the required APIs from [xrate.Limiter] to [rate.Limiter],
|
||||||
|
// which is already optimized to use [mono.Time].
|
||||||
|
type parentChildTokenBuckets struct {
|
||||||
|
parent *xrate.Limiter // parent may be nil
|
||||||
|
child *xrate.Limiter // child is always non-nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *sclient) presentFlags() derp.PeerPresentFlags {
|
func (c *sclient) presentFlags() derp.PeerPresentFlags {
|
||||||
|
|||||||
@@ -965,11 +965,28 @@ func TestPerClientRateLimit(t *testing.T) {
|
|||||||
c := &sclient{
|
c := &sclient{
|
||||||
ctx: ctx,
|
ctx: ctx,
|
||||||
}
|
}
|
||||||
c.recvLim.Store(rate.NewLimiter(rate.Limit(minRateLimitTokenBucketSize), minRateLimitTokenBucketSize))
|
lim := &parentChildTokenBuckets{
|
||||||
|
// Set parent limit to half of child to enable verification of
|
||||||
|
// rate limiting across both layers with a single sclient.
|
||||||
|
parent: rate.NewLimiter(rate.Limit(minRateLimitTokenBucketSize)/2, minRateLimitTokenBucketSize),
|
||||||
|
child: rate.NewLimiter(rate.Limit(minRateLimitTokenBucketSize), minRateLimitTokenBucketSize),
|
||||||
|
}
|
||||||
|
c.recvLim.Store(lim)
|
||||||
|
wantTokens := func(t *testing.T, wantParentTokens, wantChildTokens float64) {
|
||||||
|
t.Helper()
|
||||||
|
if lim.parent.Tokens() != wantParentTokens {
|
||||||
|
t.Fatalf("want parent tokens: %v got: %v", wantParentTokens, lim.parent.Tokens())
|
||||||
|
}
|
||||||
|
if lim.child.Tokens() != wantChildTokens {
|
||||||
|
t.Fatalf("want child tokens: %v got: %v", wantChildTokens, lim.child.Tokens())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
// First call within burst should not block.
|
// First call within burst should not block.
|
||||||
c.rateLimit(minRateLimitTokenBucketSize)
|
c.rateLimit(minRateLimitTokenBucketSize)
|
||||||
|
|
||||||
|
wantTokens(t, 0, 0)
|
||||||
|
|
||||||
// Next call exceeds burst, should block until tokens replenish.
|
// Next call exceeds burst, should block until tokens replenish.
|
||||||
done := make(chan error, 1)
|
done := make(chan error, 1)
|
||||||
go func() {
|
go func() {
|
||||||
@@ -984,7 +1001,21 @@ func TestPerClientRateLimit(t *testing.T) {
|
|||||||
default:
|
default:
|
||||||
}
|
}
|
||||||
|
|
||||||
// Advance time by 1 second
|
// Advance time by 1 second, the goroutine should still be blocked
|
||||||
|
// on the parent bucket (negative tokens).
|
||||||
|
time.Sleep(1 * time.Second)
|
||||||
|
synctest.Wait()
|
||||||
|
select {
|
||||||
|
case err := <-done:
|
||||||
|
t.Fatalf("rateLimit should have blocked, but returned: %v", err)
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify the parent bucket fills at half the rate of the child.
|
||||||
|
wantTokens(t, -(minRateLimitTokenBucketSize / 2), 0)
|
||||||
|
|
||||||
|
// Advance time by another second, parent should have enough tokens
|
||||||
|
// to unblock.
|
||||||
time.Sleep(1 * time.Second)
|
time.Sleep(1 * time.Second)
|
||||||
synctest.Wait()
|
synctest.Wait()
|
||||||
|
|
||||||
@@ -996,6 +1027,8 @@ func TestPerClientRateLimit(t *testing.T) {
|
|||||||
default:
|
default:
|
||||||
t.Fatal("rateLimit should have unblocked after 1s")
|
t.Fatal("rateLimit should have unblocked after 1s")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
wantTokens(t, 0, minRateLimitTokenBucketSize)
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
|
|
||||||
@@ -1006,7 +1039,11 @@ func TestPerClientRateLimit(t *testing.T) {
|
|||||||
c := &sclient{
|
c := &sclient{
|
||||||
ctx: ctx,
|
ctx: ctx,
|
||||||
}
|
}
|
||||||
c.recvLim.Store(rate.NewLimiter(rate.Limit(minRateLimitTokenBucketSize), minRateLimitTokenBucketSize))
|
lim := &parentChildTokenBuckets{
|
||||||
|
child: rate.NewLimiter(rate.Limit(minRateLimitTokenBucketSize), minRateLimitTokenBucketSize),
|
||||||
|
parent: rate.NewLimiter(rate.Limit(minRateLimitTokenBucketSize), minRateLimitTokenBucketSize),
|
||||||
|
}
|
||||||
|
c.recvLim.Store(lim)
|
||||||
|
|
||||||
// Exhaust burst.
|
// Exhaust burst.
|
||||||
if err := c.rateLimit(minRateLimitTokenBucketSize); err != nil {
|
if err := c.rateLimit(minRateLimitTokenBucketSize); err != nil {
|
||||||
@@ -1052,18 +1089,38 @@ func TestPerClientRateLimit(t *testing.T) {
|
|||||||
t.Run("zero_config_no_limiter", func(t *testing.T) {
|
t.Run("zero_config_no_limiter", func(t *testing.T) {
|
||||||
s := New(key.NewNode(), logger.Discard)
|
s := New(key.NewNode(), logger.Discard)
|
||||||
defer s.Close()
|
defer s.Close()
|
||||||
if s.perClientRecvBytesPerSec != 0 {
|
if !reflect.DeepEqual(s.rateConfig, RateConfig{}) {
|
||||||
t.Errorf("expected zero rate limit, got %d", s.perClientRecvBytesPerSec)
|
t.Errorf("expected zero rate limit, got %+v", s.rateConfig)
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestUpdatePerClientRateLimit(t *testing.T) {
|
func verifyLimiter(t *testing.T, lim *parentChildTokenBuckets, wantRateConfig RateConfig) {
|
||||||
|
t.Helper()
|
||||||
|
if got := lim.child.Limit(); got != rate.Limit(wantRateConfig.PerClientRateLimitBytesPerSec) {
|
||||||
|
t.Errorf("client rate limit = %v; want %d", got, wantRateConfig.PerClientRateLimitBytesPerSec)
|
||||||
|
}
|
||||||
|
if got := lim.child.Burst(); got != int(wantRateConfig.PerClientRateBurstBytes) {
|
||||||
|
t.Errorf("client burst = %v; want %d", got, wantRateConfig.PerClientRateBurstBytes)
|
||||||
|
}
|
||||||
|
if got := lim.parent.Limit(); got != rate.Limit(wantRateConfig.GlobalRateLimitBytesPerSec) {
|
||||||
|
t.Errorf("global rate limit = %v, want %d", got, wantRateConfig.GlobalRateLimitBytesPerSec)
|
||||||
|
}
|
||||||
|
if got := lim.parent.Burst(); got != int(wantRateConfig.GlobalRateBurstBytes) {
|
||||||
|
t.Errorf("global burst = %v, want %d", got, wantRateConfig.GlobalRateBurstBytes)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUpdateRateLimits(t *testing.T) {
|
||||||
const (
|
const (
|
||||||
testBurst1 = derp.MaxPacketSize * 2
|
testClientBurst1 = minRateLimitTokenBucketSize + 1
|
||||||
testRate1 = 1000
|
testClientRate1 = minRateLimitTokenBucketSize + 2
|
||||||
testBurst2 = derp.MaxPacketSize * 4
|
testClientBurst2 = minRateLimitTokenBucketSize + 3
|
||||||
testRate2 = 5000
|
testClientRate2 = minRateLimitTokenBucketSize + 4
|
||||||
|
testGlobalBurst1 = minRateLimitTokenBucketSize + 5
|
||||||
|
testGlobalRate1 = minRateLimitTokenBucketSize + 6
|
||||||
|
testGlobalBurst2 = minRateLimitTokenBucketSize + 7
|
||||||
|
testGlobalRate2 = minRateLimitTokenBucketSize + 8
|
||||||
)
|
)
|
||||||
|
|
||||||
s := New(key.NewNode(), t.Logf)
|
s := New(key.NewNode(), t.Logf)
|
||||||
@@ -1084,52 +1141,46 @@ func TestUpdatePerClientRateLimit(t *testing.T) {
|
|||||||
s.clients[clientKey] = cs
|
s.clients[clientKey] = cs
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
|
|
||||||
s.UpdatePerClientRateLimit(testRate1, testBurst1)
|
rc := RateConfig{
|
||||||
|
PerClientRateLimitBytesPerSec: testClientRate1,
|
||||||
|
PerClientRateBurstBytes: testClientBurst1,
|
||||||
|
GlobalRateLimitBytesPerSec: testGlobalRate1,
|
||||||
|
GlobalRateBurstBytes: testGlobalBurst1,
|
||||||
|
}
|
||||||
|
s.UpdateRateLimits(rc)
|
||||||
|
|
||||||
lim := c.recvLim.Load()
|
lim := c.recvLim.Load()
|
||||||
if lim == nil {
|
if lim == nil {
|
||||||
t.Fatal("expected non-nil limiter after update")
|
t.Fatal("expected non-nil limiter after update")
|
||||||
}
|
}
|
||||||
if got := lim.Limit(); got != rate.Limit(testRate1) {
|
verifyLimiter(t, lim, rc)
|
||||||
t.Errorf("rate limit = %v; want %d", got, testRate1)
|
|
||||||
}
|
|
||||||
if got := lim.Burst(); got != int(testBurst1) {
|
|
||||||
t.Errorf("burst = %v; want %d", got, testBurst1)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Verify server fields updated.
|
// Verify server fields updated.
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
if s.perClientRecvBytesPerSec != testRate1 {
|
if !reflect.DeepEqual(s.rateConfig, rc) {
|
||||||
t.Errorf("server rate = %d; want %d", s.perClientRecvBytesPerSec, testRate1)
|
t.Errorf("s.rateConfig = %+v; want %+v", s.rateConfig, rc)
|
||||||
}
|
|
||||||
if s.perClientRecvBurst != testBurst1 {
|
|
||||||
t.Errorf("server burst = %d; want %d", s.perClientRecvBurst, testBurst1)
|
|
||||||
}
|
}
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
|
|
||||||
// Update again with different nonzero values. This exercises the
|
// Update again with different nonzero values.
|
||||||
// in-place update path (existing limiter is reused, not recreated).
|
rc = RateConfig{
|
||||||
prevLim := c.recvLim.Load()
|
PerClientRateLimitBytesPerSec: testClientRate2,
|
||||||
s.UpdatePerClientRateLimit(testRate2, testBurst2)
|
PerClientRateBurstBytes: testClientBurst2,
|
||||||
|
GlobalRateLimitBytesPerSec: testGlobalRate2,
|
||||||
|
GlobalRateBurstBytes: testGlobalBurst2,
|
||||||
|
}
|
||||||
|
s.UpdateRateLimits(rc)
|
||||||
lim = c.recvLim.Load()
|
lim = c.recvLim.Load()
|
||||||
if lim == nil {
|
if lim == nil {
|
||||||
t.Fatal("expected non-nil limiter after in-place update")
|
t.Fatal("expected non-nil limiter")
|
||||||
}
|
|
||||||
if lim != prevLim {
|
|
||||||
t.Error("expected same limiter pointer after in-place update")
|
|
||||||
}
|
|
||||||
if got := lim.Limit(); got != rate.Limit(testRate2) {
|
|
||||||
t.Errorf("rate limit after in-place update = %v; want %d", got, testRate2)
|
|
||||||
}
|
|
||||||
if got := lim.Burst(); got != int(testBurst2) {
|
|
||||||
t.Errorf("burst after in-place update = %v; want %d", got, testBurst2)
|
|
||||||
}
|
}
|
||||||
|
verifyLimiter(t, lim, rc)
|
||||||
|
|
||||||
// Disable rate limiting (set to 0).
|
// Disable rate limiting (set to 0).
|
||||||
s.UpdatePerClientRateLimit(0, 0)
|
s.UpdateRateLimits(RateConfig{})
|
||||||
|
|
||||||
if got := c.recvLim.Load(); got != nil {
|
if got := c.recvLim.Load(); got != nil {
|
||||||
t.Errorf("expected nil limiter after disable, got limit=%v", got.Limit())
|
t.Errorf("expected nil limiter after disable, got limit=%v", got.child.Limit())
|
||||||
}
|
}
|
||||||
|
|
||||||
// Mesh peer should always have nil limiter regardless of update.
|
// Mesh peer should always have nil limiter regardless of update.
|
||||||
@@ -1147,22 +1198,23 @@ func TestUpdatePerClientRateLimit(t *testing.T) {
|
|||||||
s.clients[meshKey] = meshCS
|
s.clients[meshKey] = meshCS
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
|
|
||||||
s.UpdatePerClientRateLimit(testRate2, testBurst2)
|
rc = RateConfig{
|
||||||
|
PerClientRateLimitBytesPerSec: testClientRate2,
|
||||||
|
PerClientRateBurstBytes: testClientBurst2,
|
||||||
|
GlobalRateLimitBytesPerSec: testGlobalRate2,
|
||||||
|
GlobalRateBurstBytes: testGlobalBurst2,
|
||||||
|
}
|
||||||
|
s.UpdateRateLimits(rc)
|
||||||
|
|
||||||
if got := meshClient.recvLim.Load(); got != nil {
|
if got := meshClient.recvLim.Load(); got != nil {
|
||||||
t.Errorf("mesh peer should have nil limiter, got limit=%v", got.Limit())
|
t.Errorf("mesh peer should have nil limiter, got limit=%v", got.child.Limit())
|
||||||
}
|
}
|
||||||
// Non-mesh client should be updated.
|
// Non-mesh client should be updated.
|
||||||
lim = c.recvLim.Load()
|
lim = c.recvLim.Load()
|
||||||
if lim == nil {
|
if lim == nil {
|
||||||
t.Fatal("expected non-nil limiter for non-mesh client")
|
t.Fatal("expected non-nil limiter for non-mesh client")
|
||||||
}
|
}
|
||||||
if got := lim.Limit(); got != rate.Limit(testRate2) {
|
verifyLimiter(t, lim, rc)
|
||||||
t.Errorf("rate limit = %v; want %d", got, testRate2)
|
|
||||||
}
|
|
||||||
if got := lim.Burst(); got != int(testBurst2) {
|
|
||||||
t.Errorf("burst = %v; want %d", got, testBurst2)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Verify dup clients are also updated.
|
// Verify dup clients are also updated.
|
||||||
dupKey := key.NewNode().Public()
|
dupKey := key.NewNode().Public()
|
||||||
@@ -1175,32 +1227,40 @@ func TestUpdatePerClientRateLimit(t *testing.T) {
|
|||||||
s.clients[dupKey] = dupCS
|
s.clients[dupKey] = dupCS
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
|
|
||||||
s.UpdatePerClientRateLimit(testRate1, testBurst1)
|
rc = RateConfig{
|
||||||
|
GlobalRateLimitBytesPerSec: testGlobalRate1,
|
||||||
|
GlobalRateBurstBytes: testGlobalBurst1,
|
||||||
|
PerClientRateLimitBytesPerSec: testClientRate1,
|
||||||
|
PerClientRateBurstBytes: testClientBurst1,
|
||||||
|
}
|
||||||
|
s.UpdateRateLimits(rc)
|
||||||
for i, d := range []*sclient{d1, d2} {
|
for i, d := range []*sclient{d1, d2} {
|
||||||
dl := d.recvLim.Load()
|
dl := d.recvLim.Load()
|
||||||
if dl == nil {
|
if dl == nil {
|
||||||
t.Fatalf("dup client %d: expected non-nil limiter", i)
|
t.Fatalf("dup client %d: expected non-nil limiter", i)
|
||||||
}
|
}
|
||||||
if got := dl.Limit(); got != rate.Limit(testRate1) {
|
verifyLimiter(t, dl, rc)
|
||||||
t.Errorf("dup client %d: rate = %v; want %d", i, got, testRate1)
|
|
||||||
}
|
|
||||||
if got := dl.Burst(); got != int(testBurst1) {
|
|
||||||
t.Errorf("dup client %d: burst = %v; want %d", i, got, testBurst1)
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestLoadRateConfig(t *testing.T) {
|
func TestLoadRateConfig(t *testing.T) {
|
||||||
for _, tt := range []struct {
|
for _, tt := range []struct {
|
||||||
name string
|
name string
|
||||||
json string
|
json string
|
||||||
wantRate uint64
|
wantRateConfig RateConfig
|
||||||
wantBurst uint64
|
|
||||||
}{
|
}{
|
||||||
{"both_set", `{"PerClientRateLimitBytesPerSec": 1250000, "PerClientRateBurstBytes": 2500000}`, 1250000, 2500000},
|
{"all_set", `{"PerClientRateLimitBytesPerSec": 1, "PerClientRateBurstBytes": 2, "GlobalRateLimitBytesPerSec": 3, "GlobalRateBurstBytes": 4}`, RateConfig{
|
||||||
{"rate_only", `{"PerClientRateLimitBytesPerSec": 500000}`, 500000, 0},
|
PerClientRateLimitBytesPerSec: 1,
|
||||||
{"zeros", `{"PerClientRateLimitBytesPerSec": 0, "PerClientRateBurstBytes": 0}`, 0, 0},
|
PerClientRateBurstBytes: 2,
|
||||||
{"empty_json", `{}`, 0, 0},
|
GlobalRateLimitBytesPerSec: 3,
|
||||||
|
GlobalRateBurstBytes: 4,
|
||||||
|
}},
|
||||||
|
{"rate_only", `{"PerClientRateLimitBytesPerSec": 1, "GlobalRateLimitBytesPerSec": 3}`, RateConfig{
|
||||||
|
PerClientRateLimitBytesPerSec: 1,
|
||||||
|
GlobalRateLimitBytesPerSec: 3,
|
||||||
|
}},
|
||||||
|
{"zeros", `{"PerClientRateLimitBytesPerSec": 0, "PerClientRateBurstBytes": 0, "GlobalRateLimitBytesPerSec": 0, "GlobalRateBurstBytes": 0}`, RateConfig{}},
|
||||||
|
{"empty_json", `{}`, RateConfig{}},
|
||||||
} {
|
} {
|
||||||
t.Run(tt.name, func(t *testing.T) {
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
f := filepath.Join(t.TempDir(), "rate.json")
|
f := filepath.Join(t.TempDir(), "rate.json")
|
||||||
@@ -1209,13 +1269,10 @@ func TestLoadRateConfig(t *testing.T) {
|
|||||||
}
|
}
|
||||||
rc, err := LoadRateConfig(f)
|
rc, err := LoadRateConfig(f)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
t.Fatalf("unexpected error: %v", err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
if rc.PerClientRateLimitBytesPerSec != tt.wantRate {
|
if !reflect.DeepEqual(rc, tt.wantRateConfig) {
|
||||||
t.Errorf("rate = %d; want %d", rc.PerClientRateLimitBytesPerSec, tt.wantRate)
|
t.Errorf("rate config = %v want %v", rc, tt.wantRateConfig)
|
||||||
}
|
|
||||||
if rc.PerClientRateBurstBytes != tt.wantBurst {
|
|
||||||
t.Errorf("burst = %d; want %d", rc.PerClientRateBurstBytes, tt.wantBurst)
|
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -1223,7 +1280,7 @@ func TestLoadRateConfig(t *testing.T) {
|
|||||||
for _, tt := range []struct {
|
for _, tt := range []struct {
|
||||||
name string
|
name string
|
||||||
path string
|
path string
|
||||||
content string // written to path if non-empty; path used as-is if empty
|
content string // written to loaded path if non-empty; path used as-is if empty
|
||||||
}{
|
}{
|
||||||
{"empty_path", "", ""},
|
{"empty_path", "", ""},
|
||||||
{"missing_file", filepath.Join(t.TempDir(), "nonexistent.json"), ""},
|
{"missing_file", filepath.Join(t.TempDir(), "nonexistent.json"), ""},
|
||||||
@@ -1267,47 +1324,51 @@ func TestLoadAndApplyRateConfig(t *testing.T) {
|
|||||||
s.clients[clientKey] = cs
|
s.clients[clientKey] = cs
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
|
|
||||||
f := writeConfig(t, `{"PerClientRateLimitBytesPerSec": 1250000, "PerClientRateBurstBytes": 2500000}`)
|
f := writeConfig(t, fmt.Sprintf(`{"PerClientRateLimitBytesPerSec": %d, "PerClientRateBurstBytes": %d, "GlobalRateLimitBytesPerSec": %d, "GlobalRateBurstBytes": %d}`,
|
||||||
|
minRateLimitTokenBucketSize, minRateLimitTokenBucketSize+1, minRateLimitTokenBucketSize+2, minRateLimitTokenBucketSize+3))
|
||||||
if err := s.LoadAndApplyRateConfig(f); err != nil {
|
if err := s.LoadAndApplyRateConfig(f); err != nil {
|
||||||
t.Fatalf("LoadAndApplyRateConfig: %v", err)
|
t.Fatalf("LoadAndApplyRateConfig: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Verify server fields.
|
// Verify server fields.
|
||||||
|
wantRateConfig := RateConfig{
|
||||||
|
PerClientRateLimitBytesPerSec: minRateLimitTokenBucketSize,
|
||||||
|
PerClientRateBurstBytes: minRateLimitTokenBucketSize + 1,
|
||||||
|
GlobalRateLimitBytesPerSec: minRateLimitTokenBucketSize + 2,
|
||||||
|
GlobalRateBurstBytes: minRateLimitTokenBucketSize + 3,
|
||||||
|
}
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
gotRate := s.perClientRecvBytesPerSec
|
if !reflect.DeepEqual(s.rateConfig, wantRateConfig) {
|
||||||
gotBurst := s.perClientRecvBurst
|
t.Errorf("s.rateConfig = %+v; want %+v", s.rateConfig, wantRateConfig)
|
||||||
|
}
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
if gotRate != 1250000 {
|
|
||||||
t.Errorf("server rate = %d; want 1250000", gotRate)
|
|
||||||
}
|
|
||||||
if gotBurst != 2500000 {
|
|
||||||
t.Errorf("server burst = %d; want 2500000", gotBurst)
|
|
||||||
}
|
|
||||||
|
|
||||||
// Verify client limiter.
|
// Verify client limiter.
|
||||||
lim := c.recvLim.Load()
|
lim := c.recvLim.Load()
|
||||||
if lim == nil {
|
if lim == nil {
|
||||||
t.Fatal("expected non-nil limiter")
|
t.Fatal("expected non-nil limiter")
|
||||||
}
|
}
|
||||||
if got := lim.Limit(); got != rate.Limit(1250000) {
|
verifyLimiter(t, lim, wantRateConfig)
|
||||||
t.Errorf("client rate = %v; want 1250000", got)
|
|
||||||
}
|
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("burst_is_at_least_max_packet_size", func(t *testing.T) {
|
t.Run("burst_is_at_least_minRateLimitTokenBucketSize", func(t *testing.T) {
|
||||||
s := New(key.NewNode(), t.Logf)
|
s := New(key.NewNode(), t.Logf)
|
||||||
defer s.Close()
|
defer s.Close()
|
||||||
|
|
||||||
f := writeConfig(t, `{"PerClientRateLimitBytesPerSec": 1250000, "PerClientRateBurstBytes": 10}`)
|
f := writeConfig(t, `{"PerClientRateLimitBytesPerSec": 1250000, "PerClientRateBurstBytes": 10, "GlobalRateLimitBytesPerSec": 1250000, "GlobalRateBurstBytes": 10}`)
|
||||||
if err := s.LoadAndApplyRateConfig(f); err != nil {
|
if err := s.LoadAndApplyRateConfig(f); err != nil {
|
||||||
t.Fatalf("LoadAndApplyRateConfig: %v", err)
|
t.Fatalf("LoadAndApplyRateConfig: %v", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
gotBurst := s.perClientRecvBurst
|
gotClientBurst := s.rateConfig.PerClientRateBurstBytes
|
||||||
|
gotGlobalBurst := s.rateConfig.GlobalRateBurstBytes
|
||||||
s.mu.Unlock()
|
s.mu.Unlock()
|
||||||
if gotBurst != minRateLimitTokenBucketSize {
|
if gotClientBurst != minRateLimitTokenBucketSize {
|
||||||
t.Errorf("burst = %d; want at least %d", gotBurst, minRateLimitTokenBucketSize)
|
t.Errorf("client burst = %d; want %d", gotClientBurst, minRateLimitTokenBucketSize)
|
||||||
|
}
|
||||||
|
if gotGlobalBurst != minRateLimitTokenBucketSize {
|
||||||
|
t.Errorf("global burst = %d; want %d", gotGlobalBurst, minRateLimitTokenBucketSize)
|
||||||
}
|
}
|
||||||
})
|
})
|
||||||
|
|
||||||
@@ -1315,10 +1376,15 @@ func TestLoadAndApplyRateConfig(t *testing.T) {
|
|||||||
s := New(key.NewNode(), t.Logf)
|
s := New(key.NewNode(), t.Logf)
|
||||||
defer s.Close()
|
defer s.Close()
|
||||||
|
|
||||||
f := writeConfig(t, `{"PerClientRateLimitBytesPerSec": 1250000, "PerClientRateBurstBytes": 2500000}`)
|
f := writeConfig(t, `{"PerClientRateLimitBytesPerSec": 1250000, "PerClientRateBurstBytes": 2500000, "GlobalRateLimitBytesPerSec": 12500000, "GlobalRateBurstBytes": 25000000}`)
|
||||||
if err := s.LoadAndApplyRateConfig(f); err != nil {
|
if err := s.LoadAndApplyRateConfig(f); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
}
|
}
|
||||||
|
s.mu.Lock()
|
||||||
|
if reflect.DeepEqual(s.rateConfig, RateConfig{}) {
|
||||||
|
t.Error("s.rateConfig is zero val; want nonzero rates")
|
||||||
|
}
|
||||||
|
s.mu.Unlock()
|
||||||
|
|
||||||
if err := os.WriteFile(f, []byte(`{}`), 0644); err != nil {
|
if err := os.WriteFile(f, []byte(`{}`), 0644); err != nil {
|
||||||
t.Fatal(err)
|
t.Fatal(err)
|
||||||
@@ -1328,11 +1394,10 @@ func TestLoadAndApplyRateConfig(t *testing.T) {
|
|||||||
}
|
}
|
||||||
|
|
||||||
s.mu.Lock()
|
s.mu.Lock()
|
||||||
gotRate := s.perClientRecvBytesPerSec
|
if !reflect.DeepEqual(s.rateConfig, RateConfig{}) {
|
||||||
s.mu.Unlock()
|
t.Errorf("s.rateConfig = %+v; want %+v", s.rateConfig, RateConfig{})
|
||||||
if gotRate != 0 {
|
|
||||||
t.Errorf("rate = %d; want 0 (unlimited)", gotRate)
|
|
||||||
}
|
}
|
||||||
|
s.mu.Unlock()
|
||||||
})
|
})
|
||||||
|
|
||||||
t.Run("propagates_errors", func(t *testing.T) {
|
t.Run("propagates_errors", func(t *testing.T) {
|
||||||
|
|||||||
Reference in New Issue
Block a user