From 79ae52aca7b8751ef54a21b4815d3ca49bc1a39c Mon Sep 17 00:00:00 2001 From: Fernandez Ludovic Date: Wed, 3 Jan 2018 16:42:40 +0100 Subject: [PATCH] feat(kv): add rate limits configuration. --- provider/kv/filler_test.go | 17 ++++++++++ provider/kv/kv_config.go | 40 ++++++++++++++++++++++ provider/kv/kv_config_test.go | 63 +++++++++++++++++++++++++++++++++++ templates/kv.tmpl | 13 ++++++++ 4 files changed, 133 insertions(+) diff --git a/provider/kv/filler_test.go b/provider/kv/filler_test.go index 9c88aa5a6..571a2cccd 100644 --- a/provider/kv/filler_test.go +++ b/provider/kv/filler_test.go @@ -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) { expected := []*store.KVPair{ {Key: "traefik/backends/backend.with.dot.too", Value: []byte("")}, diff --git a/provider/kv/kv_config.go b/provider/kv/kv_config.go index f13833b7a..d1e73e438 100644 --- a/provider/kv/kv_config.go +++ b/provider/kv/kv_config.go @@ -8,6 +8,7 @@ import ( "text/template" "github.com/BurntSushi/ty/fun" + "github.com/containous/flaeg" "github.com/containous/traefik/log" "github.com/containous/traefik/provider/label" "github.com/containous/traefik/types" @@ -36,6 +37,7 @@ func (p *Provider) buildConfiguration() *types.Configuration { // Frontend functions "getRedirect": p.getRedirect, "getErrorPages": p.getErrorPages, + "getRateLimit": p.getRateLimit, // Backend functions "getSticky": p.getSticky, @@ -120,6 +122,44 @@ func (p *Provider) getErrorPages(rootPath string) map[string]*types.ErrorPage { 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 { serverNames := p.list(backend, pathBackendServers) return fun.Filter(p.serverFilter, serverNames).([]string) diff --git a/provider/kv/kv_config_test.go b/provider/kv/kv_config_test.go index 4b00caf12..607c4afb5 100644 --- a/provider/kv/kv_config_test.go +++ b/provider/kv/kv_config_test.go @@ -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) + }) + } +} diff --git a/templates/kv.tmpl b/templates/kv.tmpl index 0952a5388..eecb0b693 100644 --- a/templates/kv.tmpl +++ b/templates/kv.tmpl @@ -91,6 +91,19 @@ {{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/"}} [frontends."{{$frontendName}}".routes."{{Last $route}}"] rule = "{{Get "" $route "/rule"}}"