diff --git a/upload-server/server/local.go b/upload-server/server/local.go index 890dd9171..7db2740f9 100644 --- a/upload-server/server/local.go +++ b/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 diff --git a/upload-server/server/local_test.go b/upload-server/server/local_test.go index 6681b8ec2..c70091815 100644 --- a/upload-server/server/local_test.go +++ b/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) diff --git a/upload-server/server/ratelimit.go b/upload-server/server/ratelimit.go new file mode 100644 index 000000000..f28d2e312 --- /dev/null +++ b/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 +} diff --git a/upload-server/server/ratelimit_test.go b/upload-server/server/ratelimit_test.go new file mode 100644 index 000000000..3ca4f1e4f --- /dev/null +++ b/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)) +} diff --git a/upload-server/server/s3.go b/upload-server/server/s3.go index c0976acb5..ffc5df01f 100644 --- a/upload-server/server/s3.go +++ b/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 } diff --git a/upload-server/server/s3_test.go b/upload-server/server/s3_test.go index a72356409..7c732383e 100644 --- a/upload-server/server/s3_test.go +++ b/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) diff --git a/upload-server/server/server.go b/upload-server/server/server.go index 29ef72732..607c5bdad 100644 --- a/upload-server/server/server.go +++ b/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) - } else { - return configureLocalHandlers(mux) + return limiter, configureS3Handlers(mux, limiter) } + return limiter, configureLocalHandlers(mux, limiter) } func getObjectKey(w http.ResponseWriter, r *http.Request) string { diff --git a/upload-server/server/signing_test.go b/upload-server/server/signing_test.go index a65eb468c..2d33d2f9b 100644 --- a/upload-server/server/signing_test.go +++ b/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 }