150 lines
4.0 KiB
Go
150 lines
4.0 KiB
Go
// 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/>.
|
|
|
|
/*
|
|
The code in this file is adapted from MIT-licensed code in github.com/go-chi/chi. Thanks chi (thi)!
|
|
|
|
See: https://github.com/go-chi/chi/blob/e6baba61759b26ddf7b14d1e02d1da81a4d76c08/middleware/throttle.go
|
|
|
|
And: https://github.com/sponsors/pkieltyka
|
|
*/
|
|
|
|
package middleware_test
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"runtime"
|
|
"strconv"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/superseriousbusiness/gotosocial/internal/middleware"
|
|
)
|
|
|
|
func TestThrottlingMiddleware(t *testing.T) {
|
|
testThrottlingMiddleware(t, 2, time.Second*10)
|
|
testThrottlingMiddleware(t, 4, time.Second*15)
|
|
testThrottlingMiddleware(t, 8, time.Second*30)
|
|
}
|
|
|
|
func testThrottlingMiddleware(t *testing.T, cpuMulti int, retryAfter time.Duration) {
|
|
// Calculate expected request limit + queue.
|
|
limit := runtime.GOMAXPROCS(0) * cpuMulti
|
|
queueLimit := limit * cpuMulti
|
|
|
|
// Calculate expected retry-after header string.
|
|
retryAfterStr := strconv.FormatUint(uint64(retryAfter/time.Second), 10)
|
|
|
|
// Gin test http engine
|
|
// (used for ctx init).
|
|
e := gin.New()
|
|
|
|
// Add middleware to the gin engine handler stack.
|
|
middleware := middleware.Throttle(cpuMulti, retryAfter)
|
|
e.Use(middleware)
|
|
|
|
// Set the blocking gin handler.
|
|
handler := blockingHandler()
|
|
e.Handle("GET", "/", handler)
|
|
|
|
var cncls []func()
|
|
|
|
for i := 0; i < queueLimit+limit; i++ {
|
|
// Prepare a gin test context.
|
|
r := httptest.NewRequest("GET", "/", nil)
|
|
rw := httptest.NewRecorder()
|
|
|
|
// Wrap request with new cancel context.
|
|
ctx, cncl := context.WithCancel(r.Context())
|
|
r = r.WithContext(ctx)
|
|
|
|
// Pass req through
|
|
// engine handler.
|
|
go e.ServeHTTP(rw, r)
|
|
time.Sleep(time.Millisecond)
|
|
|
|
// Get http result.
|
|
res := rw.Result()
|
|
|
|
if i < queueLimit {
|
|
|
|
// Check status == 200 (default, i.e not set).
|
|
if res.StatusCode != http.StatusOK {
|
|
t.Fatalf("status code was set (%d) with queueLimit=%d and request=%d", res.StatusCode, queueLimit, i)
|
|
}
|
|
|
|
// Add cancel to func slice.
|
|
cncls = append(cncls, cncl)
|
|
|
|
} else {
|
|
|
|
// Check the returned status code is expected.
|
|
if res.StatusCode != http.StatusTooManyRequests {
|
|
t.Fatalf("did not return status 429 (%d) with queueLimit=%d and request=%d", res.StatusCode, queueLimit, i)
|
|
}
|
|
|
|
// Check the returned retry-after header is set.
|
|
if res.Header.Get("Retry-After") != retryAfterStr {
|
|
t.Fatalf("did not return retry-after %s with queueLimit=%d and request=%d", retryAfterStr, queueLimit, i)
|
|
}
|
|
|
|
// Cancel on return.
|
|
defer cncl()
|
|
|
|
}
|
|
}
|
|
|
|
// Cancel all blocked reqs.
|
|
for _, cncl := range cncls {
|
|
cncl()
|
|
}
|
|
time.Sleep(time.Second)
|
|
|
|
// Check a bunchh more requests
|
|
// can now make it through after
|
|
// previous requests were released!
|
|
for i := 0; i < limit; i++ {
|
|
|
|
// Prepare a gin test context.
|
|
r := httptest.NewRequest("GET", "/", nil)
|
|
rw := httptest.NewRecorder()
|
|
|
|
// Pass req through
|
|
// engine handler.
|
|
go e.ServeHTTP(rw, r)
|
|
time.Sleep(5 * time.Millisecond)
|
|
|
|
// Get http result.
|
|
res := rw.Result()
|
|
|
|
// Check status == 200 (default, i.e not set).
|
|
if res.StatusCode != http.StatusOK {
|
|
t.Fatalf("status code was set (%d) with queueLimit=%d and request=%d", res.StatusCode, queueLimit, i)
|
|
}
|
|
}
|
|
}
|
|
|
|
func blockingHandler() gin.HandlerFunc {
|
|
return func(ctx *gin.Context) {
|
|
<-ctx.Done()
|
|
ctx.Status(201) // specifically not 200
|
|
}
|
|
}
|