coderd/httpmw/ratelimit_test.go
Rama main · réplica de lectura del repositorio original.
package httpmw_test
import (
"fmt"
"math/rand"
"net"
"net/http"
"net/http/httptest"
"sync"
"testing"
"time"
"github.com/go-chi/chi/v5"
"github.com/stretchr/testify/require"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbgen"
"github.com/coder/coder/v2/coderd/database/dbtestutil"
"github.com/coder/coder/v2/coderd/httpmw"
"github.com/coder/coder/v2/codersdk"
"github.com/coder/coder/v2/testutil"
)
func randRemoteAddr() string {
var b [4]byte
// nolint:gosec
_, _ = rand.Read(b[:])
// nolint:gosec
return fmt.Sprintf("%s:%v", net.IP(b[:]).String(), rand.Int31()%(1<<16))
}
func TestRateLimit(t *testing.T) {
t.Parallel()
t.Run("NoUserSucceeds", func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
rtr.Use(httpmw.RateLimit(1, time.Second))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
for i := 0; i < 5; i++ {
req := httptest.NewRequest("GET", "/", nil)
rec := httptest.NewRecorder()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.Equal(t, i != 0, resp.StatusCode == http.StatusTooManyRequests)
}
})
t.Run("PathNormalizationBypass", func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
rtr.Use(httpmw.RateLimit(1, time.Second))
// A wildcard route so that requests for both the canonical path and
// its redundant-slash variants reach the same handler, mirroring
// how chi's router resolves /api/v2/users//validate-password to the
// same handler as /api/v2/users/validate-password in production.
rtr.Post("/*", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
remoteAddr := randRemoteAddr()
paths := []string{
"/api/v2/users/validate-password",
"/api/v2/users//validate-password",
"/api/v2/users///validate-password",
"/api/v2/users/validate-password",
}
for i, p := range paths {
req := httptest.NewRequest("POST", p, nil)
req.RemoteAddr = remoteAddr
rec := httptest.NewRecorder()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.Equal(t, i != 0, resp.StatusCode == http.StatusTooManyRequests, "request %d (%s)", i, p)
}
})
t.Run("DifferentAPIPrefixes", func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
rtr.Use(httpmw.RateLimit(1, time.Second))
rtr.Get("/*", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
remoteAddr := randRemoteAddr()
for _, p := range []string{
"/api/v2/chats/providers",
"/api/experimental/chats/providers",
} {
req := httptest.NewRequest("GET", p, nil)
req.RemoteAddr = remoteAddr
rec := httptest.NewRecorder()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode, p)
}
})
t.Run("PerCallerSharesBucketAcrossPaths", func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
rtr.With(httpmw.RateLimitPerCaller(1, time.Second)).Post("/chats/{chat}/workspace-files", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
remoteAddr := randRemoteAddr()
for i, p := range []string{
"/chats/00000000-0000-0000-0000-000000000001/workspace-files",
"/chats/00000000-0000-0000-0000-000000000002/workspace-files",
} {
req := httptest.NewRequest("POST", p, nil)
req.RemoteAddr = remoteAddr
rec := httptest.NewRecorder()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.Equal(t, i != 0, resp.StatusCode == http.StatusTooManyRequests, "request %d (%s)", i, p)
}
})
t.Run("RandomIPs", func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
// Because these are random IPs, the limit should never be hit!
rtr.Use(httpmw.RateLimit(1, time.Second))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
for i := 0; i < 5; i++ {
req := httptest.NewRequest("GET", "/", nil)
rec := httptest.NewRecorder()
req.RemoteAddr = randRemoteAddr()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.False(t, resp.StatusCode == http.StatusTooManyRequests)
}
})
t.Run("RegularUser", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
u := dbgen.User(t, db, database.User{})
_, key := dbgen.APIKey(t, db, database.APIKey{UserID: u.ID})
rtr := chi.NewRouter()
rtr.Use(httpmw.ExtractAPIKeyMW(httpmw.ExtractAPIKeyConfig{
DB: db,
Optional: false,
}))
rtr.Use(httpmw.RateLimit(1, time.Second))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
// Bypass must fail
req := httptest.NewRequest("GET", "/", nil)
req.Header.Set(codersdk.SessionTokenHeader, key)
req.Header.Set(codersdk.BypassRatelimitHeader, "true")
rec := httptest.NewRecorder()
// Assert we're not using IP address.
req.RemoteAddr = randRemoteAddr()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
defer resp.Body.Close()
require.Equal(t, http.StatusPreconditionRequired, resp.StatusCode)
for i := 0; i < 5; i++ {
req := httptest.NewRequest("GET", "/", nil)
req.Header.Set(codersdk.SessionTokenHeader, key)
rec := httptest.NewRecorder()
// Assert we're not using IP address.
req.RemoteAddr = randRemoteAddr()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.Equal(t, i != 0, resp.StatusCode == http.StatusTooManyRequests)
}
})
t.Run("OwnerBypass", func(t *testing.T) {
t.Parallel()
db, _ := dbtestutil.NewDB(t)
u := dbgen.User(t, db, database.User{
RBACRoles: []string{codersdk.RoleOwner},
})
_, key := dbgen.APIKey(t, db, database.APIKey{UserID: u.ID})
rtr := chi.NewRouter()
rtr.Use(httpmw.ExtractAPIKeyMW(httpmw.ExtractAPIKeyConfig{
DB: db,
Optional: false,
}))
rtr.Use(httpmw.RateLimit(1, time.Second))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
for i := 0; i < 5; i++ {
req := httptest.NewRequest("GET", "/", nil)
req.Header.Set(codersdk.SessionTokenHeader, key)
req.Header.Set(codersdk.BypassRatelimitHeader, "true")
rec := httptest.NewRecorder()
// Assert we're not using IP address.
req.RemoteAddr = randRemoteAddr()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.False(t, resp.StatusCode == http.StatusTooManyRequests)
}
})
}
func TestRateLimitByAuthToken(t *testing.T) {
t.Parallel()
t.Run("LimitsByAuthHeader", func(t *testing.T) {
t.Parallel()
tests := []struct {
name string
headerName string
headerVal string
}{
{
name: "BearerToken",
headerName: "Authorization",
headerVal: "Bearer test-token-123",
},
{
name: "XApiKey",
headerName: "X-Api-Key",
headerVal: "test-api-key-456",
},
{
name: "NoToken",
headerName: "",
headerVal: "",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
rtr.Use(httpmw.RateLimitByAuthToken(2, time.Hour))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
// Same token (or IP if no token) should be rate limited after 2 requests.
for i := 0; i < 5; i++ {
req := httptest.NewRequest("GET", "/", nil)
if tt.headerName != "" {
req.Header.Set(tt.headerName, tt.headerVal)
}
rec := httptest.NewRecorder()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
if i < 2 {
require.Equal(t, http.StatusOK, resp.StatusCode, "request %d should succeed", i)
} else {
require.Equal(t, http.StatusTooManyRequests, resp.StatusCode, "request %d should be rate limited", i)
// Verify Retry-After header is set.
require.NotEmpty(t, resp.Header.Get("Retry-After"), "Retry-After header should be set")
}
}
})
}
})
t.Run("DifferentTokensNotLimited", func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
rtr.Use(httpmw.RateLimitByAuthToken(1, time.Hour))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
// Different tokens should not be rate limited against each other.
for i := 0; i < 5; i++ {
req := httptest.NewRequest("GET", "/", nil)
req.Header.Set("Authorization", fmt.Sprintf("Bearer token-%d", i))
rec := httptest.NewRecorder()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode, "request %d should succeed", i)
}
})
t.Run("DisabledWhenZero", func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
rtr.Use(httpmw.RateLimitByAuthToken(0, time.Hour))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
// Should not be rate limited when limit is 0.
for i := 0; i < 10; i++ {
req := httptest.NewRequest("GET", "/", nil)
req.Header.Set("Authorization", "Bearer same-token")
rec := httptest.NewRecorder()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
}
})
}
func TestConcurrencyLimit(t *testing.T) {
t.Parallel()
t.Run("LimitsConcurrentRequests", func(t *testing.T) {
t.Parallel()
const maxConcurrency = 2
rtr := chi.NewRouter()
rtr.Use(httpmw.ConcurrencyLimit(maxConcurrency, "Test"))
// Use a WaitGroup as a barrier to ensure all requests are in the handler
// before any of them proceed.
var handlersReady sync.WaitGroup
handlersReady.Add(maxConcurrency)
releaseHandler := make(chan struct{})
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
handlersReady.Done()
// Wait until released.
<-releaseHandler
rw.WriteHeader(http.StatusOK)
})
server := httptest.NewServer(rtr)
defer server.Close()
ctx := testutil.Context(t, testutil.WaitShort)
// Start maxConcurrency requests that will block.
// We use channels to collect errors instead of require in goroutines.
type result struct {
statusCode int
err error
}
results := make(chan result, maxConcurrency)
var wg sync.WaitGroup
for i := 0; i < maxConcurrency; i++ {
wg.Go(func() {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+"/", nil)
if err != nil {
results <- result{err: err}
return
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
results <- result{err: err}
return
}
defer resp.Body.Close()
results <- result{statusCode: resp.StatusCode}
})
}
// Wait for all requests to enter the handler with a timeout.
handlersReadyCh := make(chan struct{})
go func() {
handlersReady.Wait()
close(handlersReadyCh)
}()
select {
case <-handlersReadyCh:
case <-ctx.Done():
t.Fatal("timed out waiting for handlers to be ready")
}
// Next request should be rejected since we're at capacity.
req, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL+"/", nil)
require.NoError(t, err)
resp, err := http.DefaultClient.Do(req)
require.NoError(t, err)
defer resp.Body.Close()
require.Equal(t, http.StatusServiceUnavailable, resp.StatusCode)
// Release all blocked requests.
close(releaseHandler)
wg.Wait()
close(results)
// Check all goroutine results.
for res := range results {
require.NoError(t, res.err)
require.Equal(t, http.StatusOK, res.statusCode)
}
})
t.Run("DisabledWhenZero", func(t *testing.T) {
t.Parallel()
rtr := chi.NewRouter()
rtr.Use(httpmw.ConcurrencyLimit(0, "Test"))
rtr.Get("/", func(rw http.ResponseWriter, r *http.Request) {
rw.WriteHeader(http.StatusOK)
})
// Should not be limited when maxConcurrency is 0.
for i := 0; i < 10; i++ {
req := httptest.NewRequest("GET", "/", nil)
rec := httptest.NewRecorder()
rtr.ServeHTTP(rec, req)
resp := rec.Result()
_ = resp.Body.Close()
require.Equal(t, http.StatusOK, resp.StatusCode)
}
})
}