Logging: rate limit fronted logging endpoint (#29272)

Co-authored-by: Arve Knudsen <arve.knudsen@gmail.com>
Co-authored-by: achatterjee-grafana <70489351+achatterjee-grafana@users.noreply.github.com>

pkg/middleware/rate_limit.go

@@ -0,0 +1,25 @@
+package middleware
+
+import (
+	"time"
+
+	"golang.org/x/time/rate"
+	"gopkg.in/macaron.v1"
+
+	"github.com/grafana/grafana/pkg/models"
+)
+
+type getTimeFn func() time.Time
+
+// RateLimit is a very basic rate limiter.
+// Will allow average of "rps" requests per second over an extended period of time, with max "burst" requests at the same time.
+// getTime should return the current time. For non-testing purposes use time.Now
+func RateLimit(rps, burst int, getTime getTimeFn) macaron.Handler {
+	l := rate.NewLimiter(rate.Limit(rps), burst)
+	return func(c *models.ReqContext) {
+		if !l.AllowN(getTime(), 1) {
+			c.JsonApiErr(429, "Rate limit reached", nil)
+			return
+		}
+	}
+}

pkg/middleware/rate_limit_test.go

@@ -0,0 +1,86 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/grafana/grafana/pkg/models"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/macaron.v1"
+)
+
+type execFunc func() *httptest.ResponseRecorder
+type advanceTimeFunc func(deltaTime time.Duration)
+type rateLimiterScenarioFunc func(c execFunc, t advanceTimeFunc)
+
+func rateLimiterScenario(t *testing.T, desc string, rps int, burst int, fn rateLimiterScenarioFunc) {
+	t.Run(desc, func(t *testing.T) {
+		defaultHandler := func(c *models.ReqContext) {
+			resp := make(map[string]interface{})
+			resp["message"] = "OK"
+			c.JSON(200, resp)
+		}
+		currentTime := time.Now()
+
+		m := macaron.New()
+		m.Use(macaron.Renderer(macaron.RenderOptions{
+			Directory: "",
+			Delims:    macaron.Delims{Left: "[[", Right: "]]"},
+		}))
+		m.Use(GetContextHandler(nil, nil, nil))
+		m.Get("/foo", RateLimit(rps, burst, func() time.Time { return currentTime }), defaultHandler)
+
+		fn(func() *httptest.ResponseRecorder {
+			resp := httptest.NewRecorder()
+			req, err := http.NewRequest("GET", "/foo", nil)
+			require.NoError(t, err)
+			m.ServeHTTP(resp, req)
+			return resp
+		}, func(deltaTime time.Duration) {
+			currentTime = currentTime.Add(deltaTime)
+		})
+	})
+}
+
+func TestRateLimitMiddleware(t *testing.T) {
+	rateLimiterScenario(t, "rate limit calls, with burst", 10, 10, func(doReq execFunc, advanceTime advanceTimeFunc) {
+		// first 10 calls succeed
+		for i := 0; i < 10; i++ {
+			resp := doReq()
+			assert.Equal(t, 200, resp.Code)
+		}
+
+		// next one fails
+		resp := doReq()
+		assert.Equal(t, 429, resp.Code)
+
+		// check that requests are accepted again in 1 sec
+		advanceTime(1 * time.Second)
+
+		for i := 0; i < 10; i++ {
+			resp := doReq()
+			assert.Equal(t, 200, resp.Code)
+		}
+	})
+
+	rateLimiterScenario(t, "rate limit calls, no burst", 10, 1, func(doReq execFunc, advanceTime advanceTimeFunc) {
+		// first calls succeeds
+		resp := doReq()
+		assert.Equal(t, 200, resp.Code)
+
+		// immediately fired next one fails
+		resp = doReq()
+		assert.Equal(t, 429, resp.Code)
+
+		// but spacing calls out works
+		for i := 0; i < 10; i++ {
+			advanceTime(100 * time.Millisecond)
+			resp := doReq()
+			assert.Equal(t, 200, resp.Code)
+		}
+	})
+}