Add rate limit for upload url requests

upload-server/server/local.go

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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)
+		return limiter, configureS3Handlers(mux, limiter)
-	} else {
-		return configureLocalHandlers(mux)
 	}
+	return limiter, configureLocalHandlers(mux, limiter)
 }
 
 func getObjectKey(w http.ResponseWriter, r *http.Request) string {

upload-server/server/signing_test.go

@@ -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
 }