mirror of
https://github.com/superseriousbusiness/gotosocial
synced 2025-06-05 21:59:39 +02:00
[feature] Add rate limit exceptions option, use ISO8601 for rate limit reset (#2151)
* start updating rate limiting, add exceptions * tests, comments, tidying up * add rate limiting exceptions to example config * envparsing * nolint * apply kimbediff * add examples
This commit is contained in:
192
internal/middleware/ratelimit_test.go
Normal file
192
internal/middleware/ratelimit_test.go
Normal file
@@ -0,0 +1,192 @@
|
||||
// GoToSocial
|
||||
// Copyright (C) GoToSocial Authors admin@gotosocial.org
|
||||
// SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful,
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
package middleware_test
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/suite"
|
||||
"github.com/superseriousbusiness/gotosocial/internal/middleware"
|
||||
"github.com/superseriousbusiness/gotosocial/internal/util"
|
||||
)
|
||||
|
||||
type RateLimitTestSuite struct {
|
||||
suite.Suite
|
||||
}
|
||||
|
||||
func (suite *RateLimitTestSuite) TestRateLimit() {
|
||||
// Suppress warnings about debug mode.
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
|
||||
const (
|
||||
trustedPlatform = "X-Test-IP"
|
||||
rlLimit = "X-RateLimit-Limit"
|
||||
rlRemaining = "X-RateLimit-Remaining"
|
||||
rlReset = "X-RateLimit-Reset"
|
||||
)
|
||||
|
||||
type rlTest struct {
|
||||
limit int
|
||||
exceptions []string
|
||||
clientIP string
|
||||
shouldPanic bool
|
||||
shouldExcept bool
|
||||
}
|
||||
|
||||
for _, test := range []rlTest{
|
||||
{
|
||||
limit: 10,
|
||||
exceptions: []string{},
|
||||
clientIP: "192.0.2.0",
|
||||
shouldPanic: false,
|
||||
shouldExcept: false,
|
||||
},
|
||||
{
|
||||
limit: 10,
|
||||
exceptions: []string{},
|
||||
clientIP: "192.0.2.0",
|
||||
shouldPanic: false,
|
||||
shouldExcept: false,
|
||||
},
|
||||
{
|
||||
limit: 10,
|
||||
exceptions: []string{"192.0.2.0/24"},
|
||||
clientIP: "192.0.2.0",
|
||||
shouldPanic: false,
|
||||
shouldExcept: true,
|
||||
},
|
||||
{
|
||||
limit: 10,
|
||||
exceptions: []string{"192.0.2.0/32"},
|
||||
clientIP: "192.0.2.1",
|
||||
shouldPanic: false,
|
||||
shouldExcept: false,
|
||||
},
|
||||
{
|
||||
limit: 10,
|
||||
exceptions: []string{"Ceci n'est pas une CIDR"},
|
||||
clientIP: "192.0.2.0",
|
||||
shouldPanic: true,
|
||||
shouldExcept: false,
|
||||
},
|
||||
} {
|
||||
if test.shouldPanic {
|
||||
// Try to trigger panic.
|
||||
suite.Panics(func() {
|
||||
_ = middleware.RateLimit(
|
||||
test.limit,
|
||||
test.exceptions,
|
||||
)
|
||||
})
|
||||
continue
|
||||
}
|
||||
|
||||
rlMiddleware := middleware.RateLimit(
|
||||
test.limit,
|
||||
test.exceptions,
|
||||
)
|
||||
|
||||
// Approximate time when this limiter will reset.
|
||||
resetAt := time.Now().Add(5 * time.Minute)
|
||||
|
||||
// Make requests up to +
|
||||
// just over the limit.
|
||||
limitedAt := test.limit + 1
|
||||
for requestsCount := 1; requestsCount < limitedAt; requestsCount++ {
|
||||
var (
|
||||
recorder = httptest.NewRecorder()
|
||||
ctx, e = gin.CreateTestContext(recorder)
|
||||
)
|
||||
|
||||
// Instruct engine to derive
|
||||
// clientIP from test header.
|
||||
e.TrustedPlatform = trustedPlatform
|
||||
ctx.Request = httptest.NewRequest(http.MethodGet, "/example", nil)
|
||||
ctx.Request.Header.Add(trustedPlatform, test.clientIP)
|
||||
|
||||
// Call the rate limiter.
|
||||
rlMiddleware(ctx)
|
||||
|
||||
// Fetch RL headers if present.
|
||||
var (
|
||||
limitStr = recorder.Header().Get(rlLimit)
|
||||
remainingStr = recorder.Header().Get(rlRemaining)
|
||||
resetStr = recorder.Header().Get(rlReset)
|
||||
)
|
||||
|
||||
if test.shouldExcept {
|
||||
// Request should be allowed through,
|
||||
// no rate-limit headers should be written.
|
||||
suite.Equal(http.StatusOK, recorder.Code)
|
||||
suite.Empty(limitStr)
|
||||
suite.Empty(remainingStr)
|
||||
suite.Empty(resetStr)
|
||||
continue
|
||||
}
|
||||
|
||||
suite.Equal(strconv.Itoa(test.limit), limitStr)
|
||||
suite.Equal(strconv.Itoa(test.limit-requestsCount), remainingStr)
|
||||
|
||||
// Ensure reset is ISO8601, and resets at
|
||||
// approximate reset time (+/- 10 seconds).
|
||||
reset, err := util.ParseISO8601(resetStr)
|
||||
if err != nil {
|
||||
suite.FailNow("", "couldn't parse %s as ISO8601: %q", resetStr, err.Error())
|
||||
}
|
||||
suite.WithinDuration(resetAt, reset, 10*time.Second)
|
||||
|
||||
if requestsCount < limitedAt {
|
||||
// Request should be allowed through.
|
||||
suite.Equal(http.StatusOK, recorder.Code)
|
||||
continue
|
||||
}
|
||||
|
||||
// Request should be denied.
|
||||
suite.Equal(http.StatusTooManyRequests, recorder.Code)
|
||||
|
||||
// Make a final request with an unrelated IP to
|
||||
// ensure it's only the one IP being limited.
|
||||
var (
|
||||
unrelatedRecorder = httptest.NewRecorder()
|
||||
unrelatedCtx, unrelatedE = gin.CreateTestContext(unrelatedRecorder)
|
||||
)
|
||||
|
||||
// Instruct engine to derive
|
||||
// clientIP from test header.
|
||||
unrelatedE.TrustedPlatform = trustedPlatform
|
||||
unrelatedCtx.Request = httptest.NewRequest(http.MethodGet, "/example", nil)
|
||||
unrelatedCtx.Request.Header.Add(trustedPlatform, "192.0.2.255")
|
||||
|
||||
// Call the rate limiter.
|
||||
rlMiddleware(unrelatedCtx)
|
||||
|
||||
// Request should be allowed through.
|
||||
suite.Equal(http.StatusOK, unrelatedRecorder.Code)
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRateLimitTestSuite(t *testing.T) {
|
||||
suite.Run(t, new(RateLimitTestSuite))
|
||||
}
|
Reference in New Issue
Block a user