Fix failing tests

Author: manuel-alvarez-alvarez

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;
@@ -11,7 +12,6 @@
 import static datadog.trace.plugin.csi.util.CallSiteConstants.METHOD_HANDLER_CLASS;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.OPCODES_FQDN;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.STACK_DUP_MODE_CLASS;
-import static datadog.trace.plugin.csi.util.CallSiteConstants.TYPE_CLASS;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.TYPE_RESOLVER;
 import static datadog.trace.plugin.csi.util.CallSiteUtils.deleteFile;
 import static datadog.trace.plugin.csi.util.JavaParserUtils.getPrimaryType;
@@ -213,7 +213,8 @@ private void addAdviceLambda(
             .setArguments(
                 new NodeList<>(
                     new FieldAccessExpr()
-                        .setScope(new TypeExpr(new ClassOrInterfaceType().setName(TYPE_CLASS)))
+                        .setScope(
+                            new TypeExpr(new ClassOrInterfaceType().setName(ADVICE_TYPE_CLASS)))
                         .setName(type),
                     new StringLiteralExpr(pointCut.getOwner().getInternalName()),
                     new StringLiteralExpr(pointCut.getMethodName()),

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

@@ -30,7 +30,7 @@ private CallSiteConstants() {}
 
   public static final String HAS_ENABLED_PROPERTY_CLASS = CALL_SITES_CLASS + ".HasEnabledProperty";
 
-  public static final String TYPE_CLASS = "Type";
+  public static final String ADVICE_TYPE_CLASS = "AdviceType";
 
   public static final String STACK_DUP_MODE_CLASS = "StackDupMode";
 

dd-java-agent/agent-tooling/src/jmh/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteBenchmarkInstrumentation.java

@@ -1,5 +1,7 @@
 package datadog.trace.agent.tooling.bytebuddy.csi;
 
+import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.AROUND;
+
 import datadog.trace.agent.tooling.csi.CallSites;
 import datadog.trace.agent.tooling.csi.InvokeAdvice;
 import datadog.trace.agent.tooling.muzzle.ReferenceMatcher;
@@ -44,6 +46,7 @@ public Iterable<CallSites> get() {
       return Collections.singletonList(
           (container -> {
             container.addAdvice(
+                AROUND,
                 "javax/servlet/ServletRequest",
                 "getParameter",
                 "(Ljava/lang/String;)Ljava/lang/String;",

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/Advices.java

@@ -370,7 +370,7 @@ void onConstantPool(
         @Nonnull TypeDescription type, @Nonnull ConstantPool pool, final byte[] classFile);
   }
 
-  private interface HasType {
+  public interface HasType {
     byte getType();
 
     static CallSiteAdvice withType(final CallSiteAdvice advice, final byte type) {

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java

@@ -1,6 +1,6 @@
 package datadog.trace.agent.tooling.bytebuddy.csi;
 
-import static datadog.trace.agent.tooling.csi.CallSiteAdvice.Type.BEFORE;
+import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.AFTER;
 import static datadog.trace.api.telemetry.LogCollector.SEND_TELEMETRY;
 import static net.bytebuddy.jar.asm.ClassWriter.COMPUTE_MAXS;
 
@@ -380,8 +380,8 @@ protected void invokeAdvice(
         final String name,
         final String descriptor,
         final boolean isInterface) {
-      if (isSuperCall && advices.typeOf(advice) == BEFORE) {
-        // TODO APPSEC-57009 calls to super are not instrumented by before callsites
+      if (isSuperCall && advices.typeOf(advice) != AFTER) {
+        // TODO APPSEC-57009 calls to super are only instrumented by after call sites
         // just ignore the advice and keep on
         mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface);
       } else {

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java

@@ -73,9 +73,9 @@ enum StackDupMode {
     APPEND_ARRAY
   }
 
-  abstract class Type {
+  abstract class AdviceType {
 
-    private Type() {}
+    private AdviceType() {}
 
     public static final byte BEFORE = -1;
     public static final byte AROUND = 0;

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy

@@ -24,18 +24,20 @@ import java.security.MessageDigest
 import java.lang.reflect.Method
 import java.security.ProtectionDomain
 
+import static datadog.trace.agent.tooling.bytebuddy.csi.Advices.HasType.withType
+import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.BEFORE
 import static net.bytebuddy.matcher.ElementMatchers.any
 import static net.bytebuddy.matcher.ElementMatchers.named
 
 @CompileDynamic
 class BaseCallSiteTest extends DDSpecification {
 
-  protected CallSites mockCallSites(final CallSiteAdvice advice, final Pointcut target, final String... helpers) {
+  protected CallSites mockCallSites(final byte type = BEFORE, final CallSiteAdvice advice, final Pointcut target, final String... helpers) {
     return Stub(CallSites) {
       accept(_ as CallSites.Container) >> {
         final container = it[0] as CallSites.Container
         container.addHelpers(helpers)
-        container.addAdvice(target.type, target.method, target.descriptor, advice)
+        container.addAdvice(type, target.type, target.method, target.descriptor, advice)
       }
     }
   }
@@ -44,10 +46,16 @@ class BaseCallSiteTest extends DDSpecification {
     final advices = [:] as Map<String, Map<String, Map<String, CallSiteAdvice>>>
     final helpers = [] as Set<String>
     final container = Stub(CallSites.Container) {
-      addAdvice(_ as String, _ as String, _ as String, _ as CallSiteAdvice) >> {
-        advices.computeIfAbsent(it[0] as String, t -> [:])
-        .computeIfAbsent(it[1] as String, m -> [:])
-        .put(it[2] as String, it[3] as CallSiteAdvice)
+      addAdvice(_, _ as String, _ as String, _ as String, _ as CallSiteAdvice) >> {
+        final type = it[0] as byte
+        final owner = it[1] as String
+        final method = it[2] as String
+        final descriptor = it[3] as String
+        final advice = it[4] as CallSiteAdvice
+        advices
+        .computeIfAbsent(owner, t -> [:])
+        .computeIfAbsent(method, m -> [:])
+        .put(descriptor, withType(advice, type))
       }
       addHelpers(_ as String[]) >> {
         Collections.addAll(helpers, it[0] as String[])

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy

@@ -9,6 +9,8 @@ import net.bytebuddy.jar.asm.Type
 import java.util.concurrent.atomic.AtomicInteger
 
 import static datadog.trace.agent.tooling.csi.CallSiteAdvice.StackDupMode.PREPEND_ARRAY_CTOR
+import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.AFTER
+import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.BEFORE
 
 class CallSiteInstrumentationTest extends BaseCallSiteTest {
 
@@ -84,7 +86,7 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest {
             )
         }
       }
-    final callSiteTransformer = new CallSiteTransformer(mockAdvices([mockCallSites(advice, pointcut)]))
+    final callSiteTransformer = new CallSiteTransformer(mockAdvices([mockCallSites(AFTER, advice, pointcut)]))
 
     when:
     final transformedClass = transformType(source, target, callSiteTransformer)
@@ -101,7 +103,7 @@ class CallSiteInstrumentationTest extends BaseCallSiteTest {
     @Override
     void accept(final Container container) {
       final pointcut = buildPointcut(String.getDeclaredMethod('concat', String))
-      container.addAdvice(pointcut.type, pointcut.method, pointcut.descriptor, new StringConcatAdvice())
+      container.addAdvice(BEFORE, pointcut.type, pointcut.method, pointcut.descriptor, new StringConcatAdvice())
     }
   }
 

dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockCallSites.java

@@ -1,4 +1,4 @@
-import static datadog.trace.agent.tooling.csi.CallSiteAdvice.Type.BEFORE;
+import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.BEFORE;
 
 import datadog.trace.agent.tooling.csi.CallSiteAdvice;
 import datadog.trace.agent.tooling.csi.CallSites;

dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockRaspCallSites.java

@@ -1,4 +1,4 @@
-import static datadog.trace.agent.tooling.csi.CallSiteAdvice.Type.BEFORE;
+import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.BEFORE;
 
 import datadog.trace.agent.tooling.csi.CallSiteAdvice;
 import datadog.trace.agent.tooling.csi.CallSites;

dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ObjectInputStreamCallSiteTest.groovy

@@ -35,7 +35,7 @@ class ObjectInputStreamCallSiteTest extends AgentTestRunner {
     new TestCustomObjectInputStream(inputStream())
 
     then:
-    // TODO APPSEC-57009 calls to super are not instrumented by before callsites
+    // TODO APPSEC-57009 calls to super are only instrumented by after callsites
     0 * iastModule.onObject(_)
   }