Add rate limit for upload url requests

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))
+}