寻求ByteBuddy拦截JDBC连接属性追加值的实现示例及建议

我刚好做过类似的ByteBuddy拦截JDBC连接属性的需求，给你分享两个可行的方案，你可以根据自己的场景选择：

我们的目标是在不修改业务代码的前提下，给JDBC连接Properties里的ClientCorrelationToken值追加后缀。ByteBuddy可以通过JavaAgent拦截两类关键点：

1. JDBC获取连接的方法（比如DriverManager.getConnection或具体驱动的connect方法），修改传入的Properties参数；
2. 直接拦截java.util.Properties的put方法，对特定key的值做追加处理。

方案一：精准拦截JDBC连接方法 #

这种方案只针对JDBC连接时传入的Properties做修改，不会影响其他地方的Properties使用。

1. JavaAgent入口类 #

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.matcher.ElementMatchers;

import java.lang.instrument.Instrumentation;
import java.util.Properties;

public class JDBCCorrelationAgent {
    public static void premain(String agentArgs, Instrumentation instrumentation) {
        new AgentBuilder.Default()
                // 拦截DriverManager中带Properties参数的getConnection方法
                .type(ElementMatchers.named("java.sql.DriverManager"))
                .transform((builder, typeDesc, classLoader, module) ->
                        builder.method(ElementMatchers.named("getConnection")
                                        .and(ElementMatchers.takesArgument(2, Properties.class)))
                                .intercept(MethodDelegation.to(DriverManagerInterceptor.class))
                )
                // 同时拦截MySQL驱动的connect方法（如果你的应用用MySQL，其他驱动类似）
                .type(ElementMatchers.named("com.mysql.cj.jdbc.NonRegisteringDriver"))
                .transform((builder, typeDesc, classLoader, module) ->
                        builder.method(ElementMatchers.named("connect")
                                        .and(ElementMatchers.takesArgument(1, Properties.class)))
                                .intercept(MethodDelegation.to(DriverConnectInterceptor.class))
                )
                .installOn(instrumentation);
    }
}

2. DriverManager拦截器 #

import java.lang.reflect.Method;
import java.util.Properties;
import net.bytebuddy.implementation.bind.annotation.AllArguments;
import net.bytebuddy.implementation.bind.annotation.Origin;
import net.bytebuddy.implementation.bind.annotation.RuntimeType;

public class DriverManagerInterceptor {
    private static final String APPENDED_SUFFIX = "appendedValue";

    @RuntimeType
    public static Object intercept(@Origin Method method,
                                   @AllArguments Object[] args) throws Exception {
        // 获取传入的Properties参数（第三个参数）
        Properties props = (Properties) args[2];
        appendCorrelationToken(props);
        // 执行原方法逻辑
        return method.invoke(null, args);
    }

    private static void appendCorrelationToken(Properties props) {
        String token = props.getProperty("ClientCorrelationToken");
        if (token != null && !token.endsWith(APPENDED_SUFFIX)) {
            props.put("ClientCorrelationToken", token + APPENDED_SUFFIX);
        }
    }
}

3. 具体驱动拦截器（以MySQL为例） #

import java.lang.reflect.Method;
import java.util.Properties;
import net.bytebuddy.implementation.bind.annotation.AllArguments;
import net.bytebuddy.implementation.bind.annotation.Origin;
import net.bytebuddy.implementation.bind.annotation.RuntimeType;
import net.bytebuddy.implementation.bind.annotation.This;

public class DriverConnectInterceptor {
    private static final String APPENDED_SUFFIX = "appendedValue";

    @RuntimeType
    public static Object intercept(@This Object driverInstance,
                                   @Origin Method method,
                                   @AllArguments Object[] args) throws Exception {
        // 获取传入的Properties参数（第二个参数）
        Properties props = (Properties) args[1];
        appendCorrelationToken(props);
        // 执行原方法逻辑
        return method.invoke(driverInstance, args);
    }

    private static void appendCorrelationToken(Properties props) {
        String token = props.getProperty("ClientCorrelationToken");
        if (token != null && !token.endsWith(APPENDED_SUFFIX)) {
            props.put("ClientCorrelationToken", token + APPENDED_SUFFIX);
        }
    }
}

方案二：全局拦截Properties的put方法 #

这种方案会拦截所有地方对Properties的put操作，只要key是ClientCorrelationToken就自动追加后缀，通用性更强。

1. 修改JavaAgent入口类 #

import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.matcher.ElementMatchers;

import java.lang.instrument.Instrumentation;
import java.util.Properties;

public class JDBCCorrelationAgent {
    public static void premain(String agentArgs, Instrumentation instrumentation) {
        new AgentBuilder.Default()
                // 拦截Properties的put方法
                .type(ElementMatchers.is(Properties.class))
                .transform((builder, typeDesc, classLoader, module) ->
                        builder.method(ElementMatchers.named("put")
                                        .and(ElementMatchers.takesArgument(0, String.class))
                                        .and(ElementMatchers.takesArgument(1, Object.class)))
                                .intercept(MethodDelegation.to(PropertiesPutInterceptor.class))
                )
                .installOn(instrumentation);
    }
}

2. Properties put拦截器 #

import java.util.Properties;
import net.bytebuddy.implementation.bind.annotation.Argument;
import net.bytebuddy.implementation.bind.annotation.RuntimeType;
import net.bytebuddy.implementation.bind.annotation.This;

public class PropertiesPutInterceptor {
    private static final String TOKEN_KEY = "ClientCorrelationToken";
    private static final String APPENDED_SUFFIX = "appendedValue";

    @RuntimeType
    public static Object intercept(@This Properties props,
                                   @Argument(0) String key,
                                   @Argument(1) Object value) {
        if (TOKEN_KEY.equals(key) && value instanceof String) {
            String originalValue = (String) value;
            // 避免重复追加后缀
            if (!originalValue.endsWith(APPENDED_SUFFIX)) {
                return props.put(key, originalValue + APPENDED_SUFFIX);
            }
        }
        // 其他key正常执行原put逻辑
        return props.put(key, value);
    }
}

打包与测试 #

1. 配置MANIFEST.MF：打包Agent jar时需要指定Premain类，示例内容：

Premain-Class: com.yourpackage.JDBCCorrelationAgent
Can-Redefine-Classes: true

2. 打包成Jar：用Maven或Gradle打包，确保ByteBuddy的依赖被包含（或者用shadowJar插件打包成fat jar）；
3. 运行测试：启动应用时加上Agent参数：

java -javaagent:./your-agent.jar -jar your-application.jar

这样当业务代码设置ClientCorrelationToken为MyToken1时，最终传入JDBC的属性值就会变成MyToken1appendedValue。

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