From 668dcebf1369f1d6a21f0e362039c7f64d72f68f Mon Sep 17 00:00:00 2001
From: Andrey Meshkov <am@adguard.com>
Date: Sun, 25 Nov 2018 23:11:36 +0300
Subject: [PATCH] Added ratelimit plugin

---
 config.go                             |  3 ++-
 coredns.go                            |  1 +
 coredns_plugin/ratelimit/ratelimit.go | 20 ++++++++++++++++++--
 3 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/config.go b/config.go
index 5997834d..42589d05 100644
--- a/config.go
+++ b/config.go
@@ -253,7 +253,8 @@ const coreDNSConfigTemplate = `.:{{.Port}} {
 		{{end}}
     }{{end}}
     {{.Pprof}}
-    hosts {
+    ratelimit 30
+	hosts {
         fallthrough
     }
     {{if .UpstreamDNS}}upstream {{range .UpstreamDNS}}{{.}} {{end}} { bootstrap {{.BootstrapDNS}} }{{end}}
diff --git a/coredns.go b/coredns.go
index a21fb986..1330e593 100644
--- a/coredns.go
+++ b/coredns.go
@@ -8,6 +8,7 @@ import (
 	"sync" // Include all plugins.
 
 	_ "github.com/AdguardTeam/AdGuardHome/coredns_plugin"
+	_ "github.com/AdguardTeam/AdGuardHome/coredns_plugin/ratelimit"
 	_ "github.com/AdguardTeam/AdGuardHome/upstream"
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/coremain"
diff --git a/coredns_plugin/ratelimit/ratelimit.go b/coredns_plugin/ratelimit/ratelimit.go
index 66c1a32b..fbfbc2f0 100644
--- a/coredns_plugin/ratelimit/ratelimit.go
+++ b/coredns_plugin/ratelimit/ratelimit.go
@@ -15,6 +15,7 @@ import (
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/plugin"
 	"github.com/coredns/coredns/plugin/metrics"
+	"github.com/coredns/coredns/plugin/pkg/dnstest"
 	"github.com/coredns/coredns/request"
 	"github.com/mholt/caddy"
 	"github.com/miekg/dns"
@@ -23,6 +24,7 @@ import (
 )
 
 const defaultRatelimit = 100
+const defaultResponseSize = 1000
 
 var (
 	tokenBuckets = cache.New(time.Hour, time.Hour)
@@ -40,7 +42,22 @@ func (p *plug) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (
 		ratelimited.Inc()
 		return 0, nil
 	}
-	return plugin.NextOrFailure(p.Name(), p.Next, ctx, w, r)
+
+	// Record response to get status code and size of the reply.
+	rw := dnstest.NewRecorder(w)
+	status, err := plugin.NextOrFailure(p.Name(), p.Next, ctx, rw, r)
+
+	size := rw.Len
+
+	if size > defaultResponseSize && state.Proto() == "udp" {
+		// For large UDP responses we call allowRequest more times
+		// The exact number of times depends on the response size
+		for i := 0; i < size/defaultResponseSize; i++ {
+			p.allowRequest(ip)
+		}
+	}
+
+	return status, err
 }
 
 func (p *plug) allowRequest(ip string) (bool, error) {
@@ -94,7 +111,6 @@ type plug struct {
 	// configuration for creating above
 	ratelimit int      // in requests per second per IP
 	whitelist []string // a list of whitelisted IP addresses
-
 }
 
 func setupPlugin(c *caddy.Controller) (*plug, error) {