feat(kv): add rate limits configuration.

This commit is contained in:
Fernandez Ludovic 2018-01-03 16:42:40 +01:00 committed by Traefiker
parent 51390aa874
commit 79ae52aca7
4 changed files with 133 additions and 0 deletions

View file

@ -93,6 +93,23 @@ func withErrorPage(name string, backend, query, status string) func(map[string]s
} }
} }
func withRateLimit(extractorFunc string, opts ...func(map[string]string)) func(map[string]string) {
return func(pairs map[string]string) {
pairs[pathFrontendRateLimitExtractorFunc] = extractorFunc
for _, opt := range opts {
opt(pairs)
}
}
}
func withLimit(name string, average, burst, period string) func(map[string]string) {
return func(pairs map[string]string) {
pairs[pathFrontendRateLimitRateSet+name+pathFrontendRateLimitAverage] = average
pairs[pathFrontendRateLimitRateSet+name+pathFrontendRateLimitBurst] = burst
pairs[pathFrontendRateLimitRateSet+name+pathFrontendRateLimitPeriod] = period
}
}
func TestFiller(t *testing.T) { func TestFiller(t *testing.T) {
expected := []*store.KVPair{ expected := []*store.KVPair{
{Key: "traefik/backends/backend.with.dot.too", Value: []byte("")}, {Key: "traefik/backends/backend.with.dot.too", Value: []byte("")},

View file

@ -8,6 +8,7 @@ import (
"text/template" "text/template"
"github.com/BurntSushi/ty/fun" "github.com/BurntSushi/ty/fun"
"github.com/containous/flaeg"
"github.com/containous/traefik/log" "github.com/containous/traefik/log"
"github.com/containous/traefik/provider/label" "github.com/containous/traefik/provider/label"
"github.com/containous/traefik/types" "github.com/containous/traefik/types"
@ -36,6 +37,7 @@ func (p *Provider) buildConfiguration() *types.Configuration {
// Frontend functions // Frontend functions
"getRedirect": p.getRedirect, "getRedirect": p.getRedirect,
"getErrorPages": p.getErrorPages, "getErrorPages": p.getErrorPages,
"getRateLimit": p.getRateLimit,
// Backend functions // Backend functions
"getSticky": p.getSticky, "getSticky": p.getSticky,
@ -120,6 +122,44 @@ func (p *Provider) getErrorPages(rootPath string) map[string]*types.ErrorPage {
return errorPages return errorPages
} }
func (p *Provider) getRateLimit(rootPath string) *types.RateLimit {
extractorFunc := p.get("", rootPath, pathFrontendRateLimitExtractorFunc)
if len(extractorFunc) == 0 {
return nil
}
var limits map[string]*types.Rate
pathRateSet := p.list(rootPath, pathFrontendRateLimitRateSet)
for _, pathLimits := range pathRateSet {
if limits == nil {
limits = make(map[string]*types.Rate)
}
rawPeriod := p.get("", pathLimits+pathFrontendRateLimitPeriod)
var period flaeg.Duration
err := period.Set(rawPeriod)
if err != nil {
log.Errorf("Invalid %q value: %q", pathLimits+pathFrontendRateLimitPeriod, rawPeriod)
continue
}
limitName := p.last(pathLimits)
limits[limitName] = &types.Rate{
Average: p.getInt64(0, pathLimits+pathFrontendRateLimitAverage),
Burst: p.getInt64(0, pathLimits+pathFrontendRateLimitBurst),
Period: period,
}
}
return &types.RateLimit{
ExtractorFunc: extractorFunc,
RateSet: limits,
}
}
func (p *Provider) listServers(backend string) []string { func (p *Provider) listServers(backend string) []string {
serverNames := p.list(backend, pathBackendServers) serverNames := p.list(backend, pathBackendServers)
return fun.Filter(p.serverFilter, serverNames).([]string) return fun.Filter(p.serverFilter, serverNames).([]string)

View file

@ -819,3 +819,66 @@ func TestProviderGetErrorPages(t *testing.T) {
}) })
} }
} }
func TestProviderGetRateLimit(t *testing.T) {
testCases := []struct {
desc string
rootPath string
kvPairs []*store.KVPair
expected *types.RateLimit
}{
{
desc: "with several limits",
rootPath: "traefik/frontends/foo",
kvPairs: filler("traefik",
frontend("foo",
withRateLimit("client.ip",
withLimit("foo", "6", "12", "18"),
withLimit("bar", "3", "6", "9")))),
expected: &types.RateLimit{
ExtractorFunc: "client.ip",
RateSet: map[string]*types.Rate{
"foo": {
Average: 6,
Burst: 12,
Period: flaeg.Duration(18 * time.Second),
},
"bar": {
Average: 3,
Burst: 6,
Period: flaeg.Duration(9 * time.Second),
},
},
},
},
{
desc: "return nil when no extractor func",
rootPath: "traefik/frontends/foo",
kvPairs: filler("traefik",
frontend("foo",
withRateLimit("",
withLimit("foo", "6", "12", "18"),
withLimit("bar", "3", "6", "9")))),
expected: nil,
},
{
desc: "return nil when no rate limit keys",
rootPath: "traefik/frontends/foo",
kvPairs: filler("traefik", frontend("foo")),
expected: nil,
},
}
for _, test := range testCases {
test := test
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
p := newProviderMock(test.kvPairs)
actual := p.getRateLimit(test.rootPath)
assert.Equal(t, test.expected, actual)
})
}
}

View file

@ -91,6 +91,19 @@
{{end}} {{end}}
{{end}} {{end}}
{{ $rateLimit := getRateLimit $frontend }}
{{ if $rateLimit }}
[frontends."{{$frontendName}}".rateLimit]
extractorFunc = "{{ $rateLimit.ExtractorFunc }}"
[frontends."{{$frontendName}}".rateLimit.rateSet]
{{ range $limitName, $rateLimit := $rateLimit.RateSet }}
[frontends."{{$frontendName}}".rateLimit.rateSet.{{ $limitName }}]
period = "{{ $rateLimit.Period }}"
average = {{ $rateLimit.Average }}
burst = {{ $rateLimit.Burst }}
{{end}}
{{end}}
{{range $route := List $frontend "/routes/"}} {{range $route := List $frontend "/routes/"}}
[frontends."{{$frontendName}}".routes."{{Last $route}}"] [frontends."{{$frontendName}}".routes."{{Last $route}}"]
rule = "{{Get "" $route "/rule"}}" rule = "{{Get "" $route "/rule"}}"