Allocation of Resources Without Limits or Throttling in flagd

@@ -36,6 +36,8 @@ type Configuration struct {
 	ContextValues              map[string]any
 	HeaderToContextKeyMappings map[string]string
 	StreamDeadline             time.Duration
+	MaxRequestHeaderBytes      int64
+	MaxRequestBodyBytes        int64
 }
 
 /*

@@ -18,6 +18,8 @@ flagd start [flags]
   -h, --help                                 help for start
   -z, --log-format string                    Set the logging format, e.g. console or json (default "console")
   -m, --management-port int32                Port for management operations (default 8014)
+  -B, --max-request-body int                 Maximum allowed request body size in bytes. Requests exceeding this are rejected with HTTP 413 (OFREP) or 429 (connect). Set to 0 to disable. WARNING: disabling this limit may allow memory exhaustion from oversized requests. (default 1000000)
+  -R, --max-request-header int               Maximum allowed request header size in bytes. Requests exceeding this are rejected with HTTP 431. Set to 0 to use Go's built-in default (1 MiB). WARNING: setting a very large or zero value may allow memory exhaustion from oversized headers. (default 1000000)
   -t, --metrics-exporter string              Set the metrics exporter. Default(if unset) is Prometheus. Can be override to otel - OpenTelemetry metric exporter. Overriding to otel require otelCollectorURI to be present
   -r, --ofrep-port int32                     ofrep service port (default 8016)
   -A, --otel-ca-path string                  tls certificate authority path to use with OpenTelemetry collector

@@ -40,6 +40,8 @@ const (
 	contextValueFlagName       = "context-value"
 	headerToContextKeyFlagName = "context-from-header"
 	streamDeadlineFlagName     = "stream-deadline"
+	maxRequestBodyFlagName     = "max-request-body"
+	maxRequestHeaderFlagName   = "max-request-header"
 )
 
 func init() {
@@ -91,6 +93,8 @@ func init() {
 		"header values to context values, where key is Header name, value is context key")
 	flags.Duration(streamDeadlineFlagName, 0, "Set a server-side deadline for flagd sync and event streams (default 0, means no deadline).")
 	flags.Bool(disableSyncMetadata, false, "Disables the getMetadata endpoint of the sync service. Defaults to false, but will default to true in later versions.")
+	flags.Int64P(maxRequestBodyFlagName, "B", 1_000_000, "Maximum allowed request body size in bytes. Requests exceeding this are rejected with HTTP 413 (OFREP) or 429 (connect). Set to 0 to disable. WARNING: disabling this limit may allow memory exhaustion from oversized requests.")
+	flags.Int64P(maxRequestHeaderFlagName, "R", 1_000_000, "Maximum allowed request header size in bytes. Requests exceeding this are rejected with HTTP 431. Set to 0 to use Go's built-in default (1 MiB). WARNING: setting a very large or zero value may allow memory exhaustion from oversized headers.")
 
 	bindFlags(flags)
 }
@@ -117,6 +121,8 @@ func bindFlags(flags *pflag.FlagSet) {
 	_ = viper.BindPFlag(headerToContextKeyFlagName, flags.Lookup(headerToContextKeyFlagName))
 	_ = viper.BindPFlag(streamDeadlineFlagName, flags.Lookup(streamDeadlineFlagName))
 	_ = viper.BindPFlag(disableSyncMetadata, flags.Lookup(disableSyncMetadata))
+	_ = viper.BindPFlag(maxRequestBodyFlagName, flags.Lookup(maxRequestBodyFlagName))
+	_ = viper.BindPFlag(maxRequestHeaderFlagName, flags.Lookup(maxRequestHeaderFlagName))
 }
 
 // startCmd represents the start command
@@ -171,6 +177,16 @@ var startCmd = &cobra.Command{
 			headerToContextKeyMappings[k] = v
 		}
 
+		// Request size limits
+		maxRequestBodyBytes := viper.GetInt64(maxRequestBodyFlagName)
+		if maxRequestBodyBytes > 0 {
+			rtLogger.Info(fmt.Sprintf("request body limit set to %d bytes", maxRequestBodyBytes))
+		}
+		maxRequestHeaderBytes := viper.GetInt64(maxRequestHeaderFlagName)
+		if maxRequestHeaderBytes > 0 {
+			rtLogger.Info(fmt.Sprintf("request header limit set to %d bytes", maxRequestHeaderBytes))
+		}
+
 		// Build Runtime -----------------------------------------------------------
 		rt, err := runtime.FromConfig(logger, Version, runtime.Config{
 			CORS:                       viper.GetStringSlice(corsFlagName),
@@ -193,6 +209,8 @@ var startCmd = &cobra.Command{
 			SyncProviders:              syncProviders,
 			ContextValues:              contextValuesToMap,
 			HeaderToContextKeyMappings: headerToContextKeyMappings,
+			MaxRequestBodyBytes:        maxRequestBodyBytes,
+			MaxRequestHeaderBytes:      maxRequestHeaderBytes,
 		})
 		if err != nil {
 			rtLogger.Fatal(err.Error())

@@ -46,6 +46,8 @@ type Config struct {
 
 	ContextValues              map[string]any
 	HeaderToContextKeyMappings map[string]string
+	MaxRequestBodyBytes        int64
+	MaxRequestHeaderBytes      int64
 }
 
 // FromConfig builds a runtime from startup configurations
@@ -105,10 +107,12 @@ func FromConfig(logger *logger.Logger, version string, config Config) (*Runtime,
 
 	// ofrep service
 	ofrepService, err := ofrep.NewOfrepService(jsonEvaluator, config.CORS, ofrep.SvcConfiguration{
-		Logger: logger.WithFields(zap.String("component", "OFREPService")),
-		Port:   config.OfrepServicePort,
-		ServiceName: svcName,
-		MetricsRecorder: recorder,
+		Logger:                logger.WithFields(zap.String("component", "OFREPService")),
+		Port:                  config.OfrepServicePort,
+		ServiceName:           svcName,
+		MetricsRecorder:       recorder,
+		MaxRequestBodyBytes:   config.MaxRequestBodyBytes,
+		MaxRequestHeaderBytes: config.MaxRequestHeaderBytes,
 	},
 		config.ContextValues,
 		config.HeaderToContextKeyMappings,
@@ -165,6 +169,8 @@ func FromConfig(logger *logger.Logger, version string, config Config) (*Runtime,
 			ContextValues:              config.ContextValues,
 			HeaderToContextKeyMappings: config.HeaderToContextKeyMappings,
 			StreamDeadline:             config.StreamDeadline,
+			MaxRequestBodyBytes:        config.MaxRequestBodyBytes,
+			MaxRequestHeaderBytes:      config.MaxRequestHeaderBytes,
 		},
 		Syncs: iSyncs,
 	}, nil

@@ -205,10 +205,16 @@ func (s *ConnectService) setupServer(svcConf service.Configuration) (net.Listene
 		v2:  v2Handler,
 	}
 
+	var svcHandler http.Handler = bs
+	if svcConf.MaxRequestBodyBytes > 0 {
+		svcHandler = http.MaxBytesHandler(svcHandler, svcConf.MaxRequestBodyBytes)
+	}
+
 	s.serverMtx.Lock()
 	s.server = &http.Server{
 		ReadHeaderTimeout: time.Second,
-		Handler:           bs,
+		Handler:           svcHandler,
+		MaxHeaderBytes:    int(svcConf.MaxRequestHeaderBytes),
 	}
 	s.serverMtx.Unlock()
 

@@ -1,11 +1,13 @@
 package service
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
 	"net/http"
 	"os"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -28,7 +30,9 @@ import (
 	"google.golang.org/protobuf/types/known/structpb"
 )
 
-func TestConnectService_UnixConnection(t *testing.T) {
+const resolveAllURLFmt = "http://localhost:%d/flagd.evaluation.v1.Service/ResolveAll"
+
+func TestConnectServiceUnixConnection(t *testing.T) {
 	type evalFields struct {
 		result   bool
 		variant  string
@@ -156,17 +160,21 @@ func TestAddMiddleware(t *testing.T) {
 	}()
 
 	require.Eventually(t, func() bool {
-		resp, err := http.Get(fmt.Sprintf("http://localhost:%d/flagd.evaluation.v1.Service/ResolveAll", port))
+		resp, err := http.Get(fmt.Sprintf(resolveAllURLFmt, port))
+		if err == nil && resp != nil {
+			resp.Body.Close()
+		}
 		// with the default http handler we should get a method not allowed (405) when attempting a GET request
 		return err == nil && resp.StatusCode == http.StatusMethodNotAllowed
 	}, 3*time.Second, 100*time.Millisecond)
 
 	svc.AddMiddleware(mwMock)
 
 	// with the injected middleware, the GET method should work
-	resp, err := http.Get(fmt.Sprintf("http://localhost:%d/flagd.evaluation.v1.Service/ResolveAll", port))
+	resp, err := http.Get(fmt.Sprintf(resolveAllURLFmt, port))
 
 	require.Nil(t, err)
+	defer resp.Body.Close()
 	// verify that the status we return in the mocked middleware
 	require.Equal(t, http.StatusOK, resp.StatusCode)
 }
@@ -307,3 +315,79 @@ func TestConnectServiceShutdown(t *testing.T) {
 		t.Error("timeout while waiting for notifications")
 	}
 }
+
+// startConnectService creates a ConnectService with a mock evaluator and metric recorder,
+// starts it in a background goroutine with the given configuration, and waits until it is ready.
+// It returns the port the service is listening on.
+func startConnectService(t *testing.T, port uint16, conf iservice.Configuration) {
+	t.Helper()
+
+	ctrl := gomock.NewController(t)
+	eval := mock.NewMockIEvaluator(ctrl)
+
+	exp := metric.NewManualReader()
+	rs := resource.NewWithAttributes("testSchema")
+	metricRecorder := telemetry.NewOTelRecorder(exp, rs, "limit-test")
+
+	svc := NewConnectService(logger.NewLogger(nil, false), eval, nil, metricRecorder)
+
+	conf.ReadinessProbe = func() bool { return true }
+	conf.Port = port
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	go func() {
+		_ = svc.Serve(ctx, conf)
+	}()
+
+	require.Eventually(t, func() bool {
+		resp, err := http.Get(fmt.Sprintf(resolveAllURLFmt, port))
+		if err == nil && resp != nil {
+			resp.Body.Close()
+		}
+		return err == nil && resp != nil
+	}, 3*time.Second, 100*time.Millisecond)
+}
+
+func TestConnectServiceRequestBodySizeLimit(t *testing.T) {
+	const port = 18291
+
+	startConnectService(t, port, iservice.Configuration{
+		MaxRequestBodyBytes: 10, // allow only 10 bytes
+	})
+
+	// Valid JSON that exceeds the 10-byte body limit, so MaxBytesReader fires mid-parse.
+	largeBody := []byte(`{"flagKey":"` + strings.Repeat("a", 100) + `"}`)
+	req, err := http.NewRequest(http.MethodPost,
+		fmt.Sprintf("http://localhost:%d/flagd.evaluation.v1.Service/ResolveBoolean", port),
+		bytes.NewReader(largeBody))
+	require.NoError(t, err)
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	// connect-go maps MaxBytesError (resource exhausted) to HTTP 429.
+	require.Equal(t, http.StatusTooManyRequests, resp.StatusCode)
+}
+
+func TestConnectServiceRequestHeaderSizeLimit(t *testing.T) {
+	const port = 18292
+
+	startConnectService(t, port, iservice.Configuration{
+		MaxRequestHeaderBytes: 100, // 10000-byte test header value easily exceeds 100 + slop
+	})
+
+	req, err := http.NewRequest(http.MethodPost,
+		fmt.Sprintf("http://localhost:%d/flagd.evaluation.v1.Service/ResolveBoolean", port),
+		bytes.NewReader([]byte("{}")))
+	require.NoError(t, err)
+	// Use valid ASCII to avoid client-side rejection; value exceeds MaxHeaderBytes + slop.
+	req.Header.Set("X-Large-Header", strings.Repeat("a", 10000))
+
+	resp, err := http.DefaultClient.Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+	require.Equal(t, http.StatusRequestHeaderFieldsTooLarge, resp.StatusCode)
+}

@@ -3,6 +3,7 @@ package ofrep
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"net/http"
 
@@ -90,8 +91,9 @@ func (h *handler) HandleFlagEvaluation(w http.ResponseWriter, r *http.Request) {
 	flagKey := vars[key]
 	request, err := extractOfrepRequest(r)
 	if err != nil {
-		h.writeJSONToResponse(http.StatusBadRequest, ofrep.ContextErrorResponseFrom(flagKey), w)
-		return
+		if h.handleExtractionError(w, err, ofrep.ContextErrorResponseFrom(flagKey)) {
+			return
+		}
 	}
 	evaluationContext := flagdContext(h.Logger, requestID, request, h.contextValues, r.Header, h.headerToContextKeyMappings)
 	selectorExpression := r.Header.Get(service.FLAGD_SELECTOR_HEADER)
@@ -113,8 +115,9 @@ func (h *handler) HandleBulkEvaluation(w http.ResponseWriter, r *http.Request) {
 
 	request, err := extractOfrepRequest(r)
 	if err != nil {
-		h.writeJSONToResponse(http.StatusBadRequest, ofrep.BulkEvaluationContextError(), w)
-		return
+		if h.handleExtractionError(w, err, ofrep.BulkEvaluationContextError()) {
+			return
+		}
 	}
 
 	evaluationContext := flagdContext(h.Logger, requestID, request, h.contextValues, r.Header, h.headerToContextKeyMappings)
@@ -152,11 +155,34 @@ func (h *handler) writeJSONToResponse(status int, payload interface{}, w http.Re
 	}
 }
 
+// handleExtractionError checks for errors from extractOfrepRequest and writes an appropriate response.
+// It returns true if an error was handled.
+func (h *handler) handleExtractionError(w http.ResponseWriter, err error, errorPayload any) bool {
+	if err == nil {
+		return false
+	}
+	var maxBytesErr *http.MaxBytesError
+	if errors.As(err, &maxBytesErr) {
+		h.writeJSONToResponse(http.StatusRequestEntityTooLarge,
+			ofrep.InternalError{ErrorDetails: "request body too large"}, w)
+		return true
+	}
+	h.writeJSONToResponse(http.StatusBadRequest, errorPayload, w)
+	return true
+}
+
 func extractOfrepRequest(req *http.Request) (ofrep.Request, error) {
 	request := ofrep.Request{}
 	err := json.NewDecoder(req.Body).Decode(&request)
-	if err != nil && err.Error() != "EOF" {
-		return request, fmt.Errorf("decode error: %w", err)
+	if err != nil {
+		// Propagate MaxBytesError so callers can return 413.
+		var maxBytesErr *http.MaxBytesError
+		if errors.As(err, &maxBytesErr) {
+			return request, err
+		}
+		if err.Error() != "EOF" {
+			return request, fmt.Errorf("decode error: %w", err)
+		}
 	}
 
 	return request, nil

@@ -20,10 +20,12 @@ type IOfrepService interface {
 }
 
 type SvcConfiguration struct {
-	Logger          *logger.Logger
-	Port            uint16
-	ServiceName     string
-	MetricsRecorder telemetry.IMetricsRecorder
+	Logger                *logger.Logger
+	Port                  uint16
+	ServiceName           string
+	MetricsRecorder       telemetry.IMetricsRecorder
+	MaxRequestBodyBytes   int64
+	MaxRequestHeaderBytes int64
 }
 
 type Service struct {
@@ -40,19 +42,24 @@ func NewOfrepService(
 		AllowedMethods: []string{http.MethodPost},
 	})
 
-	h := corsMW.Handler(NewOfrepHandler(
+	var h http.Handler = NewOfrepHandler(
 		cfg.Logger,
 		evaluator,
 		contextValues,
 		headerToContextKeyMappings,
 		cfg.MetricsRecorder,
 		cfg.ServiceName,
-	))
+	)
+	if cfg.MaxRequestBodyBytes > 0 {
+		h = http.MaxBytesHandler(h, cfg.MaxRequestBodyBytes)
+	}
+	h = corsMW.Handler(h)
 
 	server := http.Server{
 		Addr:              fmt.Sprintf(":%d", cfg.Port),
 		Handler:           h,
 		ReadHeaderTimeout: 3 * time.Second,
+		MaxHeaderBytes:    int(cfg.MaxRequestHeaderBytes),
 	}
 
 	return &Service{

@@ -5,6 +5,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"strings"
 	"testing"
 	"time"
 
@@ -17,23 +18,28 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-func Test_OfrepServiceStartStop(t *testing.T) {
+const (
+	testServiceName       = "test-service"
+	errCreateOfrepService = "error creating the ofrep service: %v"
+)
+
+func TestOfrepServiceStartStop(t *testing.T) {
 	port := 18282
 	eval := mock.NewMockIEvaluator(gomock.NewController(t))
 
 	eval.EXPECT().ResolveAllValues(gomock.Any(), gomock.Any(), gomock.Any()).
 		Return([]evaluator.AnyValue{}, model.Metadata{}, nil)
 
 	cfg := SvcConfiguration{
-		Logger: logger.NewLogger(nil, false),
-		Port:   uint16(port),
-		ServiceName:     "test-service",
+		Logger:          logger.NewLogger(nil, false),
+		Port:            uint16(port),
+		ServiceName:     testServiceName,
 		MetricsRecorder: &telemetry.NoopMetricsRecorder{},
 	}
 
 	service, err := NewOfrepService(eval, []string{"*"}, cfg, nil, nil)
 	if err != nil {
-		t.Fatalf("error creating the ofrep service: %v", err)
+		t.Fatalf(errCreateOfrepService, err)
 	}
 
 	ctx, cancelFunc := context.WithCancel(context.Background())
@@ -69,6 +75,10 @@ func Test_OfrepServiceStartStop(t *testing.T) {
 }
 
 func tryResponse(method string, uri string, payload []byte) (int, error) {
+	return tryResponseWithHeaders(method, uri, payload, nil)
+}
+
+func tryResponseWithHeaders(method string, uri string, payload []byte, headers map[string]string) (int, error) {
 	client := http.Client{
 		Timeout: 3 * time.Second,
 	}
@@ -78,9 +88,90 @@ func tryResponse(method string, uri string, payload []byte) (int, error) {
 		return 0, fmt.Errorf("error forming the request: %w", err)
 	}
 
+	for k, v := range headers {
+		request.Header.Set(k, v)
+	}
+
 	rsp, err := client.Do(request)
 	if err != nil {
 		return 0, fmt.Errorf("error from the request: %w", err)
 	}
 	return rsp.StatusCode, nil
 }
+
+func TestOfrepServiceRequestBodySizeLimit(t *testing.T) {
+	svc, port := startOfrepService(t, SvcConfiguration{
+		Logger:              logger.NewLogger(nil, false),
+		Port:                18283,
+		ServiceName:         testServiceName,
+		MetricsRecorder:     &telemetry.NoopMetricsRecorder{},
+		MaxRequestBodyBytes: 10, // allow only 10 bytes
+	})
+	_ = svc // kept alive by deferred cleanup
+
+	path := fmt.Sprintf("http://localhost:%d/ofrep/v1/evaluate/flags/myFlag", port)
+	// Valid JSON whose size exceeds the 10-byte limit, so MaxBytesReader triggers mid-parse.
+	largeBody := []byte(`{"context":{"k":"` + strings.Repeat("a", 100) + `"}}`)
+
+	status, err := tryResponse(http.MethodPost, path, largeBody)
+	if err != nil {
+		t.Fatalf("unexpected request error: %v", err)
+	}
+
+	if status != http.StatusRequestEntityTooLarge {
+		t.Errorf("expected HTTP 413, got %d", status)
+	}
+}
+
+func TestOfrepServiceRequestHeaderSizeLimit(t *testing.T) {
+	svc, port := startOfrepService(t, SvcConfiguration{
+		Logger:                logger.NewLogger(nil, false),
+		Port:                  18284,
+		ServiceName:           testServiceName,
+		MetricsRecorder:       &telemetry.NoopMetricsRecorder{},
+		MaxRequestHeaderBytes: 100, // 10000-byte test header value easily exceeds 100 + slop
+	})
+	_ = svc // kept alive by deferred cleanup
+
+	path := fmt.Sprintf("http://localhost:%d/ofrep/v1/evaluate/flags/myFlag", port)
+	// The header value must exceed MaxHeaderBytes + Go's ~4096-byte read buffer slop.
+	largeHeaderValue := string(bytes.Repeat([]byte("a"), 10000))
+
+	status, err := tryResponseWithHeaders(http.MethodPost, path, []byte{}, map[string]string{
+		"X-Large-Header": largeHeaderValue,
+	})
+	if err != nil {
+		t.Fatalf("unexpected request error: %v", err)
+	}
+
+	if status != http.StatusRequestHeaderFieldsTooLarge {
+		t.Errorf("expected HTTP 431, got %d", status)
+	}
+}
+
+// startOfrepService creates, starts and returns an OFREP service with the given config.
+// It registers cleanup to stop the service when the test finishes.
+func startOfrepService(t *testing.T, cfg SvcConfiguration) (*Service, uint16) {
+	t.Helper()
+
+	eval := mock.NewMockIEvaluator(gomock.NewController(t))
+	service, err := NewOfrepService(eval, []string{"*"}, cfg, nil, nil)
+	if err != nil {
+		t.Fatalf(errCreateOfrepService, err)
+	}
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	group, gCtx := errgroup.WithContext(ctx)
+	group.Go(func() error {
+		return service.Start(gCtx)
+	})
+	t.Cleanup(func() {
+		cancelFunc()
+		_ = group.Wait()
+	})
+
+	// wait for server startup
+	<-time.After(2 * time.Second)
+
+	return service, cfg.Port
+}