[GraphQL/Limits] Separate out directive checks (MystenLabs#18664)

## Description

Trying to do the directive checks at the same time as the query checks
complicates both implementations. Split out the directive check into its
own extension.

Also fix the directive checks not looking at directives on variable
definitions.

## Test plan

```
sui-graphql-e2e-tests$ cargo nextest run \
  --features pg_integration              \
  -- limits/directives.move
```

## Stack
- MystenLabs#18660
- MystenLabs#18661 
- MystenLabs#18662 
- MystenLabs#18663 

---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol: 
- [ ] Nodes (Validators and Full nodes): 
- [ ] Indexer: 
- [ ] JSON-RPC: 
- [x] GraphQL: The service now detects unsupported directives on query
variable definitions.
- [ ] CLI: 
- [ ] Rust SDK:

crates/sui-graphql-e2e-tests/tests/limits/directives.exp

@@ -1,4 +1,4 @@
-processed 7 tasks
+processed 8 tasks
 
 init:
 A: object(0,0)
@@ -31,7 +31,7 @@ Response: {
  "locations": [
  {
  "line": 1,
- "column": 29
+ "column": 28
  }
  ],
  "extensions": {
@@ -41,26 +41,45 @@ Response: {
  ]
 }
 
-task 3 'run-graphql'. lines 44-48:
+task 3 'run-graphql'. lines 44-50:
+Response: {
+ "data": null,
+ "errors": [
+ {
+ "message": "Directive `@deprecated` is not supported. Supported directives are `@include`, `@skip`",
+ "locations": [
+ {
+ "line": 1,
+ "column": 24
+ }
+ ],
+ "extensions": {
+ "code": "BAD_USER_INPUT"
+ }
+ }
+ ]
+}
+
+task 4 'run-graphql'. lines 52-56:
 Response: {
  "data": {}
 }
 
-task 4 'run-graphql'. lines 50-54:
+task 5 'run-graphql'. lines 58-62:
 Response: {
  "data": {
  "chainIdentifier": "3989ac57"
  }
 }
 
-task 5 'run-graphql'. lines 56-60:
+task 6 'run-graphql'. lines 64-68:
 Response: {
  "data": {
  "chainIdentifier": "3989ac57"
  }
 }
 
-task 6 'run-graphql'. lines 62-66:
+task 7 'run-graphql'. lines 70-74:
 Response: {
  "data": {}
 }

crates/sui-graphql-e2e-tests/tests/limits/directives.move

@@ -11,7 +11,7 @@
 
 //# run-graphql
 
-fragment Modules on Object  @deprecated {
+fragment Modules on Object @deprecated {
  address
  asMovePackage {
  module(name: "m") {
@@ -43,6 +43,14 @@ fragment Modules on Object @deprecated {
 
 //# run-graphql
 
+query($id: SuiAddress! @deprecated) {
+ object(id: $id) {
+ address
+ }
+}
+
+//# run-graphql
+
 {
  chainIdentifier @skip(if: true)
 }

crates/sui-graphql-rpc/src/config.rs

@@ -159,6 +159,7 @@ pub struct Experiments {
 #[GraphQLConfig]
 pub struct InternalFeatureConfig {
  pub(crate) query_limits_checker: bool,
+ pub(crate) directive_checker: bool,
  pub(crate) feature_gate: bool,
  pub(crate) logger: bool,
  pub(crate) query_timeout: bool,
@@ -459,6 +460,7 @@ impl Default for InternalFeatureConfig {
  fn default() -> Self {
  Self {
  query_limits_checker: true,
+ directive_checker: true,
  feature_gate: true,
  logger: true,
  query_timeout: true,

crates/sui-graphql-rpc/src/extensions/directive_checker.rs

@@ -0,0 +1,99 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::sync::Arc;
+
+use async_graphql::{
+ extensions::{Extension, ExtensionContext, ExtensionFactory, NextParseQuery},
+ parser::types::{Directive, ExecutableDocument, Selection},
+ Positioned, ServerResult,
+};
+use async_graphql_value::Variables;
+use async_trait::async_trait;
+
+use crate::error::{code, graphql_error_at_pos};
+
+const ALLOWED_DIRECTIVES: [&str; 2] = ["include", "skip"];
+
+/// Extension factory to add a check that all the directives used in the query are accepted and
+/// understood by the service.
+pub(crate) struct DirectiveChecker;
+
+struct DirectiveCheckerExt;
+
+impl ExtensionFactory for DirectiveChecker {
+ fn create(&self) -> Arc<dyn Extension> {
+ Arc::new(DirectiveCheckerExt)
+ }
+}
+
+#[async_trait]
+impl Extension for DirectiveCheckerExt {
+ async fn parse_query(
+ &self,
+ ctx: &ExtensionContext<'_>,
+ query: &str,
+ variables: &Variables,
+ next: NextParseQuery<'_>,
+ ) -> ServerResult<ExecutableDocument> {
+ let doc = next.run(ctx, query, variables).await?;
+
+ let mut selection_sets = vec![];
+ for fragment in doc.fragments.values() {
+ check_directives(&fragment.node.directives)?;
+ selection_sets.push(&fragment.node.selection_set);
+ }
+
+ for (_name, op) in doc.operations.iter() {
+ check_directives(&op.node.directives)?;
+
+ for var in &op.node.variable_definitions {
+ check_directives(&var.node.directives)?;
+ }
+
+ selection_sets.push(&op.node.selection_set);
+ }
+
+ while let Some(selection_set) = selection_sets.pop() {
+ for selection in &selection_set.node.items {
+ match &selection.node {
+ Selection::Field(field) => {
+ check_directives(&field.node.directives)?;
+ selection_sets.push(&field.node.selection_set);
+ }
+ Selection::FragmentSpread(spread) => {
+ check_directives(&spread.node.directives)?;
+ }
+ Selection::InlineFragment(fragment) => {
+ check_directives(&fragment.node.directives)?;
+ selection_sets.push(&fragment.node.selection_set);
+ }
+ }
+ }
+ }
+
+ Ok(doc)
+ }
+}
+
+fn check_directives(directives: &[Positioned<Directive>]) -> ServerResult<()> {
+ for directive in directives {
+ let name = &directive.node.name.node;
+ if !ALLOWED_DIRECTIVES.contains(&name.as_str()) {
+ let supported: Vec<_> = ALLOWED_DIRECTIVES
+ .iter()
+ .map(|s| format!("`@{s}`"))
+ .collect();
+
+ return Err(graphql_error_at_pos(
+ code::BAD_USER_INPUT,
+ format!(
+ "Directive `@{name}` is not supported. Supported directives are {}",
+ supported.join(", "),
+ ),
+ directive.pos,
+ ));
+ }
+ }
+ Ok(())
+}

crates/sui-graphql-rpc/src/extensions/mod.rs

@@ -1,7 +1,8 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
+pub(crate) mod directive_checker;
 pub(crate) mod feature_gate;
 pub(crate) mod logger;
-pub mod query_limits_checker;
+pub(crate) mod query_limits_checker;
 pub(crate) mod timeout;

crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs

@@ -8,13 +8,12 @@ use async_graphql::extensions::NextParseQuery;
 use async_graphql::extensions::NextRequest;
 use async_graphql::extensions::{Extension, ExtensionContext, ExtensionFactory};
 use async_graphql::parser::types::{
- Directive, ExecutableDocument, Field, FragmentDefinition, Selection, SelectionSet,
+ ExecutableDocument, Field, FragmentDefinition, Selection, SelectionSet,
 };
 use async_graphql::{value, Name, Pos, Positioned, Response, ServerResult, Value, Variables};
 use async_graphql_value::Value as GqlValue;
 use axum::http::HeaderName;
-use once_cell::sync::Lazy;
-use std::collections::{BTreeSet, HashMap, VecDeque};
+use std::collections::{HashMap, VecDeque};
 use std::net::SocketAddr;
 use std::sync::Arc;
 use std::time::Instant;
@@ -254,8 +253,6 @@ fn analyze_selection_set(
 
  match &selection.node {
  Selection::Field(f) => {
- check_directives(&f.node.directives)?;
-
  let current_count =
  estimate_output_nodes_for_curr_node(f, variables, limits.default_page_size)
  * parent_node_count;
@@ -288,7 +285,6 @@ fn analyze_selection_set(
  // TODO: this is inefficient as we might loop over same fragment multiple times
  // Ideally web should cache the costs of fragments we've seen before
  // Will do as enhancement
- check_directives(&frag_def.node.directives)?;
  for selection in frag_def.node.selection_set.node.items.iter() {
  que.push_back(ToVisit {
  selection,
@@ -300,7 +296,6 @@ fn analyze_selection_set(
  }
 
  Selection::InlineFragment(fs) => {
- check_directives(&fs.node.directives)?;
  for selection in fs.node.selection_set.node.items.iter() {
  que.push_back(ToVisit {
  selection,
@@ -383,35 +378,6 @@ fn check_limits(
  Ok(())
 }
 
-// TODO: make this configurable
-fn allowed_directives() -> &'static BTreeSet<&'static str> {
- static DIRECTIVES: Lazy<BTreeSet<&str>> =
- Lazy::new(|| BTreeSet::from_iter(["skip", "include"]));
-
- Lazy::force(&DIRECTIVES)
-}
-
-fn check_directives(directives: &[Positioned<Directive>]) -> ServerResult<()> {
- for directive in directives {
- if !allowed_directives().contains(&directive.node.name.node.as_str()) {
- return Err(graphql_error_at_pos(
- code::BAD_USER_INPUT,
- format!(
- "Directive `@{}` is not supported. Supported directives are {}",
- directive.node.name.node,
- allowed_directives()
- .iter()
- .map(|s| format!("`@{}`", s))
- .collect::<Vec<_>>()
- .join(", ")
- ),
- directive.pos,
- ));
- }
- }
- Ok(())
-}
-
 /// Given a node, estimate the number of output nodes it will produce.
 fn estimate_output_nodes_for_curr_node(
  f: &Positioned<Field>,

crates/sui-graphql-rpc/src/server/builder.rs

@@ -11,6 +11,7 @@ use crate::config::{
 };
 use crate::data::package_resolver::{DbPackageStore, PackageResolver};
 use crate::data::{DataLoader, Db};
+use crate::extensions::directive_checker::DirectiveChecker;
 use crate::metrics::Metrics;
 use crate::mutation::Mutation;
 use crate::types::datatype::IMoveDatatype;
@@ -468,18 +469,27 @@ impl ServerBuilder {
  if config.internal_features.feature_gate {
  builder = builder.extension(FeatureGate);
  }
+
  if config.internal_features.logger {
  builder = builder.extension(Logger::default());
  }
+
  if config.internal_features.query_limits_checker {
  builder = builder.extension(QueryLimitsChecker);
  }
+
+ if config.internal_features.directive_checker {
+ builder = builder.extension(DirectiveChecker);
+ }
+
  if config.internal_features.query_timeout {
  builder = builder.extension(Timeout);
  }
+
  if config.internal_features.tracing {
  builder = builder.extension(Tracing);
  }
+
  if config.internal_features.apollo_tracing {
  builder = builder.extension(ApolloTracing);
  }