feat(GraphQL): Allow more control over custom logic header names (dgr…

…aph-io#5600) * Added support for custom names in headers * Continue to support old header definition * Addressed comments * Added integration test * Moved validations to rules.go
dna2fork · Jul 18, 2020 · 3c78a1d · 3c78a1d
1 parent 8551200
commit 3c78a1d
Show file tree

Hide file tree

Showing 5 changed files with 110 additions and 8 deletions.
diff --git a/graphql/e2e/custom_logic/cmd/main.go b/graphql/e2e/custom_logic/cmd/main.go
@@ -295,6 +295,26 @@ func verifyHeadersHandler(w http.ResponseWriter, r *http.Request) {
 	check2(w.Write([]byte(`[{"id":"0x3","name":"Star Wars"}]`)))
 }
 
+func verifyCustomNameHeadersHandler(w http.ResponseWriter, r *http.Request) {
+	err := verifyRequest(r, expectedRequest{
+		method:    http.MethodGet,
+		urlSuffix: "/verifyCustomNameHeaders",
+		body:      "",
+		headers: map[string][]string{
+			"X-App-Token":      {"app-token"},
+			"X-User-Id":        {"123"},
+			"Authorization": {"random-fake-token"},
+			"Accept-Encoding":  nil,
+			"User-Agent":       nil,
+		},
+	})
+	if err != nil {
+		check2(w.Write([]byte(err.Error())))
+		return
+	}
+	check2(w.Write([]byte(`[{"id":"0x3","name":"Star Wars"}]`)))
+}
+
 func twitterFollwerHandler(w http.ResponseWriter, r *http.Request) {
 	err := verifyRequest(r, expectedRequest{
 		method:    http.MethodGet,
@@ -1114,6 +1134,7 @@ func main() {
 	http.HandleFunc("/favMovies/", getFavMoviesHandler)
 	http.HandleFunc("/favMoviesPost/", postFavMoviesHandler)
 	http.HandleFunc("/verifyHeaders", verifyHeadersHandler)
+	http.HandleFunc("/verifyCustomNameHeaders", verifyCustomNameHeadersHandler)
 	http.HandleFunc("/twitterfollowers", twitterFollwerHandler)
 
 	// for mutations

diff --git a/graphql/e2e/custom_logic/custom_logic_test.go b/graphql/e2e/custom_logic/custom_logic_test.go
@@ -168,6 +168,45 @@ func TestCustomQueryShouldForwardHeaders(t *testing.T) {
 				 secretHeaders: ["Github-Api-Token", "X-App-Token"]
 		 })
 	 }
+	 
+		 # Dgraph.Secret Github-Api-Token "random-fake-token"
+		 # Dgraph.Secret app "should-be-overriden"
+	 `
+	updateSchemaRequireNoGQLErrors(t, schema)
+	time.Sleep(2 * time.Second)
+
+	query := `
+	 query {
+		 verifyHeaders(id: "0x123") {
+			 id
+			 name
+		 }
+	 }`
+	params := &common.GraphQLParams{
+		Query: query,
+		Headers: map[string][]string{
+			"X-App-Token":   []string{"app-token"},
+			"X-User-Id":     []string{"123"},
+			"Random-header": []string{"random"},
+		},
+	}
+
+	result := params.ExecuteAsPost(t, alphaURL)
+	require.Nil(t, result.Errors)
+	expected := `{"verifyHeaders":[{"id":"0x3","name":"Star Wars"}]}`
+	require.Equal(t, expected, string(result.Data))
+}
+
+func TestCustomNameForwardHeaders(t *testing.T) {
+	schema := customTypes + `
+	 type Query {
+		 verifyHeaders(id: ID!): [Movie] @custom(http: {
+				 url: "http://mock:8888/verifyCustomNameHeaders",
+				 method: "GET",
+				 forwardHeaders: ["X-App-Token:App", "X-User-Id"],
+				 secretHeaders: ["Authorization:Github-Api-Token", "X-App-Token"]
+		 })
+	 }
 
 		 # Dgraph.Secret Github-Api-Token "random-fake-token"
 		 # Dgraph.Secret X-App-Token "should-be-overriden"
@@ -185,7 +224,7 @@ func TestCustomQueryShouldForwardHeaders(t *testing.T) {
 	params := &common.GraphQLParams{
 		Query: query,
 		Headers: map[string][]string{
-			"X-App-Token":   []string{"app-token"},
+			"App":           []string{"app-token"},
 			"X-User-Id":     []string{"123"},
 			"Random-header": []string{"random"},
 		},

diff --git a/graphql/schema/rules.go b/graphql/schema/rules.go
@@ -1604,6 +1604,32 @@ func customDirectiveValidation(sch *ast.Schema,
 		}
 	}
 
+	forwardHeaders := httpArg.Value.Children.ForName("forwardHeaders")
+	if forwardHeaders != nil {
+		for _, h := range forwardHeaders.Children {
+			key := strings.Split(h.Value.Raw, ":")
+			if len(key) > 2 {
+				return append(errs, gqlerror.ErrorPosf(graphql.Position,
+					"Type %s; Field %s; forwardHeaders in @custom directive should be of the form 'remote_headername:local_headername' or just 'headername'"+
+						", found: `%s`.",
+					typ.Name, field.Name, h.Value.Raw))
+			}
+		}
+	}
+
+	secretHeaders := httpArg.Value.Children.ForName("secretHeaders")
+	if secretHeaders != nil {
+		for _, h := range secretHeaders.Children {
+			key := strings.Split(h.Value.Raw, ":")
+			if len(key) > 2 {
+				return append(errs, gqlerror.ErrorPosf(graphql.Position,
+					"Type %s; Field %s; secretHeaders in @custom directive should be of the form 'remote_headername:local_headername' or just 'headername'"+
+						", found: `%s`.",
+					typ.Name, field.Name, h.Value.Raw))
+			}
+		}
+	}
+
 	if errs != nil {
 		return errs
 	}
@@ -1613,9 +1639,13 @@ func customDirectiveValidation(sch *ast.Schema,
 		headers := http.Header{}
 		if secretHeaders != nil {
 			for _, h := range secretHeaders.Children {
+				key := strings.Split(h.Value.Raw, ":")
+				if len(key) == 1 {
+					key = []string{h.Value.Raw, h.Value.Raw}
+				}
 				// We try and fetch the value from the stored secrets.
-				val := secrets[h.Value.Raw]
-				headers.Add(h.Value.Raw, string(val))
+				val := secrets[key[1]]
+				headers.Add(key[0], string(val))
 			}
 		}
 		if err := validateRemoteGraphql(&remoteGraphqlMetadata{

diff --git a/graphql/schema/schemagen.go b/graphql/schema/schemagen.go
@@ -259,7 +259,11 @@ func getAllowedHeaders(sch *ast.Schema, definitions []string) string {
 			return
 		}
 		for _, h := range forwardHeaders.Children {
-			headers[h.Value.Raw] = struct{}{}
+			key := strings.Split(h.Value.Raw, ":")
+			if len(key) == 1 {
+				key = []string{h.Value.Raw, h.Value.Raw}
+			}
+			headers[key[1]] = struct{}{}
 		}
 	}
 

diff --git a/graphql/schema/wrappers.go b/graphql/schema/wrappers.go
@@ -847,8 +847,12 @@ func getCustomHTTPConfig(f *field, isQueryOrMutation bool) (FieldHTTPConfig, err
 	if secretHeaders != nil {
 		hc.RLock()
 		for _, h := range secretHeaders.Children {
-			val := string(hc.secrets[h.Value.Raw])
-			fconf.ForwardHeaders.Set(h.Value.Raw, val)
+			key := strings.Split(h.Value.Raw, ":")
+			if len(key) == 1 {
+				key = []string{h.Value.Raw, h.Value.Raw}
+			}
+			val := string(hc.secrets[key[1]])
+			fconf.ForwardHeaders.Set(key[0], val)
 		}
 		hc.RUnlock()
 	}
@@ -857,8 +861,12 @@ func getCustomHTTPConfig(f *field, isQueryOrMutation bool) (FieldHTTPConfig, err
 	if forwardHeaders != nil {
 		for _, h := range forwardHeaders.Children {
 			// We would override the header if it was also specified as part of secretHeaders.
-			reqHeaderVal := f.op.header.Get(h.Value.Raw)
-			fconf.ForwardHeaders.Set(h.Value.Raw, reqHeaderVal)
+			key := strings.Split(h.Value.Raw, ":")
+			if len(key) == 1 {
+				key = []string{h.Value.Raw, h.Value.Raw}
+			}
+			reqHeaderVal := f.op.header.Get(key[1])
+			fconf.ForwardHeaders.Set(key[0], reqHeaderVal)
 		}
 	}