feat: client response expressions

router-tests/header_propagation_test.go

@@ -2,6 +2,7 @@ package integration
 
 import (
 	"encoding/json"
+	"go.uber.org/zap/zapcore"
 	"net/http"
 	"strings"
 	"testing"
@@ -1034,4 +1035,337 @@ func TestHeaderPropagation(t *testing.T) {
 			})
 		})
 	})
+
+	t.Run("Client Header Rules", func(t *testing.T) {
+		t.Parallel()
+
+		t.Run("should set client headers from static expressions", func(t *testing.T) {
+			t.Parallel()
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(config.HeaderRules{
+						Client: []*config.ClientHeaderRule{
+							{
+								Name:       "X-Static-Header",
+								Expression: `"static-value"`,
+							},
+							{
+								Name:       "X-Another-Header",
+								Expression: `"another-value"`,
+							},
+						},
+					}),
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Query: queryEmployeeWithNoHobby,
+				})
+				require.Equal(t, "static-value", res.Response.Header.Get("X-Static-Header"))
+				require.Equal(t, "another-value", res.Response.Header.Get("X-Another-Header"))
+			})
+		})
+
+		t.Run("should set client headers from request headers", func(t *testing.T) {
+			t.Parallel()
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(config.HeaderRules{
+						Client: []*config.ClientHeaderRule{
+							{
+								Name:       "X-Echo-Header",
+								Expression: `request.header.Get("X-Custom-Input")`,
+							},
+						},
+					}),
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{
+					Query: queryEmployeeWithNoHobby,
+					Header: map[string][]string{
+						"X-Custom-Input": {"input-value"},
+					},
+				})
+				require.NoError(t, err)
+				require.Equal(t, "input-value", res.Response.Header.Get("X-Echo-Header"))
+			})
+		})
+
+		t.Run("should work alongside response header propagation", func(t *testing.T) {
+			t.Parallel()
+
+			t.Run("when there is a separate header", func(t *testing.T) {
+				t.Parallel()
+
+				testenv.Run(t, &testenv.Config{
+					RouterOptions: []core.Option{
+						core.WithHeaderRules(config.HeaderRules{
+							All: &config.GlobalHeaderRule{
+								Response: []*config.ResponseHeaderRule{
+									{
+										Operation: config.HeaderRuleOperationPropagate,
+										Named:     "X-Custom-Header",
+										Algorithm: config.ResponseHeaderRuleAlgorithmFirstWrite,
+									},
+								},
+							},
+							Client: []*config.ClientHeaderRule{
+								{
+									Name:       "X-Client-Header",
+									Expression: `"client-value"`,
+								},
+							},
+						}),
+					},
+					Subgraphs: subgraphsPropagateCustomHeader,
+				}, func(t *testing.T, xEnv *testenv.Environment) {
+					res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+						Query: queryEmployeeWithHobby,
+					})
+					// Check that both client header and propagated response header are present
+					require.Equal(t, "client-value", res.Response.Header.Get("X-Client-Header"))
+					require.Equal(t, employeeVal, res.Response.Header.Get("X-Custom-Header"))
+				})
+			})
+
+			t.Run("when the same header is in use ensure client header overrides", func(t *testing.T) {
+				t.Parallel()
+
+				testenv.Run(t, &testenv.Config{
+					RouterOptions: []core.Option{
+						core.WithHeaderRules(config.HeaderRules{
+							All: &config.GlobalHeaderRule{
+								Response: []*config.ResponseHeaderRule{
+									{
+										Operation: config.HeaderRuleOperationPropagate,
+										Named:     "X-Client-Header",
+										Algorithm: config.ResponseHeaderRuleAlgorithmFirstWrite,
+									},
+								},
+							},
+							Client: []*config.ClientHeaderRule{
+								{
+									Name:       "X-Client-Header",
+									Expression: `"client-value"`,
+								},
+							},
+						}),
+					},
+					Subgraphs: subgraphsPropagateCustomHeader,
+				}, func(t *testing.T, xEnv *testenv.Environment) {
+					res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+						Query: queryEmployeeWithHobby,
+					})
+					require.Equal(t, "client-value", res.Response.Header.Get("X-Client-Header"))
+				})
+			})
+		})
+
+		t.Run("should work alongside request header propagation", func(t *testing.T) {
+			t.Parallel()
+
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(config.HeaderRules{
+						All: &config.GlobalHeaderRule{
+							Request: []*config.RequestHeaderRule{
+								{
+									Operation: config.HeaderRuleOperationPropagate,
+									Named:     "X-Request-Header",
+								},
+							},
+						},
+						Client: []*config.ClientHeaderRule{
+							{
+								Name:       "X-Client-Header",
+								Expression: `request.header.Get("X-Request-Header")`,
+							},
+						},
+					}),
+				},
+				Subgraphs: testenv.SubgraphsConfig{
+					Employees: testenv.SubgraphConfig{
+						Middleware: func(handler http.Handler) http.Handler {
+							return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+								// Verify the request header was propagated to the subgraph
+								require.Equal(t, "request-value", r.Header.Get("X-Request-Header"))
+								handler.ServeHTTP(w, r)
+							})
+						},
+					},
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{
+					Query: queryEmployeeWithNoHobby,
+					Header: map[string][]string{
+						"X-Request-Header": {"request-value"},
+					},
+				})
+				require.NoError(t, err)
+				require.Equal(t, "request-value", res.Response.Header.Get("X-Client-Header"))
+			})
+		})
+
+		t.Run("should work alongside both request and response header propagation", func(t *testing.T) {
+			t.Parallel()
+
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(config.HeaderRules{
+						All: &config.GlobalHeaderRule{
+							Request: []*config.RequestHeaderRule{
+								{
+									Operation: config.HeaderRuleOperationPropagate,
+									Named:     "X-Request-Header",
+								},
+							},
+							Response: []*config.ResponseHeaderRule{
+								{
+									Operation: config.HeaderRuleOperationPropagate,
+									Named:     "X-Verification",
+									Algorithm: config.ResponseHeaderRuleAlgorithmFirstWrite,
+								},
+							},
+						},
+						Client: []*config.ClientHeaderRule{
+							{
+								Name:       "X-Client-Header",
+								Expression: `request.header.Get("X-Request-Header")`,
+							},
+						},
+					}),
+				},
+				Subgraphs: testenv.SubgraphsConfig{
+					Employees: testenv.SubgraphConfig{
+						Middleware: func(handler http.Handler) http.Handler {
+							return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+								// Verify the request header was propagated to the subgraph
+								if r.Header.Get("X-Request-Header") == "request-value" {
+									w.Header().Set("X-Verification", "passed")
+								}
+								handler.ServeHTTP(w, r)
+							})
+						},
+					},
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res, err := xEnv.MakeGraphQLRequest(testenv.GraphQLRequest{
+					Query: queryEmployeeWithNoHobby,
+					Header: map[string][]string{
+						"X-Request-Header": {"request-value"},
+					},
+				})
+				require.NoError(t, err)
+				require.Equal(t, "passed", res.Response.Header.Get("X-Verification"))
+				require.Equal(t, "request-value", res.Response.Header.Get("X-Client-Header"))
+			})
+		})
+
+		t.Run("should ignore rules that resolve to empty string", func(t *testing.T) {
+			t.Parallel()
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(config.HeaderRules{
+						Client: []*config.ClientHeaderRule{
+							{
+								Name:       "X-Empty-Header",
+								Expression: `""`,
+							},
+							{
+								Name:       "X-Missing-Header",
+								Expression: `request.header.Get("X-Does-Not-Exist")`,
+							},
+							{
+								Name:       "X-Valid-Header",
+								Expression: `"valid-value"`,
+							},
+						},
+					}),
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Query: queryEmployeeWithNoHobby,
+				})
+				// Empty headers should not be set
+				require.Equal(t, "", res.Response.Header.Get("X-Empty-Header"))
+				require.Equal(t, "", res.Response.Header.Get("X-Missing-Header"))
+				// Valid header should be set
+				require.Equal(t, "valid-value", res.Response.Header.Get("X-Valid-Header"))
+			})
+		})
+
+		t.Run("should work with errors in response", func(t *testing.T) {
+			t.Parallel()
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(config.HeaderRules{
+						Client: []*config.ClientHeaderRule{
+							{
+								Name:       "X-Client-Header",
+								Expression: `"client-value"`,
+							},
+							{
+								Name:       "X-Error-Header",
+								Expression: `request.error != nil ? "error" : "success"`,
+							},
+						},
+					}),
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Query: `{ employee(id: 1) { id rootFieldThrowsError } }`,
+				})
+				// Client header should still be set even with errors
+				require.Equal(t, "client-value", res.Response.Header.Get("X-Client-Header"))
+				require.Equal(t, "error", res.Response.Header.Get("X-Error-Header"))
+			})
+		})
+
+		t.Run("should log errors (but not error out) when client header rule evaluation fails at runtime", func(t *testing.T) {
+			t.Parallel()
+
+			testenv.Run(t, &testenv.Config{
+				LogObservation: testenv.LogObservationConfig{
+					Enabled:  true,
+					LogLevel: zapcore.ErrorLevel,
+				},
+				RouterOptions: []core.Option{
+					core.WithHeaderRules(config.HeaderRules{
+						Client: []*config.ClientHeaderRule{
+							{
+								Name:       "X-Valid-Header",
+								Expression: `"valid-value"`,
+							},
+							{
+								Name:       "X-Invalid-Header",
+								Expression: `string(int("a"))`,
+							},
+						},
+					}),
+				},
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Query: queryEmployeeWithNoHobby,
+				})
+
+				require.Equal(t, "valid-value", res.Response.Header.Get("X-Valid-Header"))
+
+				_, headerExists := res.Response.Header["X-Invalid-Header"]
+				require.False(t, headerExists)
+
+				require.Equal(t, http.StatusOK, res.Response.StatusCode)
+				require.Contains(t, res.Body, `"data"`)
+
+				logs := xEnv.Observer()
+				require.NotNil(t, logs)
+
+				errorLogs := logs.FilterMessage("Failed to apply client response header rules").All()
+				require.Len(t, errorLogs, 1)
+
+				errorLog := errorLogs[0]
+				require.Equal(t, zapcore.ErrorLevel, errorLog.Level)
+				require.Equal(t, "Failed to apply client response header rules", errorLog.Message)
+				require.NotEmpty(t, errorLog.Context)
+			})
+		})
+	})
 }

router/core/graph_server.go

@@ -1389,6 +1389,7 @@ func (s *graphServer) buildGraphMux(
 		Log:                             s.logger,
 		EnableCacheResponseHeaders:      s.engineExecutionConfiguration.Debug.EnableCacheResponseHeaders,
 		EnableResponseHeaderPropagation: s.headerRules != nil,
+		HeaderPropagation:               s.headerPropagation,
 		EngineStats:                     s.engineStats,
 		TracerProvider:                  s.tracerProvider,
 		Authorizer:                      NewCosmoAuthorizer(authorizerOptions),

router/core/graphql_handler.go

@@ -79,6 +79,7 @@ type HandlerOptions struct {
 
 	EnableCacheResponseHeaders      bool
 	EnableResponseHeaderPropagation bool
+	HeaderPropagation               *HeaderPropagation
 
 	ApolloSubscriptionMultipartPrintBoundary bool
 }
@@ -92,6 +93,7 @@ func NewGraphQLHandler(opts HandlerOptions) *GraphQLHandler {
 		executor:                                 opts.Executor,
 		enableCacheResponseHeaders:               opts.EnableCacheResponseHeaders,
 		enableResponseHeaderPropagation:          opts.EnableResponseHeaderPropagation,
+		headerPropagation:                        opts.HeaderPropagation,
 		engineStats:                              opts.EngineStats,
 		tracer:                                   tracer,
 		authorizer:                               opts.Authorizer,
@@ -128,6 +130,7 @@ type GraphQLHandler struct {
 
 	enableCacheResponseHeaders      bool
 	enableResponseHeaderPropagation bool
+	headerPropagation               *HeaderPropagation
 
 	apolloSubscriptionMultipartPrintBoundary bool
 }
@@ -190,6 +193,13 @@ func (h *GraphQLHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			w.Header().Set("Cache-Control", "no-store, no-cache, must-revalidate")
 		}
 
+		// Apply any final client response header rules
+		if h.headerPropagation != nil {
+			if err := h.headerPropagation.ApplyClientResponseHeaderRules(w, reqCtx); err != nil {
+				reqCtx.logger.Error("Failed to apply client response header rules", zap.Error(err))
+			}
+		}
+
 		// Write contents of buf to the header propagation writer
 		hpw := HeaderPropagationWriter(w, resolveCtx.Context())
 		_, err = respBuf.WriteTo(hpw)