fix bugs

wh1t3p1g · Jan 12, 2023 · b38fbcd · b38fbcd
1 parent c66566d
commit b38fbcd
Show file tree

Hide file tree

Showing 13 changed files with 250 additions and 167 deletions.
diff --git a/config/settings.properties b/config/settings.properties
@@ -20,9 +20,9 @@ tabby.build.thread.timeout                = 2
 tabby.build.isNeedToCreateIgnoreList      = false
 
 # targets to analyse
-tabby.build.target                        = cases/emcsqlproxy.war
+tabby.build.target                        = cases/commons-collections-3.2.1.jar
 tabby.build.libraries                     = libs
-tabby.build.mode                          = web
+tabby.build.mode                          = gadget
 
 # db settings
 tabby.cache.path                          = ./cache/dev

diff --git a/src/main/java/tabby/config/GlobalConfiguration.java b/src/main/java/tabby/config/GlobalConfiguration.java
@@ -67,6 +67,7 @@ public class GlobalConfiguration {
     public static boolean IS_NEED_TO_CREATE_IGNORE_LIST = true;
 
     public static boolean isInitialed = false;
+    public static boolean isNeedStop = false;
 
     static {
         if(!FileUtils.fileExists(RULES_PATH)){

diff --git a/src/main/java/tabby/core/Analyser.java b/src/main/java/tabby/core/Analyser.java
@@ -12,7 +12,6 @@
 import tabby.core.container.RulesContainer;
 import tabby.core.scanner.CallGraphScanner;
 import tabby.core.scanner.ClassInfoScanner;
-import tabby.core.scanner.FullCallGraphScanner;
 import tabby.util.FileUtils;
 
 import java.io.File;
@@ -39,8 +38,6 @@ public class Analyser {
     @Autowired
     private CallGraphScanner callGraphScanner;
 
-    @Autowired
-    private FullCallGraphScanner fullCallGraphScanner;
     @Autowired
     private RulesContainer rulesContainer;
     @Autowired
@@ -71,12 +68,7 @@ public void run() throws IOException {
             // 收集目标
             GlobalConfiguration.rulesContainer = rulesContainer;
             if(!GlobalConfiguration.IS_JDK_ONLY){
-//                Map<String, String> files = FileUtils.getTargetDirectoryJarFiles(GlobalConfiguration.TARGET);
-                long start = System.nanoTime();
                 Map<String, String> files = fileCollector.collect(GlobalConfiguration.TARGET);
-                long time = TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - start);
-                log.info("cost {} min {} seconds."
-                        , time/60, time%60);
                 cps.putAll(files);
                 targets.putAll(files);
             }
@@ -137,11 +129,8 @@ public void runSootAnalysis(Map<String, String> targets, List<String> classpaths
             // 类信息抽取
             classInfoScanner.run(realTargets);
             // 全量函数调用图构建
-            if(GlobalConfiguration.IS_FULL_CALL_GRAPH_CONSTRUCT){
-                fullCallGraphScanner.run();
-            }else{
-                callGraphScanner.run();
-            }
+            callGraphScanner.run();
+
             rulesContainer.saveStatus();
             long time = TimeUnit.NANOSECONDS.toSeconds(System.nanoTime() - start);
             log.info("Total cost {} min {} seconds."

diff --git a/src/main/java/tabby/core/collector/CallEdgeCollector.java b/src/main/java/tabby/core/collector/CallEdgeCollector.java
@@ -0,0 +1,88 @@
+package tabby.core.collector;
+
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.scheduling.annotation.Async;
+import org.springframework.stereotype.Service;
+import soot.Modifier;
+import soot.SootMethod;
+import soot.Unit;
+import soot.jimple.InvokeExpr;
+import soot.jimple.JimpleBody;
+import soot.jimple.Stmt;
+import tabby.core.container.DataContainer;
+import tabby.core.model.DefaultInvokeModel;
+import tabby.core.switcher.Switcher;
+import tabby.dal.caching.bean.ref.MethodReference;
+import tabby.util.TickTock;
+
+/**
+ * @author wh1t3p1g
+ * @since 2023/1/12
+ */
+@Slf4j
+@Service
+public class CallEdgeCollector {
+
+    @Async("tabby-collector")
+    public void collect(MethodReference methodRef, DataContainer dataContainer, TickTock tickTock){
+        try{
+            SootMethod method = methodRef.getMethod();
+            if(method == null) {
+                tickTock.countDown();
+                return; // 提取不出内容，不分析
+            }
+
+            if(methodRef.isIgnore() || methodRef.isSink()){
+                tickTock.countDown();
+                return; // 消除后续的调用边
+            }
+
+            if(method.isStatic() && method.getParameterCount() == 0){
+                // 静态函数 且 函数入参数量为0 此类函数不影响分析
+                methodRef.setInitialed(true);
+                tickTock.countDown();
+                return;
+            }
+
+            if(method.isAbstract()
+                    || Modifier.isNative(method.getModifiers())
+                    || method.isPhantom()){
+                methodRef.setInitialed(true);
+                methodRef.setActionInitialed(true);
+                tickTock.countDown();
+                return;
+            }
+
+            JimpleBody body = (JimpleBody) Switcher.retrieveBody(method, method.getSignature());
+            if(body == null) {
+                tickTock.countDown();
+                return;
+            }
+
+            DefaultInvokeModel model = new DefaultInvokeModel();
+            for(Unit unit:body.getUnits()){
+                Stmt stmt = (Stmt) unit;
+                if(stmt.containsInvokeExpr()){
+                    InvokeExpr ie = stmt.getInvokeExpr();
+                    SootMethod targetMethod = ie.getMethod();
+                    MethodReference targetMethodRef
+                            = dataContainer.getOrAddMethodRef(ie.getMethodRef(), targetMethod);
+                    model.apply(stmt, false, methodRef, targetMethodRef, dataContainer);
+                }
+            }
+        }catch (RuntimeException e){
+//            log.error(e.getMessage());
+            log.error("Something error on call graph. "+methodRef.getSignature());
+            log.error(e.getMessage());
+//            e.printStackTrace();
+        }catch (Exception e){
+            if(e instanceof InterruptedException) {
+                log.error("Thread interrupted. " + methodRef.getSignature());
+            } else {
+                log.error("Something error on call graph. "+methodRef.getSignature());
+                e.printStackTrace();
+            }
+        }
+        tickTock.countDown();
+    }
+}
diff --git a/src/main/java/tabby/core/collector/CallGraphCollector.java b/src/main/java/tabby/core/collector/CallGraphCollector.java
@@ -10,6 +10,7 @@
 import tabby.core.switcher.Switcher;
 import tabby.core.toolkit.PollutedVarsPointsToAnalysis;
 import tabby.dal.caching.bean.ref.MethodReference;
+import tabby.util.TickTock;
 
 /**
  * @author wh1t3P1g
@@ -20,23 +21,28 @@
 @Setter
 public class CallGraphCollector {
 
-//    @Async("multiCallGraphCollector")
-    public void collect(MethodReference methodRef, DataContainer dataContainer){
+//    @Async("tabby-collector")
+    public void collect(MethodReference methodRef, DataContainer dataContainer, TickTock tickTock){
         try{
             SootMethod method = methodRef.getMethod();
-            if(method == null) return; // 提取不出内容，不分析
+            if(method == null) {
+                tickTock.countDown();
+                return; // 提取不出内容，不分析
+            }
 
             if(method.isPhantom() || methodRef.isSink()
                     || methodRef.isIgnore() || method.isAbstract()
                     || Modifier.isNative(method.getModifiers())){
                 methodRef.setInitialed(true);
+                tickTock.countDown();
                 return; // sink点为不动点，无需分析该函数内的调用情况  native/抽象函数没有具体的body
             }
 
             if(method.isStatic() && method.getParameterCount() == 0){
                 // 静态函数 且 函数入参数量为0 此类函数
                 // 对于反序列化来说 均不可控 不进行分析
                 methodRef.setInitialed(true);
+                tickTock.countDown();
                 return;
             }
 
@@ -49,10 +55,17 @@ public void collect(MethodReference methodRef, DataContainer dataContainer){
                             context, dataContainer,
                             method, methodRef);
             context.clear();
-
         }catch (RuntimeException e){
             e.printStackTrace();
+        }catch (Exception e){
+            if(e instanceof InterruptedException) {
+                log.error("Thread interrupted. " + methodRef.getSignature());
+            } else {
+                log.error("Something error on call graph. "+methodRef.getSignature());
+                e.printStackTrace();
+            }
         }
+        tickTock.countDown();
     }
 
 }
diff --git a/src/main/java/tabby/core/scanner/CallGraphScanner.java b/src/main/java/tabby/core/scanner/CallGraphScanner.java
@@ -4,10 +4,13 @@
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;
+import tabby.config.GlobalConfiguration;
+import tabby.core.collector.CallEdgeCollector;
 import tabby.core.collector.CallGraphCollector;
 import tabby.core.container.DataContainer;
 import tabby.dal.caching.bean.ref.MethodReference;
 import tabby.dal.caching.service.MethodRefService;
+import tabby.util.TickTock;
 
 import java.util.ArrayList;
 import java.util.Collection;
@@ -26,13 +29,10 @@ public class CallGraphScanner {
     public MethodRefService methodRefService;
     @Autowired
     public DataContainer dataContainer;
-
     @Autowired
-    public CallGraphCollector collector;
-
-    public static int total;
-    public static int split;
-    public static int current;
+    public CallGraphCollector callGraphCollector;
+    @Autowired
+    private CallEdgeCollector callEdgeCollector;
 
     public void run() {
         collect();
@@ -42,21 +42,16 @@ public void run() {
     public void collect() {
         Collection<MethodReference> targets =
                 new ArrayList<>(dataContainer.getSavedMethodRefs().values());
-//        log.info("Load necessary method refs.");
-//        dataContainer.loadNecessaryMethodRefs();
         log.info("Build call graph. START!");
-        total = targets.size();
-        split = total / 10;
-        split = split==0?1:split;
-        int count = 0;
+        TickTock tickTock = new TickTock(targets.size(), true);
         for (MethodReference target : targets) {
-            if(count%split == 0){
-                log.info("Status: {}%, Remain: {}", String.format("%.1f",count*0.1/total*1000), (total-count));
+            if(GlobalConfiguration.IS_FULL_CALL_GRAPH_CONSTRUCT){
+                callEdgeCollector.collect(target, dataContainer, tickTock);
+            }else{
+                callGraphCollector.collect(target, dataContainer, tickTock);
             }
-            collector.collect(target, dataContainer);
-            count++;
         }
-        log.info("Status: 100%, Remain: 0");
+        tickTock.await();
         log.info("Build call graph. DONE!");
     }
 

diff --git a/src/main/java/tabby/core/scanner/ClassInfoScanner.java b/src/main/java/tabby/core/scanner/ClassInfoScanner.java
@@ -59,14 +59,15 @@ public Map<String, CompletableFuture<ClassReference>> loadAndExtract(List<String
             moduleClasses = ModulePathSourceLocator.v().getClassUnderModulePath("jrt:/");
         }
         int size = targets.size();
-        int step = Math.max(size/10, size);
+        int step = Math.min(size/10, size);
         for (final String path : targets) {
             List<String> classes = getTargetClasses(path, moduleClasses);
             if(classes == null) continue;
 
             for (String cl : classes) {
                 try{
-                    SootClass theClass = SemanticHelper.loadClass(cl);
+//                    SootClass theClass = SemanticHelper.loadClass(cl);
+                    SootClass theClass = Scene.v().loadClassAndSupport(cl);
                     if (!theClass.isPhantom()) {
                         // 这里存在类数量不一致的情况，是因为存在重复的对象
                         results.put(cl, collector.collect(theClass));