feat: implement set header (wundergraph#1196)

router-tests/header_propagation_test.go

@@ -1,15 +1,14 @@
 package integration
 
 import (
-	"net/http"
-	"strings"
-	"testing"
-	"time"
-
 	"github.com/stretchr/testify/require"
 	"github.com/wundergraph/cosmo/router-tests/testenv"
 	"github.com/wundergraph/cosmo/router/core"
 	"github.com/wundergraph/cosmo/router/pkg/config"
+	"net/http"
+	"strings"
+	"testing"
+	"time"
 )
 
 func TestHeaderPropagation(t *testing.T) {

router-tests/header_set_test.go

@@ -0,0 +1,144 @@
+package integration
+
+import (
+	"fmt"
+	"github.com/stretchr/testify/require"
+	"github.com/wundergraph/cosmo/router-tests/testenv"
+	"github.com/wundergraph/cosmo/router/core"
+	"github.com/wundergraph/cosmo/router/pkg/config"
+	"net/http"
+	"strings"
+	"testing"
+)
+
+func TestHeaderSet(t *testing.T) {
+	const (
+		customHeader = "X-Custom-Header"
+		employeeVal  = "employee-value"
+		hobbyVal     = "hobby-value"
+	)
+
+	const queryEmployeeWithHobby = `{
+	  employee(id: 1) {
+		id
+		hobbies {
+		  ... on Gaming {
+			name
+		  }
+		}
+	  }
+	}`
+
+	t.Run("RequestSet", func(t *testing.T) {
+		getRule := func(name, val string) *config.RequestHeaderRule {
+			rule := &config.RequestHeaderRule{
+				Operation: config.HeaderRuleOperationSet,
+				Name:      name,
+				Value:     val,
+			}
+			return rule
+		}
+
+		global := func(name, defaultVal string) []core.Option {
+			return []core.Option{
+				core.WithHeaderRules(config.HeaderRules{
+					All: &config.GlobalHeaderRule{
+						Request: []*config.RequestHeaderRule{
+							getRule(name, defaultVal),
+						},
+					},
+				}),
+			}
+		}
+
+		t.Run("global request rule sets header", func(t *testing.T) {
+			t.Parallel()
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: global(customHeader, employeeVal),
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Header: http.Header{},
+					Query:  fmt.Sprintf(`query { headerValue(name:"%s") }`, customHeader),
+				})
+				require.Equal(t, fmt.Sprintf(`{"data":{"headerValue":"%s"}}`, employeeVal), res.Body)
+			})
+		})
+	})
+
+	t.Run("ResponseSet", func(t *testing.T) {
+		getRule := func(name, val string) *config.ResponseHeaderRule {
+			rule := &config.ResponseHeaderRule{
+				Operation: config.HeaderRuleOperationSet,
+				Name:      name,
+				Value:     val,
+			}
+			return rule
+		}
+
+		global := func(name, defaultVal string) []core.Option {
+			return []core.Option{
+				core.WithHeaderRules(config.HeaderRules{
+					All: &config.GlobalHeaderRule{
+						Response: []*config.ResponseHeaderRule{
+							getRule(name, defaultVal),
+						},
+					},
+				}),
+			}
+		}
+
+		partial := func(name, defaultVal string) []core.Option {
+			return []core.Option{
+				core.WithHeaderRules(config.HeaderRules{
+					Subgraphs: map[string]*config.GlobalHeaderRule{
+						"employees": {
+							Response: []*config.ResponseHeaderRule{
+								getRule(name, defaultVal),
+							},
+						},
+					},
+				}),
+			}
+		}
+
+		t.Run("no set", func(t *testing.T) {
+			t.Parallel()
+			testenv.Run(t, &testenv.Config{}, func(t *testing.T, xEnv *testenv.Environment) {
+				res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Query: queryEmployeeWithHobby,
+				})
+				ch := strings.Join(res.Response.Header.Values(customHeader), ",")
+				require.Equal(t, "", ch)
+				require.Equal(t, `{"data":{"employee":{"id":1,"hobbies":[{},{"name":"Counter Strike"},{},{},{}]}}}`, res.Body)
+			})
+		})
+
+		t.Run("global set works", func(t *testing.T) {
+			t.Parallel()
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: global(customHeader, hobbyVal),
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Query: queryEmployeeWithHobby,
+				})
+				ch := strings.Join(res.Response.Header.Values(customHeader), ",")
+				require.Equal(t, hobbyVal, ch)
+				require.Equal(t, `{"data":{"employee":{"id":1,"hobbies":[{},{"name":"Counter Strike"},{},{},{}]}}}`, res.Body)
+			})
+		})
+
+		t.Run("subgraph set works", func(t *testing.T) {
+			t.Parallel()
+			testenv.Run(t, &testenv.Config{
+				RouterOptions: partial(customHeader, employeeVal),
+			}, func(t *testing.T, xEnv *testenv.Environment) {
+				res := xEnv.MakeGraphQLRequestOK(testenv.GraphQLRequest{
+					Query: queryEmployeeWithHobby,
+				})
+				ch := strings.Join(res.Response.Header.Values(customHeader), ",")
+				require.Equal(t, employeeVal, ch)
+				require.Equal(t, `{"data":{"employee":{"id":1,"hobbies":[{},{"name":"Counter Strike"},{},{},{}]}}}`, res.Body)
+			})
+		})
+	})
+}

router/core/header_rule_engine.go

@@ -152,6 +152,7 @@ func (hf *HeaderPropagation) getAllRules() ([]*config.RequestHeaderRule, []*conf
 
 func (hf *HeaderPropagation) processRule(rule config.HeaderRule, index int) error {
 	switch rule.GetOperation() {
+	case config.HeaderRuleOperationSet:
 	case config.HeaderRuleOperationPropagate:
 		if rule.GetMatching() != "" {
 			regex, err := regexp.Compile(rule.GetMatching())
@@ -236,6 +237,11 @@ func (h *HeaderPropagation) OnOriginResponse(resp *http.Response, ctx RequestCon
 }
 
 func (h *HeaderPropagation) applyResponseRule(propagation *responseHeaderPropagation, res *http.Response, rule *config.ResponseHeaderRule) {
+	if rule.Operation == config.HeaderRuleOperationSet {
+		propagation.header.Set(rule.Name, rule.Value)
+		return
+	}
+
 	if rule.Operation != config.HeaderRuleOperationPropagate {
 		return
 	}
@@ -292,6 +298,11 @@ func (h *HeaderPropagation) applyResponseRuleKeyValue(res *http.Response, propag
 }
 
 func (h *HeaderPropagation) applyRequestRule(ctx RequestContext, request *http.Request, rule *config.RequestHeaderRule) {
+	if rule.Operation == config.HeaderRuleOperationSet {
+		request.Header.Set(rule.Name, rule.Value)
+		return
+	}
+
 	if rule.Operation != config.HeaderRuleOperationPropagate {
 		return
 	}
@@ -535,6 +546,11 @@ func FetchURLRules(rules *config.HeaderRules, subgraphs []*nodev1.Subgraph, rout
 func PropagatedHeaders(rules []*config.RequestHeaderRule) (headerNames []string, headerNameRegexps []*regexp.Regexp, err error) {
 	for _, rule := range rules {
 		switch rule.Operation {
+		case config.HeaderRuleOperationSet:
+			if rule.Name == "" || rule.Value == "" {
+				return nil, nil, fmt.Errorf("invalid header set rule %+v, no header name/value combination", rule)
+			}
+			headerNames = append(headerNames, rule.Name)
 		case config.HeaderRuleOperationPropagate:
 			if rule.Matching != "" {
 				re, err := regexp.Compile(rule.Matching)

router/pkg/config/config.go

@@ -178,6 +178,7 @@ type HeaderRuleOperation string
 
 const (
 	HeaderRuleOperationPropagate HeaderRuleOperation = "propagate"
+	HeaderRuleOperationSet       HeaderRuleOperation = "set"
 )
 
 type HeaderRule interface {
@@ -188,6 +189,7 @@ type HeaderRule interface {
 type RequestHeaderRule struct {
 	// Operation describes the header operation to perform e.g. "propagate"
 	Operation HeaderRuleOperation `yaml:"op"`
+	// Propagate options
 	// Matching is the regex to match the header name against
 	Matching string `yaml:"matching"`
 	// Named is the exact header name to match
@@ -196,6 +198,12 @@ type RequestHeaderRule struct {
 	Rename string `yaml:"rename,omitempty"`
 	// Default is the default value to set if the header is not present
 	Default string `yaml:"default"`
+
+	// Set header options
+	// Name is the name of the header to set
+	Name string `yaml:"name"`
+	// Value is the value of the header to set
+	Value string `yaml:"value"`
 }
 
 func (r *RequestHeaderRule) GetOperation() HeaderRuleOperation {
@@ -232,6 +240,12 @@ type ResponseHeaderRule struct {
 	Default string `yaml:"default"`
 	// Algorithm is the algorithm to use when multiple headers are present
 	Algorithm ResponseHeaderRuleAlgorithm `yaml:"algorithm,omitempty"`
+
+	// Set header options
+	// Name is the name of the header to set
+	Name string `yaml:"name"`
+	// Value is the value of the header to set
+	Value string `yaml:"value"`
 }
 
 func (r *ResponseHeaderRule) GetOperation() HeaderRuleOperation {

router/pkg/config/config.schema.json

@@ -959,13 +959,19 @@
             "request": {
               "type": "array",
               "items": {
-                "$ref": "#/definitions/traffic_shaping_header_rule"
+                "oneOf": [
+                  { "$ref": "#/definitions/traffic_shaping_header_rule" },
+                  { "$ref": "#/definitions/set_header_rule" }
+                ]
               }
             },
             "response": {
               "type": "array",
               "items": {
-                "$ref": "#/definitions/traffic_shaping_header_response_rule"
+                "oneOf": [
+                  { "$ref": "#/definitions/traffic_shaping_header_response_rule" },
+                  { "$ref": "#/definitions/set_header_rule" }
+                ]
               }
             }
           }
@@ -979,13 +985,19 @@
               "request": {
                 "type": "array",
                 "items": {
-                  "$ref": "#/definitions/traffic_shaping_header_rule"
+                  "oneOf": [
+                    { "$ref": "#/definitions/traffic_shaping_header_rule" },
+                    { "$ref": "#/definitions/set_header_rule" }
+                  ]
                 }
               },
               "response": {
                 "type": "array",
                 "items": {
-                  "$ref": "#/definitions/traffic_shaping_header_response_rule"
+                  "oneOf": [
+                    { "$ref": "#/definitions/traffic_shaping_header_response_rule" },
+                    { "$ref": "#/definitions/set_header_rule" }
+                  ]
                 }
               }
             }
@@ -1727,6 +1739,29 @@
           "rename": {}
         }
       }
+    },
+    "set_header_rule": {
+      "type": "object",
+      "description": "The configuration for setting headers. This is used to set specific headers in requests or responses.",
+      "additionalProperties": false,
+      "properties": {
+        "op": {
+          "type": "string",
+          "const": "set",
+          "description": "The 'set' operation is used to set a specific header value."
+        },
+        "name": {
+          "type": "string",
+          "examples": ["X-API-Key"],
+          "description": "The name of the header to set."
+        },
+        "value": {
+          "type": "string",
+          "examples": ["My-Secret-Value"],
+          "description": "The value to set for the header. This can include environment variables."
+        }
+      },
+      "required": ["op", "name", "value"]
     }
   }
 }

router/pkg/config/fixtures/full.yaml

@@ -151,6 +151,9 @@ headers:
       - op: "propagate"
         named: "X-User-Id"
         default: "123"             # Set the value when the header was not set
+      - op: "set"
+        name: "X-API-Key"
+        value: "some-secret"
     response:
       - op: "propagate"
         algorithm: "append"
@@ -164,6 +167,9 @@ headers:
       response:
         - op: "propagate"
           algorithm: "most_restrictive_cache_control"
+        - op: "set"
+          name: "X-Subgraph-Key"
+          value: "some-subgraph-secret"
 
 # Authentication and Authorization
 # See https://cosmo-docs.wundergraph.com/router/authentication-and-authorization for more information