Merge pull request #400 from Nizernizer/fix/parseCustomModel-exception

Fix/parse custom model exception

Author: lostsnow

dongtai-core/src/main/java/io/dongtai/iast/core/EngineManager.java

@@ -136,6 +136,7 @@ public static void cleanThreadState() {
         EngineManager.TAINT_HASH_CODES.remove();
         EngineManager.TAINT_RANGES_POOL.remove();
         EngineManager.SCOPE_TRACKER.remove();
+        EngineManager.ENTER_REPLAY_ENTRYPOINT.remove();
         FallbackSwitch.clearHeavyHookFallback();
         EngineManager.getFallbackManager().getHookRateLimiter().remove();
         ContextManager.getCONTEXT().remove();

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

@@ -160,51 +160,57 @@ public static void handlerCustomModel(MethodEvent event) {
      * @return Set<Object>
      */
     public static Set<Object> parseCustomModel(Object model) {
-        Set<Object> modelValues = new HashSet<Object>();
-        if (!TaintPoolUtils.isNotEmpty(model)) {
-            return modelValues;
-        }
-        Class<?> sourceClass = model.getClass();
-        if (sourceClass.getClassLoader() == null) {
-            return modelValues;
-        }
-        String className = sourceClass.getName();
-        if (className.startsWith("cn.huoxian.iast.api.") ||
-                className.startsWith("io.dongtai.api.") ||
-                className.startsWith(" org.apache.shiro.web.servlet".substring(1)) ||
-                VALUES_ENUMERATOR.equals(className) ||
-                className.startsWith(SPRING_OBJECT) ||
-                className.endsWith("RequestWrapper") ||
-                className.endsWith("ResponseWrapper")
-
-        ) {
-            return modelValues;
-        }
-        // getter methods
-        Method[] methods = sourceClass.getMethods();
-        Object itemValue = null;
-        for (Method method : methods) {
-            if (!TaintPoolUtils.isAllowTaintGetterMethod(method)) {
-                continue;
+        try{
+            Set<Object> modelValues = new HashSet<Object>();
+            if (!TaintPoolUtils.isNotEmpty(model)) {
+                return modelValues;
             }
-
-            try {
-                itemValue = method.invoke(model);
-                if (!TaintPoolUtils.isNotEmpty(itemValue)) {
+            Class<?> sourceClass = model.getClass();
+            if (sourceClass.getClassLoader() == null) {
+                return modelValues;
+            }
+            String className = sourceClass.getName();
+            if (className.startsWith("cn.huoxian.iast.api.") ||
+                    className.startsWith("io.dongtai.api.") ||
+                    className.startsWith(" org.apache.tomcat".substring(1)) ||
+                    className.startsWith(" org.apache.catalina".substring(1)) ||
+                    className.startsWith(" org.apache.shiro.web.servlet".substring(1)) ||
+                    VALUES_ENUMERATOR.equals(className) ||
+                    className.startsWith(SPRING_OBJECT) ||
+                    className.contains("RequestWrapper") ||
+                    className.contains("ResponseWrapper")
+
+            ) {
+                return modelValues;
+            }
+            // getter methods
+            Method[] methods = sourceClass.getMethods();
+            Object itemValue = null;
+            for (Method method : methods) {
+                if (!TaintPoolUtils.isAllowTaintGetterMethod(method)) {
                     continue;
                 }
-                modelValues.add(itemValue);
-                if (itemValue instanceof List) {
-                    List<?> itemValueList = (List<?>) itemValue;
-                    for (Object listValue : itemValueList) {
-                        modelValues.addAll(parseCustomModel(listValue));
+
+                try {
+                    itemValue = method.invoke(model);
+                    if (!TaintPoolUtils.isNotEmpty(itemValue)) {
+                        continue;
+                    }
+                    modelValues.add(itemValue);
+                    if (itemValue instanceof List) {
+                        List<?> itemValueList = (List<?>) itemValue;
+                        for (Object listValue : itemValueList) {
+                            modelValues.addAll(parseCustomModel(listValue));
+                        }
                     }
+                } catch (Exception e) {
+                    DongTaiLog.error(e);
                 }
-            } catch (Exception e) {
-                DongTaiLog.error(e);
             }
+            return modelValues;
+        }catch (Exception e){
+            return new HashSet<Object>();
         }
-        return modelValues;
     }
 
     private static boolean allowCall(MethodEvent event) {

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

@@ -27,7 +27,6 @@ public static void buildAndReport(Object request, Object response) {
         List<GraphNode> nodeList = build();
         String report = convertToReport(nodeList, request, response);
         ThreadPools.sendPriorityReport(ApiPath.REPORT_UPLOAD, report);
-        EngineManager.ENTER_REPLAY_ENTRYPOINT.remove();
     }
 
     /**
@@ -75,7 +74,7 @@ public static List<GraphNode> build() {
                                 event.sourceTypes
                         )
                 );
-            }catch (Exception e){
+            } catch (Exception e) {
                 DongTaiLog.debug(e);
             }
         }
@@ -99,8 +98,14 @@ public static String convertToReport(List<GraphNode> nodeList, Object request, O
         detail.put(ReportKey.METHOD, requestMeta.getOrDefault("method", ""));
         detail.put(ReportKey.SECURE, requestMeta.getOrDefault("secure", ""));
         String requestURL = requestMeta.getOrDefault("requestURL", "").toString();
+        if (null == requestURL) {
+            return null;
+        }
         detail.put(ReportKey.URL, requestURL);
         String requestURI = requestMeta.getOrDefault("requestURI", "").toString();
+        if (null == requestURI) {
+            return null;
+        }
         detail.put(ReportKey.URI, requestURI);
         setURL(requestURL);
         setURI(requestURI);
@@ -123,7 +128,6 @@ public static String convertToReport(List<GraphNode> nodeList, Object request, O
         for (GraphNode node : nodeList) {
             methodPool.put(node.toJson());
         }
-
         return report.toString();
     }
 

dongtai-core/src/main/java/io/dongtai/iast/core/utils/StackUtils.java

@@ -8,8 +8,9 @@
 public class StackUtils {
     public static StackTraceElement[] createCallStack(int stackStartPos) {
         StackTraceElement[] stackTraceElements = Thread.currentThread().getStackTrace();
-        StackTraceElement[] selfCallStack = new StackTraceElement[stackTraceElements.length - stackStartPos];
+        StackTraceElement[] selfCallStack = new StackTraceElement[0];
         if (stackTraceElements.length - stackStartPos >= 0) {
+            selfCallStack = new StackTraceElement[stackTraceElements.length - stackStartPos];
             System.arraycopy(stackTraceElements, stackStartPos, selfCallStack, 0, stackTraceElements.length - stackStartPos);
         }
         return selfCallStack;