From 45055a9b4d1af5a1d1616bdb78efef939f23dcbf Mon Sep 17 00:00:00 2001
From: CarlesDD <carles.capell@datadoghq.com>
Date: Fri, 24 May 2024 08:57:41 +0200
Subject: [PATCH] Handle waf result in graphql

---
 packages/dd-trace/src/appsec/graphql.js       |  8 ++++++--
 packages/dd-trace/test/appsec/graphql.spec.js | 12 ++++++++++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/packages/dd-trace/src/appsec/graphql.js b/packages/dd-trace/src/appsec/graphql.js
index 52a17acdf3a..bf0ba6983b7 100644
--- a/packages/dd-trace/src/appsec/graphql.js
+++ b/packages/dd-trace/src/appsec/graphql.js
@@ -32,10 +32,14 @@ function onGraphqlStartResolve ({ context, resolverInfo }) {
   if (!resolverInfo || typeof resolverInfo !== 'object') return
 
   const actions = waf.run({ ephemeral: { [addresses.HTTP_INCOMING_GRAPHQL_RESOLVER]: resolverInfo } }, req)
-  if (actions?.includes('block')) {
+  if (
+    actions &&
+    (Object.keys(actions).includes('block_request') || Object.keys(actions).includes('redirect_request'))
+  ) {
     const requestData = graphqlRequestData.get(req)
     if (requestData?.isInGraphqlRequest) {
       requestData.blocked = true
+      requestData.wafAction = actions.block_request || actions.redirect_request
       context?.abortController?.abort()
     }
   }
@@ -87,7 +91,7 @@ function beforeWriteApolloGraphqlResponse ({ abortController, abortData }) {
     const rootSpan = web.root(req)
     if (!rootSpan) return
 
-    const blockingData = getBlockingData(req, specificBlockingTypes.GRAPHQL, rootSpan)
+    const blockingData = getBlockingData(req, specificBlockingTypes.GRAPHQL, rootSpan, requestData.wafAction)
     abortData.statusCode = blockingData.statusCode
     abortData.headers = blockingData.headers
     abortData.message = blockingData.body
diff --git a/packages/dd-trace/test/appsec/graphql.spec.js b/packages/dd-trace/test/appsec/graphql.spec.js
index 827c7915d06..1f3fcec6cc2 100644
--- a/packages/dd-trace/test/appsec/graphql.spec.js
+++ b/packages/dd-trace/test/appsec/graphql.spec.js
@@ -213,9 +213,17 @@ describe('GraphQL', () => {
         user: [{ id: '1234' }]
       }
 
+      const blockParameters = {
+        status_code: '401',
+        type: 'auto',
+        grpc_status_code: '10'
+      }
+
       const abortController = context.abortController
 
-      sinon.stub(waf, 'run').returns(['block'])
+      sinon.stub(waf, 'run').returns({
+        block_request: blockParameters
+      })
       sinon.stub(web, 'root').returns({})
 
       startGraphqlResolve.publish({ context, resolverInfo })
@@ -231,7 +239,7 @@ describe('GraphQL', () => {
       const abortData = {}
       apolloChannel.asyncEnd.publish({ abortController, abortData })
 
-      expect(blocking.getBlockingData).to.have.been.calledOnceWithExactly(req, 'graphql', {})
+      expect(blocking.getBlockingData).to.have.been.calledOnceWithExactly(req, 'graphql', {}, blockParameters)
     })
   })
 })