如何用Java Agent+Javaassist拦截java.net.Socket的connect方法？

我正在开展一个项目，需要在Socket连接建立后立即向其输出流写入一些头部信息。我希望在底层拦截java.net.Socket及其public void connect(SocketAddress endpoint, int timeout)方法，为此编写了如下transform函数：

public static byte[] transform(final ClassLoader loader, final String className, final byte[] classFileBuffer) {
    loader.loadClass(...);

    final ClassPool classPool = ClassPool.getDefault();
    classPool.appendClassPath(new LoaderClassPath(loader));
    classPool.appendSystemPath();
    classPool.importPackage(...);

    final CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classFileBuffer));

    if (className.equals("java/net/Socket")) {
        final CtMethod method = ctClass.getDeclaredMethod("connect");
        final String s = "System.out.println(\"I am in Socket.connect!\");";
        method.insertAfter(s);
        return ctClass.toBytecode();
    }
    
    return classFileBuffer;
}

我已确认java.net.Socket.connect()方法被调用，但字节码注入似乎从未生效，即使在插入的代码中加入测试内容也无报错。奇怪的是，相同的逻辑对java.net.SocksSocketImpl却有效，我无法排查问题所在，恳请得到帮助。

核心原因 #

1. Socket.connect是门面方法：java.net.Socket的connect方法仅做转发，实际连接逻辑交给内部的SocketImpl实现类（如PlainSocketImpl）处理。你注入的代码在Socket.connect的末尾，但如果业务代码的连接路径未走到该方法的结束逻辑，或者被底层实现绕开，注入代码就不会执行。
2. 核心类加载限制：java.net.Socket属于JDK核心类，由Bootstrap类加载器加载。你的Agent若由System类加载器加载，修改核心类时可能遇到类池无法正确获取类定义、JVM安全机制（如Java 9+模块封装）阻止修改的问题。而SocksSocketImpl作为SocketImpl子类，加载限制更宽松，所以注入逻辑生效。

解决办法 #

方案1：拦截SocketImpl的connect方法（推荐） #

既然Socket.connect最终转发到SocketImpl，直接修改抽象类java.net.SocketImpl的connect方法，所有子类都会继承执行注入代码：

if (className.equals("java/net/SocketImpl")) {
    CtMethod method = ctClass.getDeclaredMethod("connect", new CtClass[]{
        classPool.get("java.net.SocketAddress"),
        CtClass.intType
    });
    String injectCode = "// 写入头部信息逻辑\n" +
                       "java.io.OutputStream os = this.getOutputStream();\n" +
                       "os.write(\"your-header-content\".getBytes());\n" +
                       "os.flush();";
    method.insertAfter(injectCode);
    return ctClass.toBytecode();
}

方案2：修复Socket.connect的注入逻辑 #

若坚持拦截Socket.connect，需调整以下几点：

• 移除提前加载类的代码：loader.loadClass(...)会导致类提前加载，后续字节码修改无法生效。
• 适配Bootstrap类加载器：直接添加核心类的类路径到ClassPool。
  修改后的transform函数：

public static byte[] transform(final ClassLoader loader, final String className, final byte[] classFileBuffer) {
    final ClassPool classPool = ClassPool.getDefault();
    
    if (className.equals("java/net/Socket")) {
        // 添加Bootstrap加载的Socket类路径
        classPool.appendClassPath(new ClassClassPath(java.net.Socket.class));
    }
    classPool.appendClassPath(new LoaderClassPath(loader));
    classPool.appendSystemPath();
    classPool.importPackage("java.net");

    final CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classFileBuffer));

    if (className.equals("java/net/Socket")) {
        // 指定参数类型，避免重载方法匹配错误
        CtMethod method = ctClass.getDeclaredMethod("connect", new CtClass[]{
            classPool.get("java.net.SocketAddress"),
            CtClass.intType
        });
        final String s = "System.out.println(\"I am in Socket.connect!\");";
        method.insertAfter(s);
        return ctClass.toBytecode();
    }
    
    return classFileBuffer;
}

方案3：适配Java 9+模块系统 #

若使用Java 9及以上版本，需在Agent的MANIFEST.MF中添加：

Can-Redefine-Classes: true
Can-Retransform-Classes: true

并在JVM启动时添加参数开放模块权限：

--add-opens jdk.net/java.net=ALL-UNNAMED

内容的提问来源于stack exchange，提问作者MeowCat