From 5e70a6903078c13cf7dd27252f6f780d26c33392 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Mon, 29 Jul 2024 11:13:26 +0200
Subject: [PATCH 01/10] Fix GraphQL patchers so that it does not need to be
 required

---
 lib/datadog/appsec/contrib/graphql/gateway/multiplex.rb      | 2 +-
 lib/datadog/appsec/contrib/graphql/patcher.rb                | 2 +-
 lib/datadog/tracing/contrib/graphql/unified_trace.rb         | 2 +-
 lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb | 3 ++-
 4 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/lib/datadog/appsec/contrib/graphql/gateway/multiplex.rb b/lib/datadog/appsec/contrib/graphql/gateway/multiplex.rb
index df5d82f3c12..6f3a02e5c16 100644
--- a/lib/datadog/appsec/contrib/graphql/gateway/multiplex.rb
+++ b/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'
 
diff --git a/lib/datadog/appsec/contrib/graphql/patcher.rb b/lib/datadog/appsec/contrib/graphql/patcher.rb
index 12388f0cffc..b50492e3bb5 100644
--- a/lib/datadog/appsec/contrib/graphql/patcher.rb
+++ b/lib/datadog/appsec/contrib/graphql/patcher.rb
@@ -2,6 +2,7 @@
 
 require_relative '../patcher'
 require_relative 'gateway/watcher'
+require_relative 'appsec_trace' if Gem.loaded_specs['graphql'] && Gem.loaded_specs['graphql'].version >= Gem::Version.new('2.0.19')
 
 module Datadog
   module AppSec
@@ -22,7 +23,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)
diff --git a/lib/datadog/tracing/contrib/graphql/unified_trace.rb b/lib/datadog/tracing/contrib/graphql/unified_trace.rb
index 1f4ba8f73e8..2bb739351b4 100644
--- a/lib/datadog/tracing/contrib/graphql/unified_trace.rb
+++ b/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
diff --git a/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb b/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb
index 0d3b1d33d79..5f8a479337a 100644
--- a/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb
+++ b/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative 'unified_trace' if Gem.loaded_specs['graphql'] && Gem.loaded_specs['graphql'].version >= Gem::Version.new('2.0.19')
+
 module Datadog
   module Tracing
     module Contrib
@@ -9,7 +11,6 @@ module UnifiedTracePatcher
           module_function
 
           def patch!(schemas, options)
-            require_relative 'unified_trace'
             if schemas.empty?
               ::GraphQL::Schema.trace_with(UnifiedTrace, **options)
             else

From 3956263f68485b514a9e550ddd2277120d2ae849 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Mon, 29 Jul 2024 11:25:01 +0200
Subject: [PATCH 02/10] Fix rubocop offenses on conditions for require of
 tracers

---
 lib/datadog/appsec/contrib/graphql/patcher.rb                | 5 ++++-
 lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/lib/datadog/appsec/contrib/graphql/patcher.rb b/lib/datadog/appsec/contrib/graphql/patcher.rb
index b50492e3bb5..b95d12f4ca5 100644
--- a/lib/datadog/appsec/contrib/graphql/patcher.rb
+++ b/lib/datadog/appsec/contrib/graphql/patcher.rb
@@ -2,7 +2,10 @@
 
 require_relative '../patcher'
 require_relative 'gateway/watcher'
-require_relative 'appsec_trace' if Gem.loaded_specs['graphql'] && Gem.loaded_specs['graphql'].version >= Gem::Version.new('2.0.19')
+
+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
diff --git a/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb b/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb
index 5f8a479337a..addea2950e2 100644
--- a/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb
+++ b/lib/datadog/tracing/contrib/graphql/unified_trace_patcher.rb
@@ -1,6 +1,8 @@
 # frozen_string_literal: true
 
-require_relative 'unified_trace' if Gem.loaded_specs['graphql'] && Gem.loaded_specs['graphql'].version >= Gem::Version.new('2.0.19')
+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

From 7af56c5edaee666a81ee372a745ec8b8b50ca3b8 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Mon, 29 Jul 2024 11:22:15 +0200
Subject: [PATCH 03/10] Add directives support for graphql appsec

---
 .../contrib/graphql/gateway/multiplex.rb      | 36 ++++++++++++-------
 1 file changed, 23 insertions(+), 13 deletions(-)

diff --git a/lib/datadog/appsec/contrib/graphql/gateway/multiplex.rb b/lib/datadog/appsec/contrib/graphql/gateway/multiplex.rb
index 6f3a02e5c16..db38d14f7b2 100644
--- a/lib/datadog/appsec/contrib/graphql/gateway/multiplex.rb
+++ b/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

From 193511613ab9f1d7ab3997d2b68e23def5c38b87 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Mon, 29 Jul 2024 11:35:03 +0200
Subject: [PATCH 04/10] Add sig for set_hash_with_variables

---
 sig/datadog/appsec/contrib/graphql/gateway/multiplex.rbs | 2 ++
 vendor/rbs/graphql/0/graphql.rbs                         | 3 +++
 2 files changed, 5 insertions(+)

diff --git a/sig/datadog/appsec/contrib/graphql/gateway/multiplex.rbs b/sig/datadog/appsec/contrib/graphql/gateway/multiplex.rbs
index 55e6c09e736..3465dcb1193 100644
--- a/sig/datadog/appsec/contrib/graphql/gateway/multiplex.rbs
+++ b/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
diff --git a/vendor/rbs/graphql/0/graphql.rbs b/vendor/rbs/graphql/0/graphql.rbs
index 6eb0872de50..9b7e34f137d 100644
--- a/vendor/rbs/graphql/0/graphql.rbs
+++ b/vendor/rbs/graphql/0/graphql.rbs
@@ -22,6 +22,9 @@ module GraphQL
 
       class Field
       end
+
+      class Argument
+      end
     end
   end
 

From 0ff8ccde5c3e71a7d19beca0bc73ff0eb1101156 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Mon, 29 Jul 2024 12:13:30 +0200
Subject: [PATCH 05/10] Reorganized test schema and added test directive

---
 .../contrib/graphql/appsec_trace_spec.rb      |   8 +-
 .../contrib/graphql/integration_test_spec.rb  |   2 +-
 .../contrib/graphql/support/application.rb    |  17 +-
 .../graphql/support/application_schema.rb     | 150 +++++++++++-------
 4 files changed, 107 insertions(+), 70 deletions(-)

diff --git a/spec/datadog/appsec/contrib/graphql/appsec_trace_spec.rb b/spec/datadog/appsec/contrib/graphql/appsec_trace_spec.rb
index 391a46d3865..1d9175fa81b 100644
--- a/spec/datadog/appsec/contrib/graphql/appsec_trace_spec.rb
+++ b/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' }
             }
           ]
         }
diff --git a/spec/datadog/appsec/contrib/graphql/integration_test_spec.rb b/spec/datadog/appsec/contrib/graphql/integration_test_spec.rb
index 76b46794488..d516dadf119 100644
--- a/spec/datadog/appsec/contrib/graphql/integration_test_spec.rb
+++ b/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
diff --git a/spec/datadog/tracing/contrib/graphql/support/application.rb b/spec/datadog/tracing/contrib/graphql/support/application.rb
index 55390481e48..88ab68a45fc 100644
--- a/spec/datadog/tracing/contrib/graphql/support/application.rb
+++ b/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]),
diff --git a/spec/datadog/tracing/contrib/graphql/support/application_schema.rb b/spec/datadog/tracing/contrib/graphql/support/application_schema.rb
index 23b76ebafeb..1805997047a 100644
--- a/spec/datadog/tracing/contrib/graphql/support/application_schema.rb
+++ b/spec/datadog/tracing/contrib/graphql/support/application_schema.rb
@@ -1,80 +1,116 @@
 require 'graphql'
 require 'json'
 
-class TestUserType < ::GraphQL::Schema::Object
-  field :id, ::GraphQL::Types::ID, null: false
-  field :name, ::GraphQL::Types::String, null: true
-  field :created_at, ::GraphQL::Types::String, null: false
-  field :updated_at, ::GraphQL::Types::String, null: false
-end
-
-class TestGraphQLQuery < ::GraphQL::Schema::Object
-  field :user, TestUserType, null: false, description: 'Find an user by ID' do
-    argument :id, ::GraphQL::Types::ID, required: true
-  end
-
-  def user(id:)
-    return OpenStruct.new(id: id, name: 'Caniche') if Integer(id) == 10
+module TestGraphQL
+  class Case < GraphQL::Schema::Directive
+    argument :format, String
+    locations FIELD
 
-    OpenStruct.new(id: id, name: 'Bits')
+    TRANSFORMS = [
+      "upcase",
+      "downcase",
+      # ??
+    ]
+    # Implement the Directive API
+    def self.resolve(object, arguments, context)
+      path = context.namespace(:interpreter)[:current_path]
+      return_value = yield
+      transform_name = arguments[:format]
+      if TRANSFORMS.include?(transform_name) && return_value.respond_to?(transform_name)
+        return_value = return_value.public_send(transform_name)
+        response = context.namespace(:interpreter_runtime)[:runtime].final_result
+        *keys, last = path
+        keys.each do |key|
+          if response && (response = response[key])
+            next
+          else
+            break
+          end
+        end
+        if response
+          response[last] = return_value
+        end
+        nil
+      end
+    end
   end
 
-  field :userByName, TestUserType, null: false, description: 'Find an user by name' do
-    argument :name, ::GraphQL::Types::String, required: true
+  class UserType < ::GraphQL::Schema::Object
+    field :id, ::GraphQL::Types::ID, null: false
+    field :name, ::GraphQL::Types::String, null: true
+    field :created_at, ::GraphQL::Types::String, null: false
+    field :updated_at, ::GraphQL::Types::String, null: false
   end
 
-  # rubocop:disable Naming/MethodName
-  def userByName(name:)
-    return OpenStruct.new(id: 10, name: name) if name == 'Caniche'
+  class Query < ::GraphQL::Schema::Object
+    field :user, UserType, null: false, description: 'Find an user by ID' do
+      argument :id, ::GraphQL::Types::ID, required: true
+    end
 
-    OpenStruct.new(id: 1, name: name)
-  end
+    def user(id:)
+      return OpenStruct.new(id: id, name: 'Caniche') if Integer(id) == 10
 
-  field :mutationUserByName, TestUserType, null: false, description: 'Find an user by name' do
-    argument :name, ::GraphQL::Types::String, required: true
-  end
+      OpenStruct.new(id: id, name: 'Bits')
+    end
+
+    field :userByName, UserType, null: false, description: 'Find an user by name' do
+      argument :name, ::GraphQL::Types::String, required: true
+    end
+
+    # rubocop:disable Naming/MethodName
+    def userByName(name:)
+      return OpenStruct.new(id: 10, name: name) if name == 'Caniche'
 
-  def mutationUserByName(name:)
-    if Users.users[name].nil?
-      ::GraphQL::ExecutionError.new('User not found')
-    else
-      OpenStruct.new(id: Users.users[name], name: name)
+      OpenStruct.new(id: 1, name: name)
     end
+
+    field :mutationUserByName, UserType, null: false, description: 'Find an user by name' do
+      argument :name, ::GraphQL::Types::String, required: true
+    end
+
+    def mutationUserByName(name:)
+      if Users.users[name].nil?
+        ::GraphQL::ExecutionError.new('User not found')
+      else
+        OpenStruct.new(id: Users.users[name], name: name)
+      end
+    end
+    # rubocop:enable Naming/MethodName
   end
-  # rubocop:enable Naming/MethodName
-end
 
-class Users
-  class << self
-    def users
-      @users ||= {}
+  class Users
+    class << self
+      def users
+        @users ||= {}
+      end
     end
   end
-end
 
-class TestGraphQLMutationType < ::GraphQL::Schema::Object
-  class TestGraphQLMutation < ::GraphQL::Schema::Mutation
-    argument :name, ::GraphQL::Types::String, required: true
+  class MutationType < ::GraphQL::Schema::Object
+    class Mutation < ::GraphQL::Schema::Mutation
+      argument :name, ::GraphQL::Types::String, required: true
 
-    field :user, TestUserType
-    field :errors, [String], null: false
+      field :user, UserType
+      field :errors, [String], null: false
 
-    def resolve(name:)
-      if Users.users.nil? || Users.users[name].nil?
-        Users.users ||= {}
-        item = OpenStruct.new(id: Users.users.size + 1, name: name)
-        Users.users[name] = Users.users.size + 1
-        { user: item, errors: [] }
-      else
-        ::GraphQL::ExecutionError.new('User already exists')
+      def resolve(name:)
+        if Users.users.nil? || Users.users[name].nil?
+          Users.users ||= {}
+          item = OpenStruct.new(id: Users.users.size + 1, name: name)
+          Users.users[name] = Users.users.size + 1
+          { user: item, errors: [] }
+        else
+          ::GraphQL::ExecutionError.new('User already exists')
+        end
       end
     end
-  end
 
-  field :create_user, mutation: TestGraphQLMutation
-end
+    field :create_user, mutation: Mutation
+  end
 
-class TestGraphQLSchema < ::GraphQL::Schema
-  mutation(TestGraphQLMutationType)
-  query(TestGraphQLQuery)
-end
+  class Schema < ::GraphQL::Schema
+    mutation(MutationType)
+    query(Query)
+    directive(Case)
+  end
+end
\ No newline at end of file

From 6f88b2a32a36ae0c56634139d928b9aa933e3e08 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Mon, 29 Jul 2024 12:15:42 +0200
Subject: [PATCH 06/10] Added integration test for directive

---
 .../contrib/graphql/integration_test_spec.rb  | 113 ++++++++++++++++++
 1 file changed, 113 insertions(+)

diff --git a/spec/datadog/appsec/contrib/graphql/integration_test_spec.rb b/spec/datadog/appsec/contrib/graphql/integration_test_spec.rb
index d516dadf119..5209b92355a 100644
--- a/spec/datadog/appsec/contrib/graphql/integration_test_spec.rb
+++ b/spec/datadog/appsec/contrib/graphql/integration_test_spec.rb
@@ -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

From 24cd1c2d0fd6618462f27c19128a8164f5584724 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Mon, 29 Jul 2024 12:20:26 +0200
Subject: [PATCH 07/10] Fix Rubocop offenses on test directive

---
 .../contrib/graphql/support/application_schema.rb  | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/spec/datadog/tracing/contrib/graphql/support/application_schema.rb b/spec/datadog/tracing/contrib/graphql/support/application_schema.rb
index 1805997047a..ba297c1f77b 100644
--- a/spec/datadog/tracing/contrib/graphql/support/application_schema.rb
+++ b/spec/datadog/tracing/contrib/graphql/support/application_schema.rb
@@ -7,10 +7,10 @@ class Case < GraphQL::Schema::Directive
     locations FIELD
 
     TRANSFORMS = [
-      "upcase",
-      "downcase",
-      # ??
-    ]
+      'upcase',
+      'downcase',
+    ].freeze
+
     # Implement the Directive API
     def self.resolve(object, arguments, context)
       path = context.namespace(:interpreter)[:current_path]
@@ -27,9 +27,7 @@ def self.resolve(object, arguments, context)
             break
           end
         end
-        if response
-          response[last] = return_value
-        end
+        response[last] = return_value if response
         nil
       end
     end
@@ -113,4 +111,4 @@ class Schema < ::GraphQL::Schema
     query(Query)
     directive(Case)
   end
-end
\ No newline at end of file
+end

From 94394c715c99c2b058feab4ca272f8a6174f4abc Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Thu, 1 Aug 2024 09:19:14 +0200
Subject: [PATCH 08/10] Add trace_with testing for unified_tracer

---
 .../tracing/contrib/graphql/loading_spec.rb   | 31 +++++++++++++++++++
 .../graphql/unified_trace_patcher_spec.rb     | 18 +++++++++++
 2 files changed, 49 insertions(+)
 create mode 100644 spec/datadog/tracing/contrib/graphql/loading_spec.rb

diff --git a/spec/datadog/tracing/contrib/graphql/loading_spec.rb b/spec/datadog/tracing/contrib/graphql/loading_spec.rb
new file mode 100644
index 00000000000..0dacabfbd12
--- /dev/null
+++ b/spec/datadog/tracing/contrib/graphql/loading_spec.rb
@@ -0,0 +1,31 @@
+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
diff --git a/spec/datadog/tracing/contrib/graphql/unified_trace_patcher_spec.rb b/spec/datadog/tracing/contrib/graphql/unified_trace_patcher_spec.rb
index ff94f1fbc8e..d4a82adb25b 100644
--- a/spec/datadog/tracing/contrib/graphql/unified_trace_patcher_spec.rb
+++ b/spec/datadog/tracing/contrib/graphql/unified_trace_patcher_spec.rb
@@ -23,4 +23,22 @@
         end
       end
     end
+
+    # Not specific to unified trace patcher,
+    # But this should work the same way without the need to require the tracer in the schema.
+    describe '#trace_with' do
+      context 'with schema using trace_with' do
+        it_behaves_like 'graphql instrumentation with unified naming convention trace' do
+          before do
+            # Monkey patch the schema to use the unified tracer
+            # As we're not adding a new method, we cannot use allow(...).to receive(...)
+            # rubocop:disable Lint/ConstantDefinitionInBlock, RSpec/LeakyConstantDeclaration
+            class TestGraphQLSchema
+              trace_with Datadog::Tracing::Contrib::GraphQL::UnifiedTrace
+            end
+            # rubocop:enable Lint/ConstantDefinitionInBlock, RSpec/LeakyConstantDeclaration
+          end
+        end
+      end
+    end
   end

From aa7c28565f9e074a8e75ddd62fb347a422173604 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Thu, 1 Aug 2024 16:49:02 +0200
Subject: [PATCH 09/10] Add GraphQL scenario in system tests (and force execute
 GraphQL tests)

---
 .github/workflows/system-tests.yml | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/system-tests.yml b/.github/workflows/system-tests.yml
index 3fa130f5252..79b4294f412 100644
--- a/.github/workflows/system-tests.yml
+++ b/.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

From 6f6ce3e0466d565d253762978a7ca1de6d257e42 Mon Sep 17 00:00:00 2001
From: Victor Pellan <victor.pellan@datadoghq.com>
Date: Mon, 5 Aug 2024 19:58:02 +0200
Subject: [PATCH 10/10] Added require 'shellwords' to not cause flaky tests

---
 spec/datadog/tracing/contrib/graphql/loading_spec.rb | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/spec/datadog/tracing/contrib/graphql/loading_spec.rb b/spec/datadog/tracing/contrib/graphql/loading_spec.rb
index 0dacabfbd12..2b40dfe87a7 100644
--- a/spec/datadog/tracing/contrib/graphql/loading_spec.rb
+++ b/spec/datadog/tracing/contrib/graphql/loading_spec.rb
@@ -1,3 +1,5 @@
+require 'shellwords'
+
 RSpec.describe 'loading graphql' do
   context 'then datadog' do
     let(:code) do