Merge pull request #5263 from rmosolgo/manual-analysis-timeout

Use a manual timeout implementation in Analysis to handle I/O gracefully

Author: rmosolgo

lib/graphql/analysis.rb

@@ -6,11 +6,16 @@
 require "graphql/analysis/max_query_complexity"
 require "graphql/analysis/query_depth"
 require "graphql/analysis/max_query_depth"
-require "timeout"
-
 module GraphQL
   module Analysis
     AST = self
+
+    class TimeoutError < AnalysisError
+      def initialize(...)
+        super("Timeout on validation of query")
+      end
+    end
+
     module_function
     # Analyze a multiplex, and all queries within.
     # Multiplex analyzers are ran for all queries, keeping state.
@@ -61,13 +66,11 @@ def analyze_query(query, analyzers, multiplex_analyzers: [])
           if !analyzers_to_run.empty?
             visitor = GraphQL::Analysis::Visitor.new(
               query: query,
-              analyzers: analyzers_to_run
+              analyzers: analyzers_to_run,
+              timeout: query.validate_timeout_remaining,
             )
 
-            # `nil` or `0` causes no timeout
-            Timeout::timeout(query.validate_timeout_remaining) do
-              visitor.visit
-            end
+            visitor.visit
 
             if !visitor.rescued_errors.empty?
               return visitor.rescued_errors
@@ -79,8 +82,8 @@ def analyze_query(query, analyzers, multiplex_analyzers: [])
           []
         end
       end
-    rescue Timeout::Error
-      [GraphQL::AnalysisError.new("Timeout on validation of query")]
+    rescue TimeoutError => err
+      [err]
     rescue GraphQL::UnauthorizedError, GraphQL::ExecutionError
       # This error was raised during analysis and will be returned the client before execution
       []

lib/graphql/analysis/visitor.rb

@@ -10,7 +10,7 @@ module Analysis
     #
     # @see {GraphQL::Analysis::Analyzer} AST Analyzers for queries
     class Visitor < GraphQL::Language::StaticVisitor
-      def initialize(query:, analyzers:)
+      def initialize(query:, analyzers:, timeout:)
         @analyzers = analyzers
         @path = []
         @object_types = []
@@ -24,6 +24,11 @@ def initialize(query:, analyzers:)
         @types = query.types
         @response_path = []
         @skip_stack = [false]
+        @timeout_time = if timeout
+          Process.clock_gettime(Process::CLOCK_MONOTONIC, :float_second) + timeout
+        else
+          Float::INFINITY
+        end
         super(query.selected_operation)
       end
 
@@ -72,28 +77,25 @@ def response_path
         module_eval <<-RUBY, __FILE__, __LINE__
         def call_on_enter_#{node_type}(node, parent)
           @analyzers.each do |a|
-            begin
-              a.on_enter_#{node_type}(node, parent, self)
-            rescue AnalysisError => err
-              @rescued_errors << err
-            end
+            a.on_enter_#{node_type}(node, parent, self)
+          rescue AnalysisError => err
+            @rescued_errors << err
           end
         end
 
         def call_on_leave_#{node_type}(node, parent)
           @analyzers.each do |a|
-            begin
-              a.on_leave_#{node_type}(node, parent, self)
-            rescue AnalysisError => err
-              @rescued_errors << err
-            end
+            a.on_leave_#{node_type}(node, parent, self)
+          rescue AnalysisError => err
+            @rescued_errors << err
           end
         end
 
         RUBY
       end
 
       def on_operation_definition(node, parent)
+        check_timeout
         object_type = @schema.root_type_for_operation(node.operation_type)
         @object_types.push(object_type)
         @path.push("#{node.operation_type}#{node.name ? " #{node.name}" : ""}")
@@ -104,31 +106,27 @@ def on_operation_definition(node, parent)
         @path.pop
       end
 
-      def on_fragment_definition(node, parent)
-        on_fragment_with_type(node) do
-          @path.push("fragment #{node.name}")
-          @in_fragment_def = false
-          call_on_enter_fragment_definition(node, parent)
-          super
-          @in_fragment_def = false
-          call_on_leave_fragment_definition(node, parent)
-        end
-      end
-
       def on_inline_fragment(node, parent)
-        on_fragment_with_type(node) do
-          @path.push("...#{node.type ? " on #{node.type.name}" : ""}")
-          @skipping = @skip_stack.last || skip?(node)
-          @skip_stack << @skipping
-
-          call_on_enter_inline_fragment(node, parent)
-          super
-          @skipping = @skip_stack.pop
-          call_on_leave_inline_fragment(node, parent)
+        check_timeout
+        object_type = if node.type
+          @types.type(node.type.name)
+        else
+          @object_types.last
         end
+        @object_types.push(object_type)
+        @path.push("...#{node.type ? " on #{node.type.name}" : ""}")
+        @skipping = @skip_stack.last || skip?(node)
+        @skip_stack << @skipping
+        call_on_enter_inline_fragment(node, parent)
+        super
+        @skipping = @skip_stack.pop
+        call_on_leave_inline_fragment(node, parent)
+        @object_types.pop
+        @path.pop
       end
 
       def on_field(node, parent)
+        check_timeout
         @response_path.push(node.alias || node.name)
         parent_type = @object_types.last
         # This could be nil if the previous field wasn't found:
@@ -156,6 +154,7 @@ def on_field(node, parent)
       end
 
       def on_directive(node, parent)
+        check_timeout
         directive_defn = @schema.directives[node.name]
         @directive_definitions.push(directive_defn)
         call_on_enter_directive(node, parent)
@@ -165,6 +164,7 @@ def on_directive(node, parent)
       end
 
       def on_argument(node, parent)
+        check_timeout
         argument_defn = if (arg = @argument_definitions.last)
           arg_type = arg.type.unwrap
           if arg_type.kind.input_object?
@@ -190,6 +190,7 @@ def on_argument(node, parent)
       end
 
       def on_fragment_spread(node, parent)
+        check_timeout
         @path.push("... #{node.name}")
         @skipping = @skip_stack.last || skip?(node)
         @skip_stack << @skipping
@@ -267,16 +268,10 @@ def skip?(ast_node)
         !dir.empty? && !GraphQL::Execution::DirectiveChecks.include?(dir, query)
       end
 
-      def on_fragment_with_type(node)
-        object_type = if node.type
-          @types.type(node.type.name)
-        else
-          @object_types.last
+      def check_timeout
+        if Process.clock_gettime(Process::CLOCK_MONOTONIC, :float_second) > @timeout_time
+          raise GraphQL::Analysis::TimeoutError
         end
-        @object_types.push(object_type)
-        yield(node)
-        @object_types.pop
-        @path.pop
       end
     end
   end

lib/graphql/schema.rb

@@ -821,8 +821,8 @@ def subscription_execution_strategy(new_subscription_execution_strategy = nil, d
 
       attr_writer :validate_timeout
 
-      def validate_timeout(new_validate_timeout = nil)
-        if new_validate_timeout
+      def validate_timeout(new_validate_timeout = NOT_CONFIGURED)
+        if !NOT_CONFIGURED.equal?(new_validate_timeout)
           @validate_timeout = new_validate_timeout
         elsif defined?(@validate_timeout)
           @validate_timeout

spec/graphql/analysis_spec.rb

@@ -577,25 +577,69 @@ def article(id:)
   describe ".validate_timeout" do
     class AnalysisTimeoutSchema < GraphQL::Schema
       class SlowAnalyzer < GraphQL::Analysis::Analyzer
+        def initialize(...)
+          super
+          if query.context[:initialize_sleep]
+            sleep 0.6
+          end
+        end
+
         def on_enter_field(node, parent, visitor)
+          if node.name != "__typename"
+            sleep 0.1
+          end
+          super
+        end
+
+        def on_enter_directive(...)
+          sleep 0.1
+          super
+        end
+
+        def on_enter_argument(...)
           sleep 0.1
           super
         end
 
+        def on_enter_inline_fragment(...)
+          sleep 0.1
+          super
+        end
+
+        def on_enter_fragment_spread(node, _parent, _visitor)
+          if !node.name.include?("NoSleep")
+            sleep 0.1
+          end
+          super
+        end
+
         def result
           nil
         end
       end
 
       class Query < GraphQL::Schema::Object
-        field :f1, Int
+        field :f1, Int do
+          argument :a, String, required: false
+          argument :b, String, required: false
+          argument :c, String, required: false
+          argument :d, String, required: false
+          argument :e, String, required: false
+          argument :f, String, required: false
+        end
 
-        def f1
+        def f1(...)
           context[:int] ||= 0
           context[:int] += 1
         end
       end
 
+      class Nothing < GraphQL::Schema::Directive
+        locations(GraphQL::Schema::Directive::FIELD)
+        repeatable(true)
+      end
+      directive(Nothing)
+
       query(Query)
       query_analyzer(SlowAnalyzer)
       validate_timeout 0.5
@@ -608,5 +652,56 @@ def f1
       res2 = AnalysisTimeoutSchema.execute("{ f1: f1, f2: f1, f3: f1, f4: f1, f5: f1, f6: f1}")
       assert_equal ["Timeout on validation of query"], res2["errors"].map { |e| e["message"]}
     end
+
+    it "covers directives" do
+      res = AnalysisTimeoutSchema.execute("{ f1 @nothing @nothing }")
+      assert_equal({ "f1" => 1 }, res["data"])
+
+      res2 = AnalysisTimeoutSchema.execute("{ f1 @nothing @nothing @nothing @nothing @nothing }")
+      assert_equal ["Timeout on validation of query"], res2["errors"].map { |e| e["message"]}
+    end
+
+    it "covers arguments" do
+      res = AnalysisTimeoutSchema.execute("{ f1(a: \"a\", b: \"b\")}")
+      assert_equal({ "f1" => 1 }, res["data"])
+
+      res2 = AnalysisTimeoutSchema.execute('{ f1(a: "a", b: "b", c: "c", d: "d", e: "e", f: "f") }')
+      assert_equal ["Timeout on validation of query"], res2["errors"].map { |e| e["message"]}
+    end
+
+    it "covers inline fragments" do
+      res = AnalysisTimeoutSchema.execute("{ ... { f1 } ... { f1 } }")
+      assert_equal({ "f1" => 1 }, res["data"])
+
+      res2 = AnalysisTimeoutSchema.execute("{ ... { f1 } ... { f1 } ... { f1 } ... { f1 } ... { f1 } ... { f1 } }")
+      assert_equal ["Timeout on validation of query"], res2["errors"].map { |e| e["message"]}
+    end
+
+    it "covers operation definitions" do
+      res = AnalysisTimeoutSchema.execute('query Q1 { __typename }', operation_name: "Q1")
+      assert_equal({ "__typename" => "Query" }, res["data"])
+
+      res = AnalysisTimeoutSchema.execute('query Q1 { __typename }', operation_name: "Q1", context: { initialize_sleep: true })
+      assert_equal({ "__typename" => "Query" }, res["data"])
+    end
+
+    it "covers fragment spreads" do
+      res = AnalysisTimeoutSchema.execute("{ ...F } fragment F on Query { f1 }")
+      assert_equal({ "f1" => 1 }, res["data"])
+
+      res2 = AnalysisTimeoutSchema.execute('{ ...F ...F ...F ...F ...F ...F } fragment F on Query { f1 }')
+      assert_equal ["Timeout on validation of query"], res2["errors"].map { |e| e["message"]}
+    end
+
+    it "can be ignored" do
+      no_timeout_schema = Class.new(AnalysisTimeoutSchema) do
+        validate_timeout(nil)
+      end
+      res = no_timeout_schema.execute("{ f1 @nothing @nothing @nothing @nothing @nothing }")
+      assert_equal({ "f1" => 1 }, res["data"])
+
+      res2 = no_timeout_schema.execute('{ ...F ...F ...F ...F ...F ...F } fragment F on Query { f1 }')
+      assert_equal({ "f1" => 1 }, res2["data"])
+    end
   end
 end