Skip to content

Commit 4994716

Browse files
bantonssonbantonsson
committed
Try to create a type description using loadClass if the pool fails
1 parent 8ac678f commit 4994716

9 files changed

Lines changed: 371 additions & 9 deletions

File tree

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/AgentTooling.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import datadog.trace.agent.tooling.bytebuddy.DDCachingPoolStrategy;
44
import datadog.trace.agent.tooling.bytebuddy.DDLocationStrategy;
5+
import datadog.trace.api.Config;
56
import datadog.trace.api.Platform;
67
import datadog.trace.bootstrap.WeakCache;
78
import datadog.trace.bootstrap.WeakCache.Provider;
@@ -55,7 +56,8 @@ private static Provider loadWeakCacheProvider() {
5556
private static final Provider weakCacheProvider = loadWeakCacheProvider();
5657

5758
private static final DDLocationStrategy LOCATION_STRATEGY = new DDLocationStrategy();
58-
private static final DDCachingPoolStrategy POOL_STRATEGY = new DDCachingPoolStrategy();
59+
private static final DDCachingPoolStrategy POOL_STRATEGY =
60+
new DDCachingPoolStrategy(Config.get().isResolverUseLoadClassEnabled());
5961

6062
public static <K, V> WeakCache<K, V> newWeakCache() {
6163
return newWeakCache(DEFAULT_CACHE_CAPACITY);

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/bytebuddy/DDCachingPoolStrategy.java‎

Lines changed: 78 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -72,8 +72,20 @@ public class DDCachingPoolStrategy implements PoolStrategy {
7272
.build();
7373

7474
/** Fast path for bootstrap */
75-
final SharedResolutionCacheAdapter bootstrapCacheProvider =
76-
new SharedResolutionCacheAdapter(BOOTSTRAP_HASH, null, sharedResolutionCache);
75+
final SharedResolutionCacheAdapter bootstrapCacheProvider;
76+
77+
private final boolean fallBackToLoadClass;
78+
79+
public DDCachingPoolStrategy() {
80+
this(true);
81+
}
82+
83+
public DDCachingPoolStrategy(boolean fallBackToLoadClass) {
84+
this.fallBackToLoadClass = fallBackToLoadClass;
85+
bootstrapCacheProvider =
86+
new SharedResolutionCacheAdapter(
87+
BOOTSTRAP_HASH, null, sharedResolutionCache, fallBackToLoadClass);
88+
}
7789

7890
@Override
7991
public final TypePool typePool(
@@ -98,7 +110,8 @@ public WeakReference<ClassLoader> apply(ClassLoader input) {
98110

99111
private TypePool.CacheProvider createCacheProvider(
100112
final int loaderHash, final WeakReference<ClassLoader> loaderRef) {
101-
return new SharedResolutionCacheAdapter(loaderHash, loaderRef, sharedResolutionCache);
113+
return new SharedResolutionCacheAdapter(
114+
loaderHash, loaderRef, sharedResolutionCache, fallBackToLoadClass);
102115
}
103116

104117
private TypePool createCachingTypePool(
@@ -205,14 +218,17 @@ static final class SharedResolutionCacheAdapter implements TypePool.CacheProvide
205218
private final int loaderHash;
206219
private final WeakReference<ClassLoader> loaderRef;
207220
private final ConcurrentMap<TypeCacheKey, TypePool.Resolution> sharedResolutionCache;
221+
private final boolean fallBackToLoadClass;
208222

209223
SharedResolutionCacheAdapter(
210224
final int loaderHash,
211225
final WeakReference<ClassLoader> loaderRef,
212-
final ConcurrentMap<TypeCacheKey, TypePool.Resolution> sharedResolutionCache) {
226+
final ConcurrentMap<TypeCacheKey, TypePool.Resolution> sharedResolutionCache,
227+
final boolean fallBackToLoadClass) {
213228
this.loaderHash = loaderHash;
214229
this.loaderRef = loaderRef;
215230
this.sharedResolutionCache = sharedResolutionCache;
231+
this.fallBackToLoadClass = fallBackToLoadClass;
216232
}
217233

218234
@Override
@@ -236,7 +252,15 @@ public TypePool.Resolution register(final String className, TypePool.Resolution
236252
return resolution;
237253
}
238254

239-
resolution = new CachingResolution(resolution);
255+
if (fallBackToLoadClass && resolution instanceof TypePool.Resolution.Illegal) {
256+
// If the normal pool only resolution have failed then fall back to creating the type
257+
// description from a loaded type by trying to load the class. This case is very rare and is
258+
// here to handle classes that are injected directly via calls to defineClass without
259+
// providing a way to get the class bytes.
260+
resolution = new CachingResolutionForMaybeLoadableType(loaderRef, className);
261+
} else {
262+
resolution = new CachingResolution(resolution);
263+
}
240264

241265
sharedResolutionCache.put(new TypeCacheKey(loaderHash, loaderRef, className), resolution);
242266
return resolution;
@@ -248,12 +272,60 @@ public void clear() {
248272
}
249273
}
250274

275+
private static class CachingResolutionForMaybeLoadableType implements TypePool.Resolution {
276+
private final WeakReference<ClassLoader> loaderRef;
277+
private final String className;
278+
private volatile TypeDescription typeDescription = null;
279+
private volatile boolean isResolved = false;
280+
281+
public CachingResolutionForMaybeLoadableType(
282+
WeakReference<ClassLoader> loaderRef, String className) {
283+
this.loaderRef = loaderRef;
284+
this.className = className;
285+
}
286+
287+
@Override
288+
public boolean isResolved() {
289+
return isResolved;
290+
}
291+
292+
@Override
293+
public TypeDescription resolve() {
294+
// Intentionally not "thread safe". Duplicate work deemed an acceptable trade-off.
295+
if (!isResolved) {
296+
Class<?> klass = null;
297+
ClassLoader classLoader = loaderRef.get();
298+
if (classLoader != null) {
299+
try {
300+
// Please note that by doing a loadClass, the type we are resolving will bypass
301+
// transformation since we are in the middle of a transformation. This should
302+
// be a very rare occurrence and not affect any classes we want to instrument.
303+
klass = classLoader.loadClass(className);
304+
} catch (ClassNotFoundException ignored) {
305+
}
306+
}
307+
if (klass != null) {
308+
// We managed to load the class
309+
typeDescription = TypeDescription.ForLoadedType.of(klass);
310+
log.debug(
311+
"Direct loadClass type resolution of {} from class loader {} bypass transformation",
312+
className,
313+
classLoader);
314+
}
315+
isResolved = true;
316+
}
317+
if (typeDescription == null) {
318+
throw new IllegalStateException("Cannot resolve type description for " + className);
319+
}
320+
return typeDescription;
321+
}
322+
}
323+
251324
private static class CachingResolution implements TypePool.Resolution {
252325
private final TypePool.Resolution delegate;
253326
private TypeDescription cachedResolution;
254327

255328
public CachingResolution(final TypePool.Resolution delegate) {
256-
257329
this.delegate = delegate;
258330
}
259331

‎dd-java-agent/testing/src/main/groovy/datadog/trace/agent/test/AgentTestRunner.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ public abstract class AgentTestRunner extends DDSpecification {
9696
private static final AtomicInteger INSTRUMENTATION_ERROR_COUNT = new AtomicInteger(0);
9797
private static final TestRunnerListener TEST_LISTENER = new TestRunnerListener();
9898

99-
private static final Instrumentation INSTRUMENTATION;
99+
protected static final Instrumentation INSTRUMENTATION;
100100
private static volatile ClassFileTransformer activeTransformer = null;
101101

102102
static {
@@ -229,12 +229,18 @@ public void cleanUpAfterTests() {
229229
TEST_LISTENER.deactivateTest(this);
230230
}
231231

232+
/** Override to clean up things after the agent is removed */
233+
protected void cleanupAfterAgent() {}
234+
232235
@AfterClass
233-
public static synchronized void agentCleanup() {
236+
public synchronized void agentCleanup() {
234237
if (null != activeTransformer) {
235238
INSTRUMENTATION.removeTransformer(activeTransformer);
236239
activeTransformer = null;
237240
}
241+
242+
cleanupAfterAgent();
243+
238244
// Cleanup before assertion.
239245
assert INSTRUMENTATION_ERROR_COUNT.get() == 0
240246
: INSTRUMENTATION_ERROR_COUNT.get() + " Instrumentation errors during test";

‎dd-java-agent/testing/src/test/groovy/locator/ClassInjectingForkedTest.groovy‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package locator
2+
3+
import datadog.trace.agent.test.AgentTestRunner
4+
import net.bytebuddy.agent.builder.AgentBuilder
5+
6+
import java.lang.instrument.ClassFileTransformer
7+
8+
class ClassInjectingForkedTest extends AgentTestRunner {
9+
10+
static volatile ClassFileTransformer extraTransformer = null
11+
12+
@Override
13+
protected void configurePreAgent() {
14+
super.configurePreAgent()
15+
16+
// Since this method is not at all configurePreAgent, but more like
17+
// configurePreAgentAndOhByTheWayBeforeEveryTest we need to not install
18+
// the extra transformer multiple times
19+
if (!extraTransformer) {
20+
AgentBuilder builder = new AgentBuilder.Default()
21+
builder = ClassInjectingTransformer.instrument(builder)
22+
extraTransformer = builder.installOn(INSTRUMENTATION)
23+
}
24+
}
25+
26+
@Override
27+
protected void cleanupAfterAgent() {
28+
if (extraTransformer) {
29+
INSTRUMENTATION.removeTransformer(extraTransformer)
30+
extraTransformer = null
31+
}
32+
33+
super.cleanupAfterAgent()
34+
}
35+
36+
def "should find classes injected via defineClass"() {
37+
setup:
38+
def instrumented = new ClassInjectingTestInstrumentation.ToBeInstrumented("test")
39+
40+
expect:
41+
instrumented.message == "test:instrumented:${ClassInjectingTransformer.NAME}"
42+
}
43+
}

‎dd-java-agent/testing/src/test/groovy/locator/ClassInjectingLoadClassDisabledForkedTest.groovy‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
package locator
2+
3+
import datadog.trace.agent.test.AgentTestRunner
4+
import net.bytebuddy.agent.builder.AgentBuilder
5+
import net.bytebuddy.utility.JavaModule
6+
import spock.lang.Shared
7+
8+
import java.lang.instrument.ClassFileTransformer
9+
10+
/**
11+
* This test checks that we don't fall back to loadClass when it is disabled.
12+
*/
13+
class ClassInjectingLoadClassDisabledForkedTest extends AgentTestRunner {
14+
15+
static volatile ClassFileTransformer extraTransformer = null
16+
17+
@Override
18+
protected void configurePreAgent() {
19+
super.configurePreAgent()
20+
21+
injectSysConfig("dd.resolver.use.loadclass", "false")
22+
23+
// Since this method is not at all configurePreAgent, but more like
24+
// configurePreAgentAndOhByTheWayBeforeEveryTest we need to not install
25+
// the extra transformer multiple times
26+
if (!extraTransformer) {
27+
AgentBuilder builder = new AgentBuilder.Default()
28+
builder = ClassInjectingTransformer.instrument(builder)
29+
extraTransformer = builder.installOn(INSTRUMENTATION)
30+
}
31+
}
32+
33+
@Override
34+
protected void cleanupAfterAgent() {
35+
if (extraTransformer) {
36+
INSTRUMENTATION.removeTransformer(extraTransformer)
37+
extraTransformer = null
38+
}
39+
40+
super.cleanupAfterAgent()
41+
}
42+
43+
@Override
44+
protected boolean onInstrumentationError(String typeName, ClassLoader classLoader, JavaModule module, boolean loaded, Throwable throwable) {
45+
if (typeName == "${ClassInjectingTestInstrumentation.name}\$ToBeInstrumented") {
46+
instrumentationFailure = true
47+
return false
48+
}
49+
50+
return super.onInstrumentationError(typeName, classLoader, module, loaded, throwable)
51+
}
52+
53+
@Shared
54+
volatile boolean instrumentationFailure = false
55+
56+
def "should not find classes injected via defineClass"() {
57+
setup:
58+
def instrumented = new ClassInjectingTestInstrumentation.ToBeInstrumented("test")
59+
60+
expect:
61+
instrumentationFailure
62+
instrumented.message == "test:${ClassInjectingTransformer.NAME}"
63+
}
64+
}

‎dd-java-agent/testing/src/test/java/locator/ClassInjectingTestInstrumentation.java‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package locator;
2+
3+
import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
4+
import static net.bytebuddy.matcher.ElementMatchers.named;
5+
6+
import com.google.auto.service.AutoService;
7+
import datadog.trace.agent.tooling.Instrumenter;
8+
import datadog.trace.agent.tooling.Utils;
9+
import datadog.trace.agent.tooling.bytebuddy.ExceptionHandlers;
10+
import net.bytebuddy.agent.builder.AgentBuilder;
11+
import net.bytebuddy.asm.Advice;
12+
13+
@AutoService(Instrumenter.class)
14+
public class ClassInjectingTestInstrumentation implements Instrumenter {
15+
@Override
16+
public AgentBuilder instrument(AgentBuilder agentBuilder) {
17+
return agentBuilder
18+
.type(named(getClass().getName() + "$ToBeInstrumented"))
19+
.transform(
20+
new AgentBuilder.Transformer.ForAdvice()
21+
.include(Utils.getBootstrapProxy(), Utils.getAgentClassLoader())
22+
.withExceptionHandler(ExceptionHandlers.defaultExceptionHandler())
23+
.advice(isConstructor(), getClass().getName() + "$ConstructorAdvice"));
24+
}
25+
26+
public static class ConstructorAdvice {
27+
@Advice.OnMethodEnter
28+
public static void appendToMessage(
29+
@Advice.Argument(value = 0, readOnly = false) String message) {
30+
message = message + ":instrumented";
31+
}
32+
}
33+
34+
public static final class ToBeInstrumented {
35+
private final String message;
36+
37+
public ToBeInstrumented(String message) {
38+
this.message = message;
39+
}
40+
41+
public String getMessage() {
42+
StringBuilder msg = new StringBuilder(message);
43+
for (Class<?> iface : getClass().getInterfaces()) {
44+
msg.append(":");
45+
msg.append(iface.getName());
46+
}
47+
return msg.toString();
48+
}
49+
}
50+
}

0 commit comments

Comments
 (0)