[TT-11997] Backend logic for request_headers_rewrite (#6306)

### **User description** See https://tyktech.atlassian.net/browse/TT-11997 for details of the request headers rewrite rules. ___ ### **PR Type** Enhancement, Tests ___ ### **Description** - Added tests for request headers rewrite rules in GraphQL middleware. - Integrated request headers rewrite logic into the reverse proxy handler. - Added `RequestHeadersRewrite` to `ReverseProxyParams` struct. - Passed `RequestHeadersRewrite` to the GraphQL engine transport for both v1 and v2. - Implemented logic to apply request headers rewrite rules in the GraphQL engine transport, covering three different rewrite scenarios. ___ ### **Changes walkthrough** 📝 <table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Tests </strong></td><td><table> <tr> <td> <details> <summary><strong>mw_graphql_test.go</strong><dd><code>Add tests for request headers rewrite rules in GraphQL middleware.</code></dd></summary> <hr> gateway/mw_graphql_test.go <li>Added tests for request headers rewrite rules.<br> <li> Implemented three test cases for different rewrite scenarios.<br> </details> </td> <td><a href="https://github.com/TykTechnologies/tyk/pull/6306/files#diff-72c11b8efcaceccfbbd87e75565353706443bf56420d3fdba6b5bf5f632f4b33">+137/-0</a>  </td> </tr> </table></td></tr><tr><td><strong>Enhancement </strong></td><td><table> <tr> <td> <details> <summary><strong>reverse_proxy.go</strong><dd><code>Integrate request headers rewrite logic in reverse proxy handler.</code></dd></summary> <hr> gateway/reverse_proxy.go <li>Integrated request headers rewrite logic into the reverse proxy <br>handler.<br> </details> </td> <td><a href="https://github.com/TykTechnologies/tyk/pull/6306/files#diff-e6e07722257f7e41691e471185ad6d84fd56dc9e5459526ea32e9a5e8fa1a01b">+8/-7</a>      </td> </tr> <tr> <td> <details> <summary><strong>engine.go</strong><dd><code>Add RequestHeadersRewrite to ReverseProxyParams struct.</code>    </dd></summary> <hr> internal/graphengine/engine.go - Added `RequestHeadersRewrite` to `ReverseProxyParams` struct. </details> </td> <td><a href="https://github.com/TykTechnologies/tyk/pull/6306/files#diff-9d1ef4ec1b59d708ed86c7cc348b73bd0b619ef034094b5fbd4b138e11ccde6e">+8/-7</a>      </td> </tr> <tr> <td> <details> <summary><strong>graphql_go_tools_v1.go</strong><dd><code>Pass RequestHeadersRewrite to GraphQL engine transport (v1).</code></dd></summary> <hr> internal/graphengine/graphql_go_tools_v1.go - Passed `RequestHeadersRewrite` to the GraphQL engine transport. </details> </td> <td><a href="https://github.com/TykTechnologies/tyk/pull/6306/files#diff-e592cc8ca6ac39e7574765d7f2bbf19193f173791a1b0930d4dde7f9412dc882">+1/-0</a>      </td> </tr> <tr> <td> <details> <summary><strong>graphql_go_tools_v2.go</strong><dd><code>Pass RequestHeadersRewrite to GraphQL engine transport (v2).</code></dd></summary> <hr> internal/graphengine/graphql_go_tools_v2.go - Passed `RequestHeadersRewrite` to the GraphQL engine transport. </details> </td> <td><a href="https://github.com/TykTechnologies/tyk/pull/6306/files#diff-2e41b60f046c2947fbedb8fb841b5c3c962798e3ba8211c7144f326436ffabe3">+1/-0</a>      </td> </tr> <tr> <td> <details> <summary><strong>transport.go</strong><dd><code>Implement request headers rewrite logic in GraphQL engine transport.</code></dd></summary> <hr> internal/graphengine/transport.go <li>Added logic to apply request headers rewrite rules in the GraphQL <br>engine transport.<br> <li> Implemented three rewrite rules for request headers.<br> </details> </td> <td><a href="https://github.com/TykTechnologies/tyk/pull/6306/files#diff-564061c9b29366529eb1f6f10fe39671d2ac738a4731ffd2c8b04dcc0a8cd610">+96/-1</a>    </td> </tr> </table></td></tr></tr></tbody></table> ___ > 💡 **PR-Agent usage**: >Comment `/help` on the PR to get a list of all available PR-Agent tools and their descriptions
TykTechnologies · May 28, 2024 · 00a815f · 00a815f
1 parent f7cda28
commit 00a815f
Show file tree

Hide file tree

Showing 6 changed files with 251 additions and 15 deletions.
diff --git a/gateway/mw_graphql_test.go b/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

diff --git a/gateway/reverse_proxy.go b/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

diff --git a/internal/graphengine/engine.go b/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 {

diff --git a/internal/graphengine/graphql_go_tools_v1.go b/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 {

diff --git a/internal/graphengine/graphql_go_tools_v2.go b/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 {

diff --git a/internal/graphengine/transport.go b/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] {