Merge pull request #103 in DNS/adguard-dns from feature/426 to master

* commit 'f94c63ed5b4fe028a74ce074b66afcbad7b426f4':
  Set default ratelimit to 30/sec
  Fix #426: add ratelimit to directives
  Added ratelimit plugin

config.go

@@ -253,7 +253,8 @@ const coreDNSConfigTemplate = `.:{{.Port}} {
 		{{end}}
     }{{end}}
     {{.Pprof}}
-    hosts {
+    ratelimit
+	hosts {
         fallthrough
     }
     {{if .UpstreamDNS}}upstream {{range .UpstreamDNS}}{{.}} {{end}} { bootstrap {{.BootstrapDNS}} }{{end}}

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"
@@ -64,6 +65,7 @@ var directives = []string{
 	"prometheus",
 	"errors",
 	"log",
+	"ratelimit",
 	"dnsfilter",
 	"dnstap",
 	"chaos",

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"
@@ -22,7 +23,8 @@ import (
 	"golang.org/x/net/context"
 )
 
-const defaultRatelimit = 100
+const defaultRatelimit = 30
+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) {