diff --git a/beater/api/intake/handler.go b/beater/api/intake/handler.go index 3d9a4d54937..6a29d24abfe 100644 --- a/beater/api/intake/handler.go +++ b/beater/api/intake/handler.go @@ -81,6 +81,8 @@ func Handler(handler StreamHandler, requestMetadataFunc RequestMetadataFunc, bat return } + // copy batchProcessor to avoid updating closure below + batchProcessor := batchProcessor if limiter, ok := ratelimit.FromContext(c.Request.Context()); ok { // Apply rate limiting after reading but before processing any events. batchProcessor = modelprocessor.Chained{ diff --git a/beater/api/intake/handler_test.go b/beater/api/intake/handler_test.go index ed4074520f4..1c13b1fb88e 100644 --- a/beater/api/intake/handler_test.go +++ b/beater/api/intake/handler_test.go @@ -217,6 +217,32 @@ func TestRateLimiting(t *testing.T) { } } +func TestRateLimitingRequests(t *testing.T) { + // Check that rate limiting across multiple requests is handled correctly. + // + // ratelimit.ndjson contains 19 events, and we rate limit in batches of 10 + // events. The burst of 41 should be enough for 2 iterations with one left. + limiter := rate.NewLimiter(1, 41) + processor := stream.BackendProcessor(config.DefaultConfig()) + handler := Handler(processor, emptyRequestMetadata, modelprocessor.Nop{}) + + data, err := ioutil.ReadFile("../../../testdata/intake-v2/ratelimit.ndjson") + require.NoError(t, err) + for i := 0; i < 2; i++ { + r := httptest.NewRequest("POST", "/", bytes.NewBuffer(data)) + r = r.WithContext(ratelimit.ContextWithLimiter(r.Context(), limiter)) + r.Header.Add("Content-Type", "application/x-ndjson") + + w := httptest.NewRecorder() + c := request.NewContext() + c.Reset(w, r) + handler(c) + assert.Equal(t, http.StatusAccepted, w.Code) + } + assert.True(t, limiter.Allow()) + assert.False(t, limiter.Allow()) +} + type testcaseIntakeHandler struct { c *request.Context w *httptest.ResponseRecorder diff --git a/systemtest/rum_test.go b/systemtest/rum_test.go index 4d82b0a71ef..59a4ef03764 100644 --- a/systemtest/rum_test.go +++ b/systemtest/rum_test.go @@ -184,13 +184,11 @@ func TestRUMRateLimit(t *testing.T) { // Just check that rate limiting is wired up. More specific rate limiting scenarios are unit tested. var g errgroup.Group - g.Go(func() error { return sendEvents("10.11.12.13", srv.Config.RUM.RateLimit.EventLimit) }) - g.Go(func() error { return sendEvents("10.11.12.14", srv.Config.RUM.RateLimit.EventLimit) }) + g.Go(func() error { return sendEvents("10.11.12.13", srv.Config.RUM.RateLimit.EventLimit*3) }) + g.Go(func() error { return sendEvents("10.11.12.14", srv.Config.RUM.RateLimit.EventLimit*3) }) assert.NoError(t, g.Wait()) g = errgroup.Group{} - g.Go(func() error { return sendEvents("10.11.12.13", srv.Config.RUM.RateLimit.EventLimit) }) - g.Go(func() error { return sendEvents("10.11.12.14", srv.Config.RUM.RateLimit.EventLimit) }) g.Go(func() error { return sendEvents("10.11.12.15", srv.Config.RUM.RateLimit.EventLimit) }) err = g.Wait() require.Error(t, err)