104 lines
3.3 KiB
Go
104 lines
3.3 KiB
Go
package middleware
|
|
|
|
import (
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strconv"
|
|
"testing"
|
|
"time"
|
|
|
|
"agent-desk/internal/pkg/ratelimit"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func newRateLimitTestEngine(limiter *ratelimit.Limiter) *gin.Engine {
|
|
gin.SetMode(gin.TestMode)
|
|
app := gin.New()
|
|
app.POST("/probe", RateLimit(limiter), func(ctx *gin.Context) {
|
|
ctx.JSON(http.StatusOK, gin.H{"success": true})
|
|
})
|
|
return app
|
|
}
|
|
|
|
func postProbe(app *gin.Engine, clientIP string) *httptest.ResponseRecorder {
|
|
rec := httptest.NewRecorder()
|
|
req := httptest.NewRequest(http.MethodPost, "/probe", nil)
|
|
req.RemoteAddr = clientIP + ":52000"
|
|
app.ServeHTTP(rec, req)
|
|
return rec
|
|
}
|
|
|
|
func TestRateLimitRefusesAfterTheLimitAndSetsRetryAfter(t *testing.T) {
|
|
app := newRateLimitTestEngine(ratelimit.New(2, time.Minute))
|
|
|
|
for i := 1; i <= 2; i++ {
|
|
if rec := postProbe(app, "203.0.113.7"); rec.Code != http.StatusOK {
|
|
t.Fatalf("request %d got status %d, want 200", i, rec.Code)
|
|
}
|
|
}
|
|
|
|
rec := postProbe(app, "203.0.113.7")
|
|
if rec.Code != http.StatusTooManyRequests {
|
|
t.Fatalf("the third request got status %d, want 429", rec.Code)
|
|
}
|
|
|
|
retryAfter := rec.Header().Get("Retry-After")
|
|
if retryAfter == "" {
|
|
t.Fatal("a 429 without Retry-After gives the caller nothing to act on")
|
|
}
|
|
seconds, err := strconv.Atoi(retryAfter)
|
|
if err != nil {
|
|
t.Fatalf("Retry-After = %q is not a number of seconds", retryAfter)
|
|
}
|
|
// The window is one minute and the requests are microseconds apart, so the
|
|
// answer is 60 give or take a clock tick. What must not happen is a value
|
|
// above the window, which would tell the caller to wait longer than needed, or
|
|
// zero, which would have them retry immediately and collect another 429.
|
|
if seconds < 1 || seconds > 60 {
|
|
t.Errorf("Retry-After = %d seconds, want between 1 and the 60 second window", seconds)
|
|
}
|
|
|
|
// The body has to stay a JsonResult, because web/lib/api/client.ts parses the
|
|
// payload before it looks at response.ok and surfaces payload.message.
|
|
var body struct {
|
|
Success bool `json:"success"`
|
|
Message string `json:"message"`
|
|
}
|
|
if err := json.Unmarshal(rec.Body.Bytes(), &body); err != nil {
|
|
t.Fatalf("429 body is not JSON: %v (%s)", err, rec.Body.String())
|
|
}
|
|
if body.Success {
|
|
t.Error("429 body reported success=true")
|
|
}
|
|
if body.Message == "" {
|
|
t.Error("429 body carried no message, so the caller would see a blank error")
|
|
}
|
|
}
|
|
|
|
func TestRateLimitIsKeyedByClientAddress(t *testing.T) {
|
|
app := newRateLimitTestEngine(ratelimit.New(1, time.Minute))
|
|
|
|
if rec := postProbe(app, "203.0.113.7"); rec.Code != http.StatusOK {
|
|
t.Fatalf("first request got %d, want 200", rec.Code)
|
|
}
|
|
if rec := postProbe(app, "203.0.113.7"); rec.Code != http.StatusTooManyRequests {
|
|
t.Fatalf("second request from the same address got %d, want 429", rec.Code)
|
|
}
|
|
// Somebody else must not inherit the first caller's exhaustion. Behind a NAT
|
|
// this is the difference between a working product and a shared lockout.
|
|
if rec := postProbe(app, "198.51.100.9"); rec.Code != http.StatusOK {
|
|
t.Fatalf("a different address got %d, want 200", rec.Code)
|
|
}
|
|
}
|
|
|
|
func TestRateLimitWithNilLimiterAllowsEverything(t *testing.T) {
|
|
app := newRateLimitTestEngine(nil)
|
|
for i := 0; i < 50; i++ {
|
|
if rec := postProbe(app, "203.0.113.7"); rec.Code != http.StatusOK {
|
|
t.Fatalf("request %d got %d with a nil limiter, want 200", i, rec.Code)
|
|
}
|
|
}
|
|
}
|