From 36a381cd9c4654d7b44549131fcd5d9677df9f85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Fri, 14 Mar 2025 12:29:45 +0100 Subject: [PATCH 1/2] Prevent before callsites targeting calls to super in constructors (cherry picked from commit e20ae642f83e7f96b14183371b89fb19cdeb6974) --- .../plugin/csi/impl/AdviceGeneratorImpl.java | 8 ++ .../plugin/csi/impl/ext/IastExtension.java | 8 +- .../plugin/csi/util/CallSiteConstants.java | 2 + .../csi/impl/AdviceGeneratorTest.groovy | 3 + .../csi/impl/assertion/AdviceAssert.groovy | 5 ++ .../csi/impl/assertion/AssertBuilder.groovy | 6 +- .../csi/impl/ext/IastExtensionTest.groovy | 4 +- .../agent/tooling/bytebuddy/csi/Advices.java | 78 ++++++++++++++++++- .../bytebuddy/csi/CallSiteTransformer.java | 49 +++++++++++- .../agent/tooling/csi/CallSiteAdvice.java | 9 +++ .../trace/agent/tooling/csi/CallSites.java | 13 ++-- .../src/test/java/MockCallSites.java | 4 +- .../src/test/java/MockRaspCallSites.java | 4 +- .../io/ObjectInputStreamCallSiteTest.groovy | 25 +++++- .../foo/bar/TestCustomObjectInputStream.java | 12 +++ 15 files changed, 205 insertions(+), 25 deletions(-) create mode 100644 dd-java-agent/instrumentation/java-io/src/test/java/foo/bar/TestCustomObjectInputStream.java 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..fde52b6ace9 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 @@ -11,6 +11,7 @@ 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; @@ -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,9 @@ private void addAdviceLambda( .setName("addAdvice") .setArguments( new NodeList<>( + new FieldAccessExpr() + .setScope(new TypeExpr(new ClassOrInterfaceType().setName(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..dc338eeca15 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 TYPE_CLASS = "Type"; + 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/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..ba9d69058fa 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 ((HasType) 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, HasType.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 HasType { + 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, HasType { + 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, HasType { + 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..c255a422196 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.Type.BEFORE; 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) == BEFORE) { + // TODO APPSEC-57009 calls to super are not instrumented by before callsites + // 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..f2d51334c05 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 Type { + + private Type() {} + + 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/instrumentation/iast-instrumenter/src/test/java/MockCallSites.java b/dd-java-agent/instrumentation/iast-instrumenter/src/test/java/MockCallSites.java index 8b3888a5edf..ab80fabf25b 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.Type.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..6dd054c2ff2 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.Type.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..423a7a25037 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 not instrumented by before 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); + } +} From dcedbe87cb312f811da2c2478081c18b55df53af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Manuel=20=C3=81lvarez=20=C3=81lvarez?= Date: Mon, 17 Mar 2025 12:02:34 +0100 Subject: [PATCH 2/2] Fix failing tests (cherry picked from commit 70c362cff4b33c27b973b901143b11b6a1c99462) --- .../plugin/csi/impl/AdviceGeneratorImpl.java | 5 +++-- .../plugin/csi/util/CallSiteConstants.java | 2 +- .../csi/CallSiteBenchmarkInstrumentation.java | 3 +++ .../agent/tooling/bytebuddy/csi/Advices.java | 10 +++++----- .../bytebuddy/csi/CallSiteTransformer.java | 6 +++--- .../agent/tooling/csi/CallSiteAdvice.java | 4 ++-- .../agent/tooling/csi/BaseCallSiteTest.groovy | 19 +++++++++++++------ .../csi/CallSiteInstrumentationTest.groovy | 6 ++++-- .../src/test/java/MockCallSites.java | 2 +- .../src/test/java/MockRaspCallSites.java | 2 +- .../io/ObjectInputStreamCallSiteTest.groovy | 2 +- 11 files changed, 37 insertions(+), 24 deletions(-) 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 fde52b6ace9..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; @@ -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()), 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 dc338eeca15..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,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"; 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 ba9d69058fa..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 @@ -145,7 +145,7 @@ public CallSiteAdvice findAdvice( /** Gets the type of advice we are dealing with */ public byte typeOf(final CallSiteAdvice advice) { - return ((HasType) advice).getType(); + return ((TypedAdvice) advice).getType(); } public String[] getHelpers() { @@ -193,7 +193,7 @@ public void addAdvice( final Map methodAdvices = typeAdvices.computeIfAbsent(method, k -> new HashMap<>()); final CallSiteAdvice oldAdvice = - methodAdvices.put(descriptor, HasType.withType(advice, type)); + methodAdvices.put(descriptor, TypedAdvice.withType(advice, type)); if (oldAdvice != null) { throw new UnsupportedOperationException( String.format( @@ -370,7 +370,7 @@ void onConstantPool( @Nonnull TypeDescription type, @Nonnull ConstantPool pool, final byte[] classFile); } - private interface HasType { + private interface TypedAdvice { byte getType(); static CallSiteAdvice withType(final CallSiteAdvice advice, final byte type) { @@ -382,7 +382,7 @@ static CallSiteAdvice withType(final CallSiteAdvice advice, final byte type) { } } - private static class InvokeWithType implements InvokeAdvice, HasType { + private static class InvokeWithType implements InvokeAdvice, TypedAdvice { private final InvokeAdvice advice; private final byte type; @@ -408,7 +408,7 @@ public void apply( } } - private static class InvokeDynamicWithType implements InvokeDynamicAdvice, HasType { + private static class InvokeDynamicWithType implements InvokeDynamicAdvice, TypedAdvice { private final InvokeDynamicAdvice advice; private final byte type; 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 c255a422196..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,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 { 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 f2d51334c05..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 @@ -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; 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 ab80fabf25b..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,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; 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 6dd054c2ff2..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,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; 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 423a7a25037..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 @@ -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(_) }