Cleanup ratelimiter

2 files changed, 14 insertions, 14 deletions

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