storj/private/web/ratelimiter_test.go
Jeremy Wharton 6a6cc28fc1 satellite/console,private/web: Rate limit coupon code application
Rate limits application of coupon codes by user ID to prevent
brute forcing. Refactors the rate limiter to allow limiting based
on arbitrary criteria and not just by IP.

Change-Id: I99d6749bd5b5e47d7e1aeb0314e363a8e7259dba
2021-08-23 17:00:31 +00:00

71 lines
2.1 KiB
Go

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.
package web_test
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/spf13/pflag"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"storj.io/common/testcontext"
"storj.io/private/cfgstruct"
"storj.io/storj/private/web"
)
func TestNewIPRateLimiter(t *testing.T) {
// create a rate limiter with defaults except NumLimits = 2
config := web.RateLimiterConfig{}
cfgstruct.Bind(&pflag.FlagSet{}, &config, cfgstruct.UseDevDefaults())
config.NumLimits = 2
rateLimiter := web.NewIPRateLimiter(config)
// run ratelimiter cleanup until end of test
ctx := testcontext.New(t)
defer ctx.Cleanup()
ctx2, cancel := context.WithCancel(ctx)
defer cancel()
ctx.Go(func() error {
rateLimiter.Run(ctx2)
return nil
})
// make the default HTTP handler return StatusOK
handler := rateLimiter.Limit(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
}))
// expect burst number of successes
testWithAddress(ctx, t, "192.168.1.1:5000", rateLimiter.Burst(), handler)
// expect similar results for a different IP
testWithAddress(ctx, t, "127.0.0.1:5000", rateLimiter.Burst(), handler)
// expect similar results for a different IP
testWithAddress(ctx, t, "127.0.0.100:5000", rateLimiter.Burst(), handler)
// expect original IP to work again because numLimits == 2
testWithAddress(ctx, t, "192.168.1.1:5000", rateLimiter.Burst(), handler)
}
func testWithAddress(ctx context.Context, t *testing.T, remoteAddress string, burst int, handler http.Handler) {
// create HTTP request
req, err := http.NewRequestWithContext(ctx, "GET", "", nil)
require.NoError(t, err)
req.RemoteAddr = remoteAddress
// expect burst number of successes
for x := 0; x < burst; x++ {
rr := httptest.NewRecorder()
handler.ServeHTTP(rr, req)
assert.Equal(t, rr.Code, http.StatusOK, remoteAddress)
}
// then expect failure
rr := httptest.NewRecorder()
handler.ServeHTTP(rr, req)
assert.Equal(t, rr.Code, http.StatusTooManyRequests, remoteAddress)
}