fix apollo-server v5 tests

integration-tests/appsec/graphql.spec.js

@@ -14,7 +14,7 @@ describe('graphql', () => {
   let sandbox, cwd, agent, webFile, proc
 
   before(async function () {
-    sandbox = await createSandbox(['@apollo/server@4', 'graphql', 'koalas'])
+    sandbox = await createSandbox(['@apollo/server', 'graphql', 'koalas'])
     cwd = sandbox.folder
     webFile = path.join(cwd, 'graphql/index.js')
   })
@@ -43,7 +43,8 @@ describe('graphql', () => {
       assert.propertyVal(headers, 'host', `127.0.0.1:${agent.port}`)
       assert.isArray(payload)
       assert.strictEqual(payload.length, 2)
-      assert.propertyVal(payload[1][0], 'name', 'express.request')
+      // Apollo server 5 is using Node.js http server instead of express
+      assert.propertyVal(payload[1][0], 'name', 'web.request')
       assert.propertyVal(payload[1][0].metrics, '_dd.appsec.enabled', 1)
       assert.property(payload[1][0].metrics, '_dd.appsec.waf.duration')
       assert.notProperty(payload[1][0].meta, '_dd.appsec.event')
@@ -104,7 +105,8 @@ describe('graphql', () => {
       assert.propertyVal(headers, 'host', `127.0.0.1:${agent.port}`)
       assert.isArray(payload)
       assert.strictEqual(payload.length, 2)
-      assert.propertyVal(payload[1][0], 'name', 'express.request')
+      // Apollo server 5 is using Node.js http server instead of express
+      assert.propertyVal(payload[1][0], 'name', 'web.request')
       assert.propertyVal(payload[1][0].metrics, '_dd.appsec.enabled', 1)
       assert.property(payload[1][0].metrics, '_dd.appsec.waf.duration')
       assert.propertyVal(payload[1][0].meta, 'appsec.event', 'true')

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

@@ -6,6 +6,7 @@ const { addHook } = require('./helpers/instrument')
 const shimmer = require('../../datadog-shimmer')
 
 const graphqlMiddlewareChannel = dc.tracingChannel('datadog:apollo:middleware')
+const apolloHttpServerChannel = dc.tracingChannel('datadog:apollo:httpserver')
 
 const requestChannel = dc.tracingChannel('datadog:apollo:request')
 
@@ -77,17 +78,47 @@ function apolloServerHook (apolloServer) {
   return apolloServer
 }
 
-addHook(
-  { name: '@apollo/server', file: 'dist/cjs/ApolloServer.js', versions: ['>=4.0.0 <5.0.0'] },
-  apolloServerHook
-)
+function apolloDrainHttpServerHook (drainModule) {
+  shimmer.wrap(drainModule, 'ApolloServerPluginDrainHttpServer',
+    (original) => function ApolloServerPluginDrainHttpServerWithDdTrace (options) {
+      if (options.httpServer && typeof options.httpServer.listen === 'function') {
+        shimmer.wrap(options.httpServer, 'listen', (originalListen) => function wrappedListen () {
+          shimmer.wrap(this, 'emit', (emit) => function wrappedEmit (event, req, res) {
+            if (event === 'request' && req && res && apolloHttpServerChannel.start.hasSubscribers) {
+              return apolloHttpServerChannel.traceSync(emit, { req }, this, ...arguments)
+            }
+
+            return emit.apply(this, arguments)
+          })
+
+          return originalListen.apply(this, arguments)
+        })
+      }
+
+      return original.apply(this, arguments)
+    })
+
+  return drainModule
+}
+
+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)
 
 addHook(
-  { name: '@apollo/server', file: 'dist/cjs/express4/index.js', versions: ['>=4.0.0 <5.0.0'] },
-  apolloExpress4Hook
+  { name: '@apollo/server', file: 'dist/cjs/plugin/drainHttpServer/index.js', versions: ['>=5.0.0'] },
+  apolloDrainHttpServerHook
 )
 
 addHook(
-  { name: '@apollo/server', file: 'dist/cjs/utils/HeaderMap.js', versions: ['>=4.0.0 <5.0.0'] },
-  apolloHeaderMapHook
+  { name: '@apollo/server', file: 'dist/cjs/runHttpQuery.js', versions: ['>=5.0.0'] },
+  (runHttpQueryModule) => {
+    shimmer.wrap(runHttpQueryModule, 'runHttpQuery', function wrapRunHttpQuery (originalRunHttpQuery) {
+      return wrapExecuteHTTPGraphQLRequest(originalRunHttpQuery)
+    })
+
+    return runHttpQueryModule
+  }
 )

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

@@ -6,6 +6,7 @@ const dc = require('dc-polyfill')
 module.exports = {
   apolloChannel: dc.tracingChannel('datadog:apollo:request'),
   apolloServerCoreChannel: dc.tracingChannel('datadog:apollo-server-core:request'),
+  apolloHttpServerChannel: dc.tracingChannel('datadog:apollo:httpserver'),
   bodyParser: dc.channel('datadog:body-parser:read:finish'),
   childProcessExecutionTracingChannel: dc.tracingChannel('datadog:child_process:execution'),
   cookieParser: dc.channel('datadog:cookie-parser:read:finish'),

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

@@ -14,6 +14,7 @@ const web = require('../plugins/util/web')
 const {
   startGraphqlResolve,
   graphqlMiddlewareChannel,
+  apolloHttpServerChannel,
   apolloChannel,
   apolloServerCoreChannel
 } = require('./channels')
@@ -60,6 +61,22 @@ function enterInApolloMiddleware (data) {
   })
 }
 
+function enterInApolloHttpServer (data) {
+  const req = data?.req || storage('legacy').getStore()?.req
+  if (!req) return
+
+  graphqlRequestData.set(req, {
+    inApolloHttpServer: true,
+    blocked: false
+  })
+}
+
+function exitFromApolloHttpServer (data) {
+  const req = data?.req || storage('legacy').getStore()?.req
+  const requestData = graphqlRequestData.get(req)
+  if (requestData) requestData.inApolloHttpServer = false
+}
+
 function enterInApolloServerCoreRequest () {
   const req = storage('legacy').getStore()?.req
   if (!req) return
@@ -80,7 +97,9 @@ function enterInApolloRequest () {
   const req = storage('legacy').getStore()?.req
 
   const requestData = graphqlRequestData.get(req)
-  if (requestData?.inApolloMiddleware) {
+  if (requestData) {
+    // Set isInGraphqlRequest=true since this function only runs for GraphQL requests
+    // This works for both Apollo v4 (middleware) and v5 (HTTP server) contexts
     requestData.isInGraphqlRequest = true
     addSpecificEndpoint(req.method, req.originalUrl || req.url, specificBlockingTypes.GRAPHQL)
   }
@@ -131,6 +150,11 @@ function enableApollo () {
     start: enterInApolloRequest,
     asyncEnd: beforeWriteApolloGraphqlResponse
   })
+
+  apolloHttpServerChannel.subscribe({
+    start: enterInApolloHttpServer,
+    end: exitFromApolloHttpServer
+  })
 }
 
 function disableApollo () {
@@ -148,6 +172,11 @@ function disableApollo () {
     start: enterInApolloRequest,
     asyncEnd: beforeWriteApolloGraphqlResponse
   })
+
+  apolloHttpServerChannel.unsubscribe({
+    start: enterInApolloHttpServer,
+    end: exitFromApolloHttpServer
+  })
 }
 
 function enableGraphql () {

packages/dd-trace/test/appsec/graphql.apollo-server-express.plugin.spec.js

@@ -5,7 +5,7 @@ const {
   schema,
   resolvers,
   graphqlCommonTests
-} = require('./graphq.test-utils')
+} = require('./graphql.test-utils')
 
 withVersions('apollo-server-core', 'express', '>=4', expressVersion => {
   withVersions('apollo-server-core', 'apollo-server-express', apolloServerExpressVersion => {

packages/dd-trace/test/appsec/graphql.apollo-server-fastify.plugin.spec.js

@@ -5,7 +5,7 @@ const {
   schema,
   resolvers,
   graphqlCommonTests
-} = require('./graphq.test-utils')
+} = require('./graphql.test-utils')
 
 withVersions('apollo-server-core', 'fastify', '3', fastifyVersion => {
   withVersions('apollo-server-core', 'apollo-server-fastify', apolloServerFastifyVersion => {

packages/dd-trace/test/appsec/graphql.apollo-server.plugin.spec.js

@@ -6,15 +6,15 @@ const {
   schema,
   resolvers,
   graphqlCommonTests
-} = require('./graphq.test-utils')
+} = require('./graphql.test-utils')
 
 withVersions('apollo-server', '@apollo/server', apolloServerVersion => {
   const config = {}
   let ApolloServer, startStandaloneServer
   let server
 
   before(() => {
-    return agent.load(['express', 'graphql', 'apollo-server', 'http'], { client: false })
+    return agent.load(['graphql', 'apollo-server', 'http'], { client: false })
   })
 
   before(() => {