Add directives support for graphql appsec

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

@@ -32,29 +32,39 @@ 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
-                  if selection.arguments.any?
-                    selection.arguments.each do |arg|
-                      resolver_args[arg.name] =
-                        if arg.value.is_a?(::GraphQL::Language::Nodes::VariableIdentifier)
-                          query.provided_variables[arg.value.name]
-                        else
-                          arg.value
-                        end
-                    end
+                  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)
                 end
-                unless resolver_args.empty?
-                  args[query.operation_name || "query#{index + 1}"] ||= []
-                  args[query.operation_name || "query#{index + 1}"] << resolver_args
-                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?
               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
+              end
+            end
           end
         end
       end

sig/datadog/appsec/contrib/graphql/gateway/multiplex.rbs

@@ -19,6 +19,8 @@ module Datadog
             private
 
             def create_arguments_hash: () -> Hash[String, Array[Hash[String, String]]]
+
+            def set_hash_with_variables: (Hash[String, String] resolver_hash, Array[GraphQL::Language::Nodes::Argument] arguments, Hash[String, String|Integer] provided_variables) -> void
           end
         end
       end

spec/datadog/appsec/contrib/graphql/appsec_trace_spec.rb

@@ -28,10 +28,10 @@
     expect(result.to_h['errors']).to eq(
       [
         {
-          'message' => "Field 'error' doesn't exist on type 'TestGraphQLQuery'",
+          'message' => "Field 'error' doesn't exist on type 'Query'",
           'locations' => [{ 'line' => 1, 'column' => 13 }],
           'path' => ['query test', 'error'],
-          'extensions' => { 'code' => 'undefinedField', 'typeName' => 'TestGraphQLQuery', 'fieldName' => 'error' }
+          'extensions' => { 'code' => 'undefinedField', 'typeName' => 'Query', 'fieldName' => 'error' }
         }
       ]
     )
@@ -89,10 +89,10 @@
           'errors' =>
           [
             {
-              'message' => "Field 'error' doesn't exist on type 'TestGraphQLQuery'",
+              'message' => "Field 'error' doesn't exist on type 'Query'",
               'locations' => [{ 'line' => 1, 'column' => 13 }],
               'path' => ['query test', 'error'],
-              'extensions' => { 'code' => 'undefinedField', 'typeName' => 'TestGraphQLQuery', 'fieldName' => 'error' }
+              'extensions' => { 'code' => 'undefinedField', 'typeName' => 'Query', 'fieldName' => 'error' }
             }
           ]
         }

spec/datadog/appsec/contrib/graphql/integration_test_spec.rb

@@ -287,7 +287,7 @@
             it do
               post '/graphql', query: 'mutation { createUser(name: "$testattack") { user { name, id } } }'
               expect(JSON.parse(last_response.body)['errors'][0]['title']).to eq('Blocked')
-              expect(Users.users['$testattack']).to be_nil
+              expect(TestGraphQL::Users.users['$testattack']).to be_nil
             end
           end
         end
@@ -459,5 +459,118 @@
           end
         end
       end
+
+      describe 'a query with directives' do
+        subject(:response) { post '/graphql', _json: queries }
+
+        context 'with a non-triggering multiplex' do
+          let(:appsec_ruleset) { blocking_testattack }
+          let(:queries) do
+            [
+              {
+                'query' => 'query Test($format: String!) { user(id: 1) { name @case(format: $format) } }',
+                'variables' => { 'format' => 'upcase' }
+              }
+            ]
+          end
+
+          it do
+            expect(last_response.body).to eq(
+              [
+                { 'data' => { 'user' => { 'name' => 'BITS' } } },
+              ].to_json
+            )
+            expect(spans).to include(
+              an_object_having_attributes(
+                name: 'graphql.parse',
+              ),
+              an_object_having_attributes(
+                name: 'graphql.execute_multiplex',
+              ),
+              an_object_having_attributes(
+                name: 'graphql.execute',
+              )
+            )
+          end
+
+          it_behaves_like 'a POST 200 span'
+          it_behaves_like 'a trace with AppSec tags'
+          it_behaves_like 'a trace without AppSec events'
+        end
+
+        context 'with a multiplex containing a non-blocking query' do
+          let(:appsec_ruleset) { nonblocking_testattack }
+          let(:queries) do
+            [
+              {
+                'query' => 'query Test($format: String!) { user(id: 1) { name @case(format: $format) } }',
+                'variables' => { 'format' => '$testattack' }
+              }
+            ]
+          end
+
+          it do
+            expect(last_response.body).to eq(
+              [
+                { 'data' => { 'user' => { 'name' => 'Bits' } } }
+              ].to_json
+            )
+            expect(spans).to include(
+              an_object_having_attributes(
+                name: 'graphql.parse',
+              )
+            ).once
+            expect(spans).to include(
+              an_object_having_attributes(
+                name: 'graphql.execute_multiplex',
+              )
+            ).once
+            expect(spans).to include(
+              an_object_having_attributes(
+                name: 'graphql.execute',
+              )
+            ).once
+          end
+
+          it_behaves_like 'a POST 200 span'
+          it_behaves_like 'a trace with AppSec tags'
+          it_behaves_like 'a trace with AppSec events'
+        end
+
+        context 'with a multiplex containing a blocking query' do
+          let(:appsec_ruleset) { blocking_testattack }
+          let(:queries) do
+            [
+              {
+                'query' => 'query Test($format: String!) { user(id: 1) { name @case(format: $format) } }',
+                'variables' => { 'format' => '$testattack' }
+              }
+            ]
+          end
+
+          it do
+            expect(last_response.body).to eq(
+              [
+                { 'errors' => [{ 'title' => 'Blocked', 'detail' => 'Security provided by Datadog.' }] }
+              ].to_json
+            )
+            expect(spans).to include(
+              an_object_having_attributes(
+                name: 'graphql.parse',
+              )
+            ).once
+            expect(spans).to include(
+              an_object_having_attributes(
+                name: 'graphql.execute_multiplex',
+              )
+            ).once
+            expect(spans).not_to include(
+              an_object_having_attributes(
+                name: 'graphql.execute',
+              )
+            )
+          end
+        end
+      end
     end
   end

spec/datadog/tracing/contrib/graphql/support/application.rb

@@ -31,15 +31,16 @@
   # TODO: Cleaner way to reset the schema between tests (and most likely clean ::GraphQL::Schema too)
   # stub_const is required for GraphqlController, and we cannot use variables defined in let blocks in stub_const
   before do
-    Object.send(:remove_const, :TestGraphQLSchema) if defined?(TestGraphQLSchema)
-    Object.send(:remove_const, :TestGraphQLQuery) if defined?(TestGraphQLQuery)
-    Object.send(:remove_const, :TestGraphQLMutationType) if defined?(TestGraphQLMutationType)
-    Object.send(:remove_const, :Users) if defined?(Users)
-    Object.send(:remove_const, :TestUserType) if defined?(TestUserType)
+    TestGraphQL.send(:remove_const, :Case) if defined?(TestGraphQL::Case)
+    TestGraphQL.send(:remove_const, :Schema) if defined?(TestGraphQL::Schema)
+    TestGraphQL.send(:remove_const, :Query) if defined?(TestGraphQL::Query)
+    TestGraphQL.send(:remove_const, :MutationType) if defined?(TestGraphQL::MutationType)
+    TestGraphQL.send(:remove_const, :Users) if defined?(TestGraphQL::Users)
+    TestGraphQL.send(:remove_const, :UserType) if defined?(TestGraphQL::UserType)
     load 'spec/datadog/tracing/contrib/graphql/support/application_schema.rb'
   end
   let(:operation) { Datadog::AppSec::Reactive::Operation.new('test') }
-  let(:schema) { TestGraphQLSchema }
+  let(:schema) { TestGraphQL::Schema }
 end
 
 RSpec.shared_context 'with GraphQL multiplex' do
@@ -94,9 +95,9 @@ def execute
                          context: {}
                        }
                      end
-                     TestGraphQLSchema.multiplex(queries)
+                     TestGraphQL::Schema.multiplex(queries)
                    else
-                     TestGraphQLSchema.execute(
+                     TestGraphQL::Schema.execute(
                        query: params[:query],
                        operation_name: params[:operationName],
                        variables: prepare_variables(params[:variables]),