From 5f0a91a12704e15f8c29d94ab59e8a39ff8e3fb6 Mon Sep 17 00:00:00 2001 From: Mathias Hall-Andersen Date: Sun, 11 Feb 2018 22:53:39 +0100 Subject: Moved ratelimiter to internal package --- ratelimiter.go | 139 --------------------------------------------------------- 1 file changed, 139 deletions(-) delete mode 100644 ratelimiter.go (limited to 'ratelimiter.go') diff --git a/ratelimiter.go b/ratelimiter.go deleted file mode 100644 index 6e5f005..0000000 --- a/ratelimiter.go +++ /dev/null @@ -1,139 +0,0 @@ -package main - -/* Copyright (C) 2015-2017 Jason A. Donenfeld . All Rights Reserved. */ - -/* This file contains a port of the ratelimited from 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 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 { - 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 -} -- cgit v1.2.3