Added ratelimit plugin

This commit is contained in:
Andrey Meshkov 2018-11-25 23:11:36 +03:00
parent cdd2e8ecb4
commit 668dcebf13
3 changed files with 21 additions and 3 deletions

View File

@ -253,6 +253,7 @@ const coreDNSConfigTemplate = `.:{{.Port}} {
{{end}} {{end}}
}{{end}} }{{end}}
{{.Pprof}} {{.Pprof}}
ratelimit 30
hosts { hosts {
fallthrough fallthrough
} }

View File

@ -8,6 +8,7 @@ import (
"sync" // Include all plugins. "sync" // Include all plugins.
_ "github.com/AdguardTeam/AdGuardHome/coredns_plugin" _ "github.com/AdguardTeam/AdGuardHome/coredns_plugin"
_ "github.com/AdguardTeam/AdGuardHome/coredns_plugin/ratelimit"
_ "github.com/AdguardTeam/AdGuardHome/upstream" _ "github.com/AdguardTeam/AdGuardHome/upstream"
"github.com/coredns/coredns/core/dnsserver" "github.com/coredns/coredns/core/dnsserver"
"github.com/coredns/coredns/coremain" "github.com/coredns/coredns/coremain"

View File

@ -15,6 +15,7 @@ import (
"github.com/coredns/coredns/core/dnsserver" "github.com/coredns/coredns/core/dnsserver"
"github.com/coredns/coredns/plugin" "github.com/coredns/coredns/plugin"
"github.com/coredns/coredns/plugin/metrics" "github.com/coredns/coredns/plugin/metrics"
"github.com/coredns/coredns/plugin/pkg/dnstest"
"github.com/coredns/coredns/request" "github.com/coredns/coredns/request"
"github.com/mholt/caddy" "github.com/mholt/caddy"
"github.com/miekg/dns" "github.com/miekg/dns"
@ -23,6 +24,7 @@ import (
) )
const defaultRatelimit = 100 const defaultRatelimit = 100
const defaultResponseSize = 1000
var ( var (
tokenBuckets = cache.New(time.Hour, time.Hour) 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() ratelimited.Inc()
return 0, nil 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) { func (p *plug) allowRequest(ip string) (bool, error) {
@ -94,7 +111,6 @@ type plug struct {
// configuration for creating above // configuration for creating above
ratelimit int // in requests per second per IP ratelimit int // in requests per second per IP
whitelist []string // a list of whitelisted IP addresses whitelist []string // a list of whitelisted IP addresses
} }
func setupPlugin(c *caddy.Controller) (*plug, error) { func setupPlugin(c *caddy.Controller) (*plug, error) {