github.com/lologarithm/mattermost-server@v5.3.2-0.20181002060438-c82a84ed765b+incompatible/app/ratelimit_test.go (about)

     1  // Copyright (c) 2018-present Mattermost, Inc. All Rights Reserved.
     2  // See License.txt for license information.
     3  
     4  package app
     5  
     6  import (
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"strconv"
    10  	"testing"
    11  
    12  	"github.com/mattermost/mattermost-server/model"
    13  	"github.com/stretchr/testify/require"
    14  )
    15  
    16  func genRateLimitSettings(useAuth, useIP bool, header string) *model.RateLimitSettings {
    17  	return &model.RateLimitSettings{
    18  		Enable:           model.NewBool(true),
    19  		PerSec:           model.NewInt(10),
    20  		MaxBurst:         model.NewInt(100),
    21  		MemoryStoreSize:  model.NewInt(10000),
    22  		VaryByRemoteAddr: model.NewBool(useIP),
    23  		VaryByUser:       model.NewBool(useAuth),
    24  		VaryByHeader:     header,
    25  	}
    26  }
    27  
    28  func TestNewRateLimiterSuccess(t *testing.T) {
    29  	settings := genRateLimitSettings(false, false, "")
    30  	rateLimiter, err := NewRateLimiter(settings)
    31  	require.NotNil(t, rateLimiter)
    32  	require.NoError(t, err)
    33  }
    34  
    35  func TestNewRateLimiterFailure(t *testing.T) {
    36  	invalidSettings := genRateLimitSettings(false, false, "")
    37  	invalidSettings.MaxBurst = model.NewInt(-100)
    38  	rateLimiter, err := NewRateLimiter(invalidSettings)
    39  	require.Nil(t, rateLimiter)
    40  	require.Error(t, err)
    41  }
    42  
    43  func TestGenerateKey(t *testing.T) {
    44  	cases := []struct {
    45  		useAuth         bool
    46  		useIP           bool
    47  		header          string
    48  		authTokenResult string
    49  		ipResult        string
    50  		headerResult    string
    51  		expectedKey     string
    52  	}{
    53  		{false, false, "", "", "", "", ""},
    54  		{true, false, "", "resultkey", "notme", "notme", "resultkey"},
    55  		{false, true, "", "notme", "resultkey", "notme", "resultkey"},
    56  		{false, false, "myheader", "notme", "notme", "resultkey", "resultkey"},
    57  		{true, true, "", "resultkey", "ipaddr", "notme", "resultkey"},
    58  		{true, true, "", "", "ipaddr", "notme", "ipaddr"},
    59  		{true, true, "myheader", "resultkey", "ipaddr", "hadd", "resultkeyhadd"},
    60  		{true, true, "myheader", "", "ipaddr", "hadd", "ipaddrhadd"},
    61  	}
    62  
    63  	for testnum, tc := range cases {
    64  		req := httptest.NewRequest("GET", "/", nil)
    65  		if tc.authTokenResult != "" {
    66  			req.AddCookie(&http.Cookie{
    67  				Name:  model.SESSION_COOKIE_TOKEN,
    68  				Value: tc.authTokenResult,
    69  			})
    70  		}
    71  		req.RemoteAddr = tc.ipResult + ":80"
    72  		if tc.headerResult != "" {
    73  			req.Header.Set(tc.header, tc.headerResult)
    74  		}
    75  
    76  		rateLimiter, _ := NewRateLimiter(genRateLimitSettings(tc.useAuth, tc.useIP, tc.header))
    77  
    78  		key := rateLimiter.GenerateKey(req)
    79  
    80  		require.Equal(t, tc.expectedKey, key, "Wrong key on test "+strconv.Itoa(testnum))
    81  	}
    82  }