Add back ratelimit

Author: foadnh

config/confighttp/confighttp.go

@@ -290,6 +290,9 @@ type ServerConfig struct {
 	// Auth for this receiver
 	Auth *AuthConfig `mapstructure:"auth"`
 
+	// RateLimit for this receiver
+	RateLimit *RateLimit `mapstructure:"rate_limit"`
+
 	// MaxRequestBodySize sets the maximum request body size in bytes. Default: 20MiB.
 	MaxRequestBodySize int64 `mapstructure:"max_request_body_size"`
 
@@ -441,6 +444,17 @@ func (hss *ServerConfig) ToServer(_ context.Context, host component.Host, settin
 		handler = authInterceptor(handler, server, hss.Auth.RequestParameters)
 	}
 
+	// The RateLimit interceptor should always be right after auth to ensure
+	// the request rate is within an acceptable threshold.
+	if hss.RateLimit != nil {
+		limiter, err := hss.RateLimit.rateLimiter(host.GetExtensions())
+		if err != nil {
+			return nil, err
+		}
+
+		handler = rateLimitInterceptor(handler, limiter)
+	}
+
 	if hss.CORS != nil && len(hss.CORS.AllowedOrigins) > 0 {
 		co := cors.Options{
 			AllowedOrigins:   hss.CORS.AllowedOrigins,

config/confighttp/ratelimit.go

@@ -0,0 +1,51 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package confighttp // import "go.opentelemetry.io/collector/config/confighttp"
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+
+	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/extension"
+)
+
+// RateLimit defines rate limiter settings for the receiver.
+type RateLimit struct {
+	// RateLimiterID specifies the name of the extension to use in order to rate limit the incoming data point.
+	RateLimiterID component.ID `mapstructure:"rate_limiter"`
+}
+
+type rateLimiter interface {
+	extension.Extension
+
+	Take(context.Context, string, http.Header) error
+}
+
+// rateLimiter attempts to select the appropriate rateLimiter from the list of extensions,
+// based on the component id of the extension. If a rateLimiter is not found, an error is returned.
+func (rl RateLimit) rateLimiter(extensions map[component.ID]component.Component) (rateLimiter, error) {
+	if ext, found := extensions[rl.RateLimiterID]; found {
+		if limiter, ok := ext.(rateLimiter); ok {
+			return limiter, nil
+		}
+		return nil, fmt.Errorf("extension %q is not a rate limit", rl.RateLimiterID)
+	}
+	return nil, fmt.Errorf("rate limit %q not found", rl.RateLimiterID)
+}
+
+// rateLimitInterceptor adds interceptor for rate limit check.
+// It returns TooManyRequests(429) status code if rate limiter rejects the request.
+func rateLimitInterceptor(next http.Handler, limiter rateLimiter) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		err := limiter.Take(r.Context(), r.URL.Path, r.Header)
+		if err != nil {
+			http.Error(w, http.StatusText(http.StatusTooManyRequests), http.StatusTooManyRequests)
+			return
+		}
+
+		next.ServeHTTP(w, r)
+	})
+}

config/confighttp/ratelimit_test.go

@@ -0,0 +1,112 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package confighttp
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"go.opentelemetry.io/collector/component"
+	"go.opentelemetry.io/collector/component/componenttest"
+)
+
+func TestServerRateLimit(t *testing.T) {
+	// prepare
+	hss := ServerConfig{
+		Endpoint: "localhost:0",
+		RateLimit: &RateLimit{
+			RateLimiterID: component.NewID(component.MustNewType("mock")),
+		},
+	}
+
+	limiter := &mockRateLimiter{}
+
+	host := &mockHost{
+		ext: map[component.ID]component.Component{
+			component.NewID(component.MustNewType("mock")): limiter,
+		},
+	}
+
+	var handlerCalled bool
+	handler := http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {
+		handlerCalled = true
+	})
+
+	srv, err := hss.ToServer(context.Background(), host, componenttest.NewNopTelemetrySettings(), handler)
+	require.NoError(t, err)
+
+	// test
+	srv.Handler.ServeHTTP(&httptest.ResponseRecorder{}, httptest.NewRequest("GET", "/", nil))
+
+	// verify
+	assert.True(t, handlerCalled)
+	assert.Equal(t, 1, limiter.calls)
+}
+
+func TestInvalidServerRateLimit(t *testing.T) {
+	hss := ServerConfig{
+		RateLimit: &RateLimit{
+			RateLimiterID: component.NewID(component.MustNewType("non_existing")),
+		},
+	}
+
+	srv, err := hss.ToServer(context.Background(), componenttest.NewNopHost(), componenttest.NewNopTelemetrySettings(), http.NewServeMux())
+	require.Error(t, err)
+	require.Nil(t, srv)
+}
+
+func TestRejectedServerRateLimit(t *testing.T) {
+	// prepare
+	hss := ServerConfig{
+		Endpoint: "localhost:0",
+		RateLimit: &RateLimit{
+			RateLimiterID: component.NewID(component.MustNewType("mock")),
+		},
+	}
+	host := &mockHost{
+		ext: map[component.ID]component.Component{
+			component.NewID(component.MustNewType("mock")): &mockRateLimiter{
+				err: errors.New("rate limited"),
+			},
+		},
+	}
+
+	srv, err := hss.ToServer(context.Background(), host, componenttest.NewNopTelemetrySettings(), http.HandlerFunc(func(_ http.ResponseWriter, _ *http.Request) {}))
+	require.NoError(t, err)
+
+	// test
+	response := &httptest.ResponseRecorder{}
+	srv.Handler.ServeHTTP(response, httptest.NewRequest("GET", "/", nil))
+
+	// verify
+	assert.Equal(t, response.Result().StatusCode, http.StatusTooManyRequests)
+	assert.Equal(t, response.Result().Status, fmt.Sprintf("%v %s", http.StatusTooManyRequests, http.StatusText(http.StatusTooManyRequests)))
+}
+
+// Mocks
+
+type mockRateLimiter struct {
+	calls int
+	err   error
+}
+
+func (m *mockRateLimiter) Take(context.Context, string, http.Header) error {
+	m.calls++
+	return m.err
+}
+
+func (m *mockRateLimiter) Start(_ context.Context, _ component.Host) error {
+	return nil
+}
+
+func (m *mockRateLimiter) Shutdown(_ context.Context) error {
+	return nil
+}