Provide customization to allow passing X-Geo-* headers into gRPC metadata

gateway/header.go

@@ -114,16 +114,53 @@ func handleForwardResponseTrailer(w http.ResponseWriter, md runtime.ServerMetada
 	}
 }
 
+// AtlasDefaultHeaderMatcher func used to add all headers used by atlas-app-toolkit
+// This function also passes through all the headers that runtime.DefaultHeaderMatcher handles.
+// AtlasDefaultHeaderMatcher can be used as a Incoming/Outgoing header matcher.
+func AtlasDefaultHeaderMatcher() func(string) (string, bool) {
+	//Put headers only in lower case
+	allow := map[string]struct{}{
+		//X-Geo-* headers are set of geo metadata from MaxMind DB injected on ingress nginx
+		"x-geo-org":          struct{}{},
+		"x-geo-country-code": struct{}{},
+		"x-geo-country-name": struct{}{},
+		"x-geo-region-code":  struct{}{},
+		"x-geo-region-name":  struct{}{},
+		"x-geo-city-name":    struct{}{},
+		"x-geo-postal-code":  struct{}{},
+		"x-geo-latitude":     struct{}{},
+		"x-geo-longitude":    struct{}{},
+		//request id header contains unique identifier for request
+		"request-id": struct{}{},
+		//Tracing headers
+		"x-b3-traceid":      struct{}{},
+		"x-b3-parentspanid": struct{}{},
+		"x-b3-spanid":       struct{}{},
+		"x-b3-sampled":      struct{}{},
+	}
+
+	return func(h string) (string, bool) {
+		if key, ok := runtime.DefaultHeaderMatcher(h); ok {
+			return key, ok
+		}
+
+		_, ok := allow[strings.ToLower(h)]
+		return h, ok
+	}
+}
+
 // ExtendedDefaultHeaderMatcher func is used to add custom headers to be matched
 // from incoming http requests, If this returns true the header will be added to grpc context.
-// This function also passes through all the headers that runtime.DefaultHeaderMatcher handles.
+// This function also passes through all the headers that AtlasDefaultHeaderMatcher handles.
 func ExtendedDefaultHeaderMatcher(headerNames ...string) func(string) (string, bool) {
 	customHeaders := map[string]bool{}
 	for _, name := range headerNames {
 		customHeaders[strings.ToLower(name)] = true
 	}
+
+	atlasMatcher := AtlasDefaultHeaderMatcher()
 	return func(headerName string) (string, bool) {
-		if key, ok := runtime.DefaultHeaderMatcher(headerName); ok {
+		if key, ok := atlasMatcher(headerName); ok {
 			return key, ok
 		}
 		_, ok := customHeaders[strings.ToLower(headerName)]

gateway/header_test.go

@@ -97,7 +97,7 @@ func TestExtendedDefaultHeaderMatcher(t *testing.T) {
 		{
 			name:          "custom headers in | without custom headers | failure",
 			customHeaders: []string{},
-			in:            "Request-Id",
+			in:            "CustomHeader",
 			isValid:       false,
 		},
 	}
@@ -111,3 +111,101 @@ func TestExtendedDefaultHeaderMatcher(t *testing.T) {
 		})
 	}
 }
+
+func TestAtlasDefaultHeaderMatcher(t *testing.T) {
+	var customMatcherTests = []struct {
+		name    string
+		in      string
+		isValid bool
+	}{
+		{
+			name:    "X-Geo-Org | success",
+			in:      "X-Geo-Org",
+			isValid: true,
+		},
+		{
+			name:    "X-Geo-Country-Code | success",
+			in:      "X-Geo-Country-Code",
+			isValid: true,
+		},
+		{
+			name:    "X-Geo-Country-Name | success",
+			in:      "X-Geo-Country-Name",
+			isValid: true,
+		},
+		{
+			name:    "X-Geo-Region-Code | success",
+			in:      "X-Geo-Region-Code",
+			isValid: true,
+		},
+		{
+			name:    "X-Geo-Region-Name | success",
+			in:      "X-Geo-Region-Name",
+			isValid: true,
+		},
+		{
+			name:    "X-Geo-City-Name | success",
+			in:      "X-Geo-City-Name",
+			isValid: true,
+		},
+		{
+			name:    "X-Geo-Postal-Code  | success",
+			in:      "X-Geo-Postal-Code",
+			isValid: true,
+		},
+		{
+			name:    "X-Geo-Latitude | success",
+			in:      "X-Geo-Latitude",
+			isValid: true,
+		},
+		{
+			name:    "X-Geo-Longitude | success",
+			in:      "X-Geo-Longitude",
+			isValid: true,
+		},
+		{
+			name:    "Request-Id | success",
+			in:      "Request-Id",
+			isValid: true,
+		},
+		{
+			name:    "X-B3-TraceId | success",
+			in:      "X-B3-TraceId",
+			isValid: true,
+		},
+		{
+			name:    "X-B3-ParentSpanId | success",
+			in:      "X-B3-ParentSpanId",
+			isValid: true,
+		},
+		{
+			name:    "X-B3-SpanId | success",
+			in:      "X-B3-SpanId",
+			isValid: true,
+		},
+		{
+			name:    "X-B3-Sampled | success",
+			in:      "X-B3-Sampled",
+			isValid: true,
+		},
+		{
+			name:    "x-b3-sampled | success",
+			in:      "x-b3-sampled",
+			isValid: true,
+		},
+		{
+			name:    "Failed-Header | failure",
+			in:      "Failed-Header",
+			isValid: false,
+		},
+	}
+	for _, tt := range customMatcherTests {
+		t.Run(tt.name, func(t *testing.T) {
+			f := AtlasDefaultHeaderMatcher()
+			_, ok := f(tt.in)
+			if ok != tt.isValid {
+				t.Errorf("got %v, want %v", ok, tt.isValid)
+			}
+		})
+	}
+}

server/server.go

@@ -12,6 +12,7 @@ import (
 
 	"errors"
 
+	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"github.com/infobloxopen/atlas-app-toolkit/gateway"
 	"github.com/infobloxopen/atlas-app-toolkit/health"
 	"google.golang.org/grpc"
@@ -135,14 +136,18 @@ func WithHealthChecks(checker health.Checker) Option {
 func WithGateway(options ...gateway.Option) Option {
 	return func(s *Server) error {
 		s.registrars = append(s.registrars, func(mux *http.ServeMux) error {
-			_, err := gateway.NewGateway(append(options, gateway.WithMux(mux))...)
+			_, err := gateway.NewGateway(append(options,
+				gateway.WithGatewayOptions(
+					runtime.WithIncomingHeaderMatcher(
+						gateway.AtlasDefaultHeaderMatcher())),
+				gateway.WithMux(mux))...)
 			return err
 		})
 		return nil
 	}
 }
 
-// WithMiddlewaries add opportunity to add different middleware
+// WithMiddlewares add opportunity to add different middleware
 func WithMiddlewares(middleware ...Middleware) Option {
 	return func(s *Server) error {
 		s.middlewares = append(s.middlewares, middleware...)