feat(kv): add rate limits configuration.
This commit is contained in:
parent
51390aa874
commit
79ae52aca7
4 changed files with 133 additions and 0 deletions
|
@ -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("")},
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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"}}"
|
||||
|
|
Loading…
Reference in a new issue