Merge branch 'main' into fix/mcpproxy-unframed-json-sse

Author: RegondaChandan

internal/extproc/processor_impl.go

@@ -267,12 +267,10 @@ func (r *routerProcessor[ReqT, RespT, RespChunkT, EndpointSpecT]) ProcessRequest
 		Header: &corev3.HeaderValue{Key: internalapi.ModelNameHeaderKeyDefault, RawValue: []byte(originalModel)},
 	})
 	originalPath := r.requestHeaders[":path"]
-	if r.requestHeaders[originalPathHeader] == "" {
-		r.requestHeaders[originalPathHeader] = originalPath
-		additionalHeaders = append(additionalHeaders, &corev3.HeaderValueOption{
-			Header: &corev3.HeaderValue{Key: originalPathHeader, RawValue: []byte(originalPath)},
-		})
-	}
+	r.requestHeaders[originalPathHeader] = originalPath
+	additionalHeaders = append(additionalHeaders, &corev3.HeaderValueOption{
+		Header: &corev3.HeaderValue{Key: originalPathHeader, RawValue: []byte(originalPath)},
+	})
 	if r.requestHeaders[internalapi.EnvoyOriginalPathHeader] == "" {
 		r.requestHeaders[internalapi.EnvoyOriginalPathHeader] = originalPath
 		additionalHeaders = append(additionalHeaders, &corev3.HeaderValueOption{

internal/mcpproxy/config.go

@@ -31,6 +31,7 @@ type (
 		tracer                     tracingapi.MCPTracer
 		client                     http.Client
 		logRequestHeaderAttributes map[string]string
+		maxRequestBodySize         int64 // maximum allowed POST body size in bytes
 	}
 
 	mcpProxyConfig struct {

internal/mcpproxy/handlers.go

@@ -277,8 +277,18 @@ func (m *mcpRequestContext) servePOST(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	body, err := io.ReadAll(r.Body)
+	limit := m.maxRequestBodySize
+	if limit <= 0 {
+		limit = defaultMaxRequestBodySize
+	}
+	body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, limit))
 	if err != nil {
+		var maxBytesErr *http.MaxBytesError
+		if errors.As(err, &maxBytesErr) {
+			errType = metrics.MCPErrorInternal
+			onErrorResponse(w, http.StatusRequestEntityTooLarge, "request body too large")
+			return
+		}
 		errType = metrics.MCPErrorInternal
 		onErrorResponse(w, http.StatusBadRequest, err.Error())
 		return

internal/mcpproxy/handlers_test.go

@@ -168,6 +168,20 @@ func TestServePOST_InvalidJSONRPC(t *testing.T) {
 	require.Contains(t, rr.Body.String(), "invalid JSON-RPC message")
 }
 
+func TestServePOST_OversizedBody(t *testing.T) {
+	proxy := newTestMCPProxy()
+	proxy.maxRequestBodySize = 16 // tiny limit to exercise the guard without allocating a large buffer.
+
+	body := strings.NewReader(`{"jsonrpc":"2.0","method":"initialize","id":1,"params":{}}`) // > 16 bytes
+	req := httptest.NewRequest(http.MethodPost, "/mcp", body)
+	rr := httptest.NewRecorder()
+
+	proxy.servePOST(rr, req)
+
+	require.Equal(t, http.StatusRequestEntityTooLarge, rr.Code)
+	require.Contains(t, rr.Body.String(), "request body too large")
+}
+
 func TestServePOST_InvalidSessionID(t *testing.T) {
 	proxy := newTestMCPProxy()
 	req := httptest.NewRequest(http.MethodPost, "/mcp", strings.NewReader(`{"jsonrpc":"2.0","method":"tools/call","params":{"name":"test-tool"},"id":"1"}`))

internal/mcpproxy/mcpproxy.go

@@ -14,6 +14,8 @@ import (
 	"log/slog"
 	"maps"
 	"net/http"
+	"os"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -39,6 +41,20 @@ type mcpRequestContext struct {
 	perBackendMetricsRecorded bool
 }
 
+// defaultMaxRequestBodySize is the default maximum allowed POST body size in bytes (4 MiB).
+const defaultMaxRequestBodySize = 4 * 1024 * 1024
+
+// getMaxRequestBodySize returns the configured POST body limit from the environment variable,
+// falling back to 4 MiB if the variable is unset or invalid.
+func getMaxRequestBodySize() int64 {
+	if v, ok := os.LookupEnv("MCP_PROXY_MAX_REQUEST_BODY_SIZE"); ok {
+		if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {
+			return n
+		}
+	}
+	return defaultMaxRequestBodySize
+}
+
 // NewMCPProxy creates a new MCPProxy instance.
 func NewMCPProxy(l *slog.Logger, mcpMetrics metrics.MCPMetrics, tracer tracingapi.MCPTracer, sessionCrypto SessionCrypto, logRequestHeaderAttributes map[string]string) (*ProxyConfig, *http.ServeMux, error) {
 	toolChangeSignaler := newMultiWatcherSignaler() // used to signal changes to all active sessions.
@@ -49,6 +65,7 @@ func NewMCPProxy(l *slog.Logger, mcpMetrics metrics.MCPMetrics, tracer tracingap
 		l:                          l,
 		client:                     http.Client{}, // No timeout as it's enforced at Envoy level.
 		logRequestHeaderAttributes: maps.Clone(logRequestHeaderAttributes),
+		maxRequestBodySize:         getMaxRequestBodySize(),
 	}
 	mux := http.NewServeMux()
 	mux.HandleFunc(

internal/mcpproxy/mcpproxy_test.go

@@ -60,6 +60,25 @@ func (f *fakeTracer) StartSpanAndInjectMeta(context.Context, *jsonrpc.Request, m
 
 var noopTracer = tracingapi.NoopMCPTracer{}
 
+func TestGetMaxRequestBodySize(t *testing.T) {
+	for _, tc := range []struct {
+		name     string
+		envValue string
+		want     int64
+	}{
+		{name: "default when env var not set", envValue: "", want: defaultMaxRequestBodySize},
+		{name: "custom value from env var", envValue: "1048576", want: 1048576},
+		{name: "default when env var is not a number", envValue: "not-a-number", want: defaultMaxRequestBodySize},
+		{name: "default when env var is zero", envValue: "0", want: defaultMaxRequestBodySize},
+		{name: "default when env var is negative", envValue: "-1", want: defaultMaxRequestBodySize},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Setenv("MCP_PROXY_MAX_REQUEST_BODY_SIZE", tc.envValue)
+			require.Equal(t, tc.want, getMaxRequestBodySize())
+		})
+	}
+}
+
 func TestNewMCPProxy(t *testing.T) {
 	l := slog.Default()
 	proxy, mux, err := NewMCPProxy(l, stubMetrics{}, noopTracer, NewPBKDF2AesGcmSessionCrypto("test", 100), nil)