diff --git a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java index 04c7a9be704..14efc4a42d4 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java +++ b/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()), diff --git a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/ext/IastExtension.java b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/ext/IastExtension.java index e86ca097d85..46e7c96dc89 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/ext/IastExtension.java +++ b/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 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); } diff --git a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/CallSiteConstants.java b/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/CallSiteConstants.java index 390df62d9a9..8613af8958d 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/CallSiteConstants.java +++ b/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"; diff --git a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy index 3ec2a6da5c5..2958c292cac 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy +++ b/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);', diff --git a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AdviceAssert.groovy b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AdviceAssert.groovy index 419cc49178a..d83120a6c52 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AdviceAssert.groovy +++ b/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 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 diff --git a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AssertBuilder.groovy b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AssertBuilder.groovy index c8f192ca447..f80c258dfda 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AssertBuilder.groovy +++ b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AssertBuilder.groovy @@ -83,10 +83,12 @@ class AssertBuilder { 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, diff --git a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/ext/IastExtensionTest.groovy b/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/ext/IastExtensionTest.groovy index 5ffe27e95b9..bcd28fad4e4 100644 --- a/buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/ext/IastExtensionTest.groovy +++ b/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() diff --git a/dd-java-agent/agent-tooling/src/jmh/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteBenchmarkInstrumentation.java b/dd-java-agent/agent-tooling/src/jmh/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteBenchmarkInstrumentation.java index a03865f0756..4796b15f4eb 100644 --- a/dd-java-agent/agent-tooling/src/jmh/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteBenchmarkInstrumentation.java +++ b/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 get() { return Collections.singletonList( (container -> { container.addAdvice( + AROUND, "javax/servlet/ServletRequest", "getParameter", "(Ljava/lang/String;)Ljava/lang/String;", diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/Advices.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/Advices.java index a574ab0fcb4..705bebf7a87 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/Advices.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/Advices.java @@ -6,6 +6,8 @@ import datadog.trace.agent.tooling.bytebuddy.ClassFileLocators; import datadog.trace.agent.tooling.csi.CallSiteAdvice; import datadog.trace.agent.tooling.csi.CallSites; +import datadog.trace.agent.tooling.csi.InvokeAdvice; +import datadog.trace.agent.tooling.csi.InvokeDynamicAdvice; import java.lang.reflect.Field; import java.util.Arrays; import java.util.Collections; @@ -141,6 +143,11 @@ public CallSiteAdvice findAdvice( return methodAdvices.get(descriptor); } + /** Gets the type of advice we are dealing with */ + public byte typeOf(final CallSiteAdvice advice) { + return ((TypedAdvice) advice).getType(); + } + public String[] getHelpers() { return helpers; } @@ -176,15 +183,17 @@ public void addHelpers(final String... helperClassNames) { @Override public void addAdvice( - final String type, + final byte type, + final String owner, final String method, final String descriptor, final CallSiteAdvice advice) { final Map> typeAdvices = - advices.computeIfAbsent(type, k -> new HashMap<>()); + advices.computeIfAbsent(owner, k -> new HashMap<>()); final Map methodAdvices = typeAdvices.computeIfAbsent(method, k -> new HashMap<>()); - final CallSiteAdvice oldAdvice = methodAdvices.put(descriptor, advice); + final CallSiteAdvice oldAdvice = + methodAdvices.put(descriptor, TypedAdvice.withType(advice, type)); if (oldAdvice != null) { throw new UnsupportedOperationException( String.format( @@ -360,4 +369,67 @@ public interface Listener { void onConstantPool( @Nonnull TypeDescription type, @Nonnull ConstantPool pool, final byte[] classFile); } + + private interface TypedAdvice { + byte getType(); + + static CallSiteAdvice withType(final CallSiteAdvice advice, final byte type) { + if (advice instanceof InvokeAdvice) { + return new InvokeWithType((InvokeAdvice) advice, type); + } else { + return new InvokeDynamicWithType((InvokeDynamicAdvice) advice, type); + } + } + } + + private static class InvokeWithType implements InvokeAdvice, TypedAdvice { + private final InvokeAdvice advice; + private final byte type; + + private InvokeWithType(InvokeAdvice advice, byte type) { + this.advice = advice; + this.type = type; + } + + @Override + public byte getType() { + return type; + } + + @Override + public void apply( + final MethodHandler handler, + final int opcode, + final String owner, + final String name, + final String descriptor, + final boolean isInterface) { + advice.apply(handler, opcode, owner, name, descriptor, isInterface); + } + } + + private static class InvokeDynamicWithType implements InvokeDynamicAdvice, TypedAdvice { + private final InvokeDynamicAdvice advice; + private final byte type; + + private InvokeDynamicWithType(final InvokeDynamicAdvice advice, final byte type) { + this.advice = advice; + this.type = type; + } + + @Override + public byte getType() { + return type; + } + + @Override + public void apply( + final MethodHandler handler, + final String name, + final String descriptor, + final Handle bootstrapMethodHandle, + final Object... bootstrapMethodArguments) { + advice.apply(handler, name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments); + } + } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java index 991523069c1..3b38c22891b 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteTransformer.java @@ -1,5 +1,6 @@ package datadog.trace.agent.tooling.bytebuddy.csi; +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; @@ -126,7 +127,7 @@ public MethodVisitor visitMethod( private static class CallSiteMethodVisitor extends MethodVisitor implements CallSiteAdvice.MethodHandler { - private final Advices advices; + protected final Advices advices; private CallSiteMethodVisitor( @Nonnull final Advices advices, @Nonnull final MethodVisitor delegated) { @@ -144,12 +145,22 @@ public void visitMethodInsn( CallSiteAdvice advice = advices.findAdvice(owner, name, descriptor); if (advice instanceof InvokeAdvice) { - ((InvokeAdvice) advice).apply(this, opcode, owner, name, descriptor, isInterface); + invokeAdvice((InvokeAdvice) advice, opcode, owner, name, descriptor, isInterface); } else { mv.visitMethodInsn(opcode, owner, name, descriptor, isInterface); } } + protected void invokeAdvice( + final InvokeAdvice advice, + final int opcode, + final String owner, + final String name, + final String descriptor, + final boolean isInterface) { + advice.apply(this, opcode, owner, name, descriptor, isInterface); + } + @Override public void visitInvokeDynamicInsn( final String name, @@ -158,14 +169,27 @@ public void visitInvokeDynamicInsn( final Object... bootstrapMethodArguments) { CallSiteAdvice advice = advices.findAdvice(bootstrapMethodHandle); if (advice instanceof InvokeDynamicAdvice) { - ((InvokeDynamicAdvice) advice) - .apply(this, name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments); + invokeDynamicAdvice( + (InvokeDynamicAdvice) advice, + name, + descriptor, + bootstrapMethodHandle, + bootstrapMethodArguments); } else { mv.visitInvokeDynamicInsn( name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments); } } + protected void invokeDynamicAdvice( + final InvokeDynamicAdvice advice, + final String name, + final String descriptor, + final Handle bootstrapMethodHandle, + final Object... bootstrapMethodArguments) { + advice.apply(this, name, descriptor, bootstrapMethodHandle, bootstrapMethodArguments); + } + @Override public void instruction(final int opcode) { mv.visitInsn(opcode); @@ -347,5 +371,22 @@ public void dupParameters(final String methodDescriptor, final StackDupMode mode super.dupParameters( methodDescriptor, isSuperCall ? StackDupMode.PREPEND_ARRAY_SUPER_CTOR : mode); } + + @Override + protected void invokeAdvice( + final InvokeAdvice advice, + final int opcode, + final String owner, + final String name, + final String descriptor, + final boolean isInterface) { + 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 { + super.invokeAdvice(advice, opcode, owner, name, descriptor, isInterface); + } + } } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java index 0e43f3602e8..dbe925c7b8a 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSiteAdvice.java @@ -72,4 +72,13 @@ enum StackDupMode { /** Copies the parameters in an array and appends it */ APPEND_ARRAY } + + abstract class AdviceType { + + private AdviceType() {} + + public static final byte BEFORE = -1; + public static final byte AROUND = 0; + public static final byte AFTER = 1; + } } diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSites.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSites.java index 4fbefc851f3..6bb36ff7773 100644 --- a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSites.java +++ b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/csi/CallSites.java @@ -6,22 +6,25 @@ public interface CallSites extends Consumer { interface Container { default void addAdvice( - final String type, + final byte type, + final String owner, final String method, final String descriptor, final InvokeAdvice advice) { - addAdvice(type, method, descriptor, (CallSiteAdvice) advice); + addAdvice(type, owner, method, descriptor, (CallSiteAdvice) advice); } default void addAdvice( - final String type, + final byte type, + final String owner, final String method, final String descriptor, final InvokeDynamicAdvice advice) { - addAdvice(type, method, descriptor, (CallSiteAdvice) advice); + addAdvice(type, owner, method, descriptor, (CallSiteAdvice) advice); } - void addAdvice(String type, String method, String descriptor, CallSiteAdvice advice); + void addAdvice( + byte kind, String owner, String method, String descriptor, CallSiteAdvice advice); void addHelpers(String... helperClassNames); } diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy index 416092536e9..12e242e1807 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy +++ b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy @@ -24,18 +24,19 @@ import java.security.MessageDigest import java.lang.reflect.Method import java.security.ProtectionDomain +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 +45,16 @@ class BaseCallSiteTest extends DDSpecification { final advices = [:] as Map>> final helpers = [] as Set 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, advice) } addHelpers(_ as String[]) >> { Collections.addAll(helpers, it[0] as String[]) diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy index b56b94779ad..1b37b8891e7 100644 --- a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteInstrumentationTest.groovy +++ b/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()) } } diff --git a/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockCallSites.java b/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockCallSites.java index 8b3888a5edf..b569d2d874d 100644 --- a/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockCallSites.java +++ b/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockCallSites.java @@ -1,3 +1,5 @@ +import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.BEFORE; + import datadog.trace.agent.tooling.csi.CallSiteAdvice; import datadog.trace.agent.tooling.csi.CallSites; import datadog.trace.api.iast.IastCallSites; @@ -5,6 +7,6 @@ public class MockCallSites implements IastCallSites, CallSites { @Override public void accept(final Container container) { - container.addAdvice("type", "method", "descriptor", (CallSiteAdvice) null); + container.addAdvice(BEFORE, "type", "method", "descriptor", (CallSiteAdvice) null); } } diff --git a/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockRaspCallSites.java b/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockRaspCallSites.java index 408c7850931..5db049a3383 100644 --- a/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockRaspCallSites.java +++ b/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockRaspCallSites.java @@ -1,3 +1,5 @@ +import static datadog.trace.agent.tooling.csi.CallSiteAdvice.AdviceType.BEFORE; + import datadog.trace.agent.tooling.csi.CallSiteAdvice; import datadog.trace.agent.tooling.csi.CallSites; import datadog.trace.api.appsec.RaspCallSites; @@ -5,6 +7,6 @@ public class MockRaspCallSites implements RaspCallSites, CallSites { @Override public void accept(final Container container) { - container.addAdvice("type", "method", "descriptor", (CallSiteAdvice) null); + container.addAdvice(BEFORE, "typeRasp", "methodRasp", "descriptorRasp", (CallSiteAdvice) null); } } diff --git a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ObjectInputStreamCallSiteTest.groovy b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ObjectInputStreamCallSiteTest.groovy index d7c73006b05..38c997bd5ff 100644 --- a/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ObjectInputStreamCallSiteTest.groovy +++ b/dd-java-agent/instrumentation/java-io/src/test/groovy/datadog/trace/instrumentation/java/io/ObjectInputStreamCallSiteTest.groovy @@ -5,6 +5,8 @@ import datadog.trace.api.iast.InstrumentationBridge import datadog.trace.api.iast.sink.UntrustedDeserializationModule import foo.bar.TestObjectInputStreamSuite +import foo.bar.TestCustomObjectInputStream + class ObjectInputStreamCallSiteTest extends AgentTestRunner { @Override @@ -17,12 +19,29 @@ class ObjectInputStreamCallSiteTest extends AgentTestRunner { final module = Mock(UntrustedDeserializationModule) InstrumentationBridge.registerIastModule(module) - final InputStream inputStream = new ByteArrayInputStream(new byte[0]) - when: - TestObjectInputStreamSuite.init(inputStream) + TestObjectInputStreamSuite.init(inputStream()) then: 1 * module.onObject(_) } + + void 'test super call to ObjectInputStream.'() { + given: + final iastModule = Mock(UntrustedDeserializationModule) + InstrumentationBridge.registerIastModule(iastModule) + + when: + new TestCustomObjectInputStream(inputStream()) + + then: + // TODO APPSEC-57009 calls to super are only instrumented by after callsites + 0 * iastModule.onObject(_) + } + + private static InputStream inputStream() { + final baos = new ByteArrayOutputStream() + new ObjectOutputStream(baos).writeObject([:]) + return new ByteArrayInputStream(baos.toByteArray()) + } } diff --git a/dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestCustomObjectInputStream.java b/dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestCustomObjectInputStream.java new file mode 100644 index 00000000000..0010d56fe9a --- /dev/null +++ b/dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestCustomObjectInputStream.java @@ -0,0 +1,12 @@ +package foo.bar; + +import java.io.IOException; +import java.io.InputStream; +import java.io.ObjectInputStream; + +public class TestCustomObjectInputStream extends ObjectInputStream { + + public TestCustomObjectInputStream(final InputStream in) throws IOException { + super(in); + } +}