Merge 269c9e1 into 4177398

codegen/gateway.go

@@ -251,6 +251,8 @@ type EndpointSpec struct {
 	ReqHeaders map[string]*TypedHeader `yaml:"reqHeaderMap,omitempty"`
 	// ResHeaders maps headers from client to server
 	ResHeaders map[string]*TypedHeader `yaml:"resHeaderMap,omitempty"`
+	// DefaultHeaders a slice of headers that are forwarded to downstream when available
+	DefaultHeaders []string `yaml:"-"`
 	// WorkflowType, either "httpClient" or "custom".
 	// A httpClient workflow generates a http client Caller
 	// A custom workflow just imports the custom code
@@ -417,6 +419,7 @@ func NewEndpointSpec(
 		IsClientlessEndpoint: isClientlessEndpoint,
 		ClientID:             clientID,
 		ClientMethod:         clientMethod,
+		DefaultHeaders:       h.defaultHeaders,
 	}
 
 	defaultMidSpecs, err := getOrderedDefaultMiddlewareSpecs(

codegen/module_system.go

@@ -56,6 +56,7 @@ type EndpointMeta struct {
 	ResRequiredHeadersKeys []string
 	TraceKey               string
 	DeputyReqHeader        string
+	DefaultHeaders         []string
 }
 
 // EndpointCollectionMeta saves information used to generate an initializer
@@ -1273,6 +1274,7 @@ func (g *EndpointGenerator) generateEndpointFile(e *EndpointSpec, instance *Modu
 		WorkflowPkg:            workflowPkg,
 		TraceKey:               g.packageHelper.traceKey,
 		DeputyReqHeader:        g.packageHelper.DeputyReqHeader(),
+		DefaultHeaders:         e.DefaultHeaders,
 	}
 
 	targetPath := e.TargetEndpointPath(thriftServiceName, method.Name)

codegen/package.go

@@ -58,6 +58,8 @@ type PackageHelper struct {
 	middlewareSpecs map[string]*MiddlewareSpec
 	// The default middlewares for all endpoints
 	defaultMiddlewareSpecs map[string]*MiddlewareSpec
+	// The default headers to forward with all requests when present
+	defaultHeaders []string
 	// Use deputy client when this header is set
 	deputyReqHeader string
 	// traceKey is the key for unique trace id that identifies request / response pair
@@ -95,6 +97,8 @@ type PackageHelperOptions struct {
 	CopyrightHeader string
 	// header key to redirect client requests to staging environment, defaults to "X-Zanzibar-Use-Staging"
 	StagingReqHeader string
+	// The default headers to forward with all requests when present
+	DefaultHeaders []string
 	// header key to redirect client requests to local environment, defaults to "x-deputy-forwarded"
 	DeputyReqHeader string
 	// header key to uniquely identifies request/response pair, defaults to "x-trace-id"
@@ -210,6 +214,7 @@ func NewPackageHelper(
 		traceKey:               options.traceKey(),
 		moduleSearchPaths:      options.ModuleSearchPaths,
 		defaultDependencies:    options.DefaultDependencies,
+		defaultHeaders:         options.DefaultHeaders,
 	}
 	return p, nil
 }

codegen/runner/runner.go

@@ -114,6 +114,9 @@ func main() {
 	defaultDependencies := make(map[string][]string, 0)
 	config.MustGetStruct("defaultDependencies", &defaultDependencies)
 
+	defaultHeaders := make([]string, 0)
+	config.MustGetStruct("defaultHeaders", &defaultHeaders)
+
 	options := &codegen.PackageHelperOptions{
 		RelThriftRootDir:              config.MustGetString("thriftRootDir"),
 		RelTargetGenDir:               config.MustGetString("targetGenDir"),
@@ -127,6 +130,7 @@ func main() {
 		TraceKey:                      config.MustGetString("traceKey"),
 		ModuleSearchPaths:             searchPaths,
 		DefaultDependencies:           defaultDependencies,
+		DefaultHeaders:                defaultHeaders,
 	}
 
 	packageHelper, err := codegen.NewPackageHelper(

codegen/templates/workflow.tmpl

@@ -5,6 +5,7 @@ package workflow
 {{- $endpointType := .Spec.EndpointType }}
 {{- $reqHeaderMap := .ReqHeaders }}
 {{- $reqHeaderMapKeys := .ReqHeadersKeys }}
+{{- $defaultHeaders := .DefaultHeaders }}
 {{- $reqHeaderRequiredKeys := .ReqRequiredHeadersKeys }}
 {{- $resHeaderMap := .ResHeaders }}
 {{- $resHeaderMapKeys := .ResHeadersKeys }}
@@ -133,10 +134,25 @@ func (w {{$workflowStruct}}) Handle(
 	{{end}}
 
 	clientHeaders := map[string]string{}
+	{{if (ne (len $defaultHeaders) 0) }}
+	var ok bool
+	var h string
+	var k string
+	{{range $i, $k := $defaultHeaders}}
+	k = textproto.CanonicalMIMEHeaderKey("{{$k}}")
+	h, ok = reqHeaders.Get(k)
+	if ok {
+		clientHeaders[k] = h
+	}
+	{{- end -}}
+	{{- end -}}
+
 	{{if (ne (len $reqHeaderMapKeys) 0) }}
+	{{if (eq (len $defaultHeaders) 0) }}
 	var ok bool
 	var h string
 	{{- end -}}
+	{{- end -}}
 	{{range $i, $k := $reqHeaderMapKeys}}
 	h, ok = reqHeaders.Get("{{$k}}")
 	if ok {

examples/example-gateway/build.yaml

@@ -6,6 +6,9 @@ genCodePackage: github.com/uber/zanzibar/examples/example-gateway/build/gen-code
 genMock: true
 middlewareConfig: ./middlewares
 defaultMiddlewareConfig: ./middlewares/default
+defaultHeaders:
+  - x-uber-foo
+  - x-uber-bar
 packageRoot: github.com/uber/zanzibar/examples/example-gateway
 targetGenDir: ./build
 thriftRootDir: ./idl