Skip to content

Commit 60c2bd7

Browse files
lostsnowlostsnow
committed
add ignore internal and blacklist policy
1 parent 1c0a51a commit 60c2bd7

File tree

15 files changed

+229
-61
lines changed

15 files changed

+229
-61
lines changed

‎dongtai-common/src/main/java/io/dongtai/iast/common/scope/PolicyScope.java‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ public class PolicyScope {
66
private int propagatorLevel;
77
private int propagatorSkipDepth;
88
private int sinkLevel;
9+
private int ignoreInternalLevel;
910
/**
1011
* over max method pool size
1112
*/
@@ -29,7 +30,8 @@ public void enterSource() {
2930
}
3031

3132
public boolean isValidSource() {
32-
return this.agentLevel == 0 && !this.overCapacity
33+
return this.agentLevel == 0
34+
&& this.ignoreInternalLevel == 0 && !this.overCapacity
3335
&& this.sourceLevel == 1;
3436
}
3537

@@ -45,7 +47,8 @@ public void enterPropagator(boolean skipScope) {
4547
}
4648

4749
public boolean isValidPropagator() {
48-
return this.agentLevel == 0 && !this.overCapacity && this.sourceLevel == 0
50+
return this.agentLevel == 0
51+
&& this.ignoreInternalLevel == 0 && !this.overCapacity && this.sourceLevel == 0
4952
&& (this.propagatorLevel == 1 || this.propagatorSkipDepth > 0);
5053
}
5154

@@ -61,14 +64,23 @@ public void enterSink() {
6164
}
6265

6366
public boolean isValidSink() {
64-
return this.agentLevel == 0 && !this.overCapacity && this.sourceLevel == 0
67+
return this.agentLevel == 0
68+
&& this.ignoreInternalLevel == 0 && !this.overCapacity && this.sourceLevel == 0
6569
&& this.sinkLevel == 1;
6670
}
6771

6872
public void leaveSink() {
6973
this.sinkLevel = decrement(this.sinkLevel);
7074
}
7175

76+
public void enterIgnoreInternal() {
77+
this.ignoreInternalLevel++;
78+
}
79+
80+
public void leaveIgnoreInternal() {
81+
this.ignoreInternalLevel = decrement(this.ignoreInternalLevel);
82+
}
83+
7284
public boolean isOverCapacity() {
7385
return this.overCapacity;
7486
}

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/IastClassFileTransformer.java‎

Lines changed: 40 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -144,45 +144,47 @@ public byte[] transform(final ClassLoader loader,
144144
}
145145
}
146146

147-
if (null != classBeingRedefined || configMatcher.canHook(internalClassName)) {
148-
byte[] sourceCodeBak = new byte[srcByteCodeArray.length];
149-
System.arraycopy(srcByteCodeArray, 0, sourceCodeBak, 0, srcByteCodeArray.length);
150-
final ClassReader cr = new ClassReader(sourceCodeBak);
151-
152-
ClassContext classContext = new ClassContext(cr, loader);
153-
if (Modifier.isInterface(classContext.getModifier())) {
154-
sourceCodeBak = null;
155-
return null;
156-
}
157-
final String className = classContext.getClassName();
158-
159-
Set<String> ancestors = classDiagram.getDiagram(className);
160-
if (ancestors == null) {
161-
classDiagram.setLoader(loader);
162-
classDiagram.saveAncestors(className, classContext.getSuperClassName(), classContext.getInterfaces());
163-
ancestors = classDiagram.getAncestors(className, classContext.getSuperClassName(),
164-
classContext.getInterfaces());
165-
}
166-
classContext.setAncestors(ancestors);
167-
168-
final ClassWriter cw = createClassWriter(loader, cr);
169-
ClassVisitor cv = plugins.initial(cw, classContext, policyManager);
170-
171-
if (cv instanceof AbstractClassVisitor) {
172-
cr.accept(cv, ClassReader.EXPAND_FRAMES);
173-
AbstractClassVisitor dumpClassVisitor = (AbstractClassVisitor) cv;
174-
if (dumpClassVisitor.hasTransformed()) {
175-
if (null == classBeingRedefined) {
176-
transformMap.put(className, srcByteCodeArray);
177-
} else {
178-
transformMap.put(classBeingRedefined, srcByteCodeArray);
179-
}
180-
transformCount++;
181-
return dumpClassIfNecessary(cr.getClassName(), cw.toByteArray(), srcByteCodeArray);
147+
if (null == classBeingRedefined && !configMatcher.canHook(internalClassName, this.policyManager)) {
148+
return null;
149+
}
150+
151+
byte[] sourceCodeBak = new byte[srcByteCodeArray.length];
152+
System.arraycopy(srcByteCodeArray, 0, sourceCodeBak, 0, srcByteCodeArray.length);
153+
final ClassReader cr = new ClassReader(sourceCodeBak);
154+
155+
ClassContext classContext = new ClassContext(cr, loader);
156+
if (Modifier.isInterface(classContext.getModifier())) {
157+
sourceCodeBak = null;
158+
return null;
159+
}
160+
final String className = classContext.getClassName();
161+
162+
Set<String> ancestors = classDiagram.getDiagram(className);
163+
if (ancestors == null) {
164+
classDiagram.setLoader(loader);
165+
classDiagram.saveAncestors(className, classContext.getSuperClassName(), classContext.getInterfaces());
166+
ancestors = classDiagram.getAncestors(className, classContext.getSuperClassName(),
167+
classContext.getInterfaces());
168+
}
169+
classContext.setAncestors(ancestors);
170+
171+
final ClassWriter cw = createClassWriter(loader, cr);
172+
ClassVisitor cv = plugins.initial(cw, classContext, policyManager);
173+
174+
if (cv instanceof AbstractClassVisitor) {
175+
cr.accept(cv, ClassReader.EXPAND_FRAMES);
176+
AbstractClassVisitor dumpClassVisitor = (AbstractClassVisitor) cv;
177+
if (dumpClassVisitor.hasTransformed()) {
178+
if (null == classBeingRedefined) {
179+
transformMap.put(className, srcByteCodeArray);
180+
} else {
181+
transformMap.put(classBeingRedefined, srcByteCodeArray);
182182
}
183+
transformCount++;
184+
return dumpClassIfNecessary(cr.getClassName(), cw.toByteArray(), srcByteCodeArray);
183185
}
184-
sourceCodeBak = null;
185186
}
187+
sourceCodeBak = null;
186188
} catch (Throwable throwable) {
187189
DongTaiLog.warn(ErrorCode.TRANSFORM_CLASS_FAILED, internalClassName, throwable);
188190
} finally {
@@ -271,7 +273,7 @@ public Class<?>[] findForRetransform() {
271273
continue;
272274
}
273275
try {
274-
if (!configMatcher.canHook(clazz)) {
276+
if (!configMatcher.canHook(clazz, this.policyManager)) {
275277
continue;
276278
}
277279
String className = clazz.getName();
@@ -295,7 +297,7 @@ public Class<?>[] findForRetransform() {
295297
classDiagram.setDiagram(className, diagram);
296298
}
297299
for (String clazzName : diagram) {
298-
if (PolicyManager.isHookClass(clazzName) ||
300+
if (this.policyManager.isHookClass(clazzName) ||
299301
(this.policyManager.getPolicy() != null && this.policyManager.getPolicy().isMatchClass(clazzName))) {
300302
enhanceClasses[enhanceClassSize++] = clazz;
301303
break;

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/asm/AsmMethods.java‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,14 @@ static Method getAsmMethod(final Class<?> clazz,
218218
SpyDispatcher.class,
219219
"isFirstLevelSink"
220220
);
221+
Method SPY$enterIgnoreInternal = InnerHelper.getAsmMethod(
222+
SpyDispatcher.class,
223+
"enterIgnoreInternal"
224+
);
225+
Method SPY$leaveIgnoreInternal = InnerHelper.getAsmMethod(
226+
SpyDispatcher.class,
227+
"leaveIgnoreInternal"
228+
);
221229
Method SPY$collectMethodPool = InnerHelper.getAsmMethod(
222230
SpyDispatcher.class,
223231
"collectMethodPool",

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,12 @@ public MethodVisitor visitMethod(final int access, final String name, final Stri
6060
return mv;
6161
}
6262

63+
if (this.policy.isBlacklistHooks(this.context.getClassName())
64+
&& !this.policy.isIgnoreBlacklistHooks(this.context.getClassName())
65+
&& !this.policy.isIgnoreInternalHooks(this.context.getClassName())) {
66+
return mv;
67+
}
68+
6369
MethodContext methodContext = new MethodContext(this.context, name);
6470
methodContext.setModifier(access);
6571
methodContext.setDescriptor(descriptor);

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/PropagatorAdapter.java‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodC
1818
}
1919

2020
String signature = context.toString();
21-
enterScope(adapter, signature);
21+
enterScope(adapter, signature, policyNode);
2222
}
2323
}
2424

@@ -43,11 +43,16 @@ public void onMethodExit(MethodAdviceAdapter adapter, MethodVisitor mv, int opco
4343
adapter.mark(elseLabel);
4444
adapter.mark(endLabel);
4545

46-
leaveScope(adapter, signature);
46+
leaveScope(adapter, signature, policyNode);
4747
}
4848
}
4949

50-
private void enterScope(MethodAdviceAdapter adapter, String signature) {
50+
private void enterScope(MethodAdviceAdapter adapter, String signature, PolicyNode policyNode) {
51+
if (policyNode.isIgnoreInternal()) {
52+
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
53+
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterIgnoreInternal);
54+
}
55+
5156
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
5257
if (PropagatorImpl.isSkipScope(signature)) {
5358
adapter.push(true);
@@ -57,14 +62,19 @@ private void enterScope(MethodAdviceAdapter adapter, String signature) {
5762
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterPropagator);
5863
}
5964

60-
private void leaveScope(MethodAdviceAdapter adapter, String signature) {
65+
private void leaveScope(MethodAdviceAdapter adapter, String signature, PolicyNode policyNode) {
6166
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
6267
if (PropagatorImpl.isSkipScope(signature)) {
6368
adapter.push(true);
6469
} else {
6570
adapter.push(false);
6671
}
6772
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leavePropagator);
73+
74+
if (policyNode.isIgnoreInternal()) {
75+
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
76+
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveIgnoreInternal);
77+
}
6878
}
6979

7080
private void isFirstScope(MethodAdviceAdapter adapter) {

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/SinkAdapter.java‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodC
1616
continue;
1717
}
1818

19-
enterScope(adapter);
19+
enterScope(adapter, policyNode);
2020

2121
Label elseLabel = new Label();
2222
Label endLabel = new Label();
@@ -39,18 +39,28 @@ public void onMethodExit(MethodAdviceAdapter adapter, MethodVisitor mv, int opco
3939
continue;
4040
}
4141

42-
leaveScope(adapter);
42+
leaveScope(adapter, policyNode);
4343
}
4444
}
4545

46-
private void enterScope(MethodAdviceAdapter adapter) {
46+
private void enterScope(MethodAdviceAdapter adapter, PolicyNode policyNode) {
47+
if (policyNode.isIgnoreInternal()) {
48+
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
49+
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterIgnoreInternal);
50+
}
51+
4752
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
4853
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterSink);
4954
}
5055

51-
private void leaveScope(MethodAdviceAdapter adapter) {
56+
private void leaveScope(MethodAdviceAdapter adapter, PolicyNode policyNode) {
5257
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
5358
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveSink);
59+
60+
if (policyNode.isIgnoreInternal()) {
61+
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
62+
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveIgnoreInternal);
63+
}
5464
}
5565

5666
private void isFirstScope(MethodAdviceAdapter adapter) {

‎dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/SourceAdapter.java‎

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodC
1616
continue;
1717
}
1818

19-
enterScope(adapter);
19+
enterScope(adapter, policyNode);
2020
}
2121
}
2222

@@ -39,18 +39,28 @@ public void onMethodExit(MethodAdviceAdapter adapter, MethodVisitor mv, int opco
3939
adapter.mark(elseLabel);
4040
adapter.mark(endLabel);
4141

42-
leaveScope(adapter);
42+
leaveScope(adapter, policyNode);
4343
}
4444
}
4545

46-
private void enterScope(MethodAdviceAdapter adapter) {
46+
private void enterScope(MethodAdviceAdapter adapter, PolicyNode policyNode) {
47+
if (policyNode.isIgnoreInternal()) {
48+
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
49+
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterIgnoreInternal);
50+
}
51+
4752
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
4853
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterSource);
4954
}
5055

51-
private void leaveScope(MethodAdviceAdapter adapter) {
56+
private void leaveScope(MethodAdviceAdapter adapter, PolicyNode policyNode) {
5257
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
5358
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveSource);
59+
60+
if (policyNode.isIgnoreInternal()) {
61+
adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
62+
adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveIgnoreInternal);
63+
}
5464
}
5565

5666
private void isFirstScope(MethodAdviceAdapter adapter) {

‎dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -546,6 +546,34 @@ public boolean isFirstLevelSink() {
546546
}
547547
}
548548

549+
@Override
550+
public void enterIgnoreInternal() {
551+
try {
552+
if (!EngineManager.isEngineRunning()) {
553+
return;
554+
}
555+
if (ScopeManager.SCOPE_TRACKER.inAgent() || !ScopeManager.SCOPE_TRACKER.inEnterEntry()) {
556+
return;
557+
}
558+
ScopeManager.SCOPE_TRACKER.getPolicyScope().enterIgnoreInternal();
559+
} catch (Throwable ignore) {
560+
}
561+
}
562+
563+
@Override
564+
public void leaveIgnoreInternal() {
565+
try {
566+
if (!EngineManager.isEngineRunning()) {
567+
return;
568+
}
569+
if (ScopeManager.SCOPE_TRACKER.inAgent() || !ScopeManager.SCOPE_TRACKER.inEnterEntry()) {
570+
return;
571+
}
572+
ScopeManager.SCOPE_TRACKER.getPolicyScope().leaveIgnoreInternal();
573+
} catch (Throwable ignore) {
574+
}
575+
}
576+
549577
@Override
550578
public void reportService(String category, String type, String host, String port, String handler) {
551579
// @TODO: refactor

‎dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/Policy.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@ public class Policy {
1010
private final Set<String> classHooks = new HashSet<String>();
1111
private final Set<String> ancestorClassHooks = new HashSet<String>();
1212

13+
private final Set<String> blacklistHooks = new HashSet<String>();
14+
private final Set<String> ignoreInternalHooks = new HashSet<String>();
15+
private final Set<String> ignoreBlacklistHooks = new HashSet<String>();
16+
1317
public List<SourceNode> getSources() {
1418
return sources;
1519
}
@@ -51,6 +55,12 @@ public void addPolicyNode(PolicyNode node) {
5155
methodMatcher = (SignatureMethodMatcher) node.getMethodMatcher();
5256
this.policyNodesMap.put(node.toString(), node);
5357
addHooks(methodMatcher.getSignature().getClassName(), node.getInheritable());
58+
if (node.isIgnoreInternal()) {
59+
this.ignoreInternalHooks.add(methodMatcher.getSignature().getClassName());
60+
}
61+
if (node.isIgnoreBlacklist()) {
62+
this.ignoreBlacklistHooks.add(methodMatcher.getSignature().getClassName());
63+
}
5464
}
5565
}
5666

@@ -86,4 +96,20 @@ public Set<String> getClassHooks() {
8696
public Set<String> getAncestorClassHooks() {
8797
return this.ancestorClassHooks;
8898
}
99+
100+
public void addBlacklistHooks(String className) {
101+
this.blacklistHooks.add(className);
102+
}
103+
104+
public boolean isBlacklistHooks(String className) {
105+
return this.blacklistHooks.contains(className);
106+
}
107+
108+
public boolean isIgnoreInternalHooks(String className) {
109+
return this.ignoreInternalHooks.contains(className);
110+
}
111+
112+
public boolean isIgnoreBlacklistHooks(String className) {
113+
return this.ignoreBlacklistHooks.contains(className);
114+
}
89115
}

0 commit comments

Comments
 (0)