allow skip specified class from omni enhancer

This commit is contained in:
金戟 2021-04-17 19:24:02 +08:00
parent 8dde402a8c
commit 5921ccab07
5 changed files with 56 additions and 35 deletions

View File

@ -19,7 +19,7 @@ public class PreMain {
ArgumentParser.parseArgs(agentArgs);
PropertiesParser.parseFile(ArgumentParser.configFilePath);
GlobalConfig.setupLogRootPath();
if (GlobalConfig.isEnhanceThreadLocal()) {
if (GlobalConfig.shouldEnhanceThreadLocal()) {
// add transmittable thread local transformer
TtlAgent.premain(agentArgs, inst);
}

View File

@ -33,7 +33,7 @@ public class ArgumentParser {
} else if (k.equals(DUMP_PATH)) {
GlobalConfig.setDumpPath(v);
} else if (k.equals(PKG_PREFIX)) {
GlobalConfig.setPkgPrefixes(v);
GlobalConfig.setPkgPrefixWhiteList(v);
} else if (k.equals(MOCK_SCOPE)) {
GlobalConfig.setDefaultMockScope(MockScope.of(v));
} else if (k.equals(CONFIG_FILE)) {
@ -42,7 +42,7 @@ public class ArgumentParser {
} else {
// parameter with single value
if (a.equals(USE_THREAD_POOL)) {
GlobalConfig.setEnhanceThreadLocal(true);
GlobalConfig.enableEnhanceThreadLocal(true);
}
}
}

View File

@ -20,11 +20,12 @@ public class PropertiesParser {
private static final String LOG_LEVEL = "log.level";
private static final String LOG_FILE = "log.file";
private static final String DUMP_PATH = "dump.path";
private static final String PKG_PREFIX_WHITELIST = "enhance.pkgPrefix.whiteList";
private static final String PKG_PREFIX_INCLUDES = "enhance.pkgPrefix.includes";
private static final String INNER_MOCK_CLASS_NAME = "mock.innerClass.name";
private static final String DEFAULT_MOCK_SCOPE = "mock.scope.default";
private static final String ENABLE_THREAD_POOL = "thread.pool.enhance.enable";
private static final String ENABLE_OMNI_INJECT = "omni.constructor.enhance.enable";
private static final String OMNI_INJECT_EXCLUDES = "omni.constructor.enhance.pkgPrefix.excludes";
public static void parseFile(String configFilePath) {
String path = (configFilePath == null) ? DEFAULT_CONFIG_FILE : configFilePath;
@ -56,14 +57,16 @@ public class PropertiesParser {
GlobalConfig.setLogFile(v);
} else if (k.equals(DUMP_PATH)) {
GlobalConfig.setDumpPath(v);
} else if (k.equals(PKG_PREFIX_WHITELIST)) {
GlobalConfig.setPkgPrefixes(v);
} else if (k.equals(PKG_PREFIX_INCLUDES)) {
GlobalConfig.setPkgPrefixWhiteList(v);
} else if (k.equals(OMNI_INJECT_EXCLUDES)) {
GlobalConfig.setOmniPkgPrefixBlackList(v);
} else if (k.equals(DEFAULT_MOCK_SCOPE)) {
GlobalConfig.setDefaultMockScope(MockScope.of(v));
} else if (k.equals(ENABLE_THREAD_POOL)) {
GlobalConfig.setEnhanceThreadLocal(Boolean.parseBoolean(v));
GlobalConfig.enableEnhanceThreadLocal(Boolean.parseBoolean(v));
} else if (k.equals(ENABLE_OMNI_INJECT)) {
GlobalConfig.setEnhanceOmniConstructor(Boolean.parseBoolean(v));
GlobalConfig.enableEnhanceOmniConstructor(Boolean.parseBoolean(v));
} else if (k.equals(INNER_MOCK_CLASS_NAME)) {
GlobalConfig.setInnerMockClassName(v);
}

View File

@ -50,8 +50,7 @@ public class TestableClassTransformer implements ClassFileTransformer {
return null;
}
LogUtil.verbose("Handle class: " + className);
byte[] bytes = GlobalConfig.isEnhanceOmniConstructor() ?
new OmniClassHandler().getBytes(classFileBuffer) : classFileBuffer;
byte[] bytes = shouldOmniEnhance(className) ? new OmniClassHandler().getBytes(classFileBuffer) : classFileBuffer;
ClassNode cn = ClassUtil.getClassNode(className);
if (cn != null) {
return transformMock(bytes, cn);
@ -59,6 +58,12 @@ public class TestableClassTransformer implements ClassFileTransformer {
return bytes;
}
private boolean shouldOmniEnhance(String className) {
String[] blackList = GlobalConfig.getOmniPkgPrefixBlackList();
return GlobalConfig.shouldEnhanceOmniConstructor() &&
(blackList == null || !isInPrefixList(className, blackList));
}
private byte[] transformMock(byte[] bytes, ClassNode cn) {
String className = cn.name;
try {
@ -141,20 +146,18 @@ public class TestableClassTransformer implements ClassFileTransformer {
if (null == className || className.contains(CGLIB_CLASS_PATTERN)) {
return true;
}
List<String> whitePrefixes = GlobalConfig.getPkgPrefixes();
if (!whitePrefixes.isEmpty()) {
for (String prefix : whitePrefixes) {
if (className.startsWith(prefix)) {
// Only consider package in provided list as non-system class
return false;
}
}
return true;
} else {
for (String prefix : BLACKLIST_PREFIXES) {
if (className.startsWith(prefix)) {
return true;
}
String[] pkgPrefixWhiteList = GlobalConfig.getPkgPrefixWhiteList();
if (pkgPrefixWhiteList != null) {
// Only consider package in provided list as non-system class
return !isInPrefixList(className, pkgPrefixWhiteList);
}
return isInPrefixList(className, BLACKLIST_PREFIXES);
}
private boolean isInPrefixList(String name, String[] prefixList) {
for (String prefix : prefixList) {
if (name.startsWith(prefix)) {
return true;
}
}
return false;

View File

@ -30,7 +30,8 @@ public class GlobalConfig {
private static String logFile = null;
private static String dumpPath = null;
private static List<String> pkgPrefixes = new ArrayList<String>();
private static String[] pkgPrefixWhiteList = null;
private static String[] omniPkgPrefixBlackList = null;
private static MockScope defaultMockScope = MockScope.GLOBAL;
private static boolean enhanceThreadLocal = false;
private static boolean enhanceOmniConstructor = false;
@ -61,14 +62,20 @@ public class GlobalConfig {
}
}
public static List<String> getPkgPrefixes() {
return pkgPrefixes;
public static String[] getPkgPrefixWhiteList() {
return pkgPrefixWhiteList;
}
public static void setPkgPrefixes(String prefixes) {
for (String p : prefixes.split(COMMA)) {
pkgPrefixes.add(p.endsWith(DOT) ? p : p + DOT);
}
public static void setPkgPrefixWhiteList(String prefixes) {
pkgPrefixWhiteList = parsePkgPrefixList(prefixes).toArray(new String[0]);
}
public static String[] getOmniPkgPrefixBlackList() {
return omniPkgPrefixBlackList;
}
public static void setOmniPkgPrefixBlackList(String prefixes) {
omniPkgPrefixBlackList = parsePkgPrefixList(prefixes).toArray(new String[0]);
}
public static MockScope getDefaultMockScope() {
@ -108,19 +115,19 @@ public class GlobalConfig {
}
}
public static void setEnhanceThreadLocal(boolean enabled) {
public static void enableEnhanceThreadLocal(boolean enabled) {
enhanceThreadLocal = enabled;
}
public static boolean isEnhanceThreadLocal() {
public static boolean shouldEnhanceThreadLocal() {
return enhanceThreadLocal;
}
public static void setEnhanceOmniConstructor(boolean enabled) {
public static void enableEnhanceOmniConstructor(boolean enabled) {
enhanceOmniConstructor = enabled;
}
public static boolean isEnhanceOmniConstructor() {
public static boolean shouldEnhanceOmniConstructor() {
return enhanceOmniConstructor;
}
@ -131,4 +138,12 @@ public class GlobalConfig {
public static String getInnerMockClassName() {
return innerMockClassName;
}
private static List<String> parsePkgPrefixList(String prefixes) {
List<String> whiteList = new ArrayList<String>();
for (String p : prefixes.split(COMMA)) {
whiteList.add(p.endsWith(DOT) ? p : p + DOT);
}
return whiteList;
}
}