mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-28 17:49:08 +02:00
[management] move rate limiter to shared package (#7727)
This commit is contained in:
@@ -12,7 +12,7 @@ import (
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/http/middleware"
|
||||
"github.com/netbirdio/netbird/shared/ratelimit"
|
||||
|
||||
"github.com/netbirdio/netbird/upload-server/types"
|
||||
)
|
||||
@@ -28,7 +28,7 @@ type local struct {
|
||||
signer *signer
|
||||
}
|
||||
|
||||
func configureLocalHandlers(mux *http.ServeMux, limiter *middleware.APIRateLimiter) error {
|
||||
func configureLocalHandlers(mux *http.ServeMux, limiter *ratelimit.APIRateLimiter) error {
|
||||
envURL, ok := os.LookupEnv("SERVER_URL")
|
||||
if !ok {
|
||||
return fmt.Errorf("SERVER_URL environment variable is required")
|
||||
|
||||
@@ -5,27 +5,27 @@ import (
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/http/middleware"
|
||||
"github.com/netbirdio/netbird/shared/ratelimit"
|
||||
)
|
||||
|
||||
const defaultUploadBurst = 100
|
||||
|
||||
func newRateLimiter() *middleware.APIRateLimiter {
|
||||
cfg, enabled := middleware.RateLimiterConfigFromEnv()
|
||||
if os.Getenv(middleware.RateLimitingBurstEnv) == "" {
|
||||
func newRateLimiter() *ratelimit.APIRateLimiter {
|
||||
cfg, enabled := ratelimit.RateLimiterConfigFromEnv()
|
||||
if os.Getenv(ratelimit.RateLimitingBurstEnv) == "" {
|
||||
cfg.Burst = defaultUploadBurst
|
||||
}
|
||||
|
||||
// Rate limiting is enabled by default unless explicitly disabled
|
||||
if os.Getenv(middleware.RateLimitingEnabledEnv) == "" {
|
||||
if os.Getenv(ratelimit.RateLimitingEnabledEnv) == "" {
|
||||
enabled = true
|
||||
}
|
||||
|
||||
limiter := middleware.NewAPIRateLimiter(cfg)
|
||||
limiter := ratelimit.NewAPIRateLimiter(cfg)
|
||||
limiter.SetEnabled(enabled)
|
||||
|
||||
log.Infof("Upload URL rate limiting: enabled=%t rate=%.0f/min burst=%d trusted_proxies=%q",
|
||||
limiter.Enabled(), cfg.RequestsPerMinute, cfg.Burst, os.Getenv(middleware.RateLimitingTrustedProxiesEnv))
|
||||
limiter.Enabled(), cfg.RequestsPerMinute, cfg.Burst, os.Getenv(ratelimit.RateLimitingTrustedProxiesEnv))
|
||||
|
||||
return limiter
|
||||
}
|
||||
|
||||
@@ -7,11 +7,11 @@ import (
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/http/middleware"
|
||||
"github.com/netbirdio/netbird/shared/ratelimit"
|
||||
"github.com/netbirdio/netbird/upload-server/types"
|
||||
)
|
||||
|
||||
func newTestRateLimiter(t *testing.T) *middleware.APIRateLimiter {
|
||||
func newTestRateLimiter(t *testing.T) *ratelimit.APIRateLimiter {
|
||||
t.Helper()
|
||||
|
||||
limiter := newRateLimiter()
|
||||
@@ -32,8 +32,8 @@ func getUploadURL(t *testing.T, mux *http.ServeMux) int {
|
||||
}
|
||||
|
||||
func Test_GetUploadURLIsRateLimited(t *testing.T) {
|
||||
t.Setenv(middleware.RateLimitingBurstEnv, "2")
|
||||
t.Setenv(middleware.RateLimitingRPMEnv, "1")
|
||||
t.Setenv(ratelimit.RateLimitingBurstEnv, "2")
|
||||
t.Setenv(ratelimit.RateLimitingRPMEnv, "1")
|
||||
mux, _ := newLocalMux(t)
|
||||
|
||||
require.Equal(t, http.StatusOK, getUploadURL(t, mux))
|
||||
@@ -42,8 +42,8 @@ func Test_GetUploadURLIsRateLimited(t *testing.T) {
|
||||
}
|
||||
|
||||
func Test_RateLimitingIsOnByDefault(t *testing.T) {
|
||||
t.Setenv(middleware.RateLimitingEnabledEnv, "")
|
||||
t.Setenv(middleware.RateLimitingBurstEnv, "1")
|
||||
t.Setenv(ratelimit.RateLimitingEnabledEnv, "")
|
||||
t.Setenv(ratelimit.RateLimitingBurstEnv, "1")
|
||||
mux, _ := newLocalMux(t)
|
||||
|
||||
require.Equal(t, http.StatusOK, getUploadURL(t, mux))
|
||||
@@ -51,8 +51,8 @@ func Test_RateLimitingIsOnByDefault(t *testing.T) {
|
||||
}
|
||||
|
||||
func Test_RateLimitingCanBeDisabled(t *testing.T) {
|
||||
t.Setenv(middleware.RateLimitingEnabledEnv, "false")
|
||||
t.Setenv(middleware.RateLimitingBurstEnv, "1")
|
||||
t.Setenv(ratelimit.RateLimitingEnabledEnv, "false")
|
||||
t.Setenv(ratelimit.RateLimitingBurstEnv, "1")
|
||||
mux, _ := newLocalMux(t)
|
||||
|
||||
require.Equal(t, http.StatusOK, getUploadURL(t, mux))
|
||||
|
||||
@@ -12,7 +12,7 @@ import (
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/http/middleware"
|
||||
"github.com/netbirdio/netbird/shared/ratelimit"
|
||||
|
||||
"github.com/netbirdio/netbird/upload-server/types"
|
||||
)
|
||||
@@ -23,7 +23,7 @@ type sThree struct {
|
||||
presignClient *s3.PresignClient
|
||||
}
|
||||
|
||||
func configureS3Handlers(mux *http.ServeMux, limiter *middleware.APIRateLimiter) error {
|
||||
func configureS3Handlers(mux *http.ServeMux, limiter *ratelimit.APIRateLimiter) error {
|
||||
bucket := os.Getenv(bucketVar)
|
||||
region, ok := os.LookupEnv("AWS_REGION")
|
||||
if !ok {
|
||||
|
||||
@@ -10,7 +10,7 @@ import (
|
||||
"github.com/google/uuid"
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
"github.com/netbirdio/netbird/management/server/http/middleware"
|
||||
"github.com/netbirdio/netbird/shared/ratelimit"
|
||||
"github.com/netbirdio/netbird/upload-server/types"
|
||||
)
|
||||
|
||||
@@ -21,7 +21,7 @@ const (
|
||||
|
||||
type Server struct {
|
||||
srv *http.Server
|
||||
limiter *middleware.APIRateLimiter
|
||||
limiter *ratelimit.APIRateLimiter
|
||||
}
|
||||
|
||||
func NewServer() *Server {
|
||||
@@ -63,7 +63,7 @@ func (s *Server) Stop() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func configureMux(mux *http.ServeMux) (*middleware.APIRateLimiter, error) {
|
||||
func configureMux(mux *http.ServeMux) (*ratelimit.APIRateLimiter, error) {
|
||||
limiter := newRateLimiter()
|
||||
|
||||
_, ok := os.LookupEnv(bucketVar)
|
||||
|
||||
Reference in New Issue
Block a user