Making overriding environment on config based

uber · Mar 4, 2021 · 0b49f08 · 0b49f08
1 parent 899031c
commit 0b49f08
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 5 deletions.
diff --git a/examples/example-gateway/config/production.yaml b/examples/example-gateway/config/production.yaml
@@ -90,3 +90,5 @@ grpc.clientServiceNameMapping:
   echo: echo
 router.whitelistedPaths:
   - /path/whitelisted
+service.shadow.request.header: x-shadow-request
+service.shadow.env.override.enable: true
diff --git a/examples/example-gateway/config/test.yaml b/examples/example-gateway/config/test.yaml
@@ -89,3 +89,5 @@ clients.baz.alternates:
       port: 8114
 grpc.clientServiceNameMapping:
   echo: echo
+service.shadow.request.header: x-shadow-request
+service.shadow.env.override.enable: true
diff --git a/runtime/constants.go b/runtime/constants.go
@@ -53,7 +53,6 @@ const (
 	clientTchannelUnmarshalError = "client.tchannel-unmarshal-error"
 
 	// shadow headers and environment
-	shadowHeader      = "X-Shadow-Request"
 	shadowEnvironment = "shadow"
 	environmentKey    = "env"
 )

diff --git a/runtime/router.go b/runtime/router.go
@@ -89,6 +89,7 @@ type RouterEndpoint struct {
 	contextLogger    ContextLogger
 	scope            tally.Scope
 	tracer           opentracing.Tracer
+	config           *StaticConfig
 }
 
 // NewRouterEndpoint creates an endpoint that can be registered to HTTPRouter
@@ -108,6 +109,7 @@ func NewRouterEndpoint(
 		scope:            deps.Scope,
 		tracer:           deps.Tracer,
 		JSONWrapper:      deps.JSONWrapper,
+		config:           deps.Config,
 	}
 }
 

diff --git a/runtime/server_http_request.go b/runtime/server_http_request.go
@@ -101,10 +101,12 @@ func NewServerHTTPRequest(
 		logFields = append(logFields, endpoint.contextExtractor.ExtractLogFields(ctx)...)
 	}
 
-	val := r.Header.Get(shadowHeader)
-	if val != "" {
-		scopeTags[environmentKey] = shadowEnvironment
-		logFields = append(logFields, zap.String(environmentKey, shadowEnvironment))
+	// Overriding the environment for shadow requests
+	if endpoint.config != nil {
+		if endpoint.config.ContainsKey("service.shadow.env.override.enable") && endpoint.config.MustGetBoolean("service.shadow.env.override.enable") && endpoint.config.ContainsKey("service.shadow.request.header") && r.Header.Get(endpoint.config.MustGetString("service.shadow.request.header")) != "" {
+			scopeTags[environmentKey] = shadowEnvironment
+			logFields = append(logFields, zap.String(environmentKey, shadowEnvironment))
+		}
 	}
 
 	ctx = WithScopeTags(ctx, scopeTags)

diff --git a/runtime/server_http_request_test.go b/runtime/server_http_request_test.go
@@ -2256,6 +2256,7 @@ func testIncomingHTTPRequestServerLog(t *testing.T, isShadowRequest bool, enviro
 		Logger:        bgateway.ActualGateway.Logger,
 		ContextLogger: bgateway.ActualGateway.ContextLogger,
 		Tracer:        bgateway.ActualGateway.Tracer,
+		Config:        bgateway.ActualGateway.Config,
 	}
 	err = bgateway.ActualGateway.HTTPRouter.Handle(
 		"GET", "/foo", http.HandlerFunc(zanzibar.NewRouterEndpoint(