GraphQL Blocking (#3819)

* Upload module skeleton.

* Blocking in apollo, very very first version

* Move graphql implementation to another module.

* Blocking for apollo-server-core, ugly but it works, lets find a better way

* Use real blocking data

* Set blocking to true.

* Throw before resolver execution in order to stop the operation's execution flow.

* Use HttpQueryError in apollo-server-core

* Blocking test in apollo-server-fastify

* Refactor graphql blocking.

* Remove previous implementation which only supported monitoring.
* Add new waf address in order to check the payload of every resolver.
* Use apm start resolver address instead of a new one.
* Remove mock and perform an actual call to the waf.

* Add non blocking graphql test

* Move abortController constructor to context creation.

This reduces the performance overhead due to just one instance is shared across the whole query exectution.

* Add pollo-server-express block tests

* Add unit tests.

* Add @apollo/server tests

* Update test rules for blocking by `graphql.server.resolver`

* Block with graphql templates data

* Add tests.

* Block with graphql data in graphql endpoint

* Fix tests.

* Execute @apollo/server and apollo-server-express tests

* Unify code in @apollo/server and apollo-server-core

* update comments

* Add appsec.blocked tag in blocked requests

* Add test with non graphql block response

* Tests for block with redirect

* Prevent creation of resolve span when it is blocked before the execution of the resolve code

* Refactor addResolver in order to get directives information.

* Add tests to block on directives.

* Add test for directives.

* Undo prevent creating resolve span

* Configurable graphql blocking json

* Refactor graphql

* Using resolver instead of resolvers.
* Change graphql channel name to be consistent with the others.

* Small changes in blocking

* Move resover information resolution to plugin.

* Revert "Move resover information resolution to plugin."

This reverts commit 7cc8561.

* Remove resolver information from context, pass it in a different field instead.

* Throw custom exception rather than send an empty array.

* Update packages/datadog-instrumentations/src/graphql.js

Co-authored-by: Ugaitz Urien <[email protected]>

* Change a bit apollo-server-core instrumentation

* Protect Header map, if in future version it is moved/removed, prevent breaks

* Remove some duplicated code

* Update packages/datadog-instrumentations/src/apollo-server.js

Co-authored-by: Carles Capell <[email protected]>

* Fix comments in the PR

* Fix PR comments.

* Fix some comments in the PR

* Move resolver information formatting to the plugin.

* Fix PR comments.

* Fix proper use of Promise.race.

---------

Co-authored-by: Ugaitz Urien <[email protected]>
Co-authored-by: Carles Capell <[email protected]>

.github/workflows/appsec.yml

@@ -135,6 +135,20 @@ jobs:
       - run: yarn test:appsec:plugins:ci
       - uses: codecov/codecov-action@v2
 
+  graphql:
+    runs-on: ubuntu-latest
+    env:
+      PLUGINS: apollo-server|apollo-server-express|apollo-server-fastify|apollo-server-core
+    steps:
+      - uses: actions/checkout@v2
+      - uses: ./.github/actions/node/setup
+      - run: yarn install
+      - uses: ./.github/actions/node/oldest
+      - run: yarn test:appsec:plugins:ci
+      - uses: ./.github/actions/node/latest
+      - run: yarn test:appsec:plugins:ci
+      - uses: codecov/codecov-action@v2
+
   mongodb-core:
     runs-on: ubuntu-latest
     services:

docs/test.ts

@@ -108,6 +108,7 @@ tracer.init({
     obfuscatorValueRegex: '.*',
     blockedTemplateHtml: './blocked.html',
     blockedTemplateJson: './blocked.json',
+    blockedTemplateGraphql: './blockedgraphql.json',
     eventTracking: {
       mode: 'safe'
     },

index.d.ts

@@ -567,6 +567,11 @@ export declare interface TracerOptions {
      */
     blockedTemplateJson?: string,
 
+    /**
+     * Specifies a path to a custom blocking template json file for graphql requests
+     */
+    blockedTemplateGraphql?: string,
+
     /**
      * Controls the automated user event tracking configuration
      */

integration-tests/graphql.spec.js

@@ -79,7 +79,6 @@ describe('graphql', () => {
           {
             id: 'test-rule-id-1',
             name: 'test-rule-name-1',
-            on_match: ['block'],
             tags:
             {
               category: 'attack_attempt',
@@ -92,8 +91,8 @@ describe('graphql', () => {
               operator_value: '',
               parameters: [
                 {
-                  address: 'graphql.server.all_resolvers',
-                  key_path: ['images', '0', 'category'],
+                  address: 'graphql.server.resolver',
+                  key_path: ['images', 'category'],
                   value: 'testattack',
                   highlight: ['testattack']
                 }

integration-tests/graphql/graphql-rules.json

@@ -17,6 +17,9 @@
               "inputs": [
                 {
                   "address": "graphql.server.all_resolvers"
+                },
+                {
+                  "address": "graphql.server.resolver"
                 }
               ],
               "list": [
@@ -27,7 +30,7 @@
           }
         ],
         "transformers": ["lowercase"],
-        "on_match": ["block"]
+        "on_match": []
       }
     ]
 }

packages/datadog-instrumentations/src/apollo-server-core.js

@@ -0,0 +1,41 @@
+'use strict'
+
+const { AbortController } = require('node-abort-controller')
+const { addHook } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+const dc = require('dc-polyfill')
+
+const requestChannel = dc.tracingChannel('datadog:apollo-server-core:request')
+
+addHook({ name: 'apollo-server-core', file: 'dist/runHttpQuery.js', versions: ['>3.0.0'] }, runHttpQueryModule => {
+  const HttpQueryError = runHttpQueryModule.HttpQueryError
+
+  shimmer.wrap(runHttpQueryModule, 'runHttpQuery', function wrapRunHttpQuery (originalRunHttpQuery) {
+    return async function runHttpQuery () {
+      if (!requestChannel.start.hasSubscribers) {
+        return originalRunHttpQuery.apply(this, arguments)
+      }
+
+      const abortController = new AbortController()
+      const abortData = {}
+
+      const runHttpQueryResult = requestChannel.tracePromise(
+        originalRunHttpQuery,
+        { abortController, abortData },
+        this,
+        ...arguments)
+
+      const abortPromise = new Promise((resolve, reject) => {
+        abortController.signal.addEventListener('abort', (event) => {
+          // runHttpQuery callbacks are writing the response on resolve/reject.
+          // We should return blocking data in the apollo-server-core HttpQueryError object
+          reject(new HttpQueryError(abortData.statusCode, abortData.message, true, abortData.headers))
+        }, { once: true })
+      })
+
+      return Promise.race([runHttpQueryResult, abortPromise])
+    }
+  })
+
+  return runHttpQueryModule
+})

packages/datadog-instrumentations/src/apollo-server.js

@@ -0,0 +1,83 @@
+'use strict'
+
+const { AbortController } = require('node-abort-controller')
+const dc = require('dc-polyfill')
+
+const { addHook } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const graphqlMiddlewareChannel = dc.tracingChannel('datadog:apollo:middleware')
+
+const requestChannel = dc.tracingChannel('datadog:apollo:request')
+
+let HeaderMap
+
+function wrapExecuteHTTPGraphQLRequest (originalExecuteHTTPGraphQLRequest) {
+  return async function executeHTTPGraphQLRequest () {
+    if (!HeaderMap || !requestChannel.start.hasSubscribers) {
+      return originalExecuteHTTPGraphQLRequest.apply(this, arguments)
+    }
+
+    const abortController = new AbortController()
+    const abortData = {}
+
+    const graphqlResponseData = requestChannel.tracePromise(
+      originalExecuteHTTPGraphQLRequest,
+      { abortController, abortData },
+      this,
+      ...arguments)
+
+    const abortPromise = new Promise((resolve, reject) => {
+      abortController.signal.addEventListener('abort', (event) => {
+        // This method is expected to return response data
+        // with headers, status and body
+        const headers = new HeaderMap()
+        Object.keys(abortData.headers).forEach(key => {
+          headers.set(key, abortData.headers[key])
+        })
+
+        resolve({
+          headers: headers,
+          status: abortData.statusCode,
+          body: {
+            kind: 'complete',
+            string: abortData.message
+          }
+        })
+      }, { once: true })
+    })
+
+    return Promise.race([abortPromise, graphqlResponseData])
+  }
+}
+
+function apolloExpress4Hook (express4) {
+  shimmer.wrap(express4, 'expressMiddleware', function wrapExpressMiddleware (originalExpressMiddleware) {
+    return function expressMiddleware (server, options) {
+      const originalMiddleware = originalExpressMiddleware.apply(this, arguments)
+
+      return shimmer.wrap(originalMiddleware, function (req, res, next) {
+        if (!graphqlMiddlewareChannel.start.hasSubscribers) {
+          return originalMiddleware.apply(this, arguments)
+        }
+
+        return graphqlMiddlewareChannel.traceSync(originalMiddleware, { req }, this, ...arguments)
+      })
+    }
+  })
+  return express4
+}
+
+function apolloHeaderMapHook (headerMap) {
+  HeaderMap = headerMap.HeaderMap
+  return headerMap
+}
+
+function apolloServerHook (apolloServer) {
+  shimmer.wrap(apolloServer.ApolloServer.prototype, 'executeHTTPGraphQLRequest', wrapExecuteHTTPGraphQLRequest)
+  return apolloServer
+}
+
+addHook({ name: '@apollo/server', file: 'dist/cjs/ApolloServer.js', versions: ['>=4.0.0'] }, apolloServerHook)
+addHook({ name: '@apollo/server', file: 'dist/cjs/express4/index.js', versions: ['>=4.0.0'] }, apolloExpress4Hook)
+addHook({ name: '@apollo/server', file: 'dist/cjs/utils/HeaderMap.js', versions: ['>=4.0.0'] }, apolloHeaderMapHook)

packages/datadog-instrumentations/src/graphql.js

@@ -1,5 +1,7 @@
 'use strict'
 
+const { AbortController } = require('node-abort-controller')
+
 const {
   addHook,
   channel,
@@ -37,6 +39,13 @@ const validateStartCh = channel('apm:graphql:validate:start')
 const validateFinishCh = channel('apm:graphql:validate:finish')
 const validateErrorCh = channel('apm:graphql:validate:error')
 
+class AbortError extends Error {
+  constructor (message) {
+    super(message)
+    this.name = 'AbortError'
+  }
+}
+
 function getOperation (document, operationName) {
   if (!document || !Array.isArray(document.definitions)) {
     return
@@ -175,11 +184,11 @@ function wrapExecute (execute) {
           docSource: documentSources.get(document)
         })
 
-        const context = { source, asyncResource, fields: {} }
+        const context = { source, asyncResource, fields: {}, abortController: new AbortController() }
 
         contexts.set(contextValue, context)
 
-        return callInAsyncScope(exe, asyncResource, this, arguments, (err, res) => {
+        return callInAsyncScope(exe, asyncResource, this, arguments, context.abortController, (err, res) => {
           if (finishResolveCh.hasSubscribers) finishResolvers(context)
 
           const error = err || (res && res.errors && res.errors[0])
@@ -207,7 +216,7 @@ function wrapResolve (resolve) {
 
     const field = assertField(context, info, args)
 
-    return callInAsyncScope(resolve, field.asyncResource, this, arguments, (err) => {
+    return callInAsyncScope(resolve, field.asyncResource, this, arguments, context.abortController, (err) => {
       updateFieldCh.publish({ field, info, err })
     })
   }
@@ -217,10 +226,15 @@ function wrapResolve (resolve) {
   return resolveAsync
 }
 
-function callInAsyncScope (fn, aR, thisArg, args, cb) {
+function callInAsyncScope (fn, aR, thisArg, args, abortController, cb) {
   cb = cb || (() => {})
 
   return aR.runInAsyncScope(() => {
+    if (abortController?.signal.aborted) {
+      cb(null, null)
+      throw new AbortError('Aborted')
+    }
+
     try {
       const result = fn.apply(thisArg, args)
       if (result && typeof result.then === 'function') {

packages/datadog-instrumentations/src/helpers/hooks.js

@@ -1,6 +1,8 @@
 'use strict'
 
 module.exports = {
+  '@apollo/server': () => require('../apollo-server'),
+  'apollo-server-core': () => require('../apollo-server-core'),
   '@aws-sdk/smithy-client': () => require('../aws-sdk'),
   '@cucumber/cucumber': () => require('../cucumber'),
   '@playwright/test': () => require('../playwright'),

packages/datadog-plugin-graphql/src/resolve.js

@@ -1,6 +1,7 @@
 'use strict'
 
 const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+const dc = require('dc-polyfill')
 
 const collapsedPathSym = Symbol('collapsedPaths')
 
@@ -14,8 +15,6 @@ class GraphQLResolvePlugin extends TracingPlugin {
     if (!shouldInstrument(this.config, path)) return
     const computedPathString = path.join('.')
 
-    addResolver(context, info, args)
-
     if (this.config.collapse) {
       if (!context[collapsedPathSym]) {
         context[collapsedPathSym] = {}
@@ -55,6 +54,10 @@ class GraphQLResolvePlugin extends TracingPlugin {
           span.setTag(`graphql.variables.${name}`, variables[name])
         })
     }
+
+    if (this.resolverStartCh.hasSubscribers) {
+      this.resolverStartCh.publish({ context, resolverInfo: getResolverInfo(info, args) })
+    }
   }
 
   constructor (...args) {
@@ -69,6 +72,8 @@ class GraphQLResolvePlugin extends TracingPlugin {
       field.finishTime = span._getTime ? span._getTime() : 0
       field.error = field.error || err
     })
+
+    this.resolverStartCh = dc.channel('datadog:graphql:resolver:start')
   }
 
   configure (config) {
@@ -109,28 +114,31 @@ function withCollapse (responsePathAsArray) {
   }
 }
 
-function addResolver (context, info, args) {
-  if (info.rootValue && !info.rootValue[info.fieldName]) {
-    return
-  }
+function getResolverInfo (info, args) {
+  let resolverInfo = null
+  const resolverVars = {}
 
-  if (!context.resolvers) {
-    context.resolvers = {}
+  if (args && Object.keys(args).length) {
+    Object.assign(resolverVars, args)
   }
 
-  const resolvers = context.resolvers
-
-  if (!resolvers[info.fieldName]) {
-    if (args && Object.keys(args).length) {
-      resolvers[info.fieldName] = [args]
-    } else {
-      resolvers[info.fieldName] = []
+  const directives = info.fieldNodes[0].directives
+  for (const directive of directives) {
+    const argList = {}
+    for (const argument of directive['arguments']) {
+      argList[argument.name.value] = argument.value.value
     }
-  } else {
-    if (args && Object.keys(args).length) {
-      resolvers[info.fieldName].push(args)
+
+    if (Object.keys(argList).length) {
+      resolverVars[directive.name.value] = argList
     }
   }
+
+  if (Object.keys(resolverVars).length) {
+    resolverInfo = { [info.fieldName]: resolverVars }
+  }
+
+  return resolverInfo
 }
 
 module.exports = GraphQLResolvePlugin

packages/dd-trace/src/appsec/addresses.js

@@ -13,6 +13,7 @@ module.exports = {
   HTTP_INCOMING_RESPONSE_HEADERS: 'server.response.headers.no_cookies',
   // TODO: 'server.response.trailers',
   HTTP_INCOMING_GRAPHQL_RESOLVERS: 'graphql.server.all_resolvers',
+  HTTP_INCOMING_GRAPHQL_RESOLVER: 'graphql.server.resolver',
 
   HTTP_CLIENT_IP: 'http.client_ip',