github.com/MerlinKodo/gvisor@v0.0.0-20231110090155-957f62ecf90e/pkg/tcpip/stack/icmp_rate_limit.go (about)

     1  // Copyright 2018 The gVisor Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package stack
    16  
    17  import (
    18  	"github.com/MerlinKodo/gvisor/pkg/tcpip"
    19  	"golang.org/x/time/rate"
    20  )
    21  
    22  const (
    23  	// icmpLimit is the default maximum number of ICMP messages permitted by this
    24  	// rate limiter.
    25  	icmpLimit = 1000
    26  
    27  	// icmpBurst is the default number of ICMP messages that can be sent in a single
    28  	// burst.
    29  	icmpBurst = 50
    30  )
    31  
    32  // ICMPRateLimiter is a global rate limiter that controls the generation of
    33  // ICMP messages generated by the stack.
    34  type ICMPRateLimiter struct {
    35  	limiter *rate.Limiter
    36  	clock   tcpip.Clock
    37  }
    38  
    39  // NewICMPRateLimiter returns a global rate limiter for controlling the rate
    40  // at which ICMP messages are generated by the stack. The returned limiter
    41  // does not apply limits to any ICMP types by default.
    42  func NewICMPRateLimiter(clock tcpip.Clock) *ICMPRateLimiter {
    43  	return &ICMPRateLimiter{
    44  		clock:   clock,
    45  		limiter: rate.NewLimiter(icmpLimit, icmpBurst),
    46  	}
    47  }
    48  
    49  // SetLimit sets a new Limit for the limiter.
    50  func (l *ICMPRateLimiter) SetLimit(limit rate.Limit) {
    51  	l.limiter.SetLimitAt(l.clock.Now(), limit)
    52  }
    53  
    54  // Limit returns the maximum overall event rate.
    55  func (l *ICMPRateLimiter) Limit() rate.Limit {
    56  	return l.limiter.Limit()
    57  }
    58  
    59  // SetBurst sets a new burst size for the limiter.
    60  func (l *ICMPRateLimiter) SetBurst(burst int) {
    61  	l.limiter.SetBurstAt(l.clock.Now(), burst)
    62  }
    63  
    64  // Burst returns the maximum burst size.
    65  func (l *ICMPRateLimiter) Burst() int {
    66  	return l.limiter.Burst()
    67  }
    68  
    69  // Allow reports whether one ICMP message may be sent now.
    70  func (l *ICMPRateLimiter) Allow() bool {
    71  	return l.limiter.AllowN(l.clock.Now(), 1)
    72  }