mux: forward X- headers

backend/gateway/mux/mux.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"net/http/pprof"
+	"net/textproto"
 	"net/url"
 	"path"
 	"regexp"
@@ -28,6 +29,12 @@ import (
 	awsservice "github.com/lyft/clutch/backend/service/aws"
 )
 
+const (
+	xHeader        = "X-"
+	xForwardedFor  = "X-Forwarded-For"
+	xForwardedHost = "X-Forwarded-Host"
+)
+
 var apiPattern = regexp.MustCompile(`^/v\d+/`)
 
 type assetHandler struct {
@@ -180,6 +187,19 @@ func customResponseForwarder(ctx context.Context, w http.ResponseWriter, resp pr
 	return nil
 }
 
+func customHeaderMatcher(key string) (string, bool) {
+	key = textproto.CanonicalMIMEHeaderKey(key)
+	if strings.HasPrefix(key, xHeader) {
+		// exclude handling these headers as they are looked up by grpc's annotate context flow and added to the context
+		// metadata if they're not found
+		if key != xForwardedFor && key != xForwardedHost {
+			return runtime.MetadataPrefix + key, true
+		}
+	}
+	// the the default header mapping rule
+	return runtime.DefaultHeaderMatcher(key)
+}
+
 func customErrorHandler(ctx context.Context, mux *runtime.ServeMux, m runtime.Marshaler, w http.ResponseWriter, req *http.Request, err error) {
 	//  TODO(maybe): once we have non-browser clients we probably want to avoid the redirect and directly return the error.
 	if s, ok := status.FromError(err); ok && s.Code() == codes.Unauthenticated {
@@ -219,6 +239,7 @@ func New(unaryInterceptors []grpc.UnaryServerInterceptor, assets http.FileSystem
 				UnmarshalOptions: protojson.UnmarshalOptions{},
 			},
 		),
+		runtime.WithIncomingHeaderMatcher(customHeaderMatcher),
 	)
 
 	// If there is a configured asset provider, we check to see if the service is configured before proceeding.

backend/gateway/mux/mux_test.go

@@ -66,3 +66,59 @@ func TestGetAssetProivderService(t *testing.T) {
 	_, err := getAssetProviderService(assetCfg)
 	assert.Error(t, err)
 }
+
+func TestCustomHeaderMatcher(t *testing.T) {
+	testCases := []struct {
+		key          string
+		expectedKey  string
+		expectedBool bool
+	}{
+		{
+			key:          "X-Foo-Bar",
+			expectedKey:  "grpcgateway-X-Foo-Bar",
+			expectedBool: true,
+		},
+		// testing that the headers get uppercased
+		{
+			key:          "x-foo-bar",
+			expectedKey:  "grpcgateway-X-Foo-Bar",
+			expectedBool: true,
+		},
+		// testing the default rule - isPermanentHTTPHeader group
+		{
+			key:          "Cookie",
+			expectedKey:  "grpcgateway-Cookie",
+			expectedBool: true,
+		},
+		// testing the default rule - Grpc-Metadata prefix
+		{
+			key:          "Grpc-Metadata-Foo",
+			expectedKey:  "Foo",
+			expectedBool: true,
+		},
+		// testing the prefix doesn't get applied and doesn't match default rule
+		{
+			key:          xForwardedFor,
+			expectedKey:  "",
+			expectedBool: false,
+		},
+		// testing the prefix doesn't get applied and doesn't match default rule
+		{
+			key:          xForwardedHost,
+			expectedKey:  "",
+			expectedBool: false,
+		},
+		// doesn't match custom or default rules
+		{
+			key:          "Foo-Bar",
+			expectedKey:  "",
+			expectedBool: false,
+		},
+	}
+
+	for _, test := range testCases {
+		result, ok := customHeaderMatcher(test.key)
+		assert.Equal(t, test.expectedKey, result)
+		assert.Equal(t, test.expectedBool, ok)
+	}
+}