Skip to content

Commit 8cbd216

Browse files
claudeclaude
committed
Fix thread safety, error handling, and code clarity in masked agent JAR
Introduce a masked JAR architecture where only Agent and AgentClassLoader remain as .class files; all other agent classes are renamed to .classdata at build time and loaded on demand via a custom classloader. Key changes: - Add Agent.java as the new entry point (Premain-Class/Agent-Class) - Add AgentClassLoader.java for loading masked .classdata classes - Append agent JAR to bootstrap classpath before loading Main to ensure classes required by Main's static initializers are available - Surface clear error messages on init failure; throw RuntimeException in premain so the JVM reports the failure - Pass parent classloader to super() instead of keeping a separate field - Override loadClass() for explicit delegation; use readAllBytes() - Simplify processClasspaths() since bootstrap path is handled by Agent - Add shadow JAR build pipeline with .classdata renaming - Use DefaultTask for rename, add inputs/outputs for Gradle cacheability - Remove old MANIFEST.MF (replaced by agentJar task manifest attributes) https://claude.ai/code/session_01WRpQefqudtbee9uatYakjY
1 parent f5a8ef2 commit 8cbd216

File tree

5 files changed

+175
-70
lines changed

5 files changed

+175
-70
lines changed

‎btrace-agent/build.gradle‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
plugins {
2+
alias(libs.plugins.shadow)
3+
}
4+
15
compileJava {
26
// Keep Java 8 compatibility while accessing JDK internal APIs
37
options.fork = true
@@ -30,3 +34,71 @@ dependencies {
3034
tasks.named('javadoc').configure {
3135
exclude 'org/openjdk/btrace/agent/PerfReaderImpl.java'
3236
}
37+
38+
shadowJar {
39+
include {
40+
if (it.directory) {
41+
return true
42+
}
43+
if (it.path.endsWith('.jar')) {
44+
return true
45+
}
46+
if (it.path.startsWith('org/objectweb/asm/')) {
47+
if (it.path.startsWith('org/objectweb/asm/commons/') ||
48+
it.path.startsWith('org/objectweb/asm/util/') ||
49+
it.path.startsWith('org/objectweb/asm/xml/')) {
50+
return false
51+
}
52+
return true
53+
}
54+
return it.path.startsWith('org/jctools/') ||
55+
it.path.startsWith("org/slf4j/") ||
56+
it.path.startsWith('org/openjdk/btrace/agent/') ||
57+
it.path.startsWith('org/openjdk/btrace/instr/') ||
58+
it.path.startsWith('org/openjdk/btrace/core/') ||
59+
it.path.startsWith('org/openjdk/btrace/runtime/') ||
60+
it.path.startsWith('org/openjdk/btrace/services/') ||
61+
it.path.startsWith('org/openjdk/btrace/statsd/')
62+
}
63+
64+
relocate 'org.jctools', 'org.openjdk.btrace.libs.agent.org.jctools'
65+
relocate 'org.objectweb.asm', 'org.openjdk.btrace.libs.org.objectweb.asm'
66+
relocate 'org.slf4j', 'org.openjdk.btrace.libs.org.slf4j'
67+
}
68+
69+
def unpackedJarDir = layout.buildDirectory.dir("tmp/unpackedJar")
70+
71+
task unpackJar(type: Copy) {
72+
dependsOn shadowJar
73+
from zipTree(shadowJar.archiveFile)
74+
into unpackedJarDir
75+
}
76+
77+
task renameClassFiles(type: DefaultTask) {
78+
dependsOn unpackJar
79+
inputs.dir unpackedJarDir
80+
outputs.dir unpackedJarDir
81+
doLast {
82+
fileTree(unpackedJarDir.get().asFile.path + "/org/openjdk/btrace/agent").include('**/*.class').each { File file ->
83+
if (file.name != 'Agent.class' && !file.name.startsWith('AgentClassLoader')) {
84+
file.renameTo(new File(file.parent, file.name.replace('.class', '.classdata')))
85+
}
86+
}
87+
}
88+
}
89+
90+
task agentJar(type: Jar) {
91+
dependsOn renameClassFiles
92+
from unpackedJarDir
93+
archiveFileName = project.tasks.shadowJar.archiveFileName
94+
destinationDirectory = project.tasks.shadowJar.destinationDirectory
95+
manifest {
96+
attributes(
97+
"Premain-Class": "org.openjdk.btrace.agent.Agent",
98+
"Agent-Class": "org.openjdk.btrace.agent.Agent",
99+
"Can-Redefine-Classes": true,
100+
"Can-Retransform-Classes": true,
101+
"Boot-Class-Path": "btrace-boot.jar"
102+
)
103+
}
104+
}

‎btrace-agent/src/main/java/org/openjdk/btrace/agent/Agent.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package org.openjdk.btrace.agent;
2+
3+
import java.io.File;
4+
import java.lang.instrument.Instrumentation;
5+
import java.lang.reflect.Method;
6+
import java.net.URL;
7+
import java.util.jar.JarFile;
8+
9+
public final class Agent {
10+
public static void premain(String args, Instrumentation inst) {
11+
main(args, inst, true);
12+
}
13+
14+
public static void agentmain(String args, Instrumentation inst) {
15+
main(args, inst, false);
16+
}
17+
18+
private static void main(String args, Instrumentation inst, boolean isPremain) {
19+
try {
20+
// Append to bootstrap classpath before loading Main to ensure classes
21+
// required by Main's static initializers are available
22+
URL pd = Agent.class.getProtectionDomain().getCodeSource().getLocation();
23+
if (pd.toString().endsWith(".jar")) {
24+
inst.appendToBootstrapClassLoaderSearch(new JarFile(new File(pd.toURI())));
25+
}
26+
Class<?> mainClass = AgentClassLoader.getInstance().loadClass("org.openjdk.btrace.agent.Main");
27+
Method mainMethod = mainClass.getMethod("main", String.class, Instrumentation.class);
28+
mainMethod.invoke(null, args, inst);
29+
} catch (Exception e) {
30+
String msg = "BTrace agent failed to initialize: " + e.getMessage();
31+
System.err.println(msg);
32+
e.printStackTrace();
33+
if (isPremain) {
34+
throw new RuntimeException(msg, e);
35+
}
36+
}
37+
}
38+
}

‎btrace-agent/src/main/java/org/openjdk/btrace/agent/AgentClassLoader.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
package org.openjdk.btrace.agent;
2+
3+
import java.io.IOException;
4+
import java.io.InputStream;
5+
6+
final class AgentClassLoader extends ClassLoader {
7+
private static final class Singleton {
8+
private static final AgentClassLoader INSTANCE = new AgentClassLoader();
9+
}
10+
11+
private AgentClassLoader() {
12+
super(Agent.class.getClassLoader() == null
13+
? ClassLoader.getSystemClassLoader()
14+
: Agent.class.getClassLoader());
15+
}
16+
17+
public static AgentClassLoader getInstance() {
18+
return Singleton.INSTANCE;
19+
}
20+
21+
@Override
22+
public Class<?> loadClass(String name, boolean resolve) throws ClassNotFoundException {
23+
synchronized (getClassLoadingLock(name)) {
24+
Class<?> c = findLoadedClass(name);
25+
if (c != null) {
26+
return c;
27+
}
28+
// Try loading from .classdata resources first, then delegate to parent
29+
try (InputStream is = getParent().getResourceAsStream(name.replace('.', '/') + ".classdata")) {
30+
if (is != null) {
31+
byte[] data = is.readAllBytes();
32+
c = defineClass(name, data, 0, data.length);
33+
if (resolve) {
34+
resolveClass(c);
35+
}
36+
return c;
37+
}
38+
} catch (IOException e) {
39+
throw new ClassNotFoundException("Failed to read classdata for " + name, e);
40+
}
41+
return super.loadClass(name, resolve);
42+
}
43+
}
44+
}

‎btrace-agent/src/main/java/org/openjdk/btrace/agent/Main.java‎

Lines changed: 21 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -133,15 +133,7 @@ public final class Main {
133133

134134
private static final Logger log = LoggerFactory.getLogger(Main.class);
135135

136-
public static void premain(String args, Instrumentation inst) {
137-
main(args, inst);
138-
}
139-
140-
public static void agentmain(String args, Instrumentation inst) {
141-
main(args, inst);
142-
}
143-
144-
private static synchronized void main(String args, Instrumentation inst) {
136+
public static synchronized void main(String args, Instrumentation inst) {
145137
if (AGENT_DEBUG) System.err.println("[BTrace Agent] Initialization started");
146138
if (Main.inst != null) {
147139
if (AGENT_DEBUG) System.err.println("[BTrace Agent] Agent already initialized, skipping");
@@ -719,73 +711,37 @@ private static void parseArgs() {
719711
}
720712

721713
private static void processClasspaths(String libs) {
722-
// Try to find JAR via Loader.class (unmasked bootstrap class)
723-
// Main.class won't work because it's loaded from .classdata
724-
String bootPath = null;
725-
try {
726-
Class<?> loaderClass = Class.forName("org.openjdk.btrace.boot.Loader");
727-
URL loaderResource = loaderClass.getResource("Loader.class");
728-
if (loaderResource != null) {
729-
bootPath = loaderResource.toString();
730-
if (bootPath.startsWith("jar:file:")) {
731-
// Extract JAR path from jar:file:/path/to/btrace.jar!/org/openjdk/btrace/boot/Loader.class
732-
bootPath = bootPath.substring("jar:file:".length());
733-
int idx = bootPath.indexOf("!");
734-
if (idx > -1) {
735-
bootPath = bootPath.substring(0, idx);
736-
}
737-
}
738-
}
739-
} catch (ClassNotFoundException e) {
740-
// Fall back to Main.class if Loader not found (shouldn't happen)
741-
URL agentJar = Main.class.getResource("Main.class");
742-
if (agentJar != null) {
743-
bootPath = agentJar.toString().replace("jar:file:", "");
744-
int idx = bootPath.indexOf("btrace-agent.jar");
745-
if (idx > -1) {
746-
bootPath = bootPath.substring(0, idx) + "btrace-boot.jar";
747-
}
748-
}
749-
}
750-
751714
String bootClassPath = argMap.get(BOOT_CLASS_PATH);
752-
if (bootClassPath == null && bootPath != null) {
753-
bootClassPath = bootPath;
754-
} else if (bootClassPath != null && bootPath != null) {
755-
if (".".equals(bootClassPath)) {
756-
bootClassPath = bootPath;
757-
} else {
758-
bootClassPath = bootPath + File.pathSeparator + bootClassPath;
759-
}
760-
}
761-
log.debug("Bootstrap ClassPath: {}", bootClassPath);
762715

763-
StringTokenizer tokenizer = new StringTokenizer(bootClassPath, File.pathSeparator);
764-
try {
765-
while (tokenizer.hasMoreTokens()) {
766-
String path = tokenizer.nextToken();
767-
File f = new File(path);
768-
if (!f.exists()) {
769-
log.debug("BTrace bootstrap classpath resource [{}] does not exist", path);
770-
} else {
771-
if (f.isFile() && f.getName().toLowerCase().endsWith(".jar")) {
772-
JarFile jf = asJarFile(f);
773-
log.debug("Adding jar: {}", jf);
774-
inst.appendToBootstrapClassLoaderSearch(jf);
716+
if (bootClassPath != null) {
717+
log.debug("Bootstrap ClassPath: {}", bootClassPath);
718+
StringTokenizer tokenizer = new StringTokenizer(bootClassPath, File.pathSeparator);
719+
try {
720+
while (tokenizer.hasMoreTokens()) {
721+
String path = tokenizer.nextToken();
722+
File f = new File(path);
723+
if (!f.exists()) {
724+
log.warn("BTrace bootstrap classpath resource [{}] does not exist", path);
775725
} else {
776-
log.debug("ignoring boot classpath element '{}' - only jar files allowed", path);
726+
if (f.isFile() && f.getName().toLowerCase().endsWith(".jar")) {
727+
JarFile jf = asJarFile(f);
728+
log.debug("Adding jar: {}", jf);
729+
inst.appendToBootstrapClassLoaderSearch(jf);
730+
} else {
731+
log.debug("ignoring boot classpath element '{}' - only jar files allowed", path);
732+
}
777733
}
778734
}
735+
} catch (IOException ex) {
736+
log.debug("adding to boot classpath failed!", ex);
737+
return;
779738
}
780-
} catch (IOException ex) {
781-
log.debug("adding to boot classpath failed!", ex);
782-
return;
783739
}
784740

785741
String systemClassPath = argMap.get(SYSTEM_CLASS_PATH);
786742
if (systemClassPath != null) {
787743
log.debug("System ClassPath: {}", systemClassPath);
788-
tokenizer = new StringTokenizer(systemClassPath, File.pathSeparator);
744+
StringTokenizer tokenizer = new StringTokenizer(systemClassPath, File.pathSeparator);
789745
try {
790746
while (tokenizer.hasMoreTokens()) {
791747
String path = tokenizer.nextToken();

‎btrace-agent/src/main/resources/META-INF/MANIFEST.MF‎

Lines changed: 0 additions & 5 deletions
This file was deleted.

0 commit comments

Comments
 (0)