traefik/pkg/middlewares/ratelimiter/rate_limiter_test.go
2022-12-05 11:30:05 +01:00

331 lines
8 KiB
Go

package ratelimiter
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"os"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
ptypes "github.com/traefik/paerser/types"
"github.com/traefik/traefik/v2/pkg/config/dynamic"
"github.com/traefik/traefik/v2/pkg/testhelpers"
"github.com/vulcand/oxy/v2/utils"
)
func TestNewRateLimiter(t *testing.T) {
testCases := []struct {
desc string
config dynamic.RateLimit
expectedMaxDelay time.Duration
expectedSourceIP string
requestHeader string
expectedError string
}{
{
desc: "maxDelay computation",
config: dynamic.RateLimit{
Average: 200,
Burst: 10,
},
expectedMaxDelay: 2500 * time.Microsecond,
},
{
desc: "maxDelay computation, low rate regime",
config: dynamic.RateLimit{
Average: 2,
Period: ptypes.Duration(10 * time.Second),
Burst: 10,
},
expectedMaxDelay: 500 * time.Millisecond,
},
{
desc: "default SourceMatcher is remote address ip strategy",
config: dynamic.RateLimit{
Average: 200,
Burst: 10,
},
expectedSourceIP: "127.0.0.1",
},
{
desc: "SourceCriterion in config is respected",
config: dynamic.RateLimit{
Average: 200,
Burst: 10,
SourceCriterion: &dynamic.SourceCriterion{
RequestHeaderName: "Foo",
},
},
requestHeader: "bar",
},
{
desc: "SourceCriteria are mutually exclusive",
config: dynamic.RateLimit{
Average: 200,
Burst: 10,
SourceCriterion: &dynamic.SourceCriterion{
IPStrategy: &dynamic.IPStrategy{},
RequestHeaderName: "Foo",
},
},
expectedError: "iPStrategy and RequestHeaderName are mutually exclusive",
},
}
for _, test := range testCases {
test := test
t.Run(test.desc, func(t *testing.T) {
t.Parallel()
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
h, err := New(context.Background(), next, test.config, "rate-limiter")
if test.expectedError != "" {
assert.EqualError(t, err, test.expectedError)
} else {
require.NoError(t, err)
}
rtl, _ := h.(*rateLimiter)
if test.expectedMaxDelay != 0 {
assert.Equal(t, test.expectedMaxDelay, rtl.maxDelay)
}
if test.expectedSourceIP != "" {
extractor, ok := rtl.sourceMatcher.(utils.ExtractorFunc)
require.True(t, ok, "Not an ExtractorFunc")
req := http.Request{
RemoteAddr: fmt.Sprintf("%s:1234", test.expectedSourceIP),
}
ip, _, err := extractor(&req)
assert.NoError(t, err)
assert.Equal(t, test.expectedSourceIP, ip)
}
if test.requestHeader != "" {
extractor, ok := rtl.sourceMatcher.(utils.ExtractorFunc)
require.True(t, ok, "Not an ExtractorFunc")
req := http.Request{
Header: map[string][]string{
test.config.SourceCriterion.RequestHeaderName: {test.requestHeader},
},
}
hd, _, err := extractor(&req)
assert.NoError(t, err)
assert.Equal(t, test.requestHeader, hd)
}
})
}
}
func TestRateLimit(t *testing.T) {
testCases := []struct {
desc string
config dynamic.RateLimit
loadDuration time.Duration
incomingLoad int // in reqs/s
burst int
}{
{
desc: "Average is respected",
config: dynamic.RateLimit{
Average: 100,
Burst: 1,
},
loadDuration: 2 * time.Second,
incomingLoad: 400,
},
{
desc: "burst allowed, no bursty traffic",
config: dynamic.RateLimit{
Average: 100,
Burst: 100,
},
loadDuration: 2 * time.Second,
incomingLoad: 200,
},
{
desc: "burst allowed, initial burst, under capacity",
config: dynamic.RateLimit{
Average: 100,
Burst: 100,
},
loadDuration: 2 * time.Second,
incomingLoad: 200,
burst: 50,
},
{
desc: "burst allowed, initial burst, over capacity",
config: dynamic.RateLimit{
Average: 100,
Burst: 100,
},
loadDuration: 2 * time.Second,
incomingLoad: 200,
burst: 150,
},
{
desc: "burst over average, initial burst, over capacity",
config: dynamic.RateLimit{
Average: 100,
Burst: 200,
},
loadDuration: 2 * time.Second,
incomingLoad: 200,
burst: 300,
},
{
desc: "lower than 1/s",
config: dynamic.RateLimit{
Average: 5,
Period: ptypes.Duration(10 * time.Second),
},
loadDuration: 2 * time.Second,
incomingLoad: 100,
burst: 0,
},
{
desc: "lower than 1/s, longer",
config: dynamic.RateLimit{
Average: 5,
Period: ptypes.Duration(10 * time.Second),
},
loadDuration: time.Minute,
incomingLoad: 100,
burst: 0,
},
{
desc: "lower than 1/s, longer, harsher",
config: dynamic.RateLimit{
Average: 1,
Period: ptypes.Duration(time.Minute),
},
loadDuration: time.Minute,
incomingLoad: 100,
burst: 0,
},
{
desc: "period below 1 second",
config: dynamic.RateLimit{
Average: 50,
Period: ptypes.Duration(500 * time.Millisecond),
},
loadDuration: 2 * time.Second,
incomingLoad: 300,
burst: 0,
},
// TODO Try to disambiguate when it fails if it is because of too high a load.
// {
// desc: "Zero average ==> no rate limiting",
// config: dynamic.RateLimit{
// Average: 0,
// Burst: 1,
// },
// incomingLoad: 1000,
// loadDuration: time.Second,
// },
}
for _, test := range testCases {
test := test
t.Run(test.desc, func(t *testing.T) {
if test.loadDuration >= time.Minute && testing.Short() {
t.Skip("skipping test in short mode.")
}
t.Parallel()
reqCount := 0
dropped := 0
next := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
reqCount++
})
h, err := New(context.Background(), next, test.config, "rate-limiter")
require.NoError(t, err)
loadPeriod := time.Duration(1e9 / test.incomingLoad)
start := time.Now()
end := start.Add(test.loadDuration)
ticker := time.NewTicker(loadPeriod)
defer ticker.Stop()
for {
if time.Now().After(end) {
break
}
req := testhelpers.MustNewRequest(http.MethodGet, "http://localhost", nil)
req.RemoteAddr = "127.0.0.1:1234"
w := httptest.NewRecorder()
h.ServeHTTP(w, req)
if w.Result().StatusCode != http.StatusOK {
dropped++
}
if test.burst > 0 && reqCount < test.burst {
// if a burst is defined we first hammer the server with test.burst requests as fast as possible
continue
}
<-ticker.C
}
stop := time.Now()
elapsed := stop.Sub(start)
burst := test.config.Burst
if burst < 1 {
// actual default value
burst = 1
}
period := time.Duration(test.config.Period)
if period == 0 {
period = time.Second
}
if test.config.Average == 0 {
if reqCount < 75*test.incomingLoad/100 {
t.Fatalf("we (arbitrarily) expect at least 75%% of the requests to go through with no rate limiting, and yet only %d/%d went through", reqCount, test.incomingLoad)
}
if dropped != 0 {
t.Fatalf("no request should have been dropped if rate limiting is disabled, and yet %d were", dropped)
}
return
}
// Note that even when there is no bursty traffic,
// we take into account the configured burst,
// because it also helps absorbing non-bursty traffic.
rate := float64(test.config.Average) / float64(period)
wantCount := int(int64(rate*float64(test.loadDuration)) + burst)
// Allow for a 2% leeway
maxCount := wantCount * 102 / 100
// With very high CPU loads,
// we can expect some extra delay in addition to the rate limiting we already do,
// so we allow for some extra leeway there.
// Feel free to adjust wrt to the load on e.g. the CI.
minCount := computeMinCount(wantCount)
if reqCount < minCount {
t.Fatalf("rate was slower than expected: %d requests (wanted > %d) in %v", reqCount, minCount, elapsed)
}
if reqCount > maxCount {
t.Fatalf("rate was faster than expected: %d requests (wanted < %d) in %v", reqCount, maxCount, elapsed)
}
})
}
}
func computeMinCount(wantCount int) int {
if os.Getenv("CI") != "" {
return wantCount * 60 / 100
}
return wantCount * 95 / 100
}