diff options
| -rw-r--r-- | src/device.go | 9 | ||||
| -rw-r--r-- | src/ratelimiter.go | 140 | ||||
| -rw-r--r-- | src/ratelimiter_test.go | 98 | ||||
| -rw-r--r-- | src/receive.go | 4 | ||||
| -rw-r--r-- | src/replay_test.go | 2 |
5 files changed, 248 insertions, 5 deletions
diff --git a/src/device.go b/src/device.go index 12d1ed9..2a2ad62 100644 --- a/src/device.go +++ b/src/device.go @@ -31,9 +31,10 @@ type Device struct { signal struct { stop chan struct{} } - underLoad int32 // used as an atomic bool - peers map[NoisePublicKey]*Peer - mac MACStateDevice + underLoad int32 // used as an atomic bool + ratelimiter Ratelimiter + peers map[NoisePublicKey]*Peer + mac MACStateDevice } func (device *Device) SetPrivateKey(sk NoisePrivateKey) { @@ -66,6 +67,7 @@ func NewDevice(tun TUNDevice, logLevel int) *Device { device.mtu = tun.MTU() device.peers = make(map[NoisePublicKey]*Peer) device.indices.Init() + device.ratelimiter.Init() device.routingTable.Reset() // listen @@ -99,6 +101,7 @@ func NewDevice(tun TUNDevice, logLevel int) *Device { go device.RoutineReadFromTUN(tun) go device.RoutineReceiveIncomming() go device.RoutineWriteToTUN(tun) + go device.ratelimiter.RoutineGarbageCollector(device.signal.stop) return device } diff --git a/src/ratelimiter.go b/src/ratelimiter.go new file mode 100644 index 0000000..a9f719b --- /dev/null +++ b/src/ratelimiter.go @@ -0,0 +1,140 @@ +package main + +/* Implementation of the ratelimited form the linux kernel version + * + * + * + */ + +import ( + "net" + "sync" + "time" +) + +const ( + RatelimiterPacketsPerSecond = 20 + RatelimiterPacketsBurstable = 5 + RatelimiterGarbageCollectTime = time.Second + RatelimiterPacketCost = 1000000000 / RatelimiterPacketsPerSecond + RatelimiterMaxTokens = RatelimiterPacketCost * RatelimiterPacketsBurstable +) + +type RatelimiterEntry struct { + mutex sync.Mutex + lastTime time.Time + tokens int64 +} + +type Ratelimiter struct { + mutex sync.RWMutex + lastGarbageCollect time.Time + tableIPv4 map[[net.IPv4len]byte]*RatelimiterEntry + tableIPv6 map[[net.IPv6len]byte]*RatelimiterEntry +} + +func (rate *Ratelimiter) Init() { + rate.mutex.Lock() + defer rate.mutex.Unlock() + rate.tableIPv4 = make(map[[net.IPv4len]byte]*RatelimiterEntry) + rate.tableIPv6 = make(map[[net.IPv6len]byte]*RatelimiterEntry) + rate.lastGarbageCollect = time.Now() +} + +func (rate *Ratelimiter) GarbageCollectEntries() { + rate.mutex.Lock() + + // remove unused IPv4 entries + + for key, entry := range rate.tableIPv4 { + entry.mutex.Lock() + if time.Now().Sub(entry.lastTime) > RatelimiterGarbageCollectTime { + delete(rate.tableIPv4, key) + } + entry.mutex.Unlock() + } + + // remove unused IPv6 entries + + for key, entry := range rate.tableIPv6 { + entry.mutex.Lock() + if time.Now().Sub(entry.lastTime) > RatelimiterGarbageCollectTime { + delete(rate.tableIPv6, key) + } + entry.mutex.Unlock() + } + + rate.mutex.Unlock() +} + +func (rate *Ratelimiter) RoutineGarbageCollector(stop chan struct{}) { + timer := time.NewTimer(time.Second) + for { + select { + case <-stop: + return + case <-timer.C: + rate.GarbageCollectEntries() + timer.Reset(time.Second) + } + } +} + +func (rate *Ratelimiter) Allow(ip net.IP) bool { + var entry *RatelimiterEntry + var KeyIPv4 [net.IPv4len]byte + var KeyIPv6 [net.IPv6len]byte + + // lookup entry + + IPv4 := ip.To4() + IPv6 := ip.To16() + + rate.mutex.RLock() + + if IPv4 != nil { + copy(KeyIPv4[:], IPv4) + entry = rate.tableIPv4[KeyIPv4] + } else { + copy(KeyIPv6[:], IPv6) + entry = rate.tableIPv6[KeyIPv6] + } + + rate.mutex.RUnlock() + + // make new entry if not found + + if entry == nil { + rate.mutex.Lock() + entry = new(RatelimiterEntry) + entry.tokens = RatelimiterMaxTokens - RatelimiterPacketCost + entry.lastTime = time.Now() + if IPv4 != nil { + rate.tableIPv4[KeyIPv4] = entry + } else { + rate.tableIPv6[KeyIPv6] = entry + } + rate.mutex.Unlock() + return true + } + + // add tokens to entry + + entry.mutex.Lock() + now := time.Now() + entry.tokens += now.Sub(entry.lastTime).Nanoseconds() + entry.lastTime = now + if entry.tokens > RatelimiterMaxTokens { + entry.tokens = RatelimiterMaxTokens + } + + // subtract cost of packet + + if entry.tokens > RatelimiterPacketCost { + entry.tokens -= RatelimiterPacketCost + entry.mutex.Unlock() + return true + } + entry.mutex.Unlock() + return false +} diff --git a/src/ratelimiter_test.go b/src/ratelimiter_test.go new file mode 100644 index 0000000..13b6a23 --- /dev/null +++ b/src/ratelimiter_test.go @@ -0,0 +1,98 @@ +package main + +import ( + "net" + "testing" + "time" +) + +type RatelimiterResult struct { + allowed bool + text string + wait time.Duration +} + +func TestRatelimiter(t *testing.T) { + + var ratelimiter Ratelimiter + var expectedResults []RatelimiterResult + + Nano := func(nano int64) time.Duration { + return time.Nanosecond * time.Duration(nano) + } + + Add := func(res RatelimiterResult) { + expectedResults = append( + expectedResults, + res, + ) + } + + for i := 0; i < RatelimiterPacketsBurstable; i++ { + Add(RatelimiterResult{ + allowed: true, + text: "inital burst", + }) + } + + Add(RatelimiterResult{ + allowed: false, + text: "after burst", + }) + + Add(RatelimiterResult{ + allowed: true, + wait: Nano(time.Second.Nanoseconds() / RatelimiterPacketsPerSecond), + text: "filling tokens for single packet", + }) + + Add(RatelimiterResult{ + allowed: false, + text: "not having refilled enough", + }) + + Add(RatelimiterResult{ + allowed: true, + wait: 2 * Nano(time.Second.Nanoseconds()/RatelimiterPacketsPerSecond), + text: "filling tokens for two packet burst", + }) + + Add(RatelimiterResult{ + allowed: true, + text: "second packet in 2 packet burst", + }) + + Add(RatelimiterResult{ + allowed: false, + text: "packet following 2 packet burst", + }) + + ips := []net.IP{ + net.ParseIP("127.0.0.1"), + net.ParseIP("192.168.1.1"), + net.ParseIP("172.167.2.3"), + net.ParseIP("97.231.252.215"), + net.ParseIP("248.97.91.167"), + net.ParseIP("188.208.233.47"), + net.ParseIP("104.2.183.179"), + net.ParseIP("72.129.46.120"), + net.ParseIP("2001:0db8:0a0b:12f0:0000:0000:0000:0001"), + net.ParseIP("f5c2:818f:c052:655a:9860:b136:6894:25f0"), + net.ParseIP("b2d7:15ab:48a7:b07c:a541:f144:a9fe:54fc"), + net.ParseIP("a47b:786e:1671:a22b:d6f9:4ab0:abc7:c918"), + net.ParseIP("ea1e:d155:7f7a:98fb:2bf5:9483:80f6:5445"), + net.ParseIP("3f0e:54a2:f5b4:cd19:a21d:58e1:3746:84c4"), + } + + ratelimiter.Init() + + for i, res := range expectedResults { + time.Sleep(res.wait) + for _, ip := range ips { + allowed := ratelimiter.Allow(ip) + if allowed != res.allowed { + t.Fatal("Test failed for", ip.String(), ", on:", i, "(", res.text, ")", "expected:", res.allowed, "got:", allowed) + } + } + } +} diff --git a/src/receive.go b/src/receive.go index e9c14f5..99089a9 100644 --- a/src/receive.go +++ b/src/receive.go @@ -319,6 +319,10 @@ func (device *Device) RoutineHandshake() { // ratelimit + if !device.ratelimiter.Allow(elem.source.IP) { + return + } + // handle messages switch elem.msgType { diff --git a/src/replay_test.go b/src/replay_test.go index e75c5c1..228fce6 100644 --- a/src/replay_test.go +++ b/src/replay_test.go @@ -9,8 +9,6 @@ import ( * */ -/* Copyright (C) 2015-2017 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved. */ - func TestReplay(t *testing.T) { var filter ReplayFilter |