[TT-11997] Backend logic for request_headers_rewrite

gateway/mw_graphql_test.go

@@ -616,6 +616,143 @@ func TestGraphQLMiddleware_EngineMode(t *testing.T) {
 			}...)
 		})
 
+		t.Run("apply request headers rewrite, rule one", func(t *testing.T) {
+			// Rule one:
+			//
+			// If header key/value is defined in request_headers_rewrite and remove
+			// is set to false and client sends a request with the same header key but
+			// different value, the value gets overwritten to the defined value before
+			// hitting the upstream.
+			g.Gw.BuildAndLoadAPI(func(spec *APISpec) {
+				spec.UseKeylessAccess = true
+				spec.GraphQL.Enabled = true
+				spec.GraphQL.ExecutionMode = apidef.GraphQLExecutionModeProxyOnly
+				spec.GraphQL.Version = apidef.GraphQLConfigVersion2
+				spec.GraphQL.Schema = gqlProxyUpstreamSchema
+				spec.GraphQL.Proxy.RequestHeadersRewrite = map[string]apidef.RequestHeadersRewriteConfig{
+					"X-Tyk-Test": {
+						Remove: false,
+						Value:  "value-from-rewrite-config",
+					},
+				}
+				spec.Proxy.ListenPath = "/"
+				spec.Proxy.TargetURL = testGraphQLProxyUpstream
+			})
+
+			request := gql.Request{
+				Query: `{ hello(name: "World") httpMethod }`,
+			}
+
+			_, _ = g.Run(t, []test.TestCase{
+				{
+					Data:   request,
+					Method: http.MethodPost,
+					Headers: map[string]string{
+						"X-Tyk-Key":   "tyk-value",
+						"X-Other-Key": "other-value",
+						"X-Tyk-Test":  "value-from-consumer",
+					},
+					Code:      http.StatusOK,
+					BodyMatch: `{"data":{"hello":"World","httpMethod":"POST"}}`,
+					HeadersMatch: map[string]string{
+						"X-Tyk-Key":   "tyk-value",
+						"X-Other-Key": "other-value",
+						"X-Tyk-Test":  "value-from-rewrite-config",
+					},
+				},
+			}...)
+		})
+
+		t.Run("apply request headers rewrite, rule two", func(t *testing.T) {
+			// Rule two:
+			//
+			// If header key is defined in request_headers_rewrite and remove is set
+			// to true and client sends a request with the same header key but different value,
+			// the headers gets removed completely before hitting the upstream.
+			g.Gw.BuildAndLoadAPI(func(spec *APISpec) {
+				spec.UseKeylessAccess = true
+				spec.GraphQL.Enabled = true
+				spec.GraphQL.ExecutionMode = apidef.GraphQLExecutionModeProxyOnly
+				spec.GraphQL.Version = apidef.GraphQLConfigVersion2
+				spec.GraphQL.Schema = gqlProxyUpstreamSchema
+				spec.GraphQL.Proxy.RequestHeadersRewrite = map[string]apidef.RequestHeadersRewriteConfig{
+					"X-Tyk-Test": {
+						Remove: true,
+						Value:  "value-from-rewrite-config",
+					},
+				}
+				spec.Proxy.ListenPath = "/"
+				spec.Proxy.TargetURL = testGraphQLProxyUpstream
+			})
+
+			request := gql.Request{
+				Query: `{ hello(name: "World") httpMethod }`,
+			}
+
+			_, _ = g.Run(t, []test.TestCase{
+				{
+					Data:   request,
+					Method: http.MethodPost,
+					Headers: map[string]string{
+						"X-Tyk-Key":   "tyk-value",
+						"X-Other-Key": "other-value",
+						"X-Tyk-Test":  "value-from-consumer",
+					},
+					Code:      http.StatusOK,
+					BodyMatch: `{"data":{"hello":"World","httpMethod":"POST"}}`,
+					HeadersMatch: map[string]string{
+						"X-Tyk-Key":   "tyk-value",
+						"X-Other-Key": "other-value",
+					},
+				},
+			}...)
+		})
+
+		t.Run("apply request headers rewrite, rule three", func(t *testing.T) {
+			// Rule three:
+			//
+			// If header key/value is defined in request_headers_rewrite and remove is
+			// set to false and client sends a request that does not have the same header key,
+			// the header key/value gets added before hitting the upstream.
+			g.Gw.BuildAndLoadAPI(func(spec *APISpec) {
+				spec.UseKeylessAccess = true
+				spec.GraphQL.Enabled = true
+				spec.GraphQL.ExecutionMode = apidef.GraphQLExecutionModeProxyOnly
+				spec.GraphQL.Version = apidef.GraphQLConfigVersion2
+				spec.GraphQL.Schema = gqlProxyUpstreamSchema
+				spec.GraphQL.Proxy.RequestHeadersRewrite = map[string]apidef.RequestHeadersRewriteConfig{
+					"X-Tyk-Test": {
+						Remove: false,
+						Value:  "value-from-rewrite-config",
+					},
+				}
+				spec.Proxy.ListenPath = "/"
+				spec.Proxy.TargetURL = testGraphQLProxyUpstream
+			})
+
+			request := gql.Request{
+				Query: `{ hello(name: "World") httpMethod }`,
+			}
+
+			_, _ = g.Run(t, []test.TestCase{
+				{
+					Data:   request,
+					Method: http.MethodPost,
+					Headers: map[string]string{
+						"X-Tyk-Key":   "tyk-value",
+						"X-Other-Key": "other-value",
+					},
+					Code:      http.StatusOK,
+					BodyMatch: `{"data":{"hello":"World","httpMethod":"POST"}}`,
+					HeadersMatch: map[string]string{
+						"X-Tyk-Key":   "tyk-value",
+						"X-Other-Key": "other-value",
+						"X-Tyk-Test":  "value-from-rewrite-config",
+					},
+				},
+			}...)
+		})
+
 		t.Run("proxy-only return errors from upstream", func(t *testing.T) {
 			g.Gw.BuildAndLoadAPI(func(spec *APISpec) {
 				spec.UseKeylessAccess = true

gateway/reverse_proxy.go

@@ -1011,13 +1011,14 @@ func (p *ReverseProxy) handleGraphQL(roundTripper *TykRoundTripper, outreq *http
 	needsEngine := needsGraphQLExecutionEngine(p.TykAPISpec)
 
 	res, hijacked, err = p.TykAPISpec.GraphEngine.HandleReverseProxy(graphengine.ReverseProxyParams{
-		RoundTripper:       roundTripper,
-		ResponseWriter:     w,
-		OutRequest:         outreq,
-		WebSocketUpgrader:  &p.wsUpgrader,
-		NeedsEngine:        needsEngine,
-		IsCORSPreflight:    isCORSPreflight(outreq),
-		IsWebSocketUpgrade: isWebSocketUpgrade,
+		RoundTripper:          roundTripper,
+		ResponseWriter:        w,
+		OutRequest:            outreq,
+		WebSocketUpgrader:     &p.wsUpgrader,
+		NeedsEngine:           needsEngine,
+		IsCORSPreflight:       isCORSPreflight(outreq),
+		IsWebSocketUpgrade:    isWebSocketUpgrade,
+		RequestHeadersRewrite: p.TykAPISpec.GraphQL.Proxy.RequestHeadersRewrite,
 	})
 	if err != nil {
 		return nil, hijacked, err

internal/graphengine/engine.go

@@ -85,13 +85,14 @@ type GranularAccessChecker interface {
 }
 
 type ReverseProxyParams struct {
-	RoundTripper       http.RoundTripper
-	ResponseWriter     http.ResponseWriter
-	OutRequest         *http.Request
-	WebSocketUpgrader  *websocket.Upgrader
-	NeedsEngine        bool
-	IsCORSPreflight    bool
-	IsWebSocketUpgrade bool
+	RoundTripper          http.RoundTripper
+	ResponseWriter        http.ResponseWriter
+	OutRequest            *http.Request
+	WebSocketUpgrader     *websocket.Upgrader
+	NeedsEngine           bool
+	IsCORSPreflight       bool
+	IsWebSocketUpgrade    bool
+	RequestHeadersRewrite map[string]apidef.RequestHeadersRewriteConfig
 }
 
 type ReverseProxyPreHandler interface {

internal/graphengine/graphql_go_tools_v1.go

@@ -490,6 +490,7 @@ func (r *reverseProxyPreHandlerV1) PreHandle(params ReverseProxyParams) (reverse
 		DetermineGraphQLEngineTransportType(r.apiDefinition),
 		params.RoundTripper,
 		r.newReusableBodyReadCloser,
+		params.RequestHeadersRewrite,
 	)
 
 	switch {

internal/graphengine/graphql_go_tools_v2.go

@@ -124,6 +124,7 @@ func (r *reverseProxyPreHandlerV2) PreHandle(params ReverseProxyParams) (reverse
 		DetermineGraphQLEngineTransportType(r.apiDefinition),
 		params.RoundTripper,
 		r.newReusableBodyReadCloser,
+		params.RequestHeadersRewrite,
 	)
 
 	switch {

internal/graphengine/transport.go

@@ -4,6 +4,8 @@ import (
 	"bytes"
 	"io"
 	"net/http"
+
+	"github.com/TykTechnologies/tyk/apidef"
 )
 
 type NewReusableBodyReadCloserFunc func(io.ReadCloser) (io.ReadCloser, error)
@@ -13,13 +15,20 @@ type GraphQLEngineTransport struct {
 	originalTransport         http.RoundTripper
 	transportType             GraphQLEngineTransportType
 	newReusableBodyReadCloser NewReusableBodyReadCloserFunc
+	requestHeadersRewrite     map[string]apidef.RequestHeadersRewriteConfig
 }
 
-func NewGraphQLEngineTransport(transportType GraphQLEngineTransportType, originalTransport http.RoundTripper, newReusableBodyReadCloser NewReusableBodyReadCloserFunc) *GraphQLEngineTransport {
+func NewGraphQLEngineTransport(
+	transportType GraphQLEngineTransportType,
+	originalTransport http.RoundTripper,
+	newReusableBodyReadCloser NewReusableBodyReadCloserFunc,
+	requestHeadersRewrite map[string]apidef.RequestHeadersRewriteConfig,
+) *GraphQLEngineTransport {
 	transport := &GraphQLEngineTransport{
 		originalTransport:         originalTransport,
 		transportType:             transportType,
 		newReusableBodyReadCloser: newReusableBodyReadCloser,
+		requestHeadersRewrite:     requestHeadersRewrite,
 	}
 	return transport
 }
@@ -39,6 +48,7 @@ func (g *GraphQLEngineTransport) RoundTrip(request *http.Request) (res *http.Res
 func (g *GraphQLEngineTransport) handleProxyOnly(proxyOnlyValues *GraphQLProxyOnlyContextValues, request *http.Request) (*http.Response, error) {
 	request.Method = proxyOnlyValues.forwardedRequest.Method
 	g.setProxyOnlyHeaders(proxyOnlyValues, request)
+	g.applyRequestHeadersRewriteRules(request)
 
 	response, err := g.originalTransport.RoundTrip(request)
 	if err != nil {
@@ -69,6 +79,91 @@ func (g *GraphQLEngineTransport) handleProxyOnly(proxyOnlyValues *GraphQLProxyOn
 	return response, err
 }
 
+func (g *GraphQLEngineTransport) applyRequestHeadersRewriteRules(r *http.Request) {
+	if len(g.requestHeadersRewrite) == 0 {
+		// There is no request rewrite rule, quit early.
+		return
+	}
+
+	ruleOne := func(r *http.Request, key string, values []string) bool {
+		// Rule one:
+		//
+		// If header key/value is defined in request_headers_rewrite and remove
+		// is set to false and client sends a request with the same header key but
+		// different value, the value gets overwritten to the defined value before
+		// hitting the upstream.
+
+		rewriteRule, ok := g.requestHeadersRewrite[key]
+		if !ok {
+			return false // key not exists, not apply the rule
+		}
+		if !rewriteRule.Remove {
+			if len(values) > 1 || values[0] != rewriteRule.Value {
+				// Has more than one value, so it's different.
+				// OR
+				// It has only one value, check and overwrite it if required.
+				r.Header.Del(key)
+				r.Header.Set(key, rewriteRule.Value)
+				return true // applied
+			}
+		}
+		return false // not applied
+	}
+
+	ruleTwo := func(r *http.Request, key string, values []string) bool {
+		// Rule two:
+		//
+		// If header key is defined in request_headers_rewrite and remove is set
+		// to true and client sends a request with the same header key but different value,
+		// the headers gets removed completely before hitting the upstream.
+		rewriteRule, ok := g.requestHeadersRewrite[key]
+		if !ok {
+			return false // key not exists, not apply the rule
+		}
+		if rewriteRule.Remove {
+			if len(values) > 1 || values[0] != rewriteRule.Value {
+				// Has more than one value, so it's different.
+				// OR
+				// It has only one value, check and overwrite it if required.
+				r.Header.Del(key)
+				return true // applied
+			}
+		}
+		return false // not applied
+	}
+
+	// Try to apply rule one and rule two.
+	for forwardedHeaderKey, forwardedHeaderValues := range r.Header {
+		if len(forwardedHeaderValues) == 0 {
+			// This should not be possible but this check makes the rest of code simpler.
+			continue
+		}
+
+		if ruleOne(r, forwardedHeaderKey, forwardedHeaderValues) {
+			continue
+		}
+
+		if ruleTwo(r, forwardedHeaderKey, forwardedHeaderValues) {
+			continue
+		}
+	}
+
+	// Rule three:
+	//
+	// If header key/value is defined in request_headers_rewrite and remove is
+	// set to false and client sends a request that does not have the same header key,
+	// the header key/value gets added before hitting the upstream.
+	for headerKey, rewriteRule := range g.requestHeadersRewrite {
+		if rewriteRule.Remove {
+			continue
+		}
+		existingHeaderValue := r.Header.Get(headerKey)
+		if existingHeaderValue == "" {
+			r.Header.Set(headerKey, rewriteRule.Value)
+		}
+	}
+}
+
 func (g *GraphQLEngineTransport) setProxyOnlyHeaders(proxyOnlyValues *GraphQLProxyOnlyContextValues, r *http.Request) {
 	for forwardedHeaderKey, forwardedHeaderValues := range proxyOnlyValues.forwardedRequest.Header {
 		if proxyOnlyValues.ignoreForwardedHeaders[forwardedHeaderKey] {