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); ArgumentParser.parseArgs(agentArgs);
PropertiesParser.parseFile(ArgumentParser.configFilePath); PropertiesParser.parseFile(ArgumentParser.configFilePath);
GlobalConfig.setupLogRootPath(); GlobalConfig.setupLogRootPath();
if (GlobalConfig.isEnhanceThreadLocal()) { if (GlobalConfig.shouldEnhanceThreadLocal()) {
// add transmittable thread local transformer // add transmittable thread local transformer
TtlAgent.premain(agentArgs, inst); TtlAgent.premain(agentArgs, inst);
} }

View File

@ -33,7 +33,7 @@ public class ArgumentParser {
} else if (k.equals(DUMP_PATH)) { } else if (k.equals(DUMP_PATH)) {
GlobalConfig.setDumpPath(v); GlobalConfig.setDumpPath(v);
} else if (k.equals(PKG_PREFIX)) { } else if (k.equals(PKG_PREFIX)) {
GlobalConfig.setPkgPrefixes(v); GlobalConfig.setPkgPrefixWhiteList(v);
} else if (k.equals(MOCK_SCOPE)) { } else if (k.equals(MOCK_SCOPE)) {
GlobalConfig.setDefaultMockScope(MockScope.of(v)); GlobalConfig.setDefaultMockScope(MockScope.of(v));
} else if (k.equals(CONFIG_FILE)) { } else if (k.equals(CONFIG_FILE)) {
@ -42,7 +42,7 @@ public class ArgumentParser {
} else { } else {
// parameter with single value // parameter with single value
if (a.equals(USE_THREAD_POOL)) { 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_LEVEL = "log.level";
private static final String LOG_FILE = "log.file"; private static final String LOG_FILE = "log.file";
private static final String DUMP_PATH = "dump.path"; 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 INNER_MOCK_CLASS_NAME = "mock.innerClass.name";
private static final String DEFAULT_MOCK_SCOPE = "mock.scope.default"; 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_THREAD_POOL = "thread.pool.enhance.enable";
private static final String ENABLE_OMNI_INJECT = "omni.constructor.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) { public static void parseFile(String configFilePath) {
String path = (configFilePath == null) ? DEFAULT_CONFIG_FILE : configFilePath; String path = (configFilePath == null) ? DEFAULT_CONFIG_FILE : configFilePath;
@ -56,14 +57,16 @@ public class PropertiesParser {
GlobalConfig.setLogFile(v); GlobalConfig.setLogFile(v);
} else if (k.equals(DUMP_PATH)) { } else if (k.equals(DUMP_PATH)) {
GlobalConfig.setDumpPath(v); GlobalConfig.setDumpPath(v);
} else if (k.equals(PKG_PREFIX_WHITELIST)) { } else if (k.equals(PKG_PREFIX_INCLUDES)) {
GlobalConfig.setPkgPrefixes(v); GlobalConfig.setPkgPrefixWhiteList(v);
} else if (k.equals(OMNI_INJECT_EXCLUDES)) {
GlobalConfig.setOmniPkgPrefixBlackList(v);
} else if (k.equals(DEFAULT_MOCK_SCOPE)) { } else if (k.equals(DEFAULT_MOCK_SCOPE)) {
GlobalConfig.setDefaultMockScope(MockScope.of(v)); GlobalConfig.setDefaultMockScope(MockScope.of(v));
} else if (k.equals(ENABLE_THREAD_POOL)) { } else if (k.equals(ENABLE_THREAD_POOL)) {
GlobalConfig.setEnhanceThreadLocal(Boolean.parseBoolean(v)); GlobalConfig.enableEnhanceThreadLocal(Boolean.parseBoolean(v));
} else if (k.equals(ENABLE_OMNI_INJECT)) { } 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)) { } else if (k.equals(INNER_MOCK_CLASS_NAME)) {
GlobalConfig.setInnerMockClassName(v); GlobalConfig.setInnerMockClassName(v);
} }

View File

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

View File

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