feat: add domain rate-limiter

Changelog: added

2 files changed, 121 insertions, 2 deletions

diff --git a/internal/handlers/ratelimiter.go b/internal/handlers/ratelimiter.go
index 8263f497..52281f6e 100644
--- a/internal/handlers/ratelimiter.go
+++ b/internal/handlers/ratelimiter.go
@@ -4,14 +4,16 @@ import (
 	"net/http"
 
 	"gitlab.com/gitlab-org/gitlab-pages/internal/config"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/feature"
 	"gitlab.com/gitlab-org/gitlab-pages/internal/ratelimiter"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/request"
 	"gitlab.com/gitlab-org/gitlab-pages/metrics"
 )
 
 // Ratelimiter configures the ratelimiter middleware
 // TODO: make this unexported once https://gitlab.com/gitlab-org/gitlab-pages/-/issues/670 is done
 func Ratelimiter(handler http.Handler, config *config.RateLimit) http.Handler {
-	rl := ratelimiter.New(
+	sourceIPLimiter := ratelimiter.New(
 		"source_ip",
 		ratelimiter.WithCacheMaxSize(ratelimiter.DefaultSourceIPCacheSize),
 		ratelimiter.WithCachedEntriesMetric(metrics.RateLimitSourceIPCachedEntries),
@@ -19,7 +21,22 @@ func Ratelimiter(handler http.Handler, config *config.RateLimit) http.Handler {
 		ratelimiter.WithBlockedCountMetric(metrics.RateLimitSourceIPBlockedCount),
 		ratelimiter.WithLimitPerSecond(config.SourceIPLimitPerSecond),
 		ratelimiter.WithBurstSize(config.SourceIPBurst),
+		ratelimiter.WithEnforce(feature.EnforceIPRateLimits.Enabled()),
 	)
 
-	return rl.Middleware(handler)
+	handler = sourceIPLimiter.Middleware(handler)
+
+	domainLimiter := ratelimiter.New(
+		"domain",
+		ratelimiter.WithCacheMaxSize(ratelimiter.DefaultDomainCacheSize),
+		ratelimiter.WithKeyFunc(request.GetHostWithoutPort),
+		ratelimiter.WithCachedEntriesMetric(metrics.RateLimitDomainCachedEntries),
+		ratelimiter.WithCachedRequestsMetric(metrics.RateLimitDomainCacheRequests),
+		ratelimiter.WithBlockedCountMetric(metrics.RateLimitDomainBlockedCount),
+		ratelimiter.WithLimitPerSecond(config.DomainLimitPerSecond),
+		ratelimiter.WithBurstSize(config.DomainBurst),
+		ratelimiter.WithEnforce(feature.EnforceDomainRateLimits.Enabled()),
+	)
+
+	return domainLimiter.Middleware(handler)
 }
diff --git a/internal/handlers/ratelimiter_test.go b/internal/handlers/ratelimiter_test.go
new file mode 100644
index 00000000..43acfc9a
--- /dev/null
+++ b/internal/handlers/ratelimiter_test.go
@@ -0,0 +1,102 @@
+package handlers
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-pages/internal/config"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/feature"
+	"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
+)
+
+var next = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(http.StatusNoContent)
+})
+
+func TestRatelimiter(t *testing.T) {
+	tt := map[string]struct {
+		firstRemoteAddr    string
+		firstTarget        string
+		secondRemoteAddr   string
+		secondTarget       string
+		sourceIPEnforced   bool
+		domainEnforced     bool
+		expectedSecondCode int
+	}{
+		"rejected_by_ip": {
+			firstRemoteAddr:    "10.0.0.1",
+			firstTarget:        "https://domain.gitlab.io",
+			secondRemoteAddr:   "10.0.0.1",
+			secondTarget:       "https://different.gitlab.io",
+			sourceIPEnforced:   true,
+			domainEnforced:     true,
+			expectedSecondCode: http.StatusTooManyRequests,
+		},
+		"rejected_by_domain": {
+			firstRemoteAddr:    "10.0.0.1",
+			firstTarget:        "https://domain.gitlab.io",
+			secondRemoteAddr:   "10.0.0.2",
+			secondTarget:       "https://domain.gitlab.io",
+			sourceIPEnforced:   true,
+			domainEnforced:     true,
+			expectedSecondCode: http.StatusTooManyRequests,
+		},
+		"ip_rate_limiter_disabled": {
+			firstRemoteAddr:    "10.0.0.1",
+			firstTarget:        "https://domain.gitlab.io",
+			secondRemoteAddr:   "10.0.0.1",
+			secondTarget:       "https://different.gitlab.io",
+			sourceIPEnforced:   false,
+			domainEnforced:     true,
+			expectedSecondCode: http.StatusNoContent,
+		},
+		"domain_rate_limiter_disabled": {
+			firstRemoteAddr:    "10.0.0.1",
+			firstTarget:        "https://domain.gitlab.io",
+			secondRemoteAddr:   "10.0.0.2",
+			secondTarget:       "https://domain.gitlab.io",
+			sourceIPEnforced:   true,
+			domainEnforced:     false,
+			expectedSecondCode: http.StatusNoContent,
+		},
+		"different_ip_and_domain_passes": {
+			firstRemoteAddr:    "10.0.0.1",
+			firstTarget:        "https://domain.gitlab.io",
+			secondRemoteAddr:   "10.0.0.2",
+			secondTarget:       "https://different.gitlab.io",
+			sourceIPEnforced:   true,
+			domainEnforced:     true,
+			expectedSecondCode: http.StatusNoContent,
+		},
+	}
+
+	for name, tc := range tt {
+		t.Run(name, func(t *testing.T) {
+			testhelpers.StubFeatureFlagValue(t, feature.EnforceIPRateLimits.EnvVariable, tc.sourceIPEnforced)
+			testhelpers.StubFeatureFlagValue(t, feature.EnforceDomainRateLimits.EnvVariable, tc.domainEnforced)
+
+			conf := config.RateLimit{
+				SourceIPLimitPerSecond: 0.1,
+				SourceIPBurst:          1,
+				DomainLimitPerSecond:   0.1,
+				DomainBurst:            1,
+			}
+
+			handler := Ratelimiter(next, &conf)
+
+			r1 := httptest.NewRequest(http.MethodGet, tc.firstTarget, nil)
+			r1.RemoteAddr = tc.firstRemoteAddr
+
+			firstCode, _ := testhelpers.PerformRequest(t, handler, r1)
+			require.Equal(t, http.StatusNoContent, firstCode)
+
+			r2 := httptest.NewRequest(http.MethodGet, tc.secondTarget, nil)
+			r2.RemoteAddr = tc.secondRemoteAddr
+			secondCode, _ := testhelpers.PerformRequest(t, handler, r2)
+			require.Equal(t, tc.expectedSecondCode, secondCode)
+		})
+	}
+}