mirror of
https://github.com/mattermost/mattermost.git
synced 2025-02-25 18:55:24 -06:00
034dbc07e3
Previously, an error occuring in NewRateLimiter would return a nil reference – which would be de-referenced just after, making the server crash.
83 lines
2.4 KiB
Go
83 lines
2.4 KiB
Go
// Copyright (c) 2018-present Mattermost, Inc. All Rights Reserved.
|
|
// See License.txt for license information.
|
|
|
|
package app
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strconv"
|
|
"testing"
|
|
|
|
"github.com/mattermost/mattermost-server/model"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func genRateLimitSettings(useAuth, useIP bool, header string) *model.RateLimitSettings {
|
|
return &model.RateLimitSettings{
|
|
Enable: model.NewBool(true),
|
|
PerSec: model.NewInt(10),
|
|
MaxBurst: model.NewInt(100),
|
|
MemoryStoreSize: model.NewInt(10000),
|
|
VaryByRemoteAddr: model.NewBool(useIP),
|
|
VaryByUser: model.NewBool(useAuth),
|
|
VaryByHeader: header,
|
|
}
|
|
}
|
|
|
|
func TestNewRateLimiterSuccess(t *testing.T) {
|
|
settings := genRateLimitSettings(false, false, "")
|
|
rateLimiter, err := NewRateLimiter(settings)
|
|
require.NotNil(t, rateLimiter)
|
|
require.NoError(t, err)
|
|
}
|
|
|
|
func TestNewRateLimiterFailure(t *testing.T) {
|
|
invalidSettings := genRateLimitSettings(false, false, "")
|
|
invalidSettings.MaxBurst = model.NewInt(-100)
|
|
rateLimiter, err := NewRateLimiter(invalidSettings)
|
|
require.Nil(t, rateLimiter)
|
|
require.Error(t, err)
|
|
}
|
|
|
|
func TestGenerateKey(t *testing.T) {
|
|
cases := []struct {
|
|
useAuth bool
|
|
useIP bool
|
|
header string
|
|
authTokenResult string
|
|
ipResult string
|
|
headerResult string
|
|
expectedKey string
|
|
}{
|
|
{false, false, "", "", "", "", ""},
|
|
{true, false, "", "resultkey", "notme", "notme", "resultkey"},
|
|
{false, true, "", "notme", "resultkey", "notme", "resultkey"},
|
|
{false, false, "myheader", "notme", "notme", "resultkey", "resultkey"},
|
|
{true, true, "", "resultkey", "ipaddr", "notme", "resultkey"},
|
|
{true, true, "", "", "ipaddr", "notme", "ipaddr"},
|
|
{true, true, "myheader", "resultkey", "ipaddr", "hadd", "resultkeyhadd"},
|
|
{true, true, "myheader", "", "ipaddr", "hadd", "ipaddrhadd"},
|
|
}
|
|
|
|
for testnum, tc := range cases {
|
|
req := httptest.NewRequest("GET", "/", nil)
|
|
if tc.authTokenResult != "" {
|
|
req.AddCookie(&http.Cookie{
|
|
Name: model.SESSION_COOKIE_TOKEN,
|
|
Value: tc.authTokenResult,
|
|
})
|
|
}
|
|
req.RemoteAddr = tc.ipResult + ":80"
|
|
if tc.headerResult != "" {
|
|
req.Header.Set(tc.header, tc.headerResult)
|
|
}
|
|
|
|
rateLimiter, _ := NewRateLimiter(genRateLimitSettings(tc.useAuth, tc.useIP, tc.header))
|
|
|
|
key := rateLimiter.GenerateKey(req)
|
|
|
|
require.Equal(t, tc.expectedKey, key, "Wrong key on test "+strconv.Itoa(testnum))
|
|
}
|
|
}
|