webhooker/internal/middleware/ratelimit_test.go

package middleware_test

import (
	"context"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"

	"github.com/stretchr/testify/assert"
	"sneak.berlin/go/webhooker/internal/config"
	"sneak.berlin/go/webhooker/internal/middleware"
)

func TestLoginRateLimit_AllowsGET(t *testing.T) {
	t.Parallel()

	m, _ := testMiddleware(t, config.EnvironmentDev)

	var callCount int

	handler := m.LoginRateLimit()(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			callCount++

			w.WriteHeader(http.StatusOK)
		},
	))

	// GET requests should never be rate-limited
	for i := range 20 {
		req := httptest.NewRequestWithContext(
			context.Background(),
			http.MethodGet, "/pages/login", nil,
		)
		req.RemoteAddr = "192.168.1.1:12345"

		w := httptest.NewRecorder()
		handler.ServeHTTP(w, req)

		assert.Equal(
			t, http.StatusOK, w.Code,
			"GET request %d should pass", i,
		)
	}

	assert.Equal(t, 20, callCount)
}

// runPostLimitTest exercises a POST-only rate limit middleware:
// the first limit POSTs to path from ip must pass, and the next
// one must be rejected with 429 without reaching the handler.
func runPostLimitTest(
	t *testing.T,
	mw func(http.Handler) http.Handler,
	limit int,
	path, ip string,
) {
	t.Helper()

	var callCount int

	handler := mw(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			callCount++

			w.WriteHeader(http.StatusOK)
		},
	))

	// The first limit POST requests should succeed
	for i := range limit {
		req := httptest.NewRequestWithContext(
			context.Background(),
			http.MethodPost, path, nil,
		)
		req.RemoteAddr = ip

		w := httptest.NewRecorder()
		handler.ServeHTTP(w, req)

		assert.Equal(
			t, http.StatusOK, w.Code,
			"POST request %d should pass", i,
		)
	}

	// Next POST should be rate-limited
	req := httptest.NewRequestWithContext(
		context.Background(),
		http.MethodPost, path, nil,
	)
	req.RemoteAddr = ip

	w := httptest.NewRecorder()
	handler.ServeHTTP(w, req)

	assert.Equal(
		t, http.StatusTooManyRequests, w.Code,
		"POST after limit should be 429",
	)
	assert.Equal(t, limit, callCount)
}

func TestLoginRateLimit_LimitsPOST(t *testing.T) {
	t.Parallel()

	m, _ := testMiddleware(t, config.EnvironmentDev)

	runPostLimitTest(
		t,
		m.LoginRateLimit(),
		middleware.LoginRateLimitConst,
		"/pages/login",
		"10.0.0.1:12345",
	)
}

func TestPasswordChangeRateLimit_LimitsPOST(t *testing.T) {
	t.Parallel()

	m, _ := testMiddleware(t, config.EnvironmentDev)

	runPostLimitTest(
		t,
		m.PasswordChangeRateLimit(),
		middleware.PasswordChangeRateLimitConst,
		"/user/admin/password",
		"10.0.0.2:12345",
	)
}

func TestLoginRateLimit_IndependentPerIP(t *testing.T) {
	t.Parallel()

	m, _ := testMiddleware(t, config.EnvironmentDev)

	handler := m.LoginRateLimit()(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusOK)
		},
	))

	// Exhaust limit for IP1
	for range middleware.LoginRateLimitConst {
		req := httptest.NewRequestWithContext(
			context.Background(),
			http.MethodPost, "/pages/login", nil,
		)
		req.RemoteAddr = "1.2.3.4:12345"

		w := httptest.NewRecorder()
		handler.ServeHTTP(w, req)
	}

	// IP1 should be rate-limited
	req := httptest.NewRequestWithContext(
		context.Background(),
		http.MethodPost, "/pages/login", nil,
	)
	req.RemoteAddr = "1.2.3.4:12345"

	w := httptest.NewRecorder()
	handler.ServeHTTP(w, req)

	assert.Equal(t, http.StatusTooManyRequests, w.Code)

	// IP2 should still be allowed
	req2 := httptest.NewRequestWithContext(
		context.Background(),
		http.MethodPost, "/pages/login", nil,
	)
	req2.RemoteAddr = "5.6.7.8:12345"

	w2 := httptest.NewRecorder()
	handler.ServeHTTP(w2, req2)

	assert.Equal(
		t, http.StatusOK, w2.Code,
		"different IP should not be affected",
	)
}

// receiverLimitedHandler builds a ReceiverRateLimit-wrapped
// handler with the given per-minute limit.
func receiverLimitedHandler(
	t *testing.T, limit int,
) http.Handler {
	t.Helper()

	log := slog.New(slog.NewTextHandler(
		os.Stderr,
		&slog.HandlerOptions{Level: slog.LevelDebug},
	))

	m := middleware.NewForTest(
		log,
		&config.Config{ReceiverRateLimit: limit},
		nil,
	)

	return m.ReceiverRateLimit()(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			w.WriteHeader(http.StatusOK)
		},
	))
}

// receiverPost sends one POST to the handler from the given IP
// and path and returns the recorder.
func receiverPost(
	handler http.Handler, ip, path string,
) *httptest.ResponseRecorder {
	req := httptest.NewRequestWithContext(
		context.Background(),
		http.MethodPost, path, nil,
	)
	req.RemoteAddr = ip

	w := httptest.NewRecorder()
	handler.ServeHTTP(w, req)

	return w
}

func TestReceiverRateLimit_LimitsPerIPAndPath(t *testing.T) {
	t.Parallel()

	const limit = 3

	handler := receiverLimitedHandler(t, limit)

	// The first limit requests from one IP to one entrypoint
	// pass.
	for i := range limit {
		w := receiverPost(
			handler, "9.9.9.9:1234", "/webhook/uuid-a",
		)
		assert.Equal(
			t, http.StatusOK, w.Code,
			"request %d should pass", i,
		)
	}

	// The next request over the limit is rejected with a 429
	// carrying a Retry-After header.
	w := receiverPost(
		handler, "9.9.9.9:1234", "/webhook/uuid-a",
	)
	assert.Equal(t, http.StatusTooManyRequests, w.Code)
	assert.NotEmpty(
		t, w.Header().Get("Retry-After"),
		"429 must carry a Retry-After header",
	)

	// The same IP is not limited on a different entrypoint.
	w = receiverPost(
		handler, "9.9.9.9:1234", "/webhook/uuid-b",
	)
	assert.Equal(
		t, http.StatusOK, w.Code,
		"a different entrypoint must not be affected",
	)

	// A different IP is not limited on the same entrypoint.
	w = receiverPost(
		handler, "8.8.8.8:1234", "/webhook/uuid-a",
	)
	assert.Equal(
		t, http.StatusOK, w.Code,
		"a different client IP must not be affected",
	)
}

// TestReceiverRateLimit_CountsEveryMethod proves the receiver
// limit counts non-POST requests too: a GET shares the bucket
// with a POST and is itself rejected once over the limit.
func TestReceiverRateLimit_CountsEveryMethod(t *testing.T) {
	t.Parallel()

	const (
		limit = 2
		ip    = "7.7.7.7:1234"
		path  = "/webhook/uuid-c"
	)

	handler := receiverLimitedHandler(t, limit)

	get := func() *httptest.ResponseRecorder {
		req := httptest.NewRequestWithContext(
			context.Background(), http.MethodGet, path, nil,
		)
		req.RemoteAddr = ip

		w := httptest.NewRecorder()
		handler.ServeHTTP(w, req)

		return w
	}

	// One POST plus one GET fill the bucket, so the GET must
	// have been counted.
	assert.Equal(
		t, http.StatusOK, receiverPost(handler, ip, path).Code,
	)
	assert.Equal(t, http.StatusOK, get().Code)

	assert.Equal(
		t, http.StatusTooManyRequests, get().Code,
		"a GET over the limit must be rate-limited",
	)
}