Fix building of GraphQL queries arguments hash

lib/datadog/appsec/contrib/graphql/appsec_trace.rb

@@ -28,20 +28,6 @@ def execute_multiplex(multiplex:)
 
             multiplex_return
           end
-
-          private
-
-          def active_trace
-            return unless defined?(Datadog::Tracing)
-
-            Datadog::Tracing.active_trace
-          end
-
-          def active_span
-            return unless defined?(Datadog::Tracing)
-
-            Datadog::Tracing.active_span
-          end
         end
       end
     end

lib/datadog/appsec/contrib/graphql/gateway/multiplex.rb

@@ -19,7 +19,7 @@ def initialize(multiplex)
             end
 
             def arguments
-              @arguments ||= create_arguments_hash
+              @arguments ||= build_arguments_hash
             end
 
             def queries
@@ -28,41 +28,77 @@ def queries
 
             private
 
-            def create_arguments_hash
-              args = {}
-              @multiplex.queries.each_with_index do |query, index|
-                resolver_args = {}
-                resolver_dirs = {}
-                selections = (query.selected_operation.selections.dup if query.selected_operation) || []
-                # Iterative tree traversal
-                while selections.any?
-                  selection = selections.shift
-                  set_hash_with_variables(resolver_args, selection.arguments, query.provided_variables)
-                  selection.directives.each do |dir|
-                    resolver_dirs[dir.name] ||= {}
-                    set_hash_with_variables(resolver_dirs[dir.name], dir.arguments, query.provided_variables)
-                  end
-                  selections.concat(selection.selections)
+            # This method builds an array of argument hashes for each field with arguments in the query.
+            #
+            # For example, given the following query:
+            # query ($postSlug: ID = "my-first-post", $withComments: Boolean!) {
+            #   firstPost: post(slug: $postSlug) {
+            #     title
+            #     comments @include(if: $withComments) {
+            #       author { name }
+            #       content
+            #     }
+            #   }
+            # }
+            #
+            # The result would be:
+            # {"post"=>[{"slug"=>"my-first-post"}], "comments"=>[{"include"=>{"if"=>true}}]}
+            #
+            # Note that the `comments` "include" directive is included in the arguments list
+            def build_arguments_hash
+              queries.each_with_object({}) do |query, args_hash|
+                next unless query.selected_operation
+
+                arguments_from_selections(query.selected_operation.selections, query.variables, args_hash)
+              end
+            end
+
+            def arguments_from_selections(selections, query_variables, args_hash)
+              selections.each do |selection|
+                # rubocop:disable Style/ClassEqualityComparison
+                next unless selection.class.name == Integration::AST_NODE_CLASS_NAMES[:field]
+                # rubocop:enable Style/ClassEqualityComparison
+
+                selection_name = selection.alias || selection.name
+
+                if !selection.arguments.empty? || !selection.directives.empty?
+                  args_hash[selection_name] ||= []
+                  args_hash[selection_name] <<
+                    arguments_hash(selection.arguments, query_variables).merge!(
+                      arguments_from_directives(selection.directives, query_variables)
+                    )
                 end
-                next if resolver_args.empty? && resolver_dirs.empty?
 
-                args_resolver = (args[query.operation_name || "query#{index + 1}"] ||= [])
-                # We don't want to add empty hashes so we check again if the arguments and directives are empty
-                args_resolver << resolver_args unless resolver_args.empty?
-                args_resolver << resolver_dirs unless resolver_dirs.empty?
+                arguments_from_selections(selection.selections, query_variables, args_hash)
+              end
+            end
+
+            def arguments_from_directives(directives, query_variables)
+              directives.each_with_object({}) do |directive, args_hash|
+                # rubocop:disable Style/ClassEqualityComparison
+                next unless directive.class.name == Integration::AST_NODE_CLASS_NAMES[:directive]
+                # rubocop:enable Style/ClassEqualityComparison
+
+                args_hash[directive.name] = arguments_hash(directive.arguments, query_variables)
+              end
+            end
+
+            def arguments_hash(arguments, query_variables)
+              arguments.each_with_object({}) do |argument, args_hash|
+                args_hash[argument.name] = argument_value(argument, query_variables)
               end
-              args
             end
 
-            # Set the resolver hash (resolver_args and resolver_dirs) with the arguments and provided variables
-            def set_hash_with_variables(resolver_hash, arguments, provided_variables)
-              arguments.each do |arg|
-                resolver_hash[arg.name] =
-                  if arg.value.is_a?(::GraphQL::Language::Nodes::VariableIdentifier)
-                    provided_variables[arg.value.name]
-                  else
-                    arg.value
-                  end
+            def argument_value(argument, query_variables)
+              case argument.value.class.name
+              when Integration::AST_NODE_CLASS_NAMES[:variable_identifier]
+                # we need to pass query.variables here instead of query.provided_variables,
+                # since #provided_variables don't know anything about variable default value
+                query_variables[argument.value.name]
+              when Integration::AST_NODE_CLASS_NAMES[:input_object]
+                arguments_hash(argument.value.arguments, query_variables)
+              else
+                argument.value
               end
             end
           end

lib/datadog/appsec/contrib/graphql/gateway/watcher.rb

@@ -24,27 +24,30 @@ def watch_multiplex(gateway = Instrumentation.gateway)
                 gateway.watch('graphql.multiplex', :appsec) do |stack, gateway_multiplex|
                   block = false
                   event = nil
+
                   scope = AppSec::Scope.active_scope
 
-                  AppSec::Reactive::Operation.new('graphql.multiplex') do |op|
-                    GraphQL::Reactive::Multiplex.subscribe(op, scope.processor_context) do |result|
-                      event = {
-                        waf_result: result,
-                        trace: scope.trace,
-                        span: scope.service_entry_span,
-                        multiplex: gateway_multiplex,
-                        actions: result.actions
-                      }
+                  if scope
+                    AppSec::Reactive::Operation.new('graphql.multiplex') do |op|
+                      GraphQL::Reactive::Multiplex.subscribe(op, scope.processor_context) do |result|
+                        event = {
+                          waf_result: result,
+                          trace: scope.trace,
+                          span: scope.service_entry_span,
+                          multiplex: gateway_multiplex,
+                          actions: result.actions
+                        }
+
+                        if scope.service_entry_span
+                          scope.service_entry_span.set_tag('appsec.blocked', 'true') if result.actions.include?('block')
+                          scope.service_entry_span.set_tag('appsec.event', 'true')
+                        end
 
-                      if scope.service_entry_span
-                        scope.service_entry_span.set_tag('appsec.blocked', 'true') if result.actions.include?('block')
-                        scope.service_entry_span.set_tag('appsec.event', 'true')
+                        scope.processor_context.events << event
                       end
 
-                      scope.processor_context.events << event
+                      block = GraphQL::Reactive::Multiplex.publish(op, gateway_multiplex)
                     end
-
-                    block = GraphQL::Reactive::Multiplex.publish(op, gateway_multiplex)
                   end
 
                   next [nil, [[:block, event]]] if block

lib/datadog/appsec/contrib/graphql/integration.rb

@@ -13,6 +13,13 @@ class Integration
 
           MINIMUM_VERSION = Gem::Version.new('2.0.19')
 
+          AST_NODE_CLASS_NAMES = {
+            field: 'GraphQL::Language::Nodes::Field',
+            directive: 'GraphQL::Language::Nodes::Directive',
+            variable_identifier: 'GraphQL::Language::Nodes::VariableIdentifier',
+            input_object: 'GraphQL::Language::Nodes::InputObject',
+          }.freeze
+
           register_as :graphql, auto_patch: false
 
           def self.version
@@ -24,13 +31,19 @@ def self.loaded?
           end
 
           def self.compatible?
-            super && version >= MINIMUM_VERSION
+            super && version >= MINIMUM_VERSION && ast_node_classes_defined?
           end
 
           def self.auto_instrument?
             true
           end
 
+          def self.ast_node_classes_defined?
+            AST_NODE_CLASS_NAMES.all? do |_, class_name|
+              Object.const_defined?(class_name)
+            end
+          end
+
           def patcher
             Patcher
           end