Added GraphQL monitoring support (#6375)

* Added smoke-test SpringBoot GraphQL appication

* Arguments capturing in GraphQL

* Refactor

* Fixed tests

* Fixed tests

* Fixed GraphQL resolver address

* Missing smoke-test

* Fixed tests

dd-java-agent/appsec/src/main/java/com/datadog/appsec/event/data/KnownAddresses.java

@@ -104,6 +104,9 @@ public interface KnownAddresses {
   // XXX: Not really used yet, but it's a known address and we should not treat it as unknown.
   Address<Object> GRAPHQL_SERVER_RESOLVER = new Address<>("graphql.server.resolver");
 
+  Address<Map<String, ?>> SERVER_GRAPHQL_ALL_RESOLVERS =
+      new Address<>("server.graphql.all_resolvers");
+
   Address<String> USER_ID = new Address<>("usr.id");
 
   Address<Map<String, Object>> WAF_CONTEXT_PROCESSOR = new Address<>("waf.context.processor");
@@ -158,6 +161,8 @@ static Address<?> forName(String name) {
         return GRAPHQL_SERVER_ALL_RESOLVERS;
       case "graphql.server.resolver":
         return GRAPHQL_SERVER_RESOLVER;
+      case "server.graphql.all_resolvers":
+        return SERVER_GRAPHQL_ALL_RESOLVERS;
       case "usr.id":
         return USER_ID;
       case "waf.context.processor":

dd-java-agent/appsec/src/main/java/com/datadog/appsec/gateway/GatewayBridge.java

@@ -73,6 +73,7 @@ public class GatewayBridge {
   private volatile DataSubscriberInfo pathParamsSubInfo;
   private volatile DataSubscriberInfo respDataSubInfo;
   private volatile DataSubscriberInfo grpcServerRequestMsgSubInfo;
+  private volatile DataSubscriberInfo graphqlServerRequestMsgSubInfo;
   private volatile DataSubscriberInfo requestEndSubInfo;
 
   public GatewayBridge(
@@ -391,6 +392,33 @@ public void init() {
             }
           }
         });
+
+    subscriptionService.registerCallback(
+        EVENTS.graphqlServerRequestMessage(),
+        (RequestContext ctx_, Map<String, ?> data) -> {
+          AppSecRequestContext ctx = ctx_.getData(RequestContextSlot.APPSEC);
+          if (ctx == null) {
+            return NoopFlow.INSTANCE;
+          }
+          while (true) {
+            DataSubscriberInfo subInfo = graphqlServerRequestMsgSubInfo;
+            if (subInfo == null) {
+              subInfo =
+                  producerService.getDataSubscribers(KnownAddresses.GRAPHQL_SERVER_ALL_RESOLVERS);
+              graphqlServerRequestMsgSubInfo = subInfo;
+            }
+            if (subInfo == null || subInfo.isEmpty()) {
+              return NoopFlow.INSTANCE;
+            }
+            DataBundle bundle =
+                new SingletonDataBundle<>(KnownAddresses.GRAPHQL_SERVER_ALL_RESOLVERS, data);
+            try {
+              return producerService.publishDataEvent(subInfo, ctx, bundle, true);
+            } catch (ExpiredSubscriberInfoException e) {
+              graphqlServerRequestMsgSubInfo = null;
+            }
+          }
+        });
   }
 
   public void stop() {

dd-java-agent/appsec/src/main/java/com/datadog/appsec/powerwaf/PowerWAFModule.java

@@ -384,6 +384,7 @@ private static Collection<Address<?>> getUsedAddresses(PowerwafContext ctx) {
     addressList.add(KnownAddresses.REQUEST_BODY_RAW);
     addressList.add(KnownAddresses.RESPONSE_HEADERS_NO_COOKIES);
     addressList.add(KnownAddresses.RESPONSE_BODY_OBJECT);
+    addressList.add(KnownAddresses.GRAPHQL_SERVER_ALL_RESOLVERS);
 
     return addressList;
   }

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/event/data/KnownAddressesSpecification.groovy

@@ -39,7 +39,7 @@ class KnownAddressesSpecification extends Specification {
 
   void 'number of known addresses is expected number'() {
     expect:
-    Address.instanceCount() == 26
+    Address.instanceCount() == 27
     KnownAddresses.WAF_CONTEXT_PROCESSOR.serial == Address.instanceCount() - 1
   }
 }

dd-java-agent/appsec/src/test/groovy/com/datadog/appsec/gateway/GatewayBridgeSpecification.groovy

@@ -77,6 +77,7 @@ class GatewayBridgeSpecification extends DDSpecification {
   TriConsumer<RequestContext, String, String> respHeaderCB
   Function<RequestContext, Flow<Void>> respHeadersDoneCB
   BiFunction<RequestContext, Object, Flow<Void>> grpcServerRequestMessageCB
+  BiFunction<RequestContext, Map<String, Object>, Flow<Void>> graphqlServerRequestMessageCB
 
   void setup() {
     callInitAndCaptureCBs()
@@ -410,6 +411,7 @@ class GatewayBridgeSpecification extends DDSpecification {
     1 * ig.registerCallback(EVENTS.responseHeader(), _) >> { respHeaderCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.responseHeaderDone(), _) >> { respHeadersDoneCB = it[1]; null }
     1 * ig.registerCallback(EVENTS.grpcServerRequestMessage(), _) >> { grpcServerRequestMessageCB = it[1]; null }
+    1 * ig.registerCallback(EVENTS.graphqlServerRequestMessage(), _) >> { graphqlServerRequestMessageCB = it[1]; null }
     0 * ig.registerCallback(_, _)
 
     bridge.init()

dd-java-agent/instrumentation/graphql-java-14.0/src/main/java/datadog/trace/instrumentation/graphqljava/GraphQLDecorator.java

@@ -1,10 +1,27 @@
 package datadog.trace.instrumentation.graphqljava;
 
+import static datadog.trace.api.gateway.Events.EVENTS;
+
+import datadog.trace.api.gateway.CallbackProvider;
+import datadog.trace.api.gateway.Flow;
+import datadog.trace.api.gateway.RequestContext;
+import datadog.trace.api.gateway.RequestContextSlot;
 import datadog.trace.api.naming.SpanNaming;
+import datadog.trace.bootstrap.ActiveSubsystems;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.AgentTracer;
 import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
 import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
 import datadog.trace.bootstrap.instrumentation.decorator.BaseDecorator;
+import graphql.execution.ExecutionContext;
+import graphql.language.Argument;
+import graphql.language.Field;
+import graphql.language.Selection;
+import graphql.language.StringValue;
+import graphql.language.Value;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.BiFunction;
 
 public class GraphQLDecorator extends BaseDecorator {
   public static final GraphQLDecorator DECORATE = new GraphQLDecorator();
@@ -16,6 +33,11 @@ public class GraphQLDecorator extends BaseDecorator {
       UTF8BytesString.create("graphql.validation");
   public static final CharSequence GRAPHQL_JAVA = UTF8BytesString.create("graphql-java");
 
+  // Extract this to allow for easier testing
+  protected AgentTracer.TracerAPI tracer() {
+    return AgentTracer.get();
+  }
+
   @Override
   protected String[] instrumentationNames() {
     return new String[] {"graphql-java"};
@@ -36,4 +58,52 @@ public AgentSpan afterStart(final AgentSpan span) {
     span.setMeasured(true);
     return super.afterStart(span);
   }
+
+  public AgentSpan onRequest(final AgentSpan span, final ExecutionContext context) {
+
+    if (ActiveSubsystems.APPSEC_ACTIVE) {
+
+      Map<String, Map<String, String>> resolversArgs = new HashMap<>();
+
+      for (Selection<?> selection :
+          context.getOperationDefinition().getSelectionSet().getSelections()) {
+        if (selection instanceof Field) {
+          Field field = (Field) selection;
+          String name = field.getName();
+
+          Map<String, String> arguments = new HashMap<>();
+
+          for (Argument argument : field.getArguments()) {
+            String fieldName = argument.getName();
+            Value<?> fieldValue = argument.getValue();
+            if (fieldValue instanceof StringValue) {
+              String stringValue = ((StringValue) fieldValue).getValue();
+              arguments.put(fieldName, stringValue);
+            }
+          }
+          resolversArgs.put(name, arguments);
+        }
+      }
+
+      CallbackProvider cbp = tracer().getCallbackProvider(RequestContextSlot.APPSEC);
+      RequestContext ctx = span.getRequestContext();
+      if (cbp == null || resolversArgs.isEmpty() || ctx == null) {
+        return null;
+      }
+
+      BiFunction<RequestContext, Map<String, ?>, Flow<Void>> graphqlResolverCallback =
+          cbp.getCallback(EVENTS.graphqlServerRequestMessage());
+      if (graphqlResolverCallback == null) {
+        return null;
+      }
+
+      Flow<Void> flow = graphqlResolverCallback.apply(ctx, resolversArgs);
+      if (flow.getAction() instanceof Flow.Action.RequestBlockingAction) {
+        // Blocking will be implemented in future PRs
+        // span.setRequestBlockingAction((Flow.Action.RequestBlockingAction) flow.getAction());
+      }
+    }
+
+    return span;
+  }
 }

dd-java-agent/instrumentation/graphql-java-14.0/src/main/java/datadog/trace/instrumentation/graphqljava/GraphQLInstrumentation.java

@@ -101,6 +101,7 @@ public InstrumentationContext<ExecutionResult> beginExecuteOperation(
     requestSpan.setTag("graphql.operation.name", operationName);
     String resourceName = operationName != null ? operationName : state.getQuery();
     requestSpan.setResourceName(resourceName);
+    DECORATE.onRequest(requestSpan, parameters.getExecutionContext());
     return SimpleInstrumentationContext.noOp();
   }
 

dd-smoke-tests/appsec/springboot-graphql/build.gradle

@@ -0,0 +1,45 @@
+plugins {
+  id "com.github.johnrengelman.shadow"
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+description = 'SpringBoot GraphQL Smoke Tests.'
+
+// The standard spring-boot plugin doesn't play nice with our project
+// so we'll build a fat jar instead
+jar {
+  manifest {
+    attributes('Main-Class': 'datadog.smoketest.appsec.springbootgraphql.SpringbootGraphqlApplication')
+  }
+}
+
+shadowJar {
+  mergeServiceFiles {
+    include 'META-INF/spring.*'
+  }
+}
+
+// Use Java 11 to build application
+tasks.withType(JavaCompile) {
+  setJavaVersion(delegate, 11)
+}
+
+dependencies {
+  implementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.7.0'
+  implementation group: 'org.springframework.boot', name: 'spring-boot-starter-graphql', version: '2.7.0'
+  implementation(group: 'com.fasterxml.jackson.core', name: 'jackson-databind', version: '2.15.3')
+
+  testImplementation project(':dd-smoke-tests:appsec')
+}
+
+tasks.withType(Test).configureEach {
+  dependsOn "shadowJar"
+
+  jvmArgs "-Ddatadog.smoketest.appsec.springboot-graphql.shadowJar.path=${tasks.shadowJar.archiveFile.get()}"
+}
+
+task testRuntimeActivation(type: Test) {
+  jvmArgs '-Dsmoke_test.appsec.enabled=inactive',
+    "-Ddatadog.smoketest.appsec.springboot-graphql.shadowJar.path=${tasks.shadowJar.archiveFile.get()}"
+}
+tasks['check'].dependsOn(testRuntimeActivation)

dd-smoke-tests/appsec/springboot-graphql/src/main/java/datadog/smoketest/appsec/springbootgraphql/SpringbootGraphqlApplication.java

@@ -0,0 +1,12 @@
+package datadog.smoketest.appsec.springbootgraphql;
+
+import org.springframework.boot.SpringApplication;
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+
+@SpringBootApplication
+public class SpringbootGraphqlApplication {
+
+  public static void main(String[] args) {
+    SpringApplication.run(SpringbootGraphqlApplication.class, args);
+  }
+}

dd-smoke-tests/appsec/springboot-graphql/src/main/java/datadog/smoketest/appsec/springbootgraphql/controller/BookController.java

@@ -0,0 +1,21 @@
+package datadog.smoketest.appsec.springbootgraphql.controller;
+
+import datadog.smoketest.appsec.springbootgraphql.dao.Author;
+import datadog.smoketest.appsec.springbootgraphql.dao.Book;
+import org.springframework.graphql.data.method.annotation.Argument;
+import org.springframework.graphql.data.method.annotation.QueryMapping;
+import org.springframework.graphql.data.method.annotation.SchemaMapping;
+import org.springframework.stereotype.Controller;
+
+@Controller
+public class BookController {
+  @QueryMapping
+  public Book bookById(@Argument String id) {
+    return Book.getById(id);
+  }
+
+  @SchemaMapping
+  public Author author(Book book) {
+    return Author.getById(book.getAuthorId());
+  }
+}

dd-smoke-tests/appsec/springboot-graphql/src/main/java/datadog/smoketest/appsec/springbootgraphql/dao/Author.java

@@ -0,0 +1,39 @@
+package datadog.smoketest.appsec.springbootgraphql.dao;
+
+import java.util.Arrays;
+import java.util.List;
+
+public class Author {
+
+  private final String id;
+  private final String firstName;
+  private final String lastName;
+
+  private Author(String id, String firstName, String lastName) {
+    this.id = id;
+    this.firstName = firstName;
+    this.lastName = lastName;
+  }
+
+  public String getId() {
+    return id;
+  }
+
+  public String getFirstName() {
+    return firstName;
+  }
+
+  public String getLastName() {
+    return lastName;
+  }
+
+  private static List<Author> authors =
+      Arrays.asList(
+          new Author("author-1", "Joshua", "Bloch"),
+          new Author("author-2", "Douglas", "Adams"),
+          new Author("author-3", "Bill", "Bryson"));
+
+  public static Author getById(String id) {
+    return authors.stream().filter(author -> author.id.equals(id)).findFirst().orElse(null);
+  }
+}

dd-smoke-tests/appsec/springboot-graphql/src/main/java/datadog/smoketest/appsec/springbootgraphql/dao/Book.java

@@ -0,0 +1,45 @@
+package datadog.smoketest.appsec.springbootgraphql.dao;
+
+import java.util.Arrays;
+import java.util.List;
+
+public class Book {
+
+  private final String id;
+  private final String name;
+  private final int pageCount;
+  private final String authorId;
+
+  private Book(String id, String name, int pageCount, String authorId) {
+    this.id = id;
+    this.name = name;
+    this.pageCount = pageCount;
+    this.authorId = authorId;
+  }
+
+  public String getId() {
+    return id;
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public int getPageCount() {
+    return pageCount;
+  }
+
+  public String getAuthorId() {
+    return authorId;
+  }
+
+  private static List<Book> books =
+      Arrays.asList(
+          new Book("book-1", "Effective Java", 416, "author-1"),
+          new Book("book-2", "Hitchhiker's Guide to the Galaxy", 208, "author-2"),
+          new Book("book-3", "Down Under", 436, "author-3"));
+
+  public static Book getById(String id) {
+    return books.stream().filter(book -> book.id.equals(id)).findFirst().orElse(null);
+  }
+}

dd-smoke-tests/appsec/springboot-graphql/src/main/resources/application.properties

@@ -0,0 +1 @@
+spring.graphql.graphiql.enabled=true

dd-smoke-tests/appsec/springboot-graphql/src/main/resources/graphql/schema.graphqls

@@ -0,0 +1,16 @@
+type Query {
+    bookById(id: ID): Book
+}
+
+type Book {
+    id: ID
+    name: String
+    pageCount: Int
+    author: Author
+}
+
+type Author {
+    id: ID
+    firstName: String
+    lastName: String
+}