feat(router): disallow inline argument values (ENG-9586)

router/core/context.go

@@ -651,6 +651,8 @@ type operationContext struct {
 	validationTime    time.Duration
 	planningTime      time.Duration
 	normalizationTime time.Duration
+
+	inlineArgumentsAnnotation []byte
 }
 
 func (o *operationContext) Variables() *astjson.Value {

router/core/errors.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"io"
 	"net"
 	"net/http"
 
@@ -47,19 +48,21 @@ type (
 	}
 
 	Extensions struct {
-		RateLimit     json.RawMessage `json:"rateLimit,omitempty"`
-		Authorization json.RawMessage `json:"authorization,omitempty"`
-		Trace         json.RawMessage `json:"trace,omitempty"`
-		StatusCode    int             `json:"statusCode,omitempty"`
-		Code          string          `json:"code,omitempty"`
+		RateLimit       json.RawMessage `json:"rateLimit,omitempty"`
+		Authorization   json.RawMessage `json:"authorization,omitempty"`
+		Trace           json.RawMessage `json:"trace,omitempty"`
+		InlineArguments json.RawMessage `json:"inlineArguments,omitempty"`
+		StatusCode      int             `json:"statusCode,omitempty"`
+		Code            string          `json:"code,omitempty"`
 	}
 )
 
 const (
-	ExtCodeErrPersistedQueryNotFound        = "PERSISTED_QUERY_NOT_FOUND"
-	ExtCodeErrErrorRequestCanceled          = "REQUEST_CANCELED"
-	ExtCodeErrBatchSizeExceeded             = "BATCH_LIMIT_EXCEEDED"
-	ExtCodeErrBatchSubscriptionsUnsupported = "BATCHING_SUBSCRIPTION_UNSUPPORTED"
+	ExtCodeErrPersistedQueryNotFound         = "PERSISTED_QUERY_NOT_FOUND"
+	ExtCodeErrErrorRequestCanceled           = "REQUEST_CANCELED"
+	ExtCodeErrBatchSizeExceeded              = "BATCH_LIMIT_EXCEEDED"
+	ExtCodeErrBatchSubscriptionsUnsupported  = "BATCHING_SUBSCRIPTION_UNSUPPORTED"
+	ExtCodeErrInlineArgumentValuesNotAllowed = "INLINE_ARGUMENT_VALUES_NOT_ALLOWED"
 )
 
 // isTerminalSubscriptionError reports whether the given error, when surfaced
@@ -218,112 +221,115 @@ func propagateSubgraphErrors(ctx *resolve.Context) {
 	}
 }
 
-// writeRequestErrorsParams contains parameters for writing request errors to the response.
-type writeRequestErrorsParams struct {
-	request           *http.Request
-	writer            http.ResponseWriter
-	statusCode        int
-	requestErrors     graphqlerrors.RequestErrors
-	logger            *zap.Logger
-	headerPropagation *HeaderPropagation
-}
-
-// writeRequestErrors writes the given request errors to the http.ResponseWriter.
-// It accepts a graphqlerrors.RequestErrors object and writes it to the response based on the GraphQL spec.
-func writeRequestErrors(params writeRequestErrorsParams) {
-	if params.requestErrors == nil {
-		return
-	}
-
-	params.writer.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate")
+// writeRawErrorBody routes a pre-built JSON error body ({"errors":[...]}) through the correct
+// transport framing (SSE, multipart, or plain JSON). All pre-execution error responses funnel
+// through here; use it directly when the error payload cannot be expressed as
+// graphqlerrors.RequestErrors (e.g. when the extensions object has a custom nested shape).
+func writeRawErrorBody(r *http.Request, w http.ResponseWriter, statusCode int, body []byte, logger *zap.Logger, headerPropagation *HeaderPropagation) {
+	w.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate")
 
 	// According to the tests requestContext can be nil (when called from module WriteResponseError)
 	// As such we have coded this condition defensively to be safe
-	requestContext := getRequestContext(params.request.Context())
+	requestContext := getRequestContext(r.Context())
 	isSubscription := requestContext != nil && requestContext.operation != nil && requestContext.operation.opType == "subscription"
 
 	// We only want to apply header propagation for non-subscription operations
 	// In certain cases the requestContext can be nil, e.g.:- when called from the batch handler
-	if params.headerPropagation != nil && requestContext != nil && !isSubscription {
-		err := params.headerPropagation.ApplyRouterResponseHeaderRules(params.writer, requestContext)
-		if err != nil && params.logger != nil {
-			params.logger.Error("Failed to apply router response header rules on error cases", zap.Error(err))
+	if headerPropagation != nil && requestContext != nil && !isSubscription {
+		if err := headerPropagation.ApplyRouterResponseHeaderRules(w, requestContext); err != nil && logger != nil {
+			logger.Error("Failed to apply router response header rules on error cases", zap.Error(err))
 		}
 	}
 
-	wgRequestParams := NegotiateSubscriptionParams(params.request, !isSubscription)
+	wgRequestParams := NegotiateSubscriptionParams(r, !isSubscription)
 
-	// Is subscription
 	if wgRequestParams.UseSse || wgRequestParams.UseMultipart {
-		setSubscriptionHeaders(wgRequestParams, params.request, params.writer)
+		setSubscriptionHeaders(wgRequestParams, r, w)
+	} else {
+		w.Header().Set("Content-Type", "application/json; charset=utf-8")
+	}
+	if statusCode != 0 {
+		w.WriteHeader(statusCode)
+	}
 
-		if params.statusCode != 0 {
-			params.writer.WriteHeader(params.statusCode)
+	switch {
+	case wgRequestParams.UseSse:
+		if err := writeChunks(w, []byte(GetWriterPrefix(true, false, false)), body, []byte("\n\n")); err != nil {
+			logWriteResponseError(logger, err)
 		}
-
-		if wgRequestParams.UseSse {
-			_, err := params.writer.Write([]byte("event: next\ndata: "))
-			if err != nil {
-				if params.logger != nil {
-					if rErrors.IsBrokenPipe(err) {
-						params.logger.Warn("Broken pipe, error writing response", zap.Error(err))
-						return
-					}
-					params.logger.Error("Error writing response", zap.Error(err))
-				}
-				return
-			}
-		} else if wgRequestParams.UseMultipart {
-			// Handle multipart error response
-			if err := writeMultipartError(params.writer, params.requestErrors, isSubscription); err != nil {
-				if params.logger != nil {
-					params.logger.Error("error writing multipart response", zap.Error(err))
-				}
-			}
-			return
+	case wgRequestParams.UseMultipart:
+		if err := writeMultipartError(w, body, isSubscription); err != nil && logger != nil {
+			logger.Error("error writing multipart response", zap.Error(err))
 		}
-	} else {
-		// Regular request
-		params.writer.Header().Set("Content-Type", "application/json; charset=utf-8")
-		if params.statusCode != 0 {
-			params.writer.WriteHeader(params.statusCode)
+	default:
+		if _, err := w.Write(body); err != nil {
+			logWriteResponseError(logger, err)
 		}
 	}
+}
 
-	if _, err := params.requestErrors.WriteResponse(params.writer); err != nil {
-		if params.logger != nil {
-			if rErrors.IsBrokenPipe(err) {
-				params.logger.Warn("Broken pipe, error writing response", zap.Error(err))
-				return
-			}
-			params.logger.Error("Error writing response", zap.Error(err))
+// writeChunks writes each chunk in order, stopping at the first write error.
+func writeChunks(w io.Writer, chunks ...[]byte) error {
+	for _, chunk := range chunks {
+		if _, err := w.Write(chunk); err != nil {
+			return err
 		}
 	}
+	return nil
 }
 
-// writeMultipartError writes the error response in a multipart format with proper boundaries and headers.
-func writeMultipartError(
-	w http.ResponseWriter,
-	requestErrors graphqlerrors.RequestErrors,
-	isSubscription bool,
-) error {
-	// Start with the multipart boundary
-	prefix := GetWriterPrefix(false, true, true)
-	if _, err := w.Write([]byte(prefix)); err != nil {
-		return err
+func logWriteResponseError(logger *zap.Logger, err error) {
+	if logger == nil {
+		return
 	}
+	if rErrors.IsBrokenPipe(err) {
+		logger.Warn("Broken pipe, error writing response", zap.Error(err))
+		return
+	}
+	logger.Error("Error writing response", zap.Error(err))
+}
 
-	// Write the actual error payload
-	response := graphqlerrors.Response{
-		Errors: requestErrors,
+// writeRequestErrorsParams contains parameters for writing request errors to the response.
+type writeRequestErrorsParams struct {
+	request           *http.Request
+	writer            http.ResponseWriter
+	statusCode        int
+	requestErrors     graphqlerrors.RequestErrors
+	logger            *zap.Logger
+	headerPropagation *HeaderPropagation
+}
+
+// writeRequestErrors writes the given request errors to the http.ResponseWriter.
+// It accepts a graphqlerrors.RequestErrors object and writes it to the response based on the GraphQL spec.
+func writeRequestErrors(params writeRequestErrorsParams) {
+	if params.requestErrors == nil {
+		return
 	}
 
-	responseBytes, err := response.Marshal()
+	response := graphqlerrors.Response{
+		Errors: params.requestErrors,
+	}
+	body, err := response.Marshal()
 	if err != nil {
+		if params.logger != nil {
+			params.logger.Error("Error marshalling request errors", zap.Error(err))
+		}
+		params.writer.WriteHeader(http.StatusInternalServerError)
+		return
+	}
+
+	writeRawErrorBody(params.request, params.writer, params.statusCode, body, params.logger, params.headerPropagation)
+}
+
+// writeMultipartError writes the error response in a multipart format with proper boundaries and headers.
+func writeMultipartError(w http.ResponseWriter, body []byte, isSubscription bool) error {
+	// Start with the multipart boundary
+	prefix := GetWriterPrefix(false, true, true)
+	if _, err := w.Write([]byte(prefix)); err != nil {
 		return err
 	}
 
-	resp, err := wrapMultipartMessage(responseBytes, isSubscription)
+	resp, err := wrapMultipartMessage(body, isSubscription)
 	if err != nil {
 		return err
 	}
@@ -333,7 +339,7 @@ func writeMultipartError(
 	// multipart chunks correctly. With this fix here (and in a few other places) the clients are now working.
 	resp = append(resp, []byte("\r\n--graphql--")...)
 
-	if _, err := w.Write([]byte(resp)); err != nil {
+	if _, err := w.Write(resp); err != nil {
 		return err
 	}
 
@@ -364,7 +370,10 @@ func writeOperationError(r *http.Request, w http.ResponseWriter, requestLogger *
 	var reportErr ReportError
 	var httpErr HttpError
 	var poNotFoundErr *persistedoperation.PersistentOperationNotFoundError
+	var inlineArgsErr *inlineArgumentsError
 	switch {
+	case errors.As(err, &inlineArgsErr):
+		writeInlineArgumentsError(r, w, inlineArgsErr, requestLogger, propagation)
 	case errors.Is(err, context.Canceled):
 		newErr := NewHttpGraphqlError("request canceled", ExtCodeErrErrorRequestCanceled, http.StatusOK)
 		writeRequestErrors(writeRequestErrorsParams{

router/core/graph_server.go

@@ -1817,6 +1817,8 @@ func (s *graphServer) buildGraphMux(
 		return nil, fmt.Errorf("failed to create operation blocker: %w", err)
 	}
 
+	inlineArgumentsChecker := NewInlineArgumentsChecker(s.securityConfiguration.DisallowInlineArguments)
+
 	graphqlPreHandler := NewPreHandler(&PreHandlerOptions{
 		Logger:                                 s.logger,
 		Executor:                               executor,
@@ -1825,6 +1827,7 @@ func (s *graphServer) buildGraphMux(
 		Planner:                                operationPlanner,
 		AccessController:                       s.accessController,
 		OperationBlocker:                       operationBlocker,
+		InlineArgumentsChecker:                 inlineArgumentsChecker,
 		RouterPublicKey:                        s.publicKey,
 		EnableRequestTracing:                   s.engineExecutionConfiguration.EnableRequestTracing,
 		ForceUnauthenticatedRequestTracing:     s.engineExecutionConfiguration.ForceUnauthenticatedRequestTracing,
@@ -1861,6 +1864,7 @@ func (s *graphServer) buildGraphMux(
 		wsMiddleware := NewWebsocketMiddleware(graphMuxCtx, WebsocketMiddlewareOptions{
 			OperationProcessor:        operationProcessor,
 			OperationBlocker:          operationBlocker,
+			InlineArgumentsChecker:    inlineArgumentsChecker,
 			Planner:                   operationPlanner,
 			GraphQLHandler:            graphqlHandler,
 			PreHandler:                graphqlPreHandler,

router/core/graphql_handler.go

@@ -10,11 +10,13 @@ import (
 	"net/http"
 	"strconv"
 	"strings"
+	"sync"
+
+	"github.com/buger/jsonparser"
 
 	"go.opentelemetry.io/otel/trace"
 	"go.uber.org/zap"
 
-	rErrors "github.com/wundergraph/cosmo/router/internal/errors"
 	"github.com/wundergraph/cosmo/router/pkg/config"
 	rotel "github.com/wundergraph/cosmo/router/pkg/otel"
 	"github.com/wundergraph/cosmo/router/pkg/statistics"
@@ -33,6 +35,16 @@ var (
 	errOperationPlanUnsupported = errors.New("unsupported operation plan")
 )
 
+// inlineArgsCaptureBufPool holds buffers used to capture the resolved response body
+// when the warn-mode inlineArguments extension needs to be injected before writing.
+var inlineArgsCaptureBufPool = sync.Pool{
+	New: func() any { return &bytes.Buffer{} },
+}
+
+// inlineArgsCaptureBufMaxRetainedCap bounds the capacity of buffers returned to the
+// pool so that a single very large response does not pin its memory indefinitely.
+const inlineArgsCaptureBufMaxRetainedCap = 1 << 20
+
 const (
 	ExecutionPlanCacheHeader          = "X-WG-Execution-Plan-Cache"
 	PersistedOperationCacheHeader     = "X-WG-Persisted-Operation-Cache"
@@ -243,14 +255,37 @@ func (h *GraphQLHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			}
 		}
 
-		info, err := h.executor.Resolver.ArenaResolveGraphQLResponse(resolveCtx, p.Response, hpw)
+		// Capture the response body when warn-mode inline args were detected so we can inject the extension.
+		resolveWriter := io.Writer(hpw)
+		var captureBuf *bytes.Buffer
+		if len(reqCtx.operation.inlineArgumentsAnnotation) > 0 {
+			captureBuf = inlineArgsCaptureBufPool.Get().(*bytes.Buffer)
+			captureBuf.Reset()
+			defer func() {
+				if captureBuf.Cap() <= inlineArgsCaptureBufMaxRetainedCap {
+					inlineArgsCaptureBufPool.Put(captureBuf)
+				}
+			}()
+			resolveWriter = captureBuf
+		}
+
+		info, err := h.executor.Resolver.ArenaResolveGraphQLResponse(resolveCtx, p.Response, resolveWriter)
 		reqCtx.dataSourceNames = getSubgraphNames(p.Response.DataSources)
 		if err != nil {
 			trackFinalResponseError(resolveCtx.Context(), err)
 			h.WriteError(resolveCtx, err, p.Response, w)
 			return
 		}
 
+		if captureBuf != nil {
+			body, setErr := jsonparser.Set(captureBuf.Bytes(), reqCtx.operation.inlineArgumentsAnnotation, "extensions", "inlineArguments")
+			if setErr != nil {
+				reqCtx.logger.Error("failed to inject inlineArguments annotation into response", zap.Error(setErr))
+				body = captureBuf.Bytes()
+			}
+			_, _ = hpw.Write(body)
+		}
+
 		// Compute actual cost for metrics/telemetry if not already set by the header callback
 		if !reqCtx.operation.costActualSet && resolveCtx.TypeNameStats != nil &&
 			reqCtx.operation.preparedPlan != nil && reqCtx.operation.preparedPlan.preparedPlan != nil {
@@ -547,13 +582,8 @@ func (h *GraphQLHandler) writeError(ctx *resolve.Context, err error, res *resolv
 		return
 	}
 
-	err = json.NewEncoder(w).Encode(response)
-	if err != nil {
-		if rErrors.IsBrokenPipe(err) {
-			requestLogger.Warn("Broken pipe, unable to write error response", zap.Error(err))
-		} else {
-			requestLogger.Error("Unable to write error response", zap.Error(err))
-		}
+	if encErr := json.NewEncoder(w).Encode(response); encErr != nil {
+		logWriteResponseError(requestLogger, encErr)
 	}
 }