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) {