Merge pull request #565 from Nizernizer/feature/add-validator-rule

feature: add validator handler.

Author: lostsnow

dongtai-common/src/main/java/io/dongtai/iast/common/config/ConfigBuilder.java

@@ -24,6 +24,8 @@ private ConfigBuilder() {
                 Config.<Boolean>create(ConfigKey.ENABLE_LOGGER));
         this.configMap.put(ConfigKey.LOGGER_LEVEL,
                 Config.<String>create(ConfigKey.LOGGER_LEVEL));
+        this.configMap.put(ConfigKey.VALIDATED_SINK,
+                Config.<Boolean>create(ConfigKey.VALIDATED_SINK).setDefaultValue(false));
     }
 
     public static ConfigBuilder getInstance() {
@@ -62,6 +64,7 @@ public void update(JSONObject config) {
         updateString(config, ConfigKey.JsonKey.JSON_VERSION_HEADER_KEY);
         updateBool(config, ConfigKey.JsonKey.JSON_ENABLE_LOGGER);
         updateString(config, ConfigKey.JsonKey.JSON_LOGGER_LEVEL);
+        updateBool(config, ConfigKey.JsonKey.JSON_VALIDATED_SINK);
         updateRequestDenyList(config);
     }
 

dongtai-common/src/main/java/io/dongtai/iast/common/config/ConfigKey.java

@@ -8,6 +8,7 @@ public enum ConfigKey {
     VERSION_HEADER_KEY,
     ENABLE_LOGGER,
     LOGGER_LEVEL,
+    VALIDATED_SINK,
     ;
 
     public enum JsonKey {
@@ -18,6 +19,7 @@ public enum JsonKey {
         JSON_VERSION_HEADER_KEY("version_header_name", VERSION_HEADER_KEY),
         JSON_ENABLE_LOGGER("enable_log", ENABLE_LOGGER),
         JSON_LOGGER_LEVEL("log_level", LOGGER_LEVEL),
+        JSON_VALIDATED_SINK("report_validated_sink", VALIDATED_SINK),
         ;
 
         private final String key;

dongtai-common/src/main/java/io/dongtai/iast/common/constants/PropertyConstant.java

@@ -31,5 +31,5 @@ public class PropertyConstant {
     public static final String PROPERTY_POLICY_PATH = "dongtai.policy.path";
     public static final String PROPERTY_UUID_PATH = "dongtai.uuid.path";
     public static final String PROPERTY_DISABLED_PLUGINS = "dongtai.disabled.plugins";
-    public static final String PROPERTY_DISABLED_FEATURES = "dongtai.disabled_features";
+    public static final String PROPERTY_DISABLED_FEATURES = "dongtai.disabled.features";
 }

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

@@ -219,6 +219,12 @@ static Method getAsmMethod(final Class<?> clazz,
             SpyDispatcher.class,
             "isFirstLevelSink"
     );
+
+    Method SPY$enterValidator = InnerHelper.getAsmMethod(
+            SpyDispatcher.class,
+            "enterValidator"
+    );
+
     Method SPY$enterIgnoreInternal = InnerHelper.getAsmMethod(
             SpyDispatcher.class,
             "enterIgnoreInternal"

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

@@ -54,6 +54,7 @@ public class ClassVisit extends AbstractClassVisitor {
                     new SourceAdapter(),
                     new PropagatorAdapter(),
                     new SinkAdapter(),
+                    new ValidatorAdapter(),
             };
         }
 

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

@@ -0,0 +1,54 @@
+package io.dongtai.iast.core.bytecode.enhance.plugin.core.adapter;
+
+import io.dongtai.iast.core.bytecode.enhance.MethodContext;
+import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNode;
+import io.dongtai.iast.core.handler.hookpoint.models.policy.ValidatorNode;
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.Opcodes;
+
+import java.util.Set;
+
+public class ValidatorAdapter extends MethodAdapter {
+    /**
+     * @param adapter
+     * @param mv
+     * @param context
+     * @param policyNodes
+     */
+    @Override
+    public void onMethodEnter(MethodAdviceAdapter adapter, MethodVisitor mv, MethodContext context, Set<PolicyNode> policyNodes) {
+    }
+
+    /**
+     * @param adapter
+     * @param mv
+     * @param opcode
+     * @param context
+     * @param policyNodes
+     */
+    @Override
+    public void onMethodExit(MethodAdviceAdapter adapter, MethodVisitor mv, int opcode, MethodContext context, Set<PolicyNode> policyNodes) {
+        for (PolicyNode policyNode : policyNodes) {
+            if (!(policyNode instanceof ValidatorNode)) {
+                continue;
+            }
+
+            Label elseLabel = new Label();
+            Label endLabel = new Label();
+
+            isEnterScope(adapter);
+            mv.visitJumpInsn(Opcodes.IFEQ, elseLabel);
+
+            adapter.trackMethod(opcode, policyNode, true);
+
+            adapter.mark(elseLabel);
+            adapter.mark(endLabel);
+        }
+    }
+
+    private void isEnterScope(MethodAdviceAdapter adapter) {
+        adapter.invokeStatic(ASM_TYPE_SPY_HANDLER, SPY_HANDLER$getDispatcher);
+        adapter.invokeInterface(ASM_TYPE_SPY_DISPATCHER, SPY$enterValidator);
+    }
+}

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

@@ -394,7 +394,8 @@ public void collectDubboResponse(Object result, byte status) {
             }
 
             if (!ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_REQUEST).isFirst()
-                    || !ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_ENTRY).in()) {
+                    || !ScopeManager.SCOPE_TRACKER.getScope(Scope.DUBBO_ENTRY).in()
+                    || ScopeManager.SCOPE_TRACKER.getScope(Scope.HTTP_REQUEST).in()) {
                 return;
             }
 
@@ -558,6 +559,17 @@ public void leaveSink() {
         }
     }
 
+    /**
+     * mark for enter validator entry point
+     */
+    @Override
+    public boolean enterValidator() {
+        if (!EngineManager.isEngineRunning()) {
+            return false;
+        }
+        return !ScopeManager.SCOPE_TRACKER.inAgent() && ScopeManager.SCOPE_TRACKER.inEnterEntry();
+    }
+
     /**
      * Determines whether it is a layer 1 Sink entry
      *
@@ -674,6 +686,9 @@ public boolean collectMethod(Object instance, Object[] parameters, Object retObj
             } else if ((policyNode instanceof SinkNode)) {
                 SinkImpl.solveSink(event, (SinkNode) policyNode);
                 return true;
+            } else if ((policyNode instanceof ValidatorNode)) {
+                ValidatorImpl.solveValidator(event,(ValidatorNode)policyNode, INVOKE_ID_SEQUENCER);
+                return true;
             }
 
             return false;
@@ -731,7 +746,7 @@ public boolean traceDubboInvoke(Object instance, String url, Object invocation,
     @Override
     public boolean isSkipCollectDubbo(Object invocation) {
         if (BlackUrlBypass.isBlackUrl()) {
-            Method setAttachmentMethod = null;
+            Method setAttachmentMethod;
             try {
                 setAttachmentMethod = invocation.getClass().getMethod("setAttachment", String.class, String.class);
                 setAttachmentMethod.setAccessible(true);
@@ -746,7 +761,7 @@ public boolean isSkipCollectDubbo(Object invocation) {
     @Override
     public boolean isSkipCollectFeign(Object instance) {
         if (BlackUrlBypass.isBlackUrl()) {
-            Field metadataField = null;
+            Field metadataField;
             try {
                 metadataField = instance.getClass().getDeclaredField("metadata");
                 metadataField.setAccessible(true);

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/DubboImpl.java

@@ -8,6 +8,7 @@
 import io.dongtai.iast.core.handler.context.ContextManager;
 import io.dongtai.iast.core.handler.hookpoint.IastClassLoader;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
+import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNodeType;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.SourceNode;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
 import io.dongtai.iast.core.handler.hookpoint.models.taint.range.TaintRange;
@@ -178,6 +179,7 @@ public static void collectDubboRequestSource(Object handler, Object invocation,
 
         int invokeId = invokeIdSequencer.getAndIncrement();
         event.setInvokeId(invokeId);
+        event.setPolicyType(PolicyNodeType.SOURCE.getName());
 
         event.source = true;
         event.setCallStacks(StackUtils.createCallStack(4));

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/PropagatorImpl.java

@@ -2,6 +2,7 @@
 
 import io.dongtai.iast.core.EngineManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
+import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNodeType;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.PropagatorNode;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
 import io.dongtai.iast.core.handler.hookpoint.models.taint.range.*;
@@ -63,6 +64,7 @@ private static void addPropagator(PropagatorNode propagatorNode, MethodEvent eve
         event.setCallStacks(StackUtils.createCallStack(6));
         int invokeId = invokeIdSequencer.getAndIncrement();
         event.setInvokeId(invokeId);
+        event.setPolicyType(PolicyNodeType.PROPAGATOR.getName());
         EngineManager.TRACK_MAP.get().put(invokeId, event);
     }
 

dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/controller/impl/SourceImpl.java

@@ -2,6 +2,7 @@
 
 import io.dongtai.iast.core.EngineManager;
 import io.dongtai.iast.core.handler.hookpoint.models.MethodEvent;
+import io.dongtai.iast.core.handler.hookpoint.models.policy.PolicyNodeType;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.SourceNode;
 import io.dongtai.iast.core.handler.hookpoint.models.policy.TaintPosition;
 import io.dongtai.iast.core.handler.hookpoint.models.taint.range.TaintRangesBuilder;
@@ -37,6 +38,7 @@ public static void solveSource(MethodEvent event, SourceNode sourceNode, AtomicI
 
         int invokeId = invokeIdSequencer.getAndIncrement();
         event.setInvokeId(invokeId);
+        event.setPolicyType(PolicyNodeType.SOURCE.getName());
 
         boolean valid = trackTarget(event, sourceNode);
         if (!valid) {