Moved ratelimiter to internal package

This commit is contained in:
Mathias Hall-Andersen 2018-02-11 22:53:39 +01:00
parent 1cf23c0005
commit 5f0a91a127
6 changed files with 60 additions and 24 deletions

View file

@ -1,6 +1,7 @@
package main package main
import ( import (
"git.zx2c4.com/wireguard-go/internal/ratelimiter"
"runtime" "runtime"
"sync" "sync"
"sync/atomic" "sync/atomic"
@ -50,7 +51,7 @@ type Device struct {
rate struct { rate struct {
underLoadUntil atomic.Value underLoadUntil atomic.Value
limiter Ratelimiter limiter ratelimiter.Ratelimiter
} }
pool struct { pool struct {
@ -300,7 +301,6 @@ func NewDevice(tun TUNDevice, logger *Logger) *Device {
go device.RoutineReadFromTUN() go device.RoutineReadFromTUN()
go device.RoutineTUNEventReader() go device.RoutineTUNEventReader()
go device.rate.limiter.RoutineGarbageCollector(device.signal.stop)
return device return device
} }
@ -355,6 +355,7 @@ func (device *Device) Close() {
device.BindClose() device.BindClose()
device.isUp.Set(false) device.isUp.Set(false)
device.RemoveAllPeers() device.RemoveAllPeers()
device.rate.limiter.Close()
device.log.Info.Println("Interface closed") device.log.Info.Println("Interface closed")
} }

View file

@ -1,4 +1,4 @@
package main package ratelimiter
/* Copyright (C) 2015-2017 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved. */ /* Copyright (C) 2015-2017 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved. */
@ -26,21 +26,48 @@ type RatelimiterEntry struct {
} }
type Ratelimiter struct { type Ratelimiter struct {
mutex sync.RWMutex mutex sync.RWMutex
lastGarbageCollect time.Time stop chan struct{}
tableIPv4 map[[net.IPv4len]byte]*RatelimiterEntry tableIPv4 map[[net.IPv4len]byte]*RatelimiterEntry
tableIPv6 map[[net.IPv6len]byte]*RatelimiterEntry tableIPv6 map[[net.IPv6len]byte]*RatelimiterEntry
}
func (rate *Ratelimiter) Close() {
rate.mutex.Lock()
defer rate.mutex.Unlock()
if rate.stop != nil {
close(rate.stop)
}
} }
func (rate *Ratelimiter) Init() { func (rate *Ratelimiter) Init() {
rate.mutex.Lock() rate.mutex.Lock()
defer rate.mutex.Unlock() defer rate.mutex.Unlock()
if rate.stop != nil {
close(rate.stop)
}
rate.stop = make(chan struct{})
rate.tableIPv4 = make(map[[net.IPv4len]byte]*RatelimiterEntry) rate.tableIPv4 = make(map[[net.IPv4len]byte]*RatelimiterEntry)
rate.tableIPv6 = make(map[[net.IPv6len]byte]*RatelimiterEntry) rate.tableIPv6 = make(map[[net.IPv6len]byte]*RatelimiterEntry)
rate.lastGarbageCollect = time.Now()
go func() {
timer := time.NewTimer(time.Second)
for {
select {
case <-rate.stop:
return
case <-timer.C:
rate.garbageCollectEntries()
timer.Reset(time.Second)
}
}
}()
} }
func (rate *Ratelimiter) GarbageCollectEntries() { func (rate *Ratelimiter) garbageCollectEntries() {
rate.mutex.Lock() rate.mutex.Lock()
// remove unused IPv4 entries // remove unused IPv4 entries
@ -66,19 +93,6 @@ func (rate *Ratelimiter) GarbageCollectEntries() {
rate.mutex.Unlock() rate.mutex.Unlock()
} }
func (rate *Ratelimiter) RoutineGarbageCollector(stop Signal) {
timer := time.NewTimer(time.Second)
for {
select {
case <-stop.Wait():
return
case <-timer.C:
rate.GarbageCollectEntries()
timer.Reset(time.Second)
}
}
}
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

View file

@ -1,4 +1,4 @@
package main package ratelimiter
import ( import (
"net" "net"

View file

@ -0,0 +1,21 @@
package tai64n
import (
"testing"
"time"
)
/* Testing the essential property of the timestamp
* as used by WireGuard.
*/
func TestMonotonic(t *testing.T) {
old := Now()
for i := 0; i < 10000; i++ {
time.Sleep(time.Nanosecond)
next := Now()
if !next.After(old) {
t.Error("TAI64N, not monotonically increasing on nano-second scale")
}
old = next
}
}

View file

@ -120,7 +120,7 @@ func (peer *Peer) TimerAnyAuthenticatedPacketTraversal() {
*/ */
func (peer *Peer) TimerHandshakeComplete() { func (peer *Peer) TimerHandshakeComplete() {
peer.signal.handshakeCompleted.Send() peer.signal.handshakeCompleted.Send()
peer.device.log.Info.Println("Negotiated new handshake for", peer.String()) peer.device.log.Info.Println(peer.String(), ": New handshake completed")
} }
/* Event: /* Event: