graphql: Read a file of expensive queries and return an error for them

Read the file /etc/graph-node/expensive-queries.txt (one GraphQL query per
line), parse the queries and respond to any query that has the same
ShapeHash as one of the included queries with a TooExpensive error

Author: lutter

graph/src/data/query/error.rs

@@ -54,6 +54,7 @@ pub enum QueryExecutionError {
     ScalarCoercionError(Pos, String, q::Value, String),
     TooComplex(u64, u64), // (complexity, max_complexity)
     TooDeep(u8),          // max_depth
+    TooExpensive,
     UndefinedFragment(String),
     // Using slow and prefetch query resolution yield different results
     IncorrectPrefetchResult { slow: q::Value, prefetch: q::Value },
@@ -209,6 +210,7 @@ impl fmt::Display for QueryExecutionError {
             EventStreamError => write!(f, "error in the subscription event stream"),
             FulltextQueryRequiresFilter => write!(f, "fulltext search queries can only use EntityFilter::Equal"),
             SubscriptionsDisabled => write!(f, "subscriptions are disabled"),
+            TooExpensive => write!(f, "query is too expensive")
         }
     }
 }

graphql/src/runner.rs

@@ -1,6 +1,6 @@
 use futures01::future;
 use graphql_parser::query as q;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap};
 use std::env;
 use std::str::FromStr;
 use std::sync::Arc;
@@ -9,12 +9,12 @@ use std::time::{Duration, Instant};
 use crate::prelude::{
     object, object_value, QueryExecutionOptions, StoreResolver, SubscriptionExecutionOptions,
 };
-use crate::query::execute_query;
+use crate::query::{execute_query, shape_hash::shape_hash};
 use crate::subscription::execute_prepared_subscription;
 use graph::prelude::{
     o, EthereumBlockPointer, GraphQlRunner as GraphQlRunnerTrait, Logger, Query,
     QueryExecutionError, QueryResult, QueryResultFuture, Store, StoreError, SubgraphDeploymentId,
-    SubgraphDeploymentStore, Subscription, SubscriptionResultFuture,
+    SubgraphDeploymentStore, Subscription, SubscriptionError, SubscriptionResultFuture,
 };
 
 use lazy_static::lazy_static;
@@ -23,6 +23,7 @@ use lazy_static::lazy_static;
 pub struct GraphQlRunner<S> {
     logger: Logger,
     store: Arc<S>,
+    expensive: HashMap<u64, Arc<q::Document>>,
 }
 
 lazy_static! {
@@ -53,10 +54,15 @@ where
     S: Store + SubgraphDeploymentStore,
 {
     /// Creates a new query runner.
-    pub fn new(logger: &Logger, store: Arc<S>) -> Self {
+    pub fn new(logger: &Logger, store: Arc<S>, expensive: &Vec<Arc<q::Document>>) -> Self {
+        let expensive = expensive
+            .into_iter()
+            .map(|doc| (shape_hash(&doc), doc.clone()))
+            .collect::<HashMap<_, _>>();
         GraphQlRunner {
             logger: logger.new(o!("component" => "GraphQlRunner")),
             store,
+            expensive,
         }
     }
 
@@ -138,6 +144,14 @@ where
             Ok(QueryResult::new(Some(q::Value::Object(values))))
         }
     }
+
+    pub fn check_too_expensive(&self, query: &Query) -> Result<(), Vec<QueryExecutionError>> {
+        if self.expensive.contains_key(&shape_hash(&query.document)) {
+            Err(vec![QueryExecutionError::TooExpensive])
+        } else {
+            Ok(())
+        }
+    }
 }
 
 impl<S> GraphQlRunnerTrait for GraphQlRunner<S>
@@ -160,14 +174,19 @@ where
         max_depth: Option<u8>,
         max_first: Option<u32>,
     ) -> QueryResultFuture {
-        let result = match self.execute(query, max_complexity, max_depth, max_first) {
-            Ok(result) => result,
-            Err(e) => QueryResult::from(e),
-        };
+        let result = self
+            .check_too_expensive(&query)
+            .and_then(|_| self.execute(query, max_complexity, max_depth, max_first))
+            .unwrap_or_else(|e| QueryResult::from(e));
         Box::new(future::ok(result))
     }
 
     fn run_subscription(&self, subscription: Subscription) -> SubscriptionResultFuture {
+        if let Err(errs) = self.check_too_expensive(&subscription.query) {
+            let err = SubscriptionError::GraphQLError(errs);
+            return Box::new(future::result(Err(err)));
+        }
+
         let query = match crate::execution::Query::new(
             subscription.query,
             *GRAPHQL_MAX_COMPLEXITY,

graphql/tests/query.rs

@@ -230,7 +230,7 @@ fn execute_query_document_with_variables(
     query: q::Document,
     variables: Option<QueryVariables>,
 ) -> QueryResult {
-    let runner = GraphQlRunner::new(&*LOGGER, STORE.clone());
+    let runner = GraphQlRunner::new(&*LOGGER, STORE.clone(), &vec![]);
     let query = Query::new(Arc::new(api_test_schema()), query, variables);
 
     return_err!(runner.execute(query, None, None, None))

node/src/main.rs

@@ -5,6 +5,8 @@ use lazy_static::lazy_static;
 use prometheus::Registry;
 use std::collections::HashMap;
 use std::env;
+use std::io::{BufRead, BufReader};
+use std::path::Path;
 use std::str::FromStr;
 use std::time::Duration;
 use tokio::sync::mpsc;
@@ -34,6 +36,7 @@ use graph_store_postgres::{
     ChainHeadUpdateListener as PostgresChainHeadUpdateListener, Store as DieselStore, StoreConfig,
     SubscriptionManager,
 };
+use graphql_parser::query as q;
 
 lazy_static! {
     // Default to an Ethereum reorg threshold to 50 blocks
@@ -68,6 +71,33 @@ enum ConnectionType {
     WS,
 }
 
+fn read_expensive_queries() -> Result<Vec<Arc<q::Document>>, std::io::Error> {
+    // A file with a list of expensive queries, one query per line
+    // Attempts to run these queries will return a
+    // QueryExecutionError::TooExpensive to clients
+    const EXPENSIVE_QUERIES: &str = "/etc/graph-node/expensive-queries.txt";
+    let path = Path::new(EXPENSIVE_QUERIES);
+    let mut queries = Vec::new();
+    if path.exists() {
+        let file = std::fs::File::open(path)?;
+        let reader = BufReader::new(file);
+        for line in reader.lines() {
+            let line = line?;
+            let query = graphql_parser::parse_query(&line).map_err(|e| {
+                let msg = format!(
+                    "invalid GraphQL query in {}: {}\n{}",
+                    EXPENSIVE_QUERIES,
+                    e.to_string(),
+                    line
+                );
+                std::io::Error::new(std::io::ErrorKind::InvalidData, msg)
+            })?;
+            queries.push(Arc::new(query));
+        }
+    }
+    Ok(queries)
+}
+
 // Saturating the blocking threads can cause all sorts of issues, so set a large maximum.
 // Ideally we'd use semaphores to not use more blocking threads than DB connections,
 // but for now this is necessary.
@@ -523,6 +553,8 @@ async fn main() {
         )))
     };
 
+    let expensive_queries = read_expensive_queries().unwrap();
+
     graph::spawn(
         futures::stream::FuturesOrdered::from_iter(stores_eth_adapters.into_iter().map(
             |(network_name, eth_adapter)| {
@@ -570,7 +602,11 @@ async fn main() {
         .and_then(move |stores| {
             let generic_store = stores.values().next().expect("error creating stores");
 
-            let graphql_runner = Arc::new(GraphQlRunner::new(&logger, generic_store.clone()));
+            let graphql_runner = Arc::new(GraphQlRunner::new(
+                &logger,
+                generic_store.clone(),
+                &expensive_queries,
+            ));
             let mut graphql_server = GraphQLQueryServer::new(
                 &logger_factory,
                 graphql_metrics_registry,