diff --git a/device.go b/device.go index 3db3609..99e451e 100644 --- a/device.go +++ b/device.go @@ -185,7 +185,7 @@ func (device *Device) IsUnderLoad() bool { now := time.Now() underLoad := len(device.queue.handshake) >= UnderLoadQueueSize if underLoad { - device.rate.underLoadUntil.Store(now.Add(time.Second)) + device.rate.underLoadUntil.Store(now.Add(UnderLoadAfterTime)) return true } @@ -276,7 +276,7 @@ func NewDevice(tun TUNDevice, logger *Logger) *Device { device.peers.keyMap = make(map[NoisePublicKey]*Peer) - // initialize anti-DoS / anti-scanning features + // initialize rate limiter device.rate.limiter.Init() device.rate.underLoadUntil.Store(time.Time{}) diff --git a/ratelimiter/ratelimiter.go b/ratelimiter/ratelimiter.go index 1aa6813..95b547e 100644 --- a/ratelimiter/ratelimiter.go +++ b/ratelimiter/ratelimiter.go @@ -58,12 +58,13 @@ func (rate *Ratelimiter) Init() { // start garbage collection routine go func() { - timer := time.NewTimer(time.Second) + ticker := time.NewTicker(time.Second) for { select { case <-rate.stop: + ticker.Stop() return - case <-timer.C: + case <-ticker.C: func() { rate.mutex.Lock() defer rate.mutex.Unlock() @@ -84,7 +85,6 @@ func (rate *Ratelimiter) Init() { entry.mutex.Unlock() } }() - timer.Reset(time.Second) } } }() @@ -92,8 +92,8 @@ func (rate *Ratelimiter) Init() { func (rate *Ratelimiter) Allow(ip net.IP) bool { var entry *RatelimiterEntry - var KeyIPv4 [net.IPv4len]byte - var KeyIPv6 [net.IPv6len]byte + var keyIPv4 [net.IPv4len]byte + var keyIPv6 [net.IPv6len]byte // lookup entry @@ -103,11 +103,11 @@ func (rate *Ratelimiter) Allow(ip net.IP) bool { rate.mutex.RLock() if IPv4 != nil { - copy(KeyIPv4[:], IPv4) - entry = rate.tableIPv4[KeyIPv4] + copy(keyIPv4[:], IPv4) + entry = rate.tableIPv4[keyIPv4] } else { - copy(KeyIPv6[:], IPv6) - entry = rate.tableIPv6[KeyIPv6] + copy(keyIPv6[:], IPv6) + entry = rate.tableIPv6[keyIPv6] } rate.mutex.RUnlock() @@ -115,14 +115,14 @@ func (rate *Ratelimiter) Allow(ip net.IP) bool { // make new entry if not found if entry == nil { - rate.mutex.Lock() entry = new(RatelimiterEntry) entry.tokens = maxTokens - packetCost entry.lastTime = time.Now() + rate.mutex.Lock() if IPv4 != nil { - rate.tableIPv4[KeyIPv4] = entry + rate.tableIPv4[keyIPv4] = entry } else { - rate.tableIPv6[KeyIPv6] = entry + rate.tableIPv6[keyIPv6] = entry } rate.mutex.Unlock() return true