add ignore internal and blacklist policy

HXSecurity · Mar 28, 2023 · db9f2ea · db9f2ea
1 parent 9ec09fa
commit db9f2ea
Show file tree

Hide file tree

Showing 15 changed files with 229 additions and 61 deletions.
diff --git a/dongtai-common/src/main/java/io/dongtai/iast/common/scope/PolicyScope.java b/dongtai-common/src/main/java/io/dongtai/iast/common/scope/PolicyScope.java
@@ -6,6 +6,7 @@ public class PolicyScope {
     private int propagatorLevel;
     private int propagatorSkipDepth;
     private int sinkLevel;
+    private int ignoreInternalLevel;
     /**
      * over max method pool size
      */
@@ -29,7 +30,8 @@ public void enterSource() {
     }
 
     public boolean isValidSource() {
-        return this.agentLevel == 0 && !this.overCapacity
+        return this.agentLevel == 0
+                && this.ignoreInternalLevel == 0 && !this.overCapacity
                 && this.sourceLevel == 1;
     }
 
@@ -45,7 +47,8 @@ public void enterPropagator(boolean skipScope) {
     }
 
     public boolean isValidPropagator() {
-        return this.agentLevel == 0 && !this.overCapacity && this.sourceLevel == 0
+        return this.agentLevel == 0
+                && this.ignoreInternalLevel == 0 && !this.overCapacity && this.sourceLevel == 0
                 && (this.propagatorLevel == 1 || this.propagatorSkipDepth > 0);
     }
 
@@ -61,14 +64,23 @@ public void enterSink() {
     }
 
     public boolean isValidSink() {
-        return this.agentLevel == 0 && !this.overCapacity && this.sourceLevel == 0
+        return this.agentLevel == 0
+                && this.ignoreInternalLevel == 0 && !this.overCapacity && this.sourceLevel == 0
                 && this.sinkLevel == 1;
     }
 
     public void leaveSink() {
         this.sinkLevel = decrement(this.sinkLevel);
     }
 
+    public void enterIgnoreInternal() {
+        this.ignoreInternalLevel++;
+    }
+
+    public void leaveIgnoreInternal() {
+        this.ignoreInternalLevel = decrement(this.ignoreInternalLevel);
+    }
+
     public boolean isOverCapacity() {
         return this.overCapacity;
     }

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/IastClassFileTransformer.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/IastClassFileTransformer.java
@@ -144,45 +144,47 @@ public byte[] transform(final ClassLoader loader,
                 }
             }
 
-            if (null != classBeingRedefined || configMatcher.canHook(internalClassName)) {
-                byte[] sourceCodeBak = new byte[srcByteCodeArray.length];
-                System.arraycopy(srcByteCodeArray, 0, sourceCodeBak, 0, srcByteCodeArray.length);
-                final ClassReader cr = new ClassReader(sourceCodeBak);
-
-                ClassContext classContext = new ClassContext(cr, loader);
-                if (Modifier.isInterface(classContext.getModifier())) {
-                    sourceCodeBak = null;
-                    return null;
-                }
-                final String className = classContext.getClassName();
-
-                Set<String> ancestors = classDiagram.getDiagram(className);
-                if (ancestors == null) {
-                    classDiagram.setLoader(loader);
-                    classDiagram.saveAncestors(className, classContext.getSuperClassName(), classContext.getInterfaces());
-                    ancestors = classDiagram.getAncestors(className, classContext.getSuperClassName(),
-                            classContext.getInterfaces());
-                }
-                classContext.setAncestors(ancestors);
-
-                final ClassWriter cw = createClassWriter(loader, cr);
-                ClassVisitor cv = plugins.initial(cw, classContext, policyManager);
-
-                if (cv instanceof AbstractClassVisitor) {
-                    cr.accept(cv, ClassReader.EXPAND_FRAMES);
-                    AbstractClassVisitor dumpClassVisitor = (AbstractClassVisitor) cv;
-                    if (dumpClassVisitor.hasTransformed()) {
-                        if (null == classBeingRedefined) {
-                            transformMap.put(className, srcByteCodeArray);
-                        } else {
-                            transformMap.put(classBeingRedefined, srcByteCodeArray);
-                        }
-                        transformCount++;
-                        return dumpClassIfNecessary(cr.getClassName(), cw.toByteArray(), srcByteCodeArray);
+            if (null == classBeingRedefined && !configMatcher.canHook(internalClassName, this.policyManager)) {
+                return null;
+            }
+
+            byte[] sourceCodeBak = new byte[srcByteCodeArray.length];
+            System.arraycopy(srcByteCodeArray, 0, sourceCodeBak, 0, srcByteCodeArray.length);
+            final ClassReader cr = new ClassReader(sourceCodeBak);
+
+            ClassContext classContext = new ClassContext(cr, loader);
+            if (Modifier.isInterface(classContext.getModifier())) {
+                sourceCodeBak = null;
+                return null;
+            }
+            final String className = classContext.getClassName();
+
+            Set<String> ancestors = classDiagram.getDiagram(className);
+            if (ancestors == null) {
+                classDiagram.setLoader(loader);
+                classDiagram.saveAncestors(className, classContext.getSuperClassName(), classContext.getInterfaces());
+                ancestors = classDiagram.getAncestors(className, classContext.getSuperClassName(),
+                        classContext.getInterfaces());
+            }
+            classContext.setAncestors(ancestors);
+
+            final ClassWriter cw = createClassWriter(loader, cr);
+            ClassVisitor cv = plugins.initial(cw, classContext, policyManager);
+
+            if (cv instanceof AbstractClassVisitor) {
+                cr.accept(cv, ClassReader.EXPAND_FRAMES);
+                AbstractClassVisitor dumpClassVisitor = (AbstractClassVisitor) cv;
+                if (dumpClassVisitor.hasTransformed()) {
+                    if (null == classBeingRedefined) {
+                        transformMap.put(className, srcByteCodeArray);
+                    } else {
+                        transformMap.put(classBeingRedefined, srcByteCodeArray);
                     }
+                    transformCount++;
+                    return dumpClassIfNecessary(cr.getClassName(), cw.toByteArray(), srcByteCodeArray);
                 }
-                sourceCodeBak = null;
             }
+            sourceCodeBak = null;
         } catch (Throwable throwable) {
             DongTaiLog.warn(ErrorCode.TRANSFORM_CLASS_FAILED, internalClassName, throwable);
         } finally {
@@ -271,7 +273,7 @@ public Class<?>[] findForRetransform() {
                 continue;
             }
             try {
-                if (!configMatcher.canHook(clazz)) {
+                if (!configMatcher.canHook(clazz, this.policyManager)) {
                     continue;
                 }
                 String className = clazz.getName();
@@ -295,7 +297,7 @@ public Class<?>[] findForRetransform() {
                     classDiagram.setDiagram(className, diagram);
                 }
                 for (String clazzName : diagram) {
-                    if (PolicyManager.isHookClass(clazzName) ||
+                    if (this.policyManager.isHookClass(clazzName) ||
                             (this.policyManager.getPolicy() != null && this.policyManager.getPolicy().isMatchClass(clazzName))) {
                         enhanceClasses[enhanceClassSize++] = clazz;
                         break;

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/asm/AsmMethods.java b/dongtai-core/src/main/java/io/dongtai/iast/core/bytecode/enhance/asm/AsmMethods.java
@@ -218,6 +218,14 @@ static Method getAsmMethod(final Class<?> clazz,
             SpyDispatcher.class,
             "isFirstLevelSink"
     );
+    Method SPY$enterIgnoreInternal = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "enterIgnoreInternal"
+    );
+    Method SPY$leaveIgnoreInternal = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "leaveIgnoreInternal"
+    );
     Method SPY$collectMethodPool = InnerHelper.getAsmMethod(
             SpyDispatcher.class,
             "collectMethodPool",

diff --git a/.../src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java b/.../src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/DispatchClassPlugin.java
@@ -60,6 +60,12 @@ public MethodVisitor visitMethod(final int access, final String name, final Stri
                 return mv;
             }
 
+            if (this.policy.isBlacklistHooks(this.context.getClassName())
+                    && !this.policy.isIgnoreBlacklistHooks(this.context.getClassName())
+                    && !this.policy.isIgnoreInternalHooks(this.context.getClassName())) {
+                return mv;
+            }
+
             MethodContext methodContext = new MethodContext(this.context, name);
             methodContext.setModifier(access);
             methodContext.setDescriptor(descriptor);

diff --git a/...ain/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/PropagatorAdapter.java b/...ain/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/PropagatorAdapter.java
@@ -18,7 +18,7 @@ public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodC
             }
 
             String signature = context.toString();
-            enterScope(adapter, signature);
+            enterScope(adapter, signature, policyNode);
         }
     }
 
@@ -43,11 +43,16 @@ public void onMethodExit(MethodAdviceAdapter adapter, MethodVisitor mv, int opco
             adapter.mark(elseLabel);
             adapter.mark(endLabel);
 
-            leaveScope(adapter, signature);
+            leaveScope(adapter, signature, policyNode);
         }
     }
 
-    private void enterScope(MethodAdviceAdapter adapter, String signature) {
+    private void enterScope(MethodAdviceAdapter adapter, String signature, PolicyNode policyNode) {
+        if (policyNode.isIgnoreInternal()) {
+            adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterIgnoreInternal);
+        }
+
         adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
         if (PropagatorImpl.isSkipScope(signature)) {
             adapter.push(true);
@@ -57,14 +62,19 @@ private void enterScope(MethodAdviceAdapter adapter, String signature) {
         adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterPropagator);
     }
 
-    private void leaveScope(MethodAdviceAdapter adapter, String signature) {
+    private void leaveScope(MethodAdviceAdapter adapter, String signature, PolicyNode policyNode) {
         adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
         if (PropagatorImpl.isSkipScope(signature)) {
             adapter.push(true);
         } else {
             adapter.push(false);
         }
         adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leavePropagator);
+
+        if (policyNode.isIgnoreInternal()) {
+            adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveIgnoreInternal);
+        }
     }
 
     private void isFirstScope(MethodAdviceAdapter adapter) {

diff --git a/.../src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/SinkAdapter.java b/.../src/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/SinkAdapter.java
@@ -16,7 +16,7 @@ public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodC
                 continue;
             }
 
-            enterScope(adapter);
+            enterScope(adapter, policyNode);
 
             Label elseLabel = new Label();
             Label endLabel = new Label();
@@ -39,18 +39,28 @@ public void onMethodExit(MethodAdviceAdapter adapter, MethodVisitor mv, int opco
                 continue;
             }
 
-            leaveScope(adapter);
+            leaveScope(adapter, policyNode);
         }
     }
 
-    private void enterScope(MethodAdviceAdapter adapter) {
+    private void enterScope(MethodAdviceAdapter adapter, PolicyNode policyNode) {
+        if (policyNode.isIgnoreInternal()) {
+            adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterIgnoreInternal);
+        }
+
         adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
         adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterSink);
     }
 
-    private void leaveScope(MethodAdviceAdapter adapter) {
+    private void leaveScope(MethodAdviceAdapter adapter, PolicyNode policyNode) {
         adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
         adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveSink);
+
+        if (policyNode.isIgnoreInternal()) {
+            adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveIgnoreInternal);
+        }
     }
 
     private void isFirstScope(MethodAdviceAdapter adapter) {

diff --git a/...rc/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/SourceAdapter.java b/...rc/main/java/io/dongtai/iast/core/bytecode/enhance/plugin/core/adapter/SourceAdapter.java
@@ -16,7 +16,7 @@ public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodC
                 continue;
             }
 
-            enterScope(adapter);
+            enterScope(adapter, policyNode);
         }
     }
 
@@ -39,18 +39,28 @@ public void onMethodExit(MethodAdviceAdapter adapter, MethodVisitor mv, int opco
             adapter.mark(elseLabel);
             adapter.mark(endLabel);
 
-            leaveScope(adapter);
+            leaveScope(adapter, policyNode);
         }
     }
 
-    private void enterScope(MethodAdviceAdapter adapter) {
+    private void enterScope(MethodAdviceAdapter adapter, PolicyNode policyNode) {
+        if (policyNode.isIgnoreInternal()) {
+            adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterIgnoreInternal);
+        }
+
         adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
         adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterSource);
     }
 
-    private void leaveScope(MethodAdviceAdapter adapter) {
+    private void leaveScope(MethodAdviceAdapter adapter, PolicyNode policyNode) {
         adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
         adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveSource);
+
+        if (policyNode.isIgnoreInternal()) {
+            adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+            adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$leaveIgnoreInternal);
+        }
     }
 
     private void isFirstScope(MethodAdviceAdapter adapter) {

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/SpyDispatcherImpl.java
@@ -546,6 +546,34 @@ public boolean isFirstLevelSink() {
         }
     }
 
+    @Override
+    public void enterIgnoreInternal() {
+        try {
+            if (!EngineManager.isEngineRunning()) {
+                return;
+            }
+            if (ScopeManager.SCOPE_TRACKER.inAgent() || !ScopeManager.SCOPE_TRACKER.inEnterEntry()) {
+                return;
+            }
+            ScopeManager.SCOPE_TRACKER.getPolicyScope().enterIgnoreInternal();
+        } catch (Throwable ignore) {
+        }
+    }
+
+    @Override
+    public void leaveIgnoreInternal() {
+        try {
+            if (!EngineManager.isEngineRunning()) {
+                return;
+            }
+            if (ScopeManager.SCOPE_TRACKER.inAgent() || !ScopeManager.SCOPE_TRACKER.inEnterEntry()) {
+                return;
+            }
+            ScopeManager.SCOPE_TRACKER.getPolicyScope().leaveIgnoreInternal();
+        } catch (Throwable ignore) {
+        }
+    }
+
     @Override
     public void reportService(String category, String type, String host, String port, String handler) {
         // @TODO: refactor

diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/Policy.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/policy/Policy.java
@@ -10,6 +10,10 @@ public class Policy {
     private final Set<String> classHooks = new HashSet<String>();
     private final Set<String> ancestorClassHooks = new HashSet<String>();
 
+    private final Set<String> blacklistHooks = new HashSet<String>();
+    private final Set<String> ignoreInternalHooks = new HashSet<String>();
+    private final Set<String> ignoreBlacklistHooks = new HashSet<String>();
+
     public List<SourceNode> getSources() {
         return sources;
     }
@@ -51,6 +55,12 @@ public void addPolicyNode(PolicyNode node) {
             methodMatcher = (SignatureMethodMatcher) node.getMethodMatcher();
             this.policyNodesMap.put(node.toString(), node);
             addHooks(methodMatcher.getSignature().getClassName(), node.getInheritable());
+            if (node.isIgnoreInternal()) {
+                this.ignoreInternalHooks.add(methodMatcher.getSignature().getClassName());
+            }
+            if (node.isIgnoreBlacklist()) {
+                this.ignoreBlacklistHooks.add(methodMatcher.getSignature().getClassName());
+            }
         }
     }
 
@@ -86,4 +96,20 @@ public Set<String> getClassHooks() {
     public Set<String> getAncestorClassHooks() {
         return this.ancestorClassHooks;
     }
+
+    public void addBlacklistHooks(String className) {
+        this.blacklistHooks.add(className);
+    }
+
+    public boolean isBlacklistHooks(String className) {
+        return this.blacklistHooks.contains(className);
+    }
+
+    public boolean isIgnoreInternalHooks(String className) {
+        return this.ignoreInternalHooks.contains(className);
+    }
+
+    public boolean isIgnoreBlacklistHooks(String className) {
+        return this.ignoreBlacklistHooks.contains(className);
+    }
 }