Skip to content

Commit b0770e9

Browse files
mccullsmcculls
authored
Ensure shaded helpers have unique names when injected into class-loaders (#8192)
* and delay creation of advice remapper+cache until needed
1 parent f3fd311 commit b0770e9

File tree

6 files changed

+120
-64
lines changed

6 files changed

+120
-64
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ private void prepareInstrumentation(InstrumenterModule module, int instrumentati
119119
new FieldBackedContextRequestRewriter(contextStore, module.name()))
120120
: null;
121121

122-
adviceShader = AdviceShader.with(module.adviceShading());
122+
adviceShader = AdviceShader.with(module);
123123

124124
String[] helperClassNames = module.helperClassNames();
125125
if (module.injectHelperDependencies()) {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ public ShadedAdviceLocator(ClassLoader adviceLoader, AdviceShader adviceShader)
1717
public Resolution locate(String className) throws IOException {
1818
final Resolution resolution = adviceLocator.locate(className);
1919
if (resolution.isResolved()) {
20-
return new Resolution.Explicit(adviceShader.shade(resolution.resolve()));
20+
return new Resolution.Explicit(adviceShader.shadeClass(resolution.resolve()));
2121
} else {
2222
return resolution;
2323
}

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

Lines changed: 111 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
11
package datadog.trace.agent.tooling;
22

3+
import static java.util.Arrays.asList;
4+
import static java.util.Collections.emptyList;
5+
36
import datadog.trace.api.cache.DDCache;
47
import datadog.trace.api.cache.DDCaches;
8+
import java.util.HashMap;
9+
import java.util.List;
510
import java.util.Map;
611
import net.bytebuddy.jar.asm.ClassReader;
712
import net.bytebuddy.jar.asm.ClassVisitor;
@@ -10,89 +15,139 @@
1015
import net.bytebuddy.jar.asm.commons.Remapper;
1116

1217
/** Shades advice bytecode by applying relocations to all references. */
13-
public final class AdviceShader extends Remapper {
14-
private final DDCache<String, String> cache = DDCaches.newFixedSizeCache(64);
18+
public final class AdviceShader {
19+
private final Map<String, String> relocations;
20+
private final List<String> helperNames;
1521

16-
/** Flattened sequence of old-prefix, new-prefix relocations. */
17-
private final String[] prefixes;
22+
private volatile Remapper remapper;
1823

19-
public static AdviceShader with(Map<String, String> relocations) {
20-
return relocations != null ? new AdviceShader(relocations) : null;
24+
/**
25+
* Used when installing {@link InstrumenterModule}s. Ensures any injected helpers have unique
26+
* names so the original and relocated modules can inject helpers into the same class-loader.
27+
*/
28+
public static AdviceShader with(InstrumenterModule module) {
29+
if (module.adviceShading() != null) {
30+
return new AdviceShader(module.adviceShading(), asList(module.helperClassNames()));
31+
}
32+
return null;
2133
}
2234

23-
AdviceShader(Map<String, String> relocations) {
24-
// convert relocations to a flattened sequence: old-prefix, new-prefix, etc.
25-
this.prefixes = new String[relocations.size() * 2];
26-
int i = 0;
27-
for (Map.Entry<String, String> e : relocations.entrySet()) {
28-
String oldPrefix = e.getKey();
29-
String newPrefix = e.getValue();
30-
if (oldPrefix.indexOf('.') > 0) {
31-
// accept dotted prefixes, but store them in their internal form
32-
this.prefixes[i++] = oldPrefix.replace('.', '/');
33-
this.prefixes[i++] = newPrefix.replace('.', '/');
34-
} else {
35-
this.prefixes[i++] = oldPrefix;
36-
this.prefixes[i++] = newPrefix;
37-
}
35+
/** Used to generate and check muzzle references. Only applies relocations declared in modules. */
36+
public static AdviceShader with(Map<String, String> relocations) {
37+
if (relocations != null) {
38+
return new AdviceShader(relocations, emptyList());
3839
}
40+
return null;
41+
}
42+
43+
private AdviceShader(Map<String, String> relocations, List<String> helperNames) {
44+
this.relocations = relocations;
45+
this.helperNames = helperNames;
3946
}
4047

4148
/** Applies shading before calling the given {@link ClassVisitor}. */
42-
public ClassVisitor shade(ClassVisitor cv) {
43-
return new ClassRemapper(cv, this);
49+
public ClassVisitor shadeClass(ClassVisitor cv) {
50+
if (null == remapper) {
51+
remapper = new AdviceMapper();
52+
}
53+
return new ClassRemapper(cv, remapper);
4454
}
4555

4656
/** Returns the result of shading the given bytecode. */
47-
public byte[] shade(byte[] bytecode) {
57+
public byte[] shadeClass(byte[] bytecode) {
4858
ClassReader cr = new ClassReader(bytecode);
4959
ClassWriter cw = new ClassWriter(null, 0);
50-
cr.accept(shade(cw), 0);
60+
cr.accept(shadeClass(cw), 0);
5161
return cw.toByteArray();
5262
}
5363

54-
@Override
55-
public String map(String internalName) {
56-
if (internalName.startsWith("java/")
57-
|| internalName.startsWith("datadog/")
58-
|| internalName.startsWith("net/bytebuddy/")) {
59-
return internalName; // never shade these references
64+
/** Generates a unique shaded name for the given helper. */
65+
public String uniqueHelper(String dottedName) {
66+
int packageEnd = dottedName.lastIndexOf('.');
67+
if (packageEnd > 0) {
68+
return dottedName.substring(0, packageEnd + 1) + "shaded" + dottedName.substring(packageEnd);
6069
}
61-
return cache.computeIfAbsent(internalName, this::shade);
70+
return dottedName;
6271
}
6372

64-
@Override
65-
public Object mapValue(Object value) {
66-
if (value instanceof String) {
67-
String text = (String) value;
68-
if (text.isEmpty()) {
69-
return text;
70-
} else if (text.indexOf('.') > 0) {
71-
return shadeDottedName(text);
73+
final class AdviceMapper extends Remapper {
74+
private final DDCache<String, String> mappingCache = DDCaches.newFixedSizeCache(64);
75+
76+
/** Flattened sequence of old-prefix, new-prefix relocations. */
77+
private final String[] prefixes;
78+
79+
private final Map<String, String> helperMapping;
80+
81+
AdviceMapper() {
82+
// record the unique names that we've given to injected helpers
83+
this.helperMapping = new HashMap<>(helperNames.size() + 1, 1f);
84+
for (String h : helperNames) {
85+
this.helperMapping.put(h.replace('.', '/'), uniqueHelper(h).replace('.', '/'));
86+
}
87+
// convert relocations to a flattened sequence: old-prefix, new-prefix, etc.
88+
this.prefixes = new String[relocations.size() * 2];
89+
int i = 0;
90+
for (Map.Entry<String, String> e : relocations.entrySet()) {
91+
String oldPrefix = e.getKey();
92+
String newPrefix = e.getValue();
93+
if (oldPrefix.indexOf('.') > 0) {
94+
// accept dotted prefixes, but store them in their internal form
95+
this.prefixes[i++] = oldPrefix.replace('.', '/');
96+
this.prefixes[i++] = newPrefix.replace('.', '/');
97+
} else {
98+
this.prefixes[i++] = oldPrefix;
99+
this.prefixes[i++] = newPrefix;
100+
}
101+
}
102+
}
103+
104+
@Override
105+
public String map(String internalName) {
106+
String uniqueName = helperMapping.get(internalName);
107+
if (uniqueName != null) {
108+
return uniqueName;
109+
}
110+
if (internalName.startsWith("java/")
111+
|| internalName.startsWith("datadog/")
112+
|| internalName.startsWith("net/bytebuddy/")) {
113+
return internalName; // never shade these references
114+
}
115+
return mappingCache.computeIfAbsent(internalName, this::shadeInternalName);
116+
}
117+
118+
@Override
119+
public Object mapValue(Object value) {
120+
if (value instanceof String) {
121+
String text = (String) value;
122+
if (text.isEmpty()) {
123+
return text;
124+
} else if (text.indexOf('.') > 0) {
125+
return shadeDottedName(text);
126+
} else {
127+
return shadeInternalName(text);
128+
}
72129
} else {
73-
return shade(text);
130+
return super.mapValue(value);
74131
}
75-
} else {
76-
return super.mapValue(value);
77132
}
78-
}
79133

80-
private String shade(String internalName) {
81-
for (int i = 0; i < prefixes.length; i += 2) {
82-
if (internalName.startsWith(prefixes[i])) {
83-
return prefixes[i + 1] + internalName.substring(prefixes[i].length());
134+
private String shadeInternalName(String internalName) {
135+
for (int i = 0; i < prefixes.length; i += 2) {
136+
if (internalName.startsWith(prefixes[i])) {
137+
return prefixes[i + 1] + internalName.substring(prefixes[i].length());
138+
}
84139
}
140+
return internalName;
85141
}
86-
return internalName;
87-
}
88142

89-
private String shadeDottedName(String name) {
90-
String internalName = name.replace('.', '/');
91-
for (int i = 0; i < prefixes.length; i += 2) {
92-
if (internalName.startsWith(prefixes[i])) {
93-
return prefixes[i + 1].replace('/', '.') + name.substring(prefixes[i].length());
143+
private String shadeDottedName(String name) {
144+
String internalName = name.replace('.', '/');
145+
for (int i = 0; i < prefixes.length; i += 2) {
146+
if (internalName.startsWith(prefixes[i])) {
147+
return prefixes[i + 1].replace('/', '.') + name.substring(prefixes[i].length());
148+
}
94149
}
150+
return name;
95151
}
96-
return name;
97152
}
98153
}

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -89,12 +89,13 @@ public HelperInjector(
8989
private Map<String, byte[]> getHelperMap() throws IOException {
9090
if (dynamicTypeMap.isEmpty()) {
9191
final Map<String, byte[]> classnameToBytes = new LinkedHashMap<>();
92-
for (final String helperClassName : helperClassNames) {
93-
byte[] classBytes = classFileLocator.locate(helperClassName).resolve();
92+
for (String helperName : helperClassNames) {
93+
byte[] classBytes = classFileLocator.locate(helperName).resolve();
9494
if (adviceShader != null) {
95-
classBytes = adviceShader.shade(classBytes);
95+
classBytes = adviceShader.shadeClass(classBytes);
96+
helperName = adviceShader.uniqueHelper(helperName);
9697
}
97-
classnameToBytes.put(helperClassName, classBytes);
98+
classnameToBytes.put(helperName, classBytes);
9899
}
99100

100101
return classnameToBytes;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -132,7 +132,7 @@ private static Map<String, byte[]> createHelperMap(final InstrumenterModule modu
132132
ClassFileLocator.ForClassLoader.of(module.getClass().getClassLoader());
133133
byte[] classBytes = locator.locate(helperName).resolve();
134134
if (null != adviceShader) {
135-
classBytes = adviceShader.shade(classBytes);
135+
classBytes = adviceShader.shadeClass(classBytes);
136136
}
137137
helperMap.put(helperName, classBytes);
138138
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ public static Map<String, Reference> createReferencesFrom(
7070
if (null == adviceShader) {
7171
reader.accept(cv, ClassReader.SKIP_FRAMES);
7272
} else {
73-
reader.accept(adviceShader.shade(cv), ClassReader.SKIP_FRAMES);
73+
reader.accept(adviceShader.shadeClass(cv), ClassReader.SKIP_FRAMES);
7474
}
7575

7676
final Map<String, Reference> instrumentationReferences = cv.getReferences();

0 commit comments

Comments
 (0)