test: wip adding acceptance tests627-add-cfg-and-metrics

5 files changed, 76 insertions, 27 deletions

diff --git a/internal/config/flags.go b/internal/config/flags.go
index e651cd21..55a99ae6 100644
--- a/internal/config/flags.go
+++ b/internal/config/flags.go
@@ -5,6 +5,8 @@ import (
 
 	"github.com/namsral/flag"
 
+	"gitlab.com/gitlab-org/gitlab-pages/internal/ratelimiter"
+
 	"gitlab.com/gitlab-org/gitlab-pages/internal/config/tls"
 )
 
@@ -16,7 +18,7 @@ var (
 	pagesRoot               = flag.String("pages-root", "shared/pages", "The directory where pages are stored")
 	pagesDomain             = flag.String("pages-domain", "gitlab-example.com", "The domain to serve static pages")
 	rateLimitSourceIP       = flag.Float64("rate-limit-source-ip", 0.0, "Rate limit per source IP in number of requests per second, 0 means is disabled")
-	rateLimitSourceIPBurst  = flag.Int("rate-limit-source-ip-burst", 0.0, "Rate limit per source IP maximum burst allowed per second")
+	rateLimitSourceIPBurst  = flag.Int("rate-limit-source-ip-burst", ratelimiter.DefaultSourceIPBurstSize, "Rate limit per source IP maximum burst allowed per second")
 	artifactsServer         = flag.String("artifacts-server", "", "API URL to proxy artifact requests to, e.g.: 'https://gitlab.com/api/v4'")
 	artifactsServerTimeout  = flag.Int("artifacts-server-timeout", 10, "Timeout (in seconds) for a proxied request to the artifacts server")
 	pagesStatus             = flag.String("pages-status", "", "The url path for a status page, e.g., /@status")
diff --git a/internal/ratelimiter/middleware.go b/internal/ratelimiter/middleware.go
index e53d825b..c3e1decc 100644
--- a/internal/ratelimiter/middleware.go
+++ b/internal/ratelimiter/middleware.go
@@ -7,9 +7,10 @@ import (
 
 	"github.com/sebest/xff"
 	"github.com/sirupsen/logrus"
+	"gitlab.com/gitlab-org/labkit/correlation"
+	"gitlab.com/gitlab-org/labkit/log"
 
 	"gitlab.com/gitlab-org/gitlab-pages/internal/httperrors"
-	"gitlab.com/gitlab-org/gitlab-pages/internal/logging"
 )
 
 // SourceIPLimiter middleware ensures that the originating
@@ -17,18 +18,23 @@ import (
 func (rl *RateLimiter) SourceIPLimiter(handler http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		host, ip, https := getReqDetails(r)
+		l := log.WithFields(logrus.Fields{
+			"correlation_id":                correlation.ExtractFromContext(r.Context()),
+			"host":                          r.Host,
+			"path":                          r.URL.Path,
+			"handler":                       "source_ip_rate_limiter",
+			"pages_domain":                  host,
+			"pages_https":                   https,
+			"source_ip":                     ip,
+			"rate_limiter_enabled":          rateLimiterEnabled(),
+			"rate_limiter_limit_per_second": rl.sourceIPLimitPerSecond,
+			"rate_limiter_burst_size":       rl.sourceIPBurstSize,
+		})
+		l.Debug("what is going on")
 
 		// http requests do not contain real IP information yet
 		if !rl.SourceIPAllowed(ip) && https {
-			logging.LogRequest(r).WithFields(logrus.Fields{
-				"handler":                       "source_ip_rate_limiter",
-				"pages_domain":                  host,
-				"pages_https":                   https,
-				"source_ip":                     ip,
-				"rate_limiter_enabled":          rateLimiterEnabled(),
-				"rate_limiter_limit_per_second": rl.sourceIPLimitPerSecond,
-				"rate_limiter_burst_size":       rl.sourceIPBurstSize,
-			}).Info("source IP hit rate limit")
+			l.Info("source IP hit rate limit")
 
 			// Only drop requests once FF_ENABLE_RATE_LIMITER is enabled
 			if rateLimiterEnabled() {
diff --git a/internal/ratelimiter/middleware_test.go b/internal/ratelimiter/middleware_test.go
index bfb68e5d..1b0dc75a 100644
--- a/internal/ratelimiter/middleware_test.go
+++ b/internal/ratelimiter/middleware_test.go
@@ -4,10 +4,11 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
-	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
 )
 
 var next = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -15,7 +16,7 @@ var next = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 })
 
 func TestSourceIPLimiter(t *testing.T) {
-	enableRateLimiter(t)
+	testhelpers.EnableRateLimiter(t)
 
 	for tn, tc := range sharedTestCases {
 		t.Run(tn, func(t *testing.T) {
@@ -99,7 +100,7 @@ func TestSourceIPRateLimit(t *testing.T) {
 				rr.RemoteAddr = tc.ip
 
 				if tc.enabled {
-					enableRateLimiter(t)
+					testhelpers.EnableRateLimiter(t)
 				}
 
 				handler := rl.SourceIPLimiter(next)
@@ -117,16 +118,3 @@ func TestSourceIPRateLimit(t *testing.T) {
 		})
 	}
 }
-
-func enableRateLimiter(t *testing.T) {
-	t.Helper()
-
-	orig := os.Getenv("FF_ENABLE_RATE_LIMITER")
-
-	err := os.Setenv("FF_ENABLE_RATE_LIMITER", "true")
-	require.NoError(t, err)
-
-	t.Cleanup(func() {
-		os.Setenv("FF_ENABLE_RATE_LIMITER", orig)
-	})
-}
diff --git a/internal/testhelpers/testhelpers.go b/internal/testhelpers/testhelpers.go
index 3ec97a79..04dfc3d2 100644
--- a/internal/testhelpers/testhelpers.go
+++ b/internal/testhelpers/testhelpers.go
@@ -77,3 +77,17 @@ func Getwd(t *testing.T) string {
 
 	return wd
 }
+
+// EnableRateLimiter environment variable
+func EnableRateLimiter(t *testing.T) {
+	t.Helper()
+
+	orig := os.Getenv("FF_ENABLE_RATE_LIMITER")
+
+	err := os.Setenv("FF_ENABLE_RATE_LIMITER", "true")
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		os.Setenv("FF_ENABLE_RATE_LIMITER", orig)
+	})
+}
diff --git a/test/acceptance/ratelimiter_test.go b/test/acceptance/ratelimiter_test.go
new file mode 100644
index 00000000..d67b3d04
--- /dev/null
+++ b/test/acceptance/ratelimiter_test.go
@@ -0,0 +1,39 @@
+package acceptance_test
+
+import (
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-pages/internal/testhelpers"
+)
+
+func TestRateLimitMiddleware(t *testing.T) {
+	testhelpers.EnableRateLimiter(t)
+
+	RunPagesProcess(t,
+		withListeners([]ListenSpec{httpListener}),
+		// 10 = 1 req every 100ms
+		withExtraArgument("rate-limit-source-ip", "1.0"),
+		withExtraArgument("rate-limit-source-ip-burst", "1"),
+	)
+
+	for i := 0; i < 20; i++ {
+		rsp1, err := GetPageFromListener(t, httpListener, "group.gitlab-example.com", "project/")
+		require.NoError(t, err)
+		rsp1.Body.Close()
+
+		// every other request should fail
+		//if i%2 != 0 {
+		//	require.Equal(t, http.StatusTooManyRequests, rsp1.StatusCode, "group.gitlab-example.com request: %d failed", i)
+		//	// wait for another token to become available
+		//	time.Sleep(100 * time.Millisecond)
+		//	continue
+		//}
+
+		require.Equal(t, http.StatusOK, rsp1.StatusCode, "group.gitlab-example.com request: %d failed", i)
+		time.Sleep(time.Millisecond)
+	}
+}