sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-09-12 17:59:06 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add rate limit for upload url requests

Browse Source
This commit is contained in:
bcmmbaga
2026-09-10 17:56:40 +03:00
parent cb074a3b14
commit d29cce2c82
8 changed files with 107 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
upload-server/server/local.go
+4 -2
View File
@@ -12,6 +12,8 @@ import (
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/management/server/http/middleware"
"github.com/netbirdio/netbird/upload-server/types"
)
@@ -26,7 +28,7 @@ type local struct {
signer *signer
}
func configureLocalHandlers(mux *http.ServeMux) error {
func configureLocalHandlers(mux *http.ServeMux, limiter *middleware.APIRateLimiter) error {
envURL, ok := os.LookupEnv("SERVER_URL")
if !ok {
return fmt.Errorf("SERVER_URL environment variable is required")
@@ -56,7 +58,7 @@ func configureLocalHandlers(mux *http.ServeMux) error {
dir: dir,
signer: uploadSigner,
}
mux.HandleFunc(types.GetURLPath, l.handlerGetUploadURL)
mux.Handle(types.GetURLPath, limiter.Middleware(http.HandlerFunc(l.handlerGetUploadURL)))
mux.HandleFunc(putURLPath+putHandler, l.handlePutRequest)
return nil
upload-server/server/local_test.go
+5 -5
View File
@@ -29,7 +29,7 @@ func Test_LocalHandlerGetUploadURL(t *testing.T) {
t.Setenv("STORE_DIR", t.TempDir())
mux := http.NewServeMux()
err := configureLocalHandlers(mux)
err := configureLocalHandlers(mux, newRateLimiter())
require.NoError(t, err)
req := httptest.NewRequest(http.MethodGet, types.GetURLPath+"?id=test-file", nil)
@@ -56,7 +56,7 @@ func Test_LocalHandlePutRequest(t *testing.T) {
t.Setenv(signingKeyVar, testSigningKey)
mux := http.NewServeMux()
err := configureLocalHandlers(mux)
err := configureLocalHandlers(mux, newRateLimiter())
require.NoError(t, err)
fileContent := []byte("test file content")
@@ -82,7 +82,7 @@ func Test_LocalHandlePutRequest_PathTraversal(t *testing.T) {
t.Setenv(signingKeyVar, testSigningKey)
mux := http.NewServeMux()
err := configureLocalHandlers(mux)
err := configureLocalHandlers(mux, newRateLimiter())
require.NoError(t, err)
fileContent := []byte("malicious content")
@@ -129,7 +129,7 @@ func Test_LocalHandlePutRequest_DuplicateFile(t *testing.T) {
t.Setenv(signingKeyVar, testSigningKey)
mux := http.NewServeMux()
err := configureLocalHandlers(mux)
err := configureLocalHandlers(mux, newRateLimiter())
require.NoError(t, err)
req := httptest.NewRequest(http.MethodPut,
@@ -156,7 +156,7 @@ func Test_LocalHandlePutRequest_BodyTooLarge(t *testing.T) {
t.Setenv(signingKeyVar, testSigningKey)
mux := http.NewServeMux()
err := configureLocalHandlers(mux)
err := configureLocalHandlers(mux, newRateLimiter())
require.NoError(t, err)
largeBody := make([]byte, maxUploadSize+1)
upload-server/server/ratelimit.go
+28
View File
@@ -0,0 +1,28 @@
package server
import (
"os"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/management/server/http/middleware"
)
const defaultUploadBurst = 100
func newRateLimiter() *middleware.APIRateLimiter {
cfg, _ := middleware.RateLimiterConfigFromEnv()
if os.Getenv(middleware.RateLimitingBurstEnv) == "" {
cfg.Burst = defaultUploadBurst
}
limiter := middleware.NewAPIRateLimiter(cfg)
if os.Getenv(middleware.RateLimitingEnabledEnv) == "false" {
limiter.SetEnabled(false)
}
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))
return limiter
}
upload-server/server/ratelimit_test.go
+50
View File
@@ -0,0 +1,50 @@
package server
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/require"
"github.com/netbirdio/netbird/management/server/http/middleware"
"github.com/netbirdio/netbird/upload-server/types"
)
func getUploadURL(t *testing.T, mux *http.ServeMux) int {
t.Helper()
req := httptest.NewRequest(http.MethodGet, types.GetURLPath+"?id=test-file", nil)
req.Header.Set(types.ClientHeader, types.ClientHeaderValue)
rec := httptest.NewRecorder()
mux.ServeHTTP(rec, req)
return rec.Code
}
func Test_GetUploadURLIsRateLimited(t *testing.T) {
t.Setenv(middleware.RateLimitingBurstEnv, "2")
t.Setenv(middleware.RateLimitingRPMEnv, "1")
mux, _ := newLocalMux(t)
require.Equal(t, http.StatusOK, getUploadURL(t, mux))
require.Equal(t, http.StatusOK, getUploadURL(t, mux))
require.Equal(t, http.StatusTooManyRequests, getUploadURL(t, mux))
}
func Test_RateLimitingIsOnByDefault(t *testing.T) {
t.Setenv(middleware.RateLimitingBurstEnv, "1")
mux, _ := newLocalMux(t)
require.Equal(t, http.StatusOK, getUploadURL(t, mux))
require.Equal(t, http.StatusTooManyRequests, getUploadURL(t, mux))
}
func Test_RateLimitingCanBeDisabled(t *testing.T) {
t.Setenv(middleware.RateLimitingEnabledEnv, "false")
t.Setenv(middleware.RateLimitingBurstEnv, "1")
mux, _ := newLocalMux(t)
require.Equal(t, http.StatusOK, getUploadURL(t, mux))
require.Equal(t, http.StatusOK, getUploadURL(t, mux))
}
upload-server/server/s3.go
+4 -2
View File
@@ -12,6 +12,8 @@ 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/upload-server/types"
)
@@ -21,7 +23,7 @@ type sThree struct {
presignClient *s3.PresignClient
}
func configureS3Handlers(mux *http.ServeMux) error {
func configureS3Handlers(mux *http.ServeMux, limiter *middleware.APIRateLimiter) error {
bucket := os.Getenv(bucketVar)
region, ok := os.LookupEnv("AWS_REGION")
if !ok {
@@ -40,7 +42,7 @@ func configureS3Handlers(mux *http.ServeMux) error {
bucket: bucket,
presignClient: s3.NewPresignClient(client),
}
mux.HandleFunc(types.GetURLPath, handler.handlerGetUploadURL)
mux.Handle(types.GetURLPath, limiter.Middleware(http.HandlerFunc(handler.handlerGetUploadURL)))
return nil
}
upload-server/server/s3_test.go
+1 -1
View File
@@ -90,7 +90,7 @@ func Test_S3HandlerGetUploadURL(t *testing.T) {
t.Setenv(bucketVar, bucketName)
mux := http.NewServeMux()
err = configureS3Handlers(mux)
err = configureS3Handlers(mux, newRateLimiter())
require.NoError(t, err)
req := httptest.NewRequest(http.MethodGet, types.GetURLPath+"?id=test-file", nil)
upload-server/server/server.go
+14 -7
View File
@@ -10,6 +10,7 @@ import (
"github.com/google/uuid"
log "github.com/sirupsen/logrus"
"github.com/netbirdio/netbird/management/server/http/middleware"
"github.com/netbirdio/netbird/upload-server/types"
)
@@ -19,7 +20,8 @@ const (
)
type Server struct {
srv *http.Server
srv *http.Server
limiter *middleware.APIRateLimiter
}
func NewServer() *Server {
@@ -29,7 +31,7 @@ func NewServer() *Server {
address = "0.0.0.0:8080"
}
mux := http.NewServeMux()
err := configureMux(mux)
limiter, err := configureMux(mux)
if err != nil {
log.Fatalf("Failed to configure server: %v", err)
}
@@ -38,7 +40,8 @@ func NewServer() *Server {
})
return &Server{
srv: &http.Server{Addr: address, Handler: mux},
srv: &http.Server{Addr: address, Handler: mux},
limiter: limiter,
}
}
@@ -48,6 +51,9 @@ func (s *Server) Start() error {
}
func (s *Server) Stop() error {
if s.limiter != nil {
s.limiter.Stop()
}
if s.srv != nil {
log.Infof("Stopping upload server on %s", s.srv.Addr)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
@@ -57,13 +63,14 @@ func (s *Server) Stop() error {
return nil
}
func configureMux(mux *http.ServeMux) error {
func configureMux(mux *http.ServeMux) (*middleware.APIRateLimiter, error) {
limiter := newRateLimiter()
_, ok := os.LookupEnv(bucketVar)
if ok {
return configureS3Handlers(mux)
} else {
return configureLocalHandlers(mux)
return limiter, configureS3Handlers(mux, limiter)
}
return limiter, configureLocalHandlers(mux, limiter)
}
func getObjectKey(w http.ResponseWriter, r *http.Request) string {
upload-server/server/signing_test.go
+1 -1
View File
@@ -25,7 +25,7 @@ func newLocalMux(t *testing.T) (*http.ServeMux, string) {
t.Setenv(signingKeyVar, testSigningKey)
mux := http.NewServeMux()
require.NoError(t, configureLocalHandlers(mux))
require.NoError(t, configureLocalHandlers(mux, newRateLimiter()))
return mux, mockDir
}