Merge pull request #3813 from DataDog/vpellan/graphql-unified-tracer-…

…require

.github/workflows/system-tests.yml

@@ -12,7 +12,8 @@ env:
   REGISTRY: ghcr.io
   REPO: ghcr.io/datadog/dd-trace-rb
   ST_REF: main
-  FORCE_TESTS:
+  FORCE_TESTS: -F tests/appsec/waf/test_addresses.py::Test_GraphQL -F tests/appsec/test_blocking_addresses.py::Test_BlockingGraphqlResolvers
+  FORCE_TESTS_SCENARIO: GRAPHQL_APPSEC
 
 jobs:
   build-harness:
@@ -96,6 +97,7 @@ jobs:
           - rails61
           - rails70
           - rails71
+          - graphql23
     runs-on: ubuntu-latest
     name: Build (${{ matrix.app }})
     steps:
@@ -236,6 +238,9 @@ jobs:
           - library: ruby
             app: rack
             scenario: TELEMETRY_METRIC_GENERATION_DISABLED
+          - library: ruby
+            app: graphql23
+            scenario: GRAPHQL_APPSEC
     runs-on: ubuntu-latest
     needs:
       - build-harness
@@ -268,7 +273,7 @@ jobs:
           docker image list
       - name: Run scenario
         run: |
-          ./run.sh ++docker ${{ matrix.scenario }} ${{ env.FORCE_TESTS }}
+          ./run.sh ++docker ${{ matrix.scenario }} ${{matrix.scenario == env.FORCE_TESTS_SCENARIO && env.FORCE_TESTS || ''}}
         env:
           DD_API_KEY: ${{ secrets.DD_APPSEC_SYSTEM_TESTS_API_KEY }}
       - name: Archive logs
@@ -296,6 +301,7 @@ jobs:
           - rails61
           - rails70
           - rails71
+          - graphql23
     runs-on: ubuntu-latest
     needs:
       - test
@@ -339,6 +345,7 @@ jobs:
           - weblog-rails60
           - weblog-rails61
           - weblog-rails70
+          - weblog-graphql23
     runs-on: ubuntu-latest
     needs:
       - test

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

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'graphql/language/nodes'
+require 'graphql'
 
 require_relative '../../../instrumentation/gateway/argument'
 
@@ -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

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

@@ -3,6 +3,10 @@
 require_relative '../patcher'
 require_relative 'gateway/watcher'
 
+if Gem.loaded_specs['graphql'] && Gem.loaded_specs['graphql'].version >= Gem::Version.new('2.0.19')
+  require_relative 'appsec_trace'
+end
+
 module Datadog
   module AppSec
     module Contrib
@@ -22,7 +26,6 @@ def target_version
           end
 
           def patch
-            require_relative 'appsec_trace'
             Gateway::Watcher.watch
             ::GraphQL::Schema.trace_with(AppSecTrace)
             Patcher.instance_variable_set(:@patched, true)

lib/datadog/tracing/contrib/graphql/unified_trace.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'graphql/tracing'
+require 'graphql'
 
 module Datadog
   module Tracing

lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb

@@ -1,5 +1,9 @@
 # frozen_string_literal: true
 
+if Gem.loaded_specs['graphql'] && Gem.loaded_specs['graphql'].version >= Gem::Version.new('2.0.19')
+  require_relative 'unified_trace'
+end
+
 module Datadog
   module Tracing
     module Contrib
@@ -9,7 +13,6 @@ module UnifiedTracePatcher
           module_function
 
           def patch!(schemas, options)
-            require_relative 'unified_trace'
             if schemas.empty?
               ::GraphQL::Schema.trace_with(UnifiedTrace, **options)
             else

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/loading_spec.rb

@@ -0,0 +1,33 @@
+require 'shellwords'
+
+RSpec.describe 'loading graphql' do
+  context 'then datadog' do
+    let(:code) do
+      <<-E
+        require "graphql"
+        require "datadog"
+        exit 0
+      E
+    end
+
+    it 'loads successfully by itself' do
+      rv = system("ruby -e #{Shellwords.shellescape(code)}")
+      expect(rv).to be true
+    end
+  end
+
+  context 'after datadog' do
+    let(:code) do
+      <<-E
+        require "datadog"
+        require "graphql"
+        exit 0
+      E
+    end
+
+    it 'loads successfully by itself' do
+      rv = system("ruby -e #{Shellwords.shellescape(code)}")
+      expect(rv).to be true
+    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]),