Hesham/wireguard-go
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Cleanup ratelimiter

Browse Source
This commit is contained in:
Jason A. Donenfeld 2018-05-13 18:42:06 +02:00
parent 2c27ab205c
commit e94185681f
2 changed files with 14 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
device.go
View File

@ -185,7 +185,7 @@ func (device *Device) IsUnderLoad() bool {
now := time.Now() now := time.Now()
underLoad := len(device.queue.handshake) >= UnderLoadQueueSize underLoad := len(device.queue.handshake) >= UnderLoadQueueSize
if underLoad { if underLoad {
device.rate.underLoadUntil.Store(now.Add(time.Second)) device.rate.underLoadUntil.Store(now.Add(UnderLoadAfterTime))
return true return true
} }
@ -276,7 +276,7 @@ func NewDevice(tun TUNDevice, logger *Logger) *Device {
device.peers.keyMap = make(map[NoisePublicKey]*Peer) device.peers.keyMap = make(map[NoisePublicKey]*Peer)
// initialize anti-DoS / anti-scanning features // initialize rate limiter
device.rate.limiter.Init() device.rate.limiter.Init()
device.rate.underLoadUntil.Store(time.Time{}) device.rate.underLoadUntil.Store(time.Time{})

24
ratelimiter/ratelimiter.go
View File

@ -58,12 +58,13 @@ func (rate *Ratelimiter) Init() {
// start garbage collection routine // start garbage collection routine
go func() { go func() {
timer := time.NewTimer(time.Second) ticker := time.NewTicker(time.Second)
for { for {
select { select {
case <-rate.stop: case <-rate.stop:
ticker.Stop()
return return
case <-timer.C: case <-ticker.C:
func() { func() {
rate.mutex.Lock() rate.mutex.Lock()
defer rate.mutex.Unlock() defer rate.mutex.Unlock()
@ -84,7 +85,6 @@ func (rate *Ratelimiter) Init() {
entry.mutex.Unlock() entry.mutex.Unlock()
} }
}() }()
timer.Reset(time.Second)
} }
} }
}() }()
@ -92,8 +92,8 @@ func (rate *Ratelimiter) Init() {
func (rate *Ratelimiter) Allow(ip net.IP) bool { func (rate *Ratelimiter) Allow(ip net.IP) bool {
var entry *RatelimiterEntry var entry *RatelimiterEntry
var KeyIPv4 [net.IPv4len]byte var keyIPv4 [net.IPv4len]byte
var KeyIPv6 [net.IPv6len]byte var keyIPv6 [net.IPv6len]byte
// lookup entry // lookup entry
@ -103,11 +103,11 @@ func (rate *Ratelimiter) Allow(ip net.IP) bool {
rate.mutex.RLock() rate.mutex.RLock()
if IPv4 != nil { if IPv4 != nil {
copy(KeyIPv4[:], IPv4) copy(keyIPv4[:], IPv4)
entry = rate.tableIPv4[KeyIPv4] entry = rate.tableIPv4[keyIPv4]
} else { } else {
copy(KeyIPv6[:], IPv6) copy(keyIPv6[:], IPv6)
entry = rate.tableIPv6[KeyIPv6] entry = rate.tableIPv6[keyIPv6]
} }
rate.mutex.RUnlock() rate.mutex.RUnlock()
@ -115,14 +115,14 @@ func (rate *Ratelimiter) Allow(ip net.IP) bool {
// make new entry if not found // make new entry if not found
if entry == nil { if entry == nil {
rate.mutex.Lock()
entry = new(RatelimiterEntry) entry = new(RatelimiterEntry)
entry.tokens = maxTokens - packetCost entry.tokens = maxTokens - packetCost
entry.lastTime = time.Now() entry.lastTime = time.Now()
rate.mutex.Lock()
if IPv4 != nil { if IPv4 != nil {
rate.tableIPv4[KeyIPv4] = entry rate.tableIPv4[keyIPv4] = entry
} else { } else {
rate.tableIPv6[KeyIPv6] = entry rate.tableIPv6[keyIPv6] = entry
} }
rate.mutex.Unlock() rate.mutex.Unlock()
return true return true