code.vegaprotocol.io/vega@v0.79.0/libs/http/rate_limit_test.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package http_test
    17  
    18  import (
    19  	"context"
    20  	"testing"
    21  	"time"
    22  
    23  	"code.vegaprotocol.io/vega/libs/config/encoding"
    24  	vghttp "code.vegaprotocol.io/vega/libs/http"
    25  
    26  	"github.com/stretchr/testify/assert"
    27  )
    28  
    29  func TestRateLimit(t *testing.T) {
    30  	ctx := context.Background()
    31  	rl, err := vghttp.NewRateLimit(
    32  		ctx,
    33  		vghttp.RateLimitConfig{
    34  			CoolDown:  encoding.Duration{Duration: 1 * time.Minute},
    35  			AllowList: []string{"1.2.3.4/32", "2.3.4.252/30", "fe80::/10"},
    36  		},
    37  	)
    38  	assert.NoError(t, err)
    39  	if assert.NotNil(t, rl) {
    40  		// IP addresses in the allow list
    41  		for _, ip := range []string{"1.2.3.4", "2.3.4.254", "fe80::abcd"} {
    42  			for i := 0; i < 10; i++ {
    43  				err = rl.NewRequest("someprefix", ip)
    44  				assert.NoError(t, err)
    45  			}
    46  		}
    47  
    48  		// IP address not in the allow list
    49  		err = rl.NewRequest("someprefix", "2.2.2.2")
    50  		assert.NoError(t, err)
    51  		err = rl.NewRequest("someprefix", "2.2.2.2")
    52  		assert.Error(t, err)
    53  	}
    54  }