dnsforward -- implement ratelimit and refuseany

dnsforward/ratelimit.go

@@ -0,0 +1,80 @@
+package dnsforward
+
+import (
+	"log"
+	"sort"
+	"time"
+
+	"github.com/beefsack/go-rate"
+	gocache "github.com/patrickmn/go-cache"
+)
+
+func (s *Server) limiterForIP(ip string) interface{} {
+	if s.ratelimitBuckets == nil {
+		s.ratelimitBuckets = gocache.New(time.Hour, time.Hour)
+	}
+
+	// check if ratelimiter for that IP already exists, if not, create
+	value, found := s.ratelimitBuckets.Get(ip)
+	if !found {
+		value = rate.New(s.Ratelimit, time.Second)
+		s.ratelimitBuckets.Set(ip, value, time.Hour)
+	}
+
+	return value
+}
+
+func (s *Server) isRatelimited(ip string) bool {
+	if s.Ratelimit == 0 { // 0 -- disabled
+		return false
+	}
+	if len(s.RatelimitWhitelist) > 0 {
+		i := sort.SearchStrings(s.RatelimitWhitelist, ip)
+
+		if i < len(s.RatelimitWhitelist) && s.RatelimitWhitelist[i] == ip {
+			// found, don't ratelimit
+			return false
+		}
+	}
+
+	value := s.limiterForIP(ip)
+	rl, ok := value.(*rate.RateLimiter)
+	if !ok {
+		log.Println("SHOULD NOT HAPPEN: non-bool entry found in safebrowsing lookup cache")
+		return false
+	}
+
+	allow, _ := rl.Try()
+	return !allow
+}
+
+func (s *Server) isRatelimitedForReply(ip string, size int) bool {
+	if s.Ratelimit == 0 { // 0 -- disabled
+		return false
+	}
+	if len(s.RatelimitWhitelist) > 0 {
+		i := sort.SearchStrings(s.RatelimitWhitelist, ip)
+
+		if i < len(s.RatelimitWhitelist) && s.RatelimitWhitelist[i] == ip {
+			// found, don't ratelimit
+			return false
+		}
+	}
+
+	value := s.limiterForIP(ip)
+	rl, ok := value.(*rate.RateLimiter)
+	if !ok {
+		log.Println("SHOULD NOT HAPPEN: non-bool entry found in safebrowsing lookup cache")
+		return false
+	}
+
+	// For large UDP responses we try more times, effectively limiting per bandwidth
+	// The exact number of times depends on the response size
+	for i := 0; i < size/1000; i++ {
+		allow, _ := rl.Try()
+		if !allow { // not allowed -> ratelimited
+			return true
+		}
+	}
+	return false
+}