From 05d86f45728bd496ab063c41a76e4f2768b37f3c Mon Sep 17 00:00:00 2001 From: Ashok Menon Date: Sat, 13 Jul 2024 17:06:13 +0100 Subject: [PATCH] [GraphQL/Limits] Separate out directive checks ## 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 ``` --- .../tests/limits/directives.exp | 31 ++++-- .../tests/limits/directives.move | 10 +- crates/sui-graphql-rpc/src/config.rs | 2 + .../src/extensions/directive_checker.rs | 99 +++++++++++++++++++ crates/sui-graphql-rpc/src/extensions/mod.rs | 3 +- .../src/extensions/query_limits_checker.rs | 38 +------ crates/sui-graphql-rpc/src/server/builder.rs | 10 ++ 7 files changed, 149 insertions(+), 44 deletions(-) create mode 100644 crates/sui-graphql-rpc/src/extensions/directive_checker.rs diff --git a/crates/sui-graphql-e2e-tests/tests/limits/directives.exp b/crates/sui-graphql-e2e-tests/tests/limits/directives.exp index 0550bdebcc88d2..0b7bd9e4ab84be 100644 --- a/crates/sui-graphql-e2e-tests/tests/limits/directives.exp +++ b/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": "8f58ad28" } } -task 5 'run-graphql'. lines 56-60: +task 6 'run-graphql'. lines 64-68: Response: { "data": { "chainIdentifier": "8f58ad28" } } -task 6 'run-graphql'. lines 62-66: +task 7 'run-graphql'. lines 70-74: Response: { "data": {} } diff --git a/crates/sui-graphql-e2e-tests/tests/limits/directives.move b/crates/sui-graphql-e2e-tests/tests/limits/directives.move index 13631cd5b6a743..4a20b9f4a18ea2 100644 --- a/crates/sui-graphql-e2e-tests/tests/limits/directives.move +++ b/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) } diff --git a/crates/sui-graphql-rpc/src/config.rs b/crates/sui-graphql-rpc/src/config.rs index 320653bcbaa3de..71590683d308ff 100644 --- a/crates/sui-graphql-rpc/src/config.rs +++ b/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, diff --git a/crates/sui-graphql-rpc/src/extensions/directive_checker.rs b/crates/sui-graphql-rpc/src/extensions/directive_checker.rs new file mode 100644 index 00000000000000..11e6ce6e80044e --- /dev/null +++ b/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 { + Arc::new(DirectiveCheckerExt) + } +} + +#[async_trait] +impl Extension for DirectiveCheckerExt { + async fn parse_query( + &self, + ctx: &ExtensionContext<'_>, + query: &str, + variables: &Variables, + next: NextParseQuery<'_>, + ) -> ServerResult { + 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]) -> 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(()) +} diff --git a/crates/sui-graphql-rpc/src/extensions/mod.rs b/crates/sui-graphql-rpc/src/extensions/mod.rs index 9c282b5ae958e3..f3b8aa792a0fed 100644 --- a/crates/sui-graphql-rpc/src/extensions/mod.rs +++ b/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; diff --git a/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs b/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs index 715241653542ba..79054f833e93ad 100644 --- a/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs +++ b/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> = - Lazy::new(|| BTreeSet::from_iter(["skip", "include"])); - - Lazy::force(&DIRECTIVES) -} - -fn check_directives(directives: &[Positioned]) -> 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::>() - .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, diff --git a/crates/sui-graphql-rpc/src/server/builder.rs b/crates/sui-graphql-rpc/src/server/builder.rs index c2429bac1f41d7..01f70dd20d390c 100644 --- a/crates/sui-graphql-rpc/src/server/builder.rs +++ b/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); }