Merge branch 'master' into vzakharov/api_sec_sampling_new

Author: smola

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java

@@ -2,6 +2,7 @@
 
 import static com.github.javaparser.ast.Modifier.Keyword.PUBLIC;
 import static datadog.trace.plugin.csi.impl.CallSiteFactory.typeResolver;
+import static datadog.trace.plugin.csi.util.CallSiteConstants.ADVICE_TYPE_CLASS;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.AUTO_SERVICE_FQDN;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.CALL_SITES_CLASS;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.CALL_SITES_FQCN;
@@ -185,20 +186,24 @@ private void addAdviceLambda(
     final MethodType pointCut = spec.getPointcut();
     final BlockStmt adviceBody = new BlockStmt();
     final Expression advice;
+    final String type;
     if (spec.isInvokeDynamic()) {
       advice = invokeDynamicAdviceSignature(adviceBody);
     } else {
       advice = invokeAdviceSignature(adviceBody);
     }
     if (spec instanceof BeforeSpecification) {
+      type = "BEFORE";
       writeStackOperations(spec, adviceBody);
       writeAdviceMethodCall(spec, adviceBody);
       writeOriginalMethodCall(spec, adviceBody);
     } else if (spec instanceof AfterSpecification) {
+      type = "AFTER";
       writeStackOperations(spec, adviceBody);
       writeOriginalMethodCall(spec, adviceBody);
       writeAdviceMethodCall(spec, adviceBody);
     } else {
+      type = "AROUND";
       writeAdviceMethodCall(spec, adviceBody);
     }
     body.addStatement(
@@ -207,6 +212,10 @@ private void addAdviceLambda(
             .setName("addAdvice")
             .setArguments(
                 new NodeList<>(
+                    new FieldAccessExpr()
+                        .setScope(
+                            new TypeExpr(new ClassOrInterfaceType().setName(ADVICE_TYPE_CLASS)))
+                        .setName(type),
                     new StringLiteralExpr(pointCut.getOwner().getInternalName()),
                     new StringLiteralExpr(pointCut.getMethodName()),
                     new StringLiteralExpr(pointCut.getMethodType().getDescriptor()),

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/ext/IastExtension.java

@@ -364,19 +364,19 @@ private static LambdaExpr findAdviceLambda(
     final MethodType pointcut = spec.getPointcut();
     for (final MethodCallExpr add : addAdvices) {
       final NodeList<Expression> arguments = add.getArguments();
-      final String owner = arguments.get(0).asStringLiteralExpr().asString();
+      final String owner = arguments.get(1).asStringLiteralExpr().asString();
       if (!owner.equals(pointcut.getOwner().getInternalName())) {
         continue;
       }
-      final String method = arguments.get(1).asStringLiteralExpr().asString();
+      final String method = arguments.get(2).asStringLiteralExpr().asString();
       if (!method.equals(pointcut.getMethodName())) {
         continue;
       }
-      final String description = arguments.get(2).asStringLiteralExpr().asString();
+      final String description = arguments.get(3).asStringLiteralExpr().asString();
       if (!description.equals(pointcut.getMethodType().getDescriptor())) {
         continue;
       }
-      return arguments.get(3).asLambdaExpr();
+      return arguments.get(4).asLambdaExpr();
     }
     throw new IllegalArgumentException("Cannot find lambda expression for pointcut " + pointcut);
   }

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/CallSiteConstants.java

@@ -30,6 +30,8 @@ private CallSiteConstants() {}
 
   public static final String HAS_ENABLED_PROPERTY_CLASS = CALL_SITES_CLASS + ".HasEnabledProperty";
 
+  public static final String ADVICE_TYPE_CLASS = "AdviceType";
+
   public static final String STACK_DUP_MODE_CLASS = "StackDupMode";
 
   public static final String METHOD_HANDLER_CLASS = "MethodHandler";

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy

@@ -44,6 +44,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(BeforeAdvice)
       advices(0) {
+        type("BEFORE")
         pointcut('java/security/MessageDigest', 'getInstance', '(Ljava/lang/String;)Ljava/security/MessageDigest;')
         statements(
           'handler.dupParameters(descriptor, StackDupMode.COPY);',
@@ -76,6 +77,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(AroundAdvice)
       advices(0) {
+        type("AROUND")
         pointcut('java/lang/String', 'replaceAll', '(Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;')
         statements(
           'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AroundAdvice", "around", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;");'
@@ -106,6 +108,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(AfterAdvice)
       advices(0) {
+        type("AFTER")
         pointcut('java/lang/String', 'concat', '(Ljava/lang/String;)Ljava/lang/String;')
         statements(
           'handler.dupInvoke(owner, descriptor, StackDupMode.COPY);',

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AdviceAssert.groovy

@@ -1,11 +1,16 @@
 package datadog.trace.plugin.csi.impl.assertion
 
 class AdviceAssert {
+  protected String type
   protected String owner
   protected String method
   protected String descriptor
   protected Collection<String> statements
 
+  void type(String type) {
+    assert type == this.type
+  }
+
   void pointcut(String owner, String method, String descriptor) {
     assert owner == this.owner
     assert method == this.method

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AssertBuilder.groovy

@@ -83,10 +83,12 @@ class AssertBuilder<C extends CallSiteAssert> {
     return getMethodCalls(acceptMethod).findAll {
       it.nameAsString == 'addAdvice'
     }.collect {
-      def (owner, method, descriptor) = it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
-      final handlerLambda = it.arguments[3].asLambdaExpr()
+      final adviceType = it.arguments.get(0).asFieldAccessExpr().getName()
+      def (owner, method, descriptor) = it.arguments.subList(1, 4)*.asStringLiteralExpr()*.asString()
+      final handlerLambda = it.arguments[4].asLambdaExpr()
       final advice = handlerLambda.body.asBlockStmt().statements*.toString()
       return new AdviceAssert([
+        type      : adviceType,
         owner     : owner,
         method    : method,
         descriptor: descriptor,

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/ext/IastExtensionTest.groovy

@@ -217,8 +217,8 @@ class IastExtensionTest extends BaseCsiPluginTest {
       return getMethodCalls(acceptMethod).findAll {
         it.nameAsString == 'addAdvice'
       }.collect {
-        def (owner, method, descriptor) = it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
-        final handlerLambda = it.arguments[3].asLambdaExpr()
+        def (owner, method, descriptor) = it.arguments.subList(1, 4)*.asStringLiteralExpr()*.asString()
+        final handlerLambda = it.arguments[4].asLambdaExpr()
         final statements = handlerLambda.body.asBlockStmt().statements
         final instrumentedStmt = statements.get(0).asIfStmt()
         final executedStmt = statements.get(1).asIfStmt()

components/context/src/main/java/datadog/context/Context.java

@@ -36,6 +36,15 @@
  * @see ContextKey
  */
 public interface Context {
+  /**
+   * Returns the empty context.
+   *
+   * @return the context containing no values at all.
+   */
+  static Context empty() {
+    return EmptyContext.INSTANCE;
+  }
+
   /**
    * Returns the root context.
    *

components/context/src/test/java/datadog/context/ContextTest.java

@@ -1,5 +1,6 @@
 package datadog.context;
 
+import static datadog.context.Context.empty;
 import static datadog.context.Context.root;
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -21,6 +22,17 @@ class ContextTest {
   static final ContextKey<Float> FLOAT_KEY = ContextKey.named("float-key");
   static final ContextKey<Long> LONG_KEY = ContextKey.named("long-key");
 
+  @Test
+  void testEmpty() {
+    // Test empty is always the same
+    Context empty = empty();
+    assertEquals(empty, empty(), "Empty context should be consistent");
+    // Test empty is not mutated
+    String stringValue = "value";
+    empty.with(STRING_KEY, stringValue);
+    assertEquals(empty, empty(), "Empty context should be immutable");
+  }
+
   @Test
   void testRoot() {
     // Test root is always the same

dd-java-agent/agent-ci-visibility/src/main/java/datadog/trace/civisibility/CiVisibilityRepoServices.java

@@ -9,6 +9,7 @@
 import datadog.trace.civisibility.ci.CIProviderInfo;
 import datadog.trace.civisibility.ci.CITagsProvider;
 import datadog.trace.civisibility.ci.PullRequestInfo;
+import datadog.trace.civisibility.ci.env.CiEnvironment;
 import datadog.trace.civisibility.codeowners.Codeowners;
 import datadog.trace.civisibility.codeowners.CodeownersProvider;
 import datadog.trace.civisibility.codeowners.NoCodeowners;
@@ -39,6 +40,7 @@
 import java.nio.file.Paths;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -64,7 +66,7 @@ public class CiVisibilityRepoServices {
     ciProvider = ciProviderInfo.getProvider();
 
     CIInfo ciInfo = ciProviderInfo.buildCIInfo();
-    PullRequestInfo pullRequestInfo = ciProviderInfo.buildPullRequestInfo();
+    PullRequestInfo pullRequestInfo = buildPullRequestInfo(services.environment, ciProviderInfo);
 
     if (pullRequestInfo.isNotEmpty()) {
       LOGGER.info("PR detected: {}", pullRequestInfo);
@@ -107,6 +109,22 @@ public class CiVisibilityRepoServices {
     }
   }
 
+  @Nonnull
+  private static PullRequestInfo buildPullRequestInfo(
+      CiEnvironment environment, CIProviderInfo ciProviderInfo) {
+    PullRequestInfo ciProviderPrInfo = ciProviderInfo.buildPullRequestInfo();
+    if (ciProviderPrInfo.isNotEmpty()) {
+      return ciProviderPrInfo;
+    }
+
+    // could not get PR info from CI provider,
+    // check if it was set manually
+    return new PullRequestInfo(
+        null,
+        environment.get(Constants.DDCI_PULL_REQUEST_TARGET_SHA),
+        environment.get(Constants.DDCI_PULL_REQUEST_SOURCE_SHA));
+  }
+
   private static String getRepoRoot(CIInfo ciInfo, GitClient.Factory gitClientFactory) {
     String ciWorkspace = ciInfo.getNormalizedCiWorkspace();
     if (Strings.isNotBlank(ciWorkspace)) {