From d9baf52bcacd3961dcb9969eae4ac194a8780585 Mon Sep 17 00:00:00 2001
From: dekkku <161748598+dekkku@users.noreply.github.com>
Date: Sat, 16 Nov 2024 14:05:31 +0530
Subject: [PATCH] fix(2911): add validation in federation resolvers (#3102)

---
 src/core/config/transformer/subgraph.rs       | 89 ++++++++++++++++---
 ...apollo-federation-validation.md_error.snap | 27 ++++++
 .../execution/apollo-federation-validation.md | 46 ++++++++++
 3 files changed, 152 insertions(+), 10 deletions(-)
 create mode 100644 tests/core/snapshots/apollo-federation-validation.md_error.snap
 create mode 100644 tests/execution/apollo-federation-validation.md

diff --git a/src/core/config/transformer/subgraph.rs b/src/core/config/transformer/subgraph.rs
index ed1bf950d8..b95dc2d5a3 100644
--- a/src/core/config/transformer/subgraph.rs
+++ b/src/core/config/transformer/subgraph.rs
@@ -39,22 +39,24 @@ impl Transform for Subgraph {
             // if federation is disabled don't process the config
             return Valid::succeed(config);
         }
-
+        let config_types = config.types.clone();
         let mut resolver_by_type = BTreeMap::new();
 
         let valid = Valid::from_iter(config.types.iter_mut(), |(type_name, ty)| {
             if let Some(resolver) = &ty.resolver {
                 resolver_by_type.insert(type_name.clone(), resolver.clone());
 
-                KeysExtractor::extract_keys(resolver).and_then(|fields| match fields {
-                    Some(fields) => {
-                        let key = Key { fields };
-
-                        to_directive(key.to_directive()).map(|directive| {
-                            ty.directives.push(directive);
-                        })
-                    }
-                    None => Valid::succeed(()),
+                KeysExtractor::validate(&config_types, resolver, type_name).and_then(|_| {
+                    KeysExtractor::extract_keys(resolver).and_then(|fields| match fields {
+                        Some(fields) => {
+                            let key = Key { fields };
+
+                            to_directive(key.to_directive()).map(|directive| {
+                                ty.directives.push(directive);
+                            })
+                        }
+                        None => Valid::succeed(()),
+                    })
                 })
             } else {
                 Valid::succeed(())
@@ -192,6 +194,72 @@ fn combine_keys(v: Vec<Keys>) -> Keys {
 struct KeysExtractor;
 
 impl KeysExtractor {
+    fn validate_expressions<'a>(
+        type_name: &str,
+        type_map: &BTreeMap<String, config::Type>,
+        expr_iter: impl Iterator<Item = &'a Segment>,
+    ) -> Valid<(), String> {
+        Valid::from_iter(expr_iter, |segment| {
+            if let Segment::Expression(expr) = segment {
+                if expr.len() > 1 && expr[0].as_str() == "value" {
+                    Self::validate_iter(type_map, type_name, expr.iter().skip(1))
+                } else {
+                    Valid::succeed(())
+                }
+            } else {
+                Valid::succeed(())
+            }
+        })
+        .unit()
+    }
+
+    fn validate_iter<'a>(
+        type_map: &BTreeMap<String, config::Type>,
+        current_type: &str,
+        fields_iter: impl Iterator<Item = &'a String>,
+    ) -> Valid<(), String> {
+        let mut current_type = current_type;
+        Valid::from_iter(fields_iter.enumerate(), |(index, key)| {
+            if let Some(type_def) = type_map.get(current_type) {
+                if !type_def.fields.contains_key(key) {
+                    return Valid::fail(format!(
+                        "Invalid key at index {}: '{}' is not a field of '{}'",
+                        index, key, current_type
+                    ));
+                }
+                current_type = type_def.fields[key].type_of.name();
+            } else {
+                return Valid::fail(format!("Type '{}' not found in config", current_type));
+            }
+            Valid::succeed(())
+        })
+        .unit()
+    }
+
+    fn validate(
+        type_map: &BTreeMap<String, config::Type>,
+        resolver: &Resolver,
+        type_name: &str,
+    ) -> Valid<(), String> {
+        if let Resolver::Http(http) = resolver {
+            Valid::from_iter(http.query.iter(), |q| {
+                Self::validate_expressions(
+                    type_name,
+                    type_map,
+                    Mustache::parse(&q.value).segments().iter(),
+                )
+            })
+            .and(Self::validate_expressions(
+                type_name,
+                type_map,
+                Mustache::parse(&http.url).segments().iter(),
+            ))
+            .unit()
+        } else {
+            Valid::succeed(())
+        }
+    }
+
     fn extract_keys(resolver: &Resolver) -> Valid<Option<String>, String> {
         // TODO: add validation for available fields from the type
         match resolver {
@@ -376,6 +444,7 @@ mod tests {
         }
     }
 
+    #[cfg(test)]
     mod extractor {
         use insta::assert_debug_snapshot;
         use serde_json::json;
diff --git a/tests/core/snapshots/apollo-federation-validation.md_error.snap b/tests/core/snapshots/apollo-federation-validation.md_error.snap
new file mode 100644
index 0000000000..fb321eb5ae
--- /dev/null
+++ b/tests/core/snapshots/apollo-federation-validation.md_error.snap
@@ -0,0 +1,27 @@
+---
+source: tests/core/spec.rs
+expression: errors
+---
+[
+  {
+    "message": "Type 'AccountType' not found in config",
+    "trace": [
+      "Account"
+    ],
+    "description": null
+  },
+  {
+    "message": "Invalid key at index 0: 'blogId' is not a field of 'Blog'",
+    "trace": [
+      "Blog"
+    ],
+    "description": null
+  },
+  {
+    "message": "Invalid key at index 1: 'userId' is not a field of 'Blog'",
+    "trace": [
+      "User"
+    ],
+    "description": null
+  }
+]
diff --git a/tests/execution/apollo-federation-validation.md b/tests/execution/apollo-federation-validation.md
new file mode 100644
index 0000000000..b02d77dd88
--- /dev/null
+++ b/tests/execution/apollo-federation-validation.md
@@ -0,0 +1,46 @@
+---
+error: true
+---
+
+# Apollo federation validation
+
+```graphql @config
+schema @server(port: 8000, enableFederation: true) @upstream(httpCache: 42, batch: {delay: 100}) {
+  query: Query
+}
+
+type Query {
+  post(id: Int!): Post @http(url: "http://jsonplaceholder.typicode.com/posts/{{.args.id}}")
+}
+
+type User @http(url: "http://jsonplaceholder.typicode.com/users/{{.value.blog.userId}}") {
+  id: Int!
+  username: String!
+  blog: Blog!
+}
+
+type Post
+  @http(
+    url: "http://jsonplaceholder.typicode.com/posts"
+    query: [{key: "id", value: "{{.value.id}}"}]
+    batchKey: ["blog", "blogId"]
+  ) {
+  id: Int!
+  blog: Blog!
+}
+
+type Account
+  @http(
+    url: "http://jsonplaceholder.typicode.com/posts"
+    query: [{key: "id", value: "{{.value.type.id}}"}]
+    batchKey: ["blog", "blogId"]
+  ) {
+  id: Int!
+  balance: Blog!
+  type: AccountType
+}
+
+type Blog @http(url: "http://jsonplaceholder.typicode.com/posts", query: [{key: "id", value: "{{.value.blogId}}"}]) {
+  id: Int!
+}
+```