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 1/3] 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) {

From e4998651fe65c3166264c11c3ac123189cfb92eb Mon Sep 17 00:00:00 2001
From: Andrey Meshkov <am@adguard.com>
Date: Sun, 25 Nov 2018 23:17:49 +0300
Subject: [PATCH 2/3] Fix #426: add ratelimit to directives

---
 coredns.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/coredns.go b/coredns.go
index 1330e593..c56a8f00 100644
--- a/coredns.go
+++ b/coredns.go
@@ -66,6 +66,7 @@ var directives = []string{
 	"errors",
 	"log",
 	"dnsfilter",
+	"ratelimit",
 	"dnstap",
 	"chaos",
 	"loadbalance",

From f94c63ed5b4fe028a74ce074b66afcbad7b426f4 Mon Sep 17 00:00:00 2001
From: Andrey Meshkov <am@adguard.com>
Date: Sun, 25 Nov 2018 23:32:28 +0300
Subject: [PATCH 3/3] Set default ratelimit to 30/sec

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

diff --git a/config.go b/config.go
index 42589d05..2db48d38 100644
--- a/config.go
+++ b/config.go
@@ -253,7 +253,7 @@ const coreDNSConfigTemplate = `.:{{.Port}} {
 		{{end}}
     }{{end}}
     {{.Pprof}}
-    ratelimit 30
+    ratelimit
 	hosts {
         fallthrough
     }
diff --git a/coredns.go b/coredns.go
index c56a8f00..45854056 100644
--- a/coredns.go
+++ b/coredns.go
@@ -65,8 +65,8 @@ var directives = []string{
 	"prometheus",
 	"errors",
 	"log",
-	"dnsfilter",
 	"ratelimit",
+	"dnsfilter",
 	"dnstap",
 	"chaos",
 	"loadbalance",
diff --git a/coredns_plugin/ratelimit/ratelimit.go b/coredns_plugin/ratelimit/ratelimit.go
index fbfbc2f0..5d8abc70 100644
--- a/coredns_plugin/ratelimit/ratelimit.go
+++ b/coredns_plugin/ratelimit/ratelimit.go
@@ -23,7 +23,7 @@ import (
 	"golang.org/x/net/context"
 )
 
-const defaultRatelimit = 100
+const defaultRatelimit = 30
 const defaultResponseSize = 1000
 
 var (